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

LIANMING HUANG, Shangyu Wu, Yufei Cui, Ying Xiong, Haibo Hu, Xue Liu, Tei-Wei Kuo, Nan Guan, Chun Jason Xue

Keywords: Early Exit; Retrieval Augmentation; Large Language Model

TL;DR: This paper introduces RAEE, a retrieval-augmented early exit framework that accelerates inference while enhancing performance.

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.

Primary Area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Submission Number: 24600