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

## Abstract

Abstract Deploying large language model inference remains challenging due to their high computational overhead.
Early exiting optimizes model inference by adaptively reducing the number of inference layers.
Existing methods typically train internal classifiers to determine whether to exit at intermediate layers.
However, such classifier-based early exiting frameworks require significant effort to train the classifiers while can only achieve comparable performance at best.
To address these limitations, this paper proposes RAEE, a robust R etrieval- A ugmented E arly E xiting framework for efficient inference.
First, this paper demonstrates that the early exiting problem can be modeled as a distribution prediction problem, where the distribution is approximated using similar data’s exiting information.
Then, this paper details the process of collecting exiting information to build the retrieval database.
Finally, based on the pre-built retrieval database, RAEE leverages the retrieved similar data’s exiting information to guide the backbone model to exit at the layer, which is predicted by the approximated distribution.
Experimental results demonstrate that the proposed RAEE can significantly accelerate inference.
More importantly, RAEE can also achieve a robust zero-shot performance on 8 downstream tasks.

## 1 Introduction

Large language models have been widely used in various application scenarios due to their excellent performance (Thoppilan et al., ; Touvron et al., ; Scao et al., ).
However, deploying large language models on resource-constrained devices is still challenging due to the high computational overheads of performing model inference (Dao et al., ; Liu et al., ).
Model pruning, as an advanced technique, provides a new direction for efficient inference (Valicenti et al., ; Ma et al., ).
It selectively removes less important weights or connections from the neural network to reduce complexity and computational requirements without significantly degrading performance.
One popular model pruning method is the early exiting technique, which speeds up inference by adaptively reducing the number of inference layers.

Most early exiting frameworks (Liu et al., ; Zhu, ; Fan et al., ) leverage classifiers to predict the exiting layer and then stop the inference at the predicted exiting layer.
Those classifier-based early exiting frameworks can be categorized into three types according to the training strategies.
One type is the training-based early exiting (Zhu, ; Zhou et al., ; Zhu et al., ), which requires training the classifiers along with the inference model.
These methods introduce significant training and fine-tuning overheads, particularly when applied to large language models.
Another branch of classifier-based early exiting can be concluded as semi-training-based early exiting (Fan et al., ).
The backbone models in this type of early exiting framework would not be updated, and only classifiers would be fitted to predict the exiting layer.
These methods may not capture the patterns between inputs and exiting layers well, requiring significant human effort in feature engineering.
The last one is training-free early exiting (Sun et al., ), which requires no parameter updates and uses heuristics to determine the exiting layer.
These methods lack the generalization ability in predicting the exiting layer and often fail to achieve optimal or sub-optimal solutions.
Moreover, most existing early exiting frameworks sacrifice the model performance for acceleration (Fan et al., ; Sun et al., ; Schuster et al., ; Bae et al., ).

To address the above limitations, this paper first shows that the exiting layer predictions can be solved by predicting from an exiting distribution.
Then, this paper presents the observations that similar data’s exiting information can be used to approximate the exiting distribution.
Based on these observations, this paper proposes the RAEE, a robust retrieval-augmented early exiting framework for efficient inference.
RAEE collects exiting information from training data and builds the indexing to retrieve similar data’s exiting information.
During the inference, RAEE predicts the exiting layer based on the top-k nearest neighbors’ exiting information and stops the model forwarding at the predicted exiting layer.

We conduct comprehensive experiments to evaluate the proposed RAEE and various comparison methods on 8 downstream tasks.
Experimental results demonstrate that RAEE can accelerate the model inference while achieving robust model performance.
Codes are available at (^1^11https://anonymous.4open.science/r/RAEE-D724).

The main contributions of this paper are:

- •
We model the early exiting problem as a distribution prediction problem and demonstrate that the exiting distribution can be approximated by the exiting information of similar data.
- •
We propose a robust retrieval-augmented early exiting framework, named RAEE, which leverages the external database to guide the early exiting.
- •
Experimental results show that the proposed RAEE can not only accelerate the model inference but also achieve a robust performance.

Figure: Figure 1: The overview of retrieval-augmented early exiting framework.
Refer to caption: x1.png

## 2 Motivations

In this section, we demonstrate why using retrieval-based techniques is a simple yet effective way to augment the early exiting framework during the inference stage.

Problem Statement.
Formally, early exiting can be defined as follows: Given a backbone model $\mathcal{M}$ with $m$ layers and an input $x$,
The early exiting framework aims to design an exiting function or classifier $l=f(x)$ to determine whether to exit at the layer $l$.
The final prediction $y$ is then transformed from the intermediate output states $h_{l}$ of the $l$-th layer.
And the final prediction probability can be formulated as,

$$ $P(y\;|\;x)=P(y\;|\;h_{f(x)}).$ (1) $$

, where $f(x)$ is trained or built on the downstream tasks’ training data $\mathcal{D}$.

Limitations of Existing Early Exiting.
Existing early exiting frameworks can be categorized into training-based, semi-training-based, and training-free exiting methods.
For training-based exiting, most works jointly optimize the exiting classifiers and the backbone model (Zhu et al., , ; Zhu, ).
These works fine-tune the pre-trained backbone model with exiting mechanisms over downstream tasks’ training data and update both the backbone model and classifier parameters.
Although these approaches could adapt the model with the early exiting framework and accelerate the inference while maintaining performance, the fine-tuning overhead is still a non-negligible amount, especially when fine-tuning modern large language models.

Instead of training both classifiers and the backbone model, other works (Fan et al., ) only fit the classifiers based on features extracted from the backbone model, with parameters in the backbone model fixed.
These works can significantly reduce the training overhead of early exiting framework, as they aim to address the early exiting problem using efficient traditional machine learning approaches for classifier fitting, e.g., techniques for fitting support vector machine.
However, such methods that only train classifiers may not effectively capture the patterns between inputs and exiting layers.
The reasons may lie in that
1) the classifiers only focus on distinguishing the extract features at each layer, rather than recognize the patterns between input embeddings and exiting layers;
2) the feature extraction requires human effort to design, where classifiers only indirectly learn where to exit.

Training-free early exiting frameworks refer to computing the exiting layer with heuristics, e.g., hashing functions (Sun et al., ).
These techniques basically introduce no or few training overhead and are easy to set up.
Although these works can significantly accelerate the building process as well as the classifying process, they lack generalization in exiting layer predictions.
The predictions made by these heuristics may also be sensitive to different inputs, leading to unstable performance of the early exiting framework.

Figure: Figure 2: The probability of different exiting layers on SST-2 test data and corresponding top-8 nearest neighbor SST-2 training data.
Refer to caption: x2.png

Motivations of Retrieval-Augmented Early Exiting.
To address the above limitations, this paper aims to leverage retrieval-based techniques to guide the early exiting, which requires no parameters to update and has a generalization capability.
Existing retrieval databases basically use clustering and product quantization to build the retrieval indexing over millions of embeddings, which can perform efficient approximate nearest neighbor searching.
The building process of the indexing is not resource-constrained, which can run on either GPUs or CPUs.
Besides, since the retrieval database stores the original data, every retrieval in the retrieval-augmented early exiting can be regarded as a generalization over several semantically similar data, while those original data may be used to train classifiers or build hashing functions in other early exiting frameworks.
The retrieval-augmented techniques also exhibit strong adaptability to new data.
By simply adding new data to the retrieval database, the index can retrieve up-to-date information.

