RAEE: A Robust Retrieval-Augmented Early Exit Framework for Efficient Inference

TL;DR

RAEE introduces a retrieval-augmented framework for early exit in large language models, improving inference efficiency and robustness by leveraging similar data for better exit decisions without significant training overhead.

Contribution

The paper proposes RAEE, a novel retrieval-augmented early exit method that enhances model performance and robustness during inference without extensive retraining.

Findings

Accelerates inference across multiple tasks

Maintains robust zero-shot performance

Reduces training overhead for early exit methods

Abstract

Deploying large language model inference remains challenging due to their high computational overhead. Early exit optimizes model inference by adaptively reducing the number of inference layers. Current methods typically train internal classifiers or use heuristic methods to determine the exit layer. However, those methods either introduce significant training overheads or lead to performance degradation. To address these limitations, this paper proposes RAEE, a robust Retrieval-Augmented Early Exit framework that not only enables early exit but also enhances model performance through corrective exit information at intermediate layers. This paper first demonstrates that the early exit problem can be effectively modeled as a distribution prediction problem, in which the distribution can be further approximated through the exit information of similar data. Subsequently, this paper introduces the process of collecting exit information of correct predictions and the steps to construct the retrieval database. Finally, leveraging the pre-constructed retrieval database, RAEE utilizes the exit information from retrieved similar data to guide the backbone model's exit. Experimental results demonstrate that RAEE can not only accelerate inference while achieving robust zero-shot performance across eight downstream tasks.