To further demonstrate the efficacy of retrieval-augmented early exiting, this paper conducts some analysis experiments in Figure 2 to show the probability of different exiting layers on two SST-2 test data and corresponding top-8 nearest neighbor SST-2 training data.
The probability in Figure 2 is calculated as the normalized logits of the answer label token.
Note that it is impossible to obtain the probability of each exiting layer during the real inference as no labels are provided.
The blue line with dots in Figure 2 shows the probability of correct predictions exiting at each layer.
For example, in Figure 2 (a), when exiting at layer 18, the probability of correct predictions is about 0.68, while exiting at layer 9 only has a probability of 0.01.
The box plot in Figure 2 presents the probability statistics of the top-8 nearest neighbor training data, indicating the probability of exiting at each layer.
For instance, when exiting at layer 18, the maximum exit probability is approximately 0.81, the minimum is around 0.39, and the box plot indicates that half of the neighbors exhibit a probability larger than 0.54.

The blue line in Figure 2 represents the target distribution for the classifier to learn.
The experimental results in Figure 2 (a) and (b) reveal that the exiting probabilities exhibit a similar pattern to those of the nearest neighbors.
Additionally, the exiting layer varies across different inputs.
Consequently, the experimental results lead to two key conclusions:

- 1.
The exiting probability can be approximated by the probability of the top-k nearest neighbors.
- 2.
Different inputs exhibit distinct probabilities for exiting layers.

These observations motivate us to propose a retrieval-augmented early exiting framework.

## 3 Methodology

In this section, this paper describes the proposed retrieval-augmented early exiting framework in detail.
First, this paper introduces the details of how to build the retrieval database for the early exiting.
Then, this paper presents the retrieval-augmented early exiting baseline called RAEE.

### 3.1 Building the Retrieval Database for Early Exiting

The critical factor in constructing the retrieval database is selecting appropriate keys and values.
To avoid introducing too much retrieving overheads, this paper only retrieves once at the beginning of the backbone model.
Consider the training data $\mathcal{D}=\{(x^{train}_{1},y^{train}_{1}),\ldots,(x^{train}_{|\mathcal{D}|},
y^{train}_{|\mathcal{D}|})\}$ and a backbone model $\mathcal{M}$ with $m$ layers $\{\mathcal{L}_{1},\ldots,\mathcal{L}_{m}\}$. In this context, as shown in the top part of Figure 1, the keys $\mathcal{K}$ are input embeddings of training data, which can be obtained from an extra encoder model $\mathcal{E}$ , such as BERT (Devlin et al., ), or the outputs of embedding layers in the backbone model $\mathcal{M}_{emb}$,

$$ $\mathcal{K}=\{e_{i}\}^{|\mathcal{D}|}_{i=1}=\{\mathcal{E}(x^{train}_{i})\}^{| \mathcal{D}|}_{i=1}.$ (2) $$

For the values, this paper collects a set of possible exiting layers $l_{i}$ and corresponding probabilities $p_{i}$ for each embedding $e_{i}$, i.e., $v_{i}=\{(l^{j}_{i},p^{j}_{i})\}^{m_{i}}_{j=1}$, where $m_{i}$ indicates the number of possible exiting layers for the embedding $e_{i}$.
The layer $l$ chosen as the exiting layer is determined by whether the outputs of this layer $h_{l}$ can be used to make the right predictions $\hat{y}$ compared to the training labels $y^{train}$.
Then, the values $\mathcal{V}$ are all sets of possible exiting layers,

$$ $\mathcal{V}=\{v_{i}\}^{|D|}_{i=1}=\left\{\{(l^{j}_{i},p^{j}_{i})\}^{m_{i}}_{j= 1}\right\}^{|D|}_{i=1}.$ (3) $$

We follow the same dataset splitting used in the LM-BFF (Gao et al., ), the collecting process requires no parameters to update, only model inference is performed.
After collecting keys and values for the retrieval databases, this paper uses state-of-the-art approximate nearest neighbor search indexing, such as FAISS (Johnson et al., ), and efficient key-value stores to build the retrieval database.
More details can be found in Algorithm 1 of Appendix B.

### 3.2 RAEE

In this section, this paper then presents a retrieval-augmented early exiting framework named RAEE to optimize the model inference.
RAEE regards the exiting layer as a random variable $z$, taking values in the set of $\{1,\ldots,m\}$, where $m$ is the total number of layers in the backbone model $\mathcal{M}$.
The probability mass function $P(z=l)$ represents the probability of the case that the backbone model exits at the layer $l$.
With the gold label, we can observe that the random variable $z$ follows an unknown discrete distribution $F$.
Then, this paper shows how to leverage the retrieval database to approximate the distribution $F$.

Given an input $x$, RAEE first retrieves top-$k$ nearest neighbors $\{v_{1},\ldots,v_{k}\}$, where each neighbor $v_{i}$ has $m_{i}$ possible exiting layers.
Naturally, we can approximate the distribution $F$ by estimating the probability function $P(z=l)$,

$$ $P(z=l\;|\;x)=\sum^{k}_{i=1}P(v_{i}\;|\;x)\cdot\sum^{m_{i}}_{j}\mathbbm{1}\left (any(l^{j}_{i}=l\;\&\&\;p^{j}_{i}\geq\tau)\right)\cdot p^{j}_{i}$ (4) $$

, where $\mathbbm{1}$ is the indicator function that returns 1 if the condition is true and 0 otherwise, $any(\cdot)$ is the function that returns true if one condition is true and false otherwise, $\tau$ is the threshold for filtering the layers with extremely low probability, the inner loop only count once since there is at most one possible exiting layer of neighbor $i$ that is equalt to $l$.
Since different neighbors should have different contributions to the probability function $P(z=l)$, RAEE uses the reciprocal of the scaled distance between each neighbor and the query to estimate the contribution,

$$ $P(v_{i}\;|\;x)=\frac{\min\left(\{distance(v_{j},x)\}^{k}_{j=1}\right)}{ distance(v_{i},x)}$ (5) $$

Then, RAEE designs a function $f(x)$ to determine the exiting layer, which selects the layer that maximizes the probability function $P(z=l)$,

$$ $f(x)=\arg\max_{l}P(z=l\;|\;x)$ (6) $$

Notably, when there are multiple exiting layers with the same maximal probability, RAEE chooses the earliest exiting layer.

The bottom part of Figure 1 shows the inference workflow of RAEE.
Specifically, RAEE first simultaneously feeds the inputs into both the backbone model for the label predictions and the same encoder used in the building process for the query embeddings.
Then, the retriever in RAEE retrieves the top-$k$ nearest neighbors in the retrieval databases based on the query embeddings.
After obtaining all possible exiting layers of $k$ nearest neighbors, RAEE computes the exiting layers based on the Equations 4-6.
Finally, RAEE stops the forwarding at the calculated exiting layer, and passes the intermediate outputs of the exiting layer to the final prediction layer, e.g., LM Head in language models, to obtain the final predictions (Equation 1).
More details can be found in Algorithm 2 of Appendix B.

**Table 1: Zero-shot performance of different methods on 8 downstream tasks. Methods with an asterisk superscript ‘*’ are evaluated on NVIDIA A100 GPU. The results in bold are the best results on the task. ‘D-BERT’ and ‘D-RoBERTa’ refer to DistilBERT and DistilRoBERTa, respectively.**
| Methods | SST-2 | SST-5 | MR | CR | MPQA | SUBJ | TREC | CoLA | Avg |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Pretrained Models |  |  |  |  |  |  |  |  |  |
| RoBERTa-Large | 83.60 | 34.98 | 80.80 | 79.55 | 67.60 | 51.45 | 32.40 | 2.03 | 54.05 |
| ElasticBERT | 51.15 | 27.78 | 50.10 | 50.00 | 47.05 | 50.00 | 18.60 | 6.02 | 37.59 |
| T5-Large^* | 49.31 | 23.12 | 50.40 | 50.90 | 45.40 | 52.75 | 27.60 | -4.64 | 36.86 |
| Llama-3-8B^* | 62.84 | 26.06 | 59.65 | 72.90 | 51.75 | 52.80 | 8.40 | 0.00 | 41.80 |
| Gemma-7B^* | 49.08 | 28.64 | 50.05 | 50.10 | 50.00 | 48.05 | 18.00 | -0.79 | 36.64 |
| Static Models |  |  |  |  |  |  |  |  |  |
| D-BERT | 71.90 | 25.61 | 67.65 | 77.10 | 64.85 | 70.35 | 28.40 | -0.70 | 50.65 |
| D-RoBERTa | 81.77 | 28.37 | 77.95 | 84.20 | 69.50 | 54.35 | 32.00 | -1.18 | 53.37 |
| Dynamic Models |  |  |  |  |  |  |  |  |  |
| HashEE | 52.75 | 18.64 | 48.95 | 52.50 | 55.85 | 45.35 | 23.80 | -3.86 | 36.75 |
| AdaInfer | 50.92 | 34.98 | 50.45 | 49.60 | 60.90 | 50.65 | 32.40 | -1.62 | 41.04 |
| DeeBERT^* | 52.29 | 18.05 | 50.60 | 50.00 | 75.95 | 80.85 | 16.20 | 0.00 | 42.99 |
| CALM^* | 51.72 | 23.17 | 49.25 | 50.55 | 49.80 | 49.90 | 18.00 | 0.00 | 36.55 |
| SLEB^* | 51.38 | 16.61 | 50.05 | 50.65 | 50.10 | 47.55 | 25.40 | 0.00 | 36.47 |
| RAEE-RoBERTa | 84.06 | 33.44 | 81.75 | 67.90 | 79.25 | 83.60 | 60.60 | 15.55 | 63.27 |
| RAEE-T5^* | 53.21 | 27.19 | 50.60 | 51.55 | 55.90 | 49.90 | 39.80 | 12.55 | 42.59 |
| RAEE-Llama^* | 75.57 | 32.04 | 67.90 | 72.75 | 76.10 | 89.95 | 46.00 | 7.84 | 58.52 |
| RAEE-Gemma^* | 73.17 | 32.35 | 67.35 | 56.85 | 75.80 | 89.90 | 28.40 | 11.43 | 54.41 |

## 4 Experiments

In this section, this paper first introduces the dataset and the experimental setting.
Then, this paper presents the main results of 8 downstream tasks.
This paper also conducts an ablation study and analysis of RAEE to show the impact of these factors on model performance.

### 4.1 Dataset and Experimental Setup

Datasets We conduct comprehensive experiments across 8 downstream tasks from GLUE benchmark (Wang et al., ).
These tasks cover sentiment analysis, opinion polarity analysis, grammatical judgment, natural language inference, paraphrasing, etc.

Experimental Settings
The proposed RAEE was implemented using the PyTorch framework and Transformer.
We evaluated RoBERTa, ElasticBERT (Liu et al., ), D-BERT (Sanh et al., ), D-RoBERTa (Sanh et al., ), HashEE (Sun et al., ), AdaInfer (Fan et al., ), and RAEE-RoBERTa on one NVIDIA Quadro RTX5000 GPU with 16GB GPU memory, while others on one NVIDIA A100 GPU with 80GB GPU memory.
The experiments were conducted in the zero-shot setting. For dynamic models with classifiers, we only train the classifiers on downstream tasks.
The evaluation metric is accuracy, except for CoLA, which is measured by the Matthew correlation coefficient.
We evaluate RAEE on various backbone models, such as RoBERTa-large, T5-Large (Raffel et al., ), Llama-3-8B (Dubey et al., ), and Gemma-7B (Team et al., ).
The number of retrieved nearest neighbors is set to 12 in this paper.

To validate the effectiveness, we compared RAEE with three types of methods.
Pretrained Models:
1) RoBERTa-Large (Gao et al., ), a state-of-the-art encoder model, where the prompt-based version is used;
2) ElasticBERT (Liu et al., ), a pre-trained multi-exit transformer model, where the large version is used in this paper;
3) T5-Large (Raffel et al., ), a versatile transformer-based model for various NLP tasks;
4) Llama-3-8B (Dubey et al., ), a pre-trained model with strength in specific language scenarios;
5) Gemma-7B (Team et al., ), a model with potential for outstanding performance in specific settings.
Static Models:
1) DistilBERT (Sanh et al., ), a distilled version of BERT-base model;
2) DistilRoBERTa (Sanh et al., ), a distilled version of the RoBERTa-base model.
Dynamic Models:
1) HashEE (Sun et al., ), a hash-based early exiting approach with ElasticBERT-large as its backbone model;
2) AdaInfer (Fan et al., ), an SVM-based early exiting method with our reproduced version on RoBERTa-large;
3) DeeBERT (Xin et al., ), a classical entropy-thresholding-based early exiting method with RoBERTa-Large as its backbone model;
4) CALM (Schuster et al., ), a classical entropy-thresholding-based early exiting method with T5-Large (Raffel et al., ) as its backbone model;
5) SLEB (Song et al., ), a method that tackles the limitation of early exiting methods by eliminating redundant transformer blocks with Llama-3-8b (Dubey et al., ) as its backbone model.
The templates are listed in the Appendix A.
More details about the experimental setup can be found in Appendix E.

### 4.2 Main Results

Table 1 presents the main results, comparing the performance of the RAEE method against three other method types across eight downstream tasks.
From the experimental results, the proposed RAEE can achieve the best zero-shot performance on average.
Compared to the backbone model, such as Roberta-Large, T5-Large, Llama-3-8B, and Gemma-7B without early exiting, RAEE outperforms them on most tasks.
Besides, RAEE also significantly outperforms existing early exiting frameworks on all tasks.
The above results demonstrate the robustness of the proposed RAEE by consistently achieving strong performance across various tasks and backbone models.

Figure: (a) SST-2
Refer to caption: x4.png

Figure 3 shows the inference latency of RAEE and comparisons.
Due to space limits, we only show the inference latency on the first four tasks, and more results can be found in Appendix C.
Experimental results show that the proposed RAEE can consistently accelerate the model inference over various backbone models.
Although static models can achieve a fast inference speed due to the small model architecture, those models require great efforts to fine-tune or distill from backbone models.
The exceptionally long inference latency of current dynamic models is primarily due to the internal classifiers’ poor performance in zero-shot scenarios. The inadequate performance of these classifiers leads to late exits, which in turn causes high inference overheads.

### 4.3 Reasons for Significant Performance Improvement

The main results demonstrate that the proposed RAEE can significantly outperform backbone models, which is not an easy-understanding and intuitive phenomenon compared to previous early exiting methods.
The reason of such phenomenon lies in that the retrieval database in RAEE also acts as an error corrector, which contains the exiting information of examples that are correctly predicted by intermediate layers, but backbone models without early exiting fail to predict.

**Table 2: Model performance and inference latency of different methods with RoBERTa-Large.**
| Models | SST-2 | SST-5 | MR | CR | MPQA | SubJ | TREC | CoLA | Avg |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Performance $\uparrow$ |  |  |  |  |  |  |  |  |  |
| RoBERTa-Large | 83.60 | 34.98 | 80.80 | 79.55 | 67.60 | 51.45 | 32.40 | 2.03 | 54.05 |
| RAEE w/o | 85.32 | 34.75 | 81.05 | 71.55 | 75.80 | 63.90 | 42.60 | -2.10 | 56.61 |
| RAEE | 85.21 | 33.44 | 81.75 | 67.70 | 78.90 | 83.80 | 62.60 | 13.75 | 63.39 |
| Latency (ms) $\downarrow$ |  |  |  |  |  |  |  |  |  |
| RoBERTa-Large | 49.40 | 49.63 | 49.63 | 49.39 | 49.02 | 49.44 | 49.65 | 49.11 | 49.41 |
| RAEE w/o | 48.18 | 52.07 | 47.87 | 41.73 | 45.97 | 42.66 | 43.50 | 46.28 | 46.03 |
| RAEE | 47.27 | 36.10 | 45.82 | 38.49 | 41.84 | 33.76 | 34.47 | 31.74 | 38.69 |

To better support the above claims, we also conducted an analysis experiment using the retrieval database, which only contained exit information based on examples that backbone models correctly predicted without early exiting.
As shown in Table 2, RAEE w/o refers to the method built on only correctly predicted examples.
As expected, RAEE w/o achieves comparable performance to baselines but accelerates the inference process.
This is because the test data that is correctly predicted by RAEE w/o can also be correctly predicted by backbone models.
However, due to a lack of exiting information on examples where backbone models fail to predict, RAEE w/o also fails to predict on the test data where backbone models fail.
Therefore, when providing the exiting information based on examples where backbone models fail to predict but intermediate outputs succeed in predicting, RAEE can make correct predictions and exit earlier.

### 4.4 Ablation Study of Top- k

Table 3 illustrates the impact of varying the number of retrievals on the distribution approximation.
As $k$ increases, the proposed RAEE-RoBERTa improves the overall performance from 61.22 to 63.27.
This suggests that more retrieved exiting information can help enhance the approximation performance.
However, when $k$ beyond 12, the overall performance degrades from 63.27 to 62.83.
The reasons may lie in that providing the exiting information of the retrievals that are not quite related to the query would introduce noise, thus misleading the final predictions.
This also implies that only a few exiting information is enough to approximate the exiting distribution, and the time cost of the retrieving process can be saved.

**Table 3: The impact of retrieval number $k$ on the distribution approximation.**
| $k$ | SST-2 | SST-5 | MR | CR | MPQA | SUBJ | TREC | CoLA | Avg |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 2 | 78.21 | 32.94 | 77.35 | 68.50 | 78.25 | 80.05 | 60.20 | 14.26 | 61.22 |
| 4 | 81.08 | 33.48 | 79.40 | 67.90 | 79.45 | 82.95 | 61.20 | 14.78 | 62.53 |
| 8 | 83.60 | 34.12 | 81.00 | 68.05 | 79.75 | 83.65 | 61.20 | 12.28 | 62.96 |
| 12 | 84.06 | 33.44 | 81.75 | 67.90 | 79.25 | 83.60 | 60.60 | 15.55 | 63.27 |
| 16 | 83.72 | 32.71 | 82.05 | 69.10 | 78.70 | 83.35 | 61.20 | 15.24 | 63.26 |
| 20 | 83.14 | 33.53 | 81.85 | 69.70 | 78.45 | 83.05 | 61.00 | 11.94 | 62.83 |

### 4.5 Ablation Study on the Retrieval Database Size

Table 4 shows the performance of RAEE-RoBERTa with different sizes of retrieval databases.
The size of the retrieval databases implies how similar embeddings would be retrieved, thus impacting the confidence of the provided exiting information.
As the database size increases, the performance of RAEE-RoBERTa increases significantly from 60.85 to 63.27 on average.
This demonstrates that collecting more data can improve the generalization of RAEE, thus approximating the exiting distribution more accurately.

**Table 4: The impact of retrieval database size on the distribution approximation. The percentage refers to the amount of training data that is used to build the retrieval database.**
| Database Size | SST-2 | SST-5 | MR | CR | MPQA | SUBJ | TREC | CoLA | Avg |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 20% | 83.14 | 32.35 | 81.40 | 67.80 | 75.55 | 77.85 | 57.20 | 11.47 | 60.85 |
| 50% | 82.45 | 32.31 | 81.30 | 65.35 | 76.65 | 82.65 | 58.20 | 16.83 | 61.97 |
| 100% | 84.06 | 33.44 | 81.75 | 67.90 | 79.25 | 83.60 | 60.60 | 15.55 | 63.27 |

## 5 Conclusion

This paper models the early exiting problem as a distribution approximation problem and observes that similar data’s exiting information can be used to approximate.
Then, this paper proposes a retrieval-augmented early exiting framework named RAEE.
Experimental results show that RAEE can accelerate the model inference while significantly improving the model performance.
For future work, exploring the performance of applying the retrieval-augmented early exiting framework in fine-tuning scenarios is worthwhile.

## Appendix A Templates on All Tasks

Table 5 provides an overview of the manual templates and selected label words used for each dataset in this paper. These templates and label words were created following LM-BFF (Gao et al., ).

**Table 5: Templates and label words used in this paper.**
| Task | Prompts | Label word |
| --- | --- | --- |
| SST-2 | [CLS] $x$ It was [MASK]. [SEP] | “0”:“terrible”, “1”:“great” |
| SST-5 | [CLS] $x$ It was [MASK]. [SEP] | “0”:“terrible”,“1”: “bad”, |
|  |  | “2”: “okay”,“3”: “good”,“4”: “great” |
| MR | [CLS] $x$ It was [MASK]. [SEP] | “0”:“terrible”, “1”:“great” |
| CR | [CLS] $x$ It was [MASK]. [SEP] | “0”:“terrible”, “1”:“great” |
| MPQA | [CLS] $x$ It was [MASK]. [SEP] | “0”:“terrible”, “1”:“great” |
| SUBJ | [CLS] $x$ This is [MASK]. [SEP] | “0”:“subjective”, “1”:“objective” |
| TREC | [CLS] [MASK] $x$ [SEP] | “0”:“Description”,“1”:“Entity”,“2”:“Expression”, |
|  |  | “3”:“Human”,“4”:“Location”,“5”:“Number” |
| CoLA | [CLS] $x$ It was [MASK]. [SEP] | “0”:“incorrect”, “1”:“correct” |

## Appendix B Detailed Algorithms of RAEE

Figure: Algorithm 1 Collect the keys and values for building the retrieval database.

Algorithm 1 collects keys and values for building the retrieval database.
For each sample in training data $\mathcal{D}$, the backbone model $\mathcal{M}$ is traversed layer by layer to compute the hidden state $h_{j}$ and corresponding $logits$.
If a prediction $\hat{y}$ at a certain layer $j$ matches the sample’s true label $y^{train}_{i}$, the exiting information, including the layer $j$ and the probability $p^{j}_{i}$, is added to the sample’s value list (Line 2-12).
When the encoder $\mathcal{E}$ is unavailable (Line 16), RAEE utilizes the hidden states from the backbone model $\mathcal{M}$ as embeddings for indexing.
The specific layer from which the hidden states are extracted is treated as a hyperparameter that the user can define.

Figure: Algorithm 2 Model inference with the synchronized retrieval-augmented early exiting.

Algorithm 2 performs model inference with retrieval-augmented early exiting.
When the encoder $\mathcal{E}$ is unavailable (Line 5), RAEE utilizes the hidden states from the backbone model $\mathcal{M}$ as embeddings for querying. The specific layer from which the hidden states are extracted is treated as a hyperparameter that the user can define.
The inference process (Lines 9-11) and the retrieving process (Lines 2-8) can be executed in parallel for more efficient implementations.
Specifically, we can use two threads to execute each process.
During the inference process with asynchronized RAEE, the calculated intermediate results would be collected for early exiting.

Figure: (a) MPQA
Refer to caption: x8.png

## Appendix C Inference Latency

Figure 4 shows the inference latency of different methods of the last four downstream tasks.
Experimental results demonstrate the consistent performance of the proposed RAEE compared to Figure 3.

## Appendix D Exiting Layers

Table 6 compares the average exiting layers of the RAEE method against two other method types across eight downstream tasks.
Experimental results show that the RAEE method can exit earlier than the two other method types, thus reducing computational overhead during model inference for better model efficiency.
This result also aligns with the expectations in the motivation example.
This suggests that the RAEE method can accurately approximate the gold exiting layer distribution by using the retrieval database.
Although AdaInfer exits earlier than the RAEE method, it exhibits quite poor performance, as shown in Table 1.
The reason may be that during the zero-shot inference scenario, the collected features can only provide limited information for the SVM, thus resulting in unstable prediction performance.

**Table 6: Exiting layers of different methods with RoBERTa-Large, T5-Large, Llama-3-8B, and Gemma-7B on 8 downstream tasks. The sum of the number of layers in the encoder and the decoder counts the number of layers for T5-large (Raffel et al., ). Due to the strong code coupling, it is hard to collect the exiting layer information of CALM(Schuster et al., ). Methods with an asterisk superscript ‘*’ are evaluated on NVIDIA A100 GPU with 80G GPU memory, while others are evaluated on NVIDIA RTX 5000.**
| Model | SST-2 $\downarrow$ | SST-5 $\downarrow$ | MR $\downarrow$ | CR $\downarrow$ | MPQA $\downarrow$ | Subj $\downarrow$ | TREC $\downarrow$ | CoLA $\downarrow$ | Avg $\downarrow$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Pretrained Models |  |  |  |  |  |  |  |  |  |
| RoBERTa-Large^* | 24.00 | 24.00 | 24.00 | 24.00 | 24.00 | 24.00 | 24.00 | 24.00 | 24.00 |
| T5-Large^* | 48.00 | 48.00 | 48.00 | 48.00 | 48.00 | 48.00 | 48.00 | 48.00 | 48.00 |
| Llama-3-8B^* | 32.00 | 32.00 | 32.00 | 32.00 | 32.00 | 32.00 | 32.00 | 32.00 | 32.00 |
| Gemma-7B^* | 28.00 | 28.00 | 28.00 | 28.00 | 28.00 | 28.00 | 28.00 | 28.00 | 28.00 |
| Dynamic Models |  |  |  |  |  |  |  |  |  |
| AdaInfer | 1.00 | 24.00 | 0.00 | 0.47 | 18.00 | 0.08 | 24.00 | 4.00 | 8.94 |
| DeeBERT^* | 22.95 | 24.00 | 23.33 | 8.98 | 15.90 | 10.36 | 24.00 | 18.31 | 18.48 |
| SLEB^* | 28.00 | 28.00 | 28.00 | 28.00 | 28.00 | 28.00 | 28.00 | 28.00 | 28.00 |
| RAEE-RoBERTa | 18.61 | 15.87 | 18.69 | 15.38 | 17.32 | 14.00 | 15.25 | 12.70 | 15.98 |
| RAEE-T5^* | 22.05 | 18.44 | 21.74 | 26.91 | 17.77 | 18.71 | 27.34 | 18.72 | 21.46 |
| RAEE-Llama^* | 29.86 | 27.08 | 29.72 | 31.16 | 28.82 | 21.45 | 24.24 | 30.97 | 27.91 |
| RAEE-Gemma^* | 10.97 | 17.72 | 11.86 | 3.27 | 14.75 | 0.52 | 12.97 | 20.08 | 11.52 |

## Appendix E Implementation Details

This section lists the implementation details.

- •
For DeeBERT(Xin et al., ), we use RoBERTa-Large as its backbone model.
Since DeeBERT(Xin et al., ) is a classical entropy-thresholding-based early-exit method, it requires first fine-tuning the backbone model on the downstream task and then updating all but the last off-ramp, for a fair comparison, we only update the off-ramp in DeeBERT on each downstream task.
We also use RoBERTa-large as the backbone model and train all off-ramps for 50 epochs (much larger than the default setting of 10 epochs).
Other experimental settings for DeeBERT(Xin et al., ) remain as default.
- •
For CALM (Schuster et al., ), we use T5-Large (Raffel et al., ) as its backbone model.
CALM (Schuster et al., ) is also a classical entropy-thresholding-based early-exit method, and we evaluate it under the zero-shot setting.
- •
For SLEB(Song et al., ), we use Llama-3-8b (Dubey et al., ) as its backbone model.
SLEB(Song et al., ) tackles the limitation of early exit methods by eliminating redundant transformer blocks.
Since the proposed RAEE exits at 27.91 layers, for a fair comparison, we also set the hyper-parameter num_remove_blocks of SLEB(Song et al., ) as 4 for comparable efficiency.

## Appendix F Cost of Building the Retrieval Database

We collect the statistics of building the retrieval database in Table  7, such as database size, index size, and time cost for the retrieval database for RoBERTa-Large.
The average of building the retrieval database is less than 8 minutes on RTX 5000, an acceptable overhead compared to the time cost of fine-tuning.
The index size and database size are quite small, which can be ignored compared to the backbone model size.

**Table 7: The cost of building the retrieval database.**
| Model | SST-2 | SST-5 | MR | CR | MPQA | Subj | TREC | CoLA | Avg |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Building Time (second) | 430.81 | 514.46 | 545.90 | 151.97 | 531.17 | 495.42 | 349.59 | 526.55 | 443.23 |
| # Entries | 6920 | 8544 | 8662 | 1775 | 8606 | 8000 | 5452 | 8551 | 7063.75 |
| Index Size (MB) | 3.5 | 3.9 | 3.9 | 2.0 | 3.9 | 3.8 | 3.2 | 3.9 | 3.5 |
| Database Size (MB) | 2.6 | 2.1 | 3.1 | 0.8 | 2.6 | 2.7 | 1.0 | 2.7 | 2.2 |

## Appendix G Comparisons to Fine-tuning-based Methods

We have evaluated CALM(Schuster et al., ) and FREE (Bae et al., ) on the given 8 downstream tasks in table 8 and table 9.
We finetune them with the default setting on each task for one epoch. Experiments are conducted on 1 NVIDIA A100 GPU.
Since the experimental results of CALM(Schuster et al., ) and FREE(Bae et al., ) under the zero-shot setting are the same, we only report the zero-shot results of CALM(Schuster et al., ). Although the fine-tuned CALM(Schuster et al., ) and FREE(Bae et al., ) can achieve much better performance, they require large computational resources and time for the fine-tuning process, which is NOT the main research goal of this paper.
Under the zero-shot setting, the proposed RAEE can still outperform CALM(Schuster et al., ) or FREE(Bae et al., ).

**Table 8: Model performance of different methods with T5-Large under different settings on 8 downstream tasks.**
| Model | SST-2 $\uparrow$ | SST-5 $\uparrow$ | MR $\uparrow$ | CR $\uparrow$ | MPQA $\uparrow$ | Subj $\uparrow$ | TREC $\uparrow$ | CoLA $\uparrow$ | Avg $\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Zero-Shot |  |  |  |  |  |  |  |  |  |
| RAEE-T5 | 53.21 | 27.19 | 50.60 | 51.55 | 55.90 | 49.90 | 39.80 | 12.55 | 42.59 |
| CALM | 51.72 | 23.17 | 49.25 | 50.55 | 49.80 | 49.90 | 18.00 | 0.00 | 36.55 |
| Fine-Tuning |  |  |  |  |  |  |  |  |  |
| CALM | 88.76 | 28.96 | 86.90 | 49.25 | 85.05 | 82.85 | 71.40 | -2.61 | 61.32 |
| FREE | 87.96 | 42.81 | 85.30 | 7.40 | 85.80 | 78.25 | 64.80 | 0.04 | 56.55 |

**Table 9: Inference latency(ms) of different methods with T5-Large under different settings on 8 downstream tasks.**
| Model | SST-2 $\downarrow$ | SST-5 $\downarrow$ | MR $\downarrow$ | CR $\downarrow$ | MPQA $\downarrow$ | Subj $\downarrow$ | TREC $\downarrow$ | CoLA $\downarrow$ | Avg $\downarrow$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Zero-Shot |  |  |  |  |  |  |  |  |  |
| RAEE-T5 | 14.21 | 11.56 | 11.98 | 10.33 | 9.93 | 13.60 | 18.60 | 15.93 | 13.27 |
| CALM | 74.29 | 73.87 | 71.16 | 74.57 | 116.15 | 87.27 | 69.49 | 65.60 | 79.05 |
| Fine-Tuning |  |  |  |  |  |  |  |  |  |
| CALM | 99.72 | 100.39 | 101.45 | 90.98 | 102.97 | 108.79 | 84.58 | 87.14 | 97.00 |
| FREE | 85.69 | 95.79 | 95.24 | 95.21 | 94.87 | 90.91 | 92.11 | 89.45 | 92.41 |

## Appendix H Related Work

### H.1 Early Exiting Framework

Model inference with early exit has been a popular pruning method to reduce both computation and memory overhead on text classification or generation tasks.
Most existing works (Bae et al., ; Kong et al., ; Ji et al., ; Wolczyk et al., ; Hooper et al., ) introduce classifiers in each layer to determine whether the inference should continue.
(Xin et al., ; Liu et al., ; Zhou et al., ; He et al., ) train the classifier by minimizing the differences between each layer’s outputs and final outputs, then perform early-exit-based inference according to a threshold.
(Liao et al., ) incorporates the past and future information to predict the early exit layer.
Instead of training a neural network as classifiers, (Fan et al., ) only fit the machine learning classifiers on the extracted features for the early exiting.
(Zhu, ; Zhu et al., , ; Zhang et al., ) focus on designing novel loss functions for training a more robust classifier.
(Li et al., ) incorporates sentence-level features as well as token-level features to predict the early exiting.
Different from those works, our method does not require training the classifier.
(Sun et al., ) proposed a hash-based early exiting method that uses the hashing functions to map tokens to exiting layers.
Our method predicts exit layers using a pre-built database, resulting in better generalization.
Other works (Jazbec et al., ; Li et al., ) focus on designing early exiting frameworks for image classification tasks.
They are optimizations in the different domains compared to our works.

### H.2 Retrieval-based Augmentations

Retrieval-based augmentations (Li et al., ; Wang et al., ; Xiong et al., ; Cui et al., ; Wu et al., , ) have been widely used in various natural language processing (NLP) tasks and achieved remarkable performance.
Existing works mostly leverage external knowledge databases to augment generator models on various text-generation tasks, such as language modeling (Khandelwal et al., ; Lewis et al., ; Borgeaud et al., ), question-answering (Guu et al., ; Izacard & Grave, ), machine translation (Khandelwal et al., ; Wang et al., ), dialogue system (Cheng et al., ).
Those works focus on improving the model’s generation quality, while our work aims to use the retrieval knowledge to accelerate the model’s inference.
Additionally, other works have improved model generation efficiency using external retrieval databases.
(Zhang et al., ) stores the generation trajectories of diffusion models and retrieves similar trajectories to skip several intermediate sampling steps.
These works are out of the scope of the research problems in this paper.

## Appendix I Retrieved Examples of RAEE

We show two examples from the SST-2 task and their retrieved top-k data samples.
As shown in Table 10 and Table 11, the retrieved samples are semantically similar to the query sentence, demonstrating the proposed RAEE’s efficacy.

**Table 10: Examples of data and corresponding retrieved data.**
| Query/Top-K | Sentence | Label |
| --- | --- | --- |
| Query | although laced with humor and a few fanciful touches, the film is a refreshingly serious look at young women. | 1 |
| Top-1 | the film is hard to dismiss – moody, thoughtful, and lit by flashes of mordant humor. | 1 |
| Top-2 | the movie enters a realm where few non-porn films venture, and comes across as darkly funny, energetic, and surprisingly gentle. | 1 |
| Top-3 | the movie, despite its rough edges and a tendency to sag in certain places, is wry and engrossing. | 1 |
| Top-4 | metaphors abound, but it is easy to take this film at face value and enjoy its slightly humorous and tender story. | 1 |
| Top-5 | it may not be particularly innovative, but the film’s crisp, unaffected style and air of gentle longing make it unexpectedly rewarding. | 1 |
| Top-6 | it has its faults, but it is a kind, unapologetic, sweetheart of a movie, and mandy moore leaves a positive impression. | 1 |
| Top-7 | although frailty fits into a classic genre, in its script and execution it is a remarkably original work. | 1 |
| Top-8 | unlike lots of hollywood fluff, this has layered, well-developed characters and some surprises. | 1 |
| Top-9 | as broad and cartoonish as the screenplay is, there is an accuracy of observation in the work of the director, frank novak, that keeps the film grounded in an undeniable social realism. | 1 |
| Top-10 | though its rather routine script is loaded with familiar situations, the movie has a cinematic fluidity and sense of intelligence that makes it work more than it probably should. | 1 |
| Top-11 | it tends to remind one of a really solid woody allen film, with its excellent use of new york locales and sharp writing. | 1 |
| Top-12 | though a touch too arthouse 101 in its poetic symbolism, heaven proves to be a good match of the sensibilities of two directors. | 1 |

**Table 11: Examples of data and corresponding retrieved data (Cond).**
| Query/Top-K | Sentence | Label |
| --- | --- | --- |
| Query | … a boring parade of talking heads and technical gibberish that will do little to advance the linux cause. | 0 |
| Top-1 | a vile, incoherent mess… a scummy ripoff of david cronenberg’s brilliant ‘videodrome. | 0 |
| Top-2 | completely creatively stillborn and executed in a manner that i’m not sure could be a single iota worse… a soulless hunk of exploitative garbage. | 0 |
| Top-3 | contrived, maudlin and cliche-ridden… if this sappy script was the best the contest received, those rejected must have been astronomically bad. | 0 |
| Top-4 | could as easily have been called ‘ under siege 3: in alcatraz ’… a cinematic corpse that never springs to life. | 0 |
| Top-5 | little more than a stylish exercise in revisionism whose point…is no doubt true, but serves as a rather thin moral to such a knowing fable. | 0 |
| Top-6 | a thoroughly awful movie – dumb, narratively chaotic, visually sloppy…a weird amalgam of ‘the thing’ and a geriatric scream. | 0 |
| Top-7 | on a cutting room floor somewhere lies…footage that might have made no such thing a trenchant, ironic cultural satire instead of a frustrating misfire. | 0 |
| Top-8 | …while certainly clever in spots, this too-long, spoofy update of shakespeare’s macbeth does n’t sustain a high enough level of invention. | 0 |
| Top-9 | worthless, from its pseudo-rock-video opening to the idiocy of its last frames. | 0 |
| Top-10 | comes across as a relic from a bygone era, and its convolutions… feel silly rather than plausible. | 0 |
| Top-11 | a tired, unnecessary retread…a stale copy of a picture that was n’t all that great to begin with. | 0 |
| Top-12 | (less a movie than) an appalling, odoriferous thing…so rotten in almost every single facet of production that you’ll want to crawl up your own in embarrassment. | 0 |

## Appendix J Limitations

The limitations of RAEE may lie in the following aspects.
Building Overheads. Although building the retrieval database is an offline process that can be completed in a few hours, there is still a trade-off between the generalization capability and the building overheads.
Analysis experiments have demonstrated that more data can significantly improve the RAEE’s generalization capability.
Retrieving Quality. Various types of indexing can be chosen for building the retrieval databases.
Different retrievers have different retrieving qualities.
The indexing used in RAEE may not be the optimal choice across different tasks.
Addressing these limitations is not the main focus of this paper, while future optimizations on those topics can be combined with the proposed RAEE.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 6
presentation: 2
questions: See potential experiments or analyses mentioned in weaknesses
soundness: 3
strengths: - Good empirical results on downstream tasks when compared to prior work (table 1)
- Compares inference times as well (figure 3)
- Provides a clear overview and motivation for the problem
summary: This work uses retrieval to improve performance on training-free early exit frameworks. The motivation for doing so is observing that similar data should have similar early exit patterns. Experimental results show that this method is significantly better than existing training-free early exit methods.
weaknesses: - While the paper mentions that out-of-domain performance is out of scope, I think this is a very important problem because many models today train on non-public data and in the real world we do not always have accompanying train sets to user inference. While it may be out of scope to completely understand out-of-domain performance, I would like to see the authors do some analysis, such as examining performance changes as a function of distance between test example and the nearest neighbors.
- Another interesting experiment for the above is using the LM train set (e.g. C4 for T5) instead of GLUE train sets for the database.
- Give that you have the test labels, can you do additional analysis to compare RAEE with other methods to see how often it exits at the correct layer? This can also clarify questions about inference times
- Figure captions could be improved

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 2
questions: 1. It has been observed that modeling the early exit problem as a distribution prediction issue is not a novel approach, as similar concepts have been explored in prior works. Could the authors elaborate on the specific novelties of their proposed RAEE framework compared to existing methods?(eg. Predictive Exit: Prediction of Fine-Grained Early Exits for Computation- and
Energy-Efficient Inference)
2. Table 4 indicates that the performance of RAEE improves with a larger retrieval database. How does the authors plan to balance the trade-off between database size, storage requirements, and inference efficiency, especially for resource-constrained environments?
3. In Table 1, why were only 2-3 other methods compared for a specific single backbone(why weren't some methods compared) ?
4. The paper uses a threshold of 0.9, but it's unclear how sensitive RAEE is to it. An analysis of how changes in these parameters affect performance would be useful.
soundness: 3
strengths: 1. The paper provides a novel combined method to decide the exit layer.
2. The paper is well written and easy to follow.
3. The proposed method is reasonable. Experiments support their claims. The performance can be improved in most scenarios.
summary: This paper aims to enhancing the efficiency of large language model inference by adaptively exiting the model at earlier layers. The authors model the early exit problem as a distribution prediction issue, and then use exit information from similar data to approximate the distribution. They outline the methodology for constructing a retrieval database with exit information and propose the RAEE framework, which leverages the pre-built retrieval database to predict the exit layer based on the exit information from the top-k nearest neighbors. Experimental results across eight downstream tasks demonstrate that RAEE improves inference speed while maintaining robust zero-shot performance, outperforming other early exit frameworks.
weaknesses: 1. The paper does not clearly position itself with respect to existing retrieval-augmented methods that used to accelerate the model’s inference. A more thorough literature review is needed to highlight how RAEE differs from and improves upon prior work.
2. While the data presented in figure3 is comprehensive, I noticed that the visual presentation, specifically the subscripts, could be enhanced for better readability and aesthetic appeal.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 1
questions: - If you could assess the concerns above.
- When two or more layers have the same maximal probability RAEE selects the earliest L255. I'm wondering if this is a good option. How many cases exist with the same maximal probability in each dataset? How many times RAEE choose the earliest and make mistakes? If it would use a later, could it get a correct performance?
- How much memory is added if we want to deploy this method on resource-constrained devices?
soundness: 2
strengths: **Method**

- The method is interesting, and there is a clear performance improvement. Creating a database to decide what layer to use sounds promising. I would push that idea to create a more general dataset rather than tailor it to each specific target dataset.
summary: This paper addresses the challenge of high computational demands in large language model inference. It focuses on reducing the number of inference layers required by exploiting early exits. Instead of early-exiting by training internal classifiers to decide if the model can exit after fewer layers, the authors propose RAEE, a Retrieval-Augmented Early Exit framework. RAEE treats early exit as a distribution prediction problem, where exit decisions are informed by similar data examples stored in a retrieval database (Cache). This approach allows the model to decide on exiting based on prior information from similar data, leading to faster and more efficient inference.
weaknesses: **Motivation.**

- Although I understand that the inference time could be reduced by early exiting, one of the main reasons for the high computational demand for Transformer LLMs is the KV cache. In this regard, solutions like PagedAttention [1] show a 3.5x and 24x higher throughput using LLaMA. My concern is that this paper claims to accelerate the model inference. However, it doesn't compare with those types of approaches. If the motivation of this work is to deploy LLM on resource-constrained devices, the memory aspect is key. Moreover, Figure 3 shows that the method did not always get better inference latency compared to the base method T5-L, RoBERTa/ElasticBERT or the selected baselines (AdaInfer, SLEB).

- In my opinion, the work is more interesting regarding performance improvement than latency reduction (Table 2). The results show a clear diversity in the layers. It seems that some layers are better than others for some datasets. Thus, section 4.3 needs to be explored deeply. I expected to see some statistics about the layers used the most for each dataset and what type of tasks/questions/data are better answered for the early or later layer. This raises the question of whether the retriever database could be more general to the input data type.

---
**Experiments.**
- Even though the GLUE benchmark is well-known by the NLP community, it would be good to explain it further, emphasizing the evaluation metrics and clarifying if the tasks are classification or generation.

- To better contextualize the method RAEE with the LLMs used in the comparison, it would be great to see the performance of HashEE/CALM/SLEB on each LLM. (Figure3)

- Section 4.3 requires much more analysis to understand the sentence L407-408 better. 
> - What layers are the most used? 
> - Is there some correlation between the type of data vs early exit or type of task vs early exit?
 

---
 **Presentation.**

- Some parts of the text are difficult to follow. 
 
- Figures and plots require better captions. Figure 1 is quite intuitive, but explaining each component in the caption would be ideal. Figure 2 is confusing: the caption doesn't explain the difference between (a) and (b). Also, the explanation between L131 and L142 is confusing. In general, captions for figures and tables are very superficial.

- Citations. Line 124 could cite Tip-Adapter as a training-free adaptation that uses a cache for image classification. In the same line (124), what are the existing retrieval databases that use clustering and/or product quantization?

[1] Woosuk Kwon et.al, Efficient Memory Management for Large Language Model Serving with PagedAttention, 2023

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: See weaknesses:

Also, Early Exits are not pruning based methods, instead they fall into class of dynamic inference methods as pruning reduces the weights. Here none of the weights are reduced instead the model decides which layer it should not use although it has an option to use them which is not the case with pruning.
soundness: 2
strengths: 1. The paper is well written.

2. Paper uses multiple backbone models making the claims strong.

3. The paper claims to reduce the latency overheads which is an important problem for the community.
summary: The paper proposes a new early exit approach that predicts a probability distribution over the set of layers in the model in which distribution is approximated using similar data. Specifically, it utilizes the embedding of the incoming samples and creates an embedding space where the number of spaces is equal to number of layers in the model. The decision of which space an incoming sample belongs is made based on  the possible layers it can make an exit from. During inference, an incoming sample is first checked against the top-k nearest neighbors and the sample is assigned an exiting layer based on the estimated probability distribution using the top-k neighbors.
weaknesses: 1. **Novelty:**  The paper lacks novelty as the main claim is it is learning the distribution of incoming samples using the embeddings of the samples. This is already done in DIMEE [1] work, however the objectives of both the works are different, the method to solve the problem are very similar. 

2. **Lack of proper baselines:** There are multiple existing works [2], [3]. [4], [5] that also learn the probability distribution over the exit points but the paper have neither cited them nor compared against them that reduces the overall impact of the paper. Since the final objective of the paper is to learn a distribution, there should be a comparison with existing distribution predicting methods.

3. **Overclaim:** The paper has a major claim that it can outperform under zero-shot setting but it requires the labels to create the retrieval database. Also, there will be large impact if the domain of the test dataset changes which is not explored in this paper.

4. **Lack of explanation:** As per as I believe, the better results of RAEE are due to the fact that it uses learned final layer classifier to map the hidden representations from intermediate layers to class probabilities similar to [5], however, other baselines use different classifiers at each layer that are randomly initialized hence the loss in performance. This is an apple to orange comparison, even the baselines should be tested with learned final layer classifier for fair comparison.

5. **Clustering figure:** I believe there should be an additional t-SNE plot in the paper showing that how the clusters are formed based on the exit points as shown in the DIMEE paper.

**Missing references:** There are a lot of missing references

1. DIMEE: https://arxiv.org/abs/2410.05338

2. JEI-DNN: https://openreview.net/pdf?id=jX2DT7qDam (ICLR 2024)

3. ZTW: https://proceedings.neurips.cc/paper/2021/file/149ef6419512be56a93169cd5e6fa8fd-Paper.pdf (NIPS 2021)

4. PALBERT: https://proceedings.neurips.cc/paper_files/paper/2022/file/5a9c1af5f76da0bd37903b6f23e96c74-Paper-Conference.pdf (NIPS)

5. MSDNet: https://arxiv.org/abs/1703.09844

6. CeeBERT: https://aclanthology.org/2024.findings-acl.101/ (ACL 2024)

7. ETFEE Yixin Ji, Jikai Wang, Juntao Li, Qiang Chen, Wenliang Chen, and Min Zhang. 2023. Early exit with disentangled representation and equiangular tight frame. In Findings of the Association for Computational Linguistics: ACL 2023, pages 14128–14142.

Meta-review: This paper received mixed reviews. The reviewers recognized the new, interesting, and reasonable method, its strong performance, and extensive experiments. At the same time, they raised concerns with unclear motivation and inappropriate positioning of the paper (qFBa), no comparisons with other approaches aiming at inference acceleration, e.g., PagedAttention (qFBa), lack of essential analysis supporting the main arguments (qFBa), experiments limited to in-distribution retrieval (xZ1Y, LKLV), space-time complexity caused by the use of retrieval database (Jjgv, LKLV), lack of proper baselines (LKLV), missing analysis on the sensitivity to the threshold hyperparameter (Jjgv), lack of discussion about existing retrieval-augmented methods (Jjgv), and presentation issues (xZ1Y, Jjgv, qFBa).

The authors' rebuttal and subsequent responses in the discussion period address some of these concerns but failed to fully assuage all of them: after the discussion perior, the reviewers still pointed out the issues on the motivation and positioning of the paper (qFBa), missing comparisons with PagedAttention (qFBa), concerns with the out-of-distribution experiments (xZ1Y), the complexity issue (Jjgv, LKLV), and no comparisons with relevant baselines (LKLV). As a result, two reviewers voted to reject, and a reviewer who leaned borderline toward accept still has concerns about the OOD experiment. 

Putting these together, the AC found that the remaining concerns outweigh the positive comments and the rebuttal, and thus regrets to recommend rejection. The authors are encouraged to revise the paper with the comments by the reviewers and the AC, and submit to an upcoming conference.

Decision: Reject