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

Authors: LIANMING HUANG, Shangyu Wu, Yufei Cui, Ying Xiong, Xue Liu, Tei-Wei Kuo, Nan Guan, Chun Jason Xue
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: Early Exit; Retrieval Augmentation; Large Language Model

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 to determine whether to exit at intermediate layers.
However, such classifier-based early exit 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 Retrieval-Augmented Early Exit framework for efficient inference.
This paper first demonstrates that the early exit problem can be effectively modeled as a distribution prediction problem, in which the distribution is approximated through the exit information of similar data. 
Subsequently, it outlines the methodology for collecting exit information 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 at the layer. 
Experimental results demonstrate that RAEE significantly accelerates inference while achieving robust zero-shot performance across eight downstream tasks.

Decision: reject
Reviewer scores: 6, 6, 3, 3

Reviews and discussion:
<thread>
creation date:
1730109355581

thread_type:
official_review

<thread_object>
creator: reviewer_lklv

<content>
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.

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.

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730445724227

thread_type:
official_review

<thread_object>
creator: reviewer_qfba

<content>
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.

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.

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

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730642124438

thread_type:
official_review

<thread_object>
creator: reviewer_jjgv

<content>
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.

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.

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730690223011

thread_type:
official_review

<thread_object>
creator: reviewer_xz1y

<content>
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.

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

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

questions:
See potential experiments or analyses mentioned in weaknesses
</content>
</thread_object>

</thread>
<thread>
creation date:
1732387187228

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response to Out-of-Domain Issues-2

comment:
3. **The data domain also matters but is not necessary**. Since using GLUE tasks’ training dataset to build the retrieval database has demonstrated that their data qualities are better than wikitext, we conducted experiments using different GLUE tasks’ training datasets to build the retrieval database. As shown in Rebuttal-Table 5, RAEE can achieve better performance with out-of-domain databases on SST-2, SST-5, MR, and CR tasks. But for the rest tasks, RAEE can achieve the best performance with only in-domain databases. For exit layers in Rebuttal-Table 6, similar conclusions can also be drawn that in-domain databases are not always the best choice.
    
    Rebuttal-Table 5. Performance of RAEE using different domain retrieval databases across eight classification tasks.
    
    | Metrics | SST-2(acc) | SST-5(acc) | MR(acc) | CR(acc) | MPQA(acc) | Subj(acc) | TREC(acc) | CoLA(mcc) |
    | --- | --- | --- | --- | --- | --- | --- | --- | --- |
    | RAEE (RB-L-'SST-2')  | 84.63 | 34.39 | **84.80** | 80.75 | 59.60 | 50.65 | 28.20 | 7.86 |
    | RAEE (RB-L-'SST-5')   | 75.34 | 33.57 | 73.85 | 69.20 | 62.85 | 50.35 | 18.60 | 3.22 |
    | RAEE (RB-L-'MR')   | **90.48** | **36.11** | 81.55 | 81.80 | 66.05 | 51.10 | 24.60 | -3.83 |
    | RAEE (RB-L-'CR')  | 74.66 | 32.67 | 75.40 | 68.05 | 55.60 | 49.00 | 16.60 | 2.58 |
    | RAEE (RB-L-'MPQA')  | 86.12 | 34.16 | 83.30 | **83.50** | **78.55** | 50.00 | 29.40 | -1.56 |
    | RAEE (RB-L-'Subj') | 79.70 | 32.22 | 75.90 | 72.55 | 54.70 | **84.05** | 11.80 | 8.06 |
    | RAEE (RB-L-'TREC')  | 61.01 | 26.97 | 56.55 | 75.50 | 56.30 | 49.65 | **62.40** | -1.39 |
    | RAEE (RB-L-'CoLA') | 76.72 | 33.85 | 73.95 | 68.05 | 53.25 | 57.20 | 18.00 | **14.48** |
    
    Rebuttal-Table 6. Exit layers of RAEE using different domain retrieval databases across eight classification tasks.
    
    | Layers | SST-2 | SST-5 | MR | CR | MPQA | Subj | TREC | CoLA |
    | --- | --- | --- | --- | --- | --- | --- | --- | --- |
    | RAEE (RB-L-'SST-2')  | 18.55 | 18.57 | 18.60 | 19.09 | 17.90 | 17.95 | 19.41 | 18.56 |
    | RAEE (RB-L-'SST-5')   | 14.23 | 13.93 | 13.67 | **14.09** | 18.03 | 12.37 | 15.34 | 15.47 |
    | RAEE (RB-L-'MR')   | 18.52 | 19.06 | 18.71 | 19.19 | 17.39 | 18.20 | 18.25 | 19.39 |
    | RAEE (RB-L-'CR')  | 17.10 | 17.24 | 17.59 | 15.35 | 13.18 | 18.36 | 18.38 | 16.32 |
    | RAEE (RB-L-'MPQA')  | 21.71 | 21.64 | 21.66 | 21.45 | 17.20 | 22.09 | 21.79 | 20.34 |
    | RAEE (RB-L-'Subj') | 18.48 | 18.02 | 18.06 | 16.75 | 7.67 | 13.59 | **6.94** | **11.09** |
    | RAEE (RB-L-'TREC')  | **12.29** | **12.39** | **11.74** | 16.37 | 13.60 | **12.04** | 12.82 | 16.12 |
    | RAEE (RB-L-'CoLA') | 16.36 | 15.73 | 16.00 | 14.47 | **5.28** | 14.72 | 8.35 | 12.48 |

In conclusion, although RAEE with the wikitext-based retrieval database cannot achieve as good performance as that with the in-domain retrieval database, the above analysis still demonstrates the efficacy of the proposed RAEE framework.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732387204701

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response to Out-of-Domain Issues-1

comment:
Thanks for all the reviewers’ insightful suggestions. The main concern of this paper is that RAEE cannot achieve good performance with the out-of-domain retrieval database. To analyze the case of out-of-domain issues, we have conducted experiments using the retrieval database built on wikitext-2-v1. Since there is no gold label on the text dataset, we follow the next-token prediction task setting, where the input sentence's next token is treated as the gold label. 

Specifically, we first split the whole text dataset into sentences, avoiding breaking semantics. Then, according to the backbone model’s max input length, for each sentence whose length is smaller than the max input length, we regard the last meaningful token of the sentence as the gold label; for the other sentences, we use a sliding window of the max input length of the backbone model and regard the last meaningful token of the window as the gold label. Finally, we collect the exit information in the way used in this paper. To best demonstrate the efficacy, we choose llama-3-8b as the backbone model, which is pre-trained with the next-token prediction task.

Rebuttal-Table 1. Performance of the Llama-3-8b and RAEE (Llama-wiki) across eight classification tasks.

| Metrics | SST-2(acc) | SST-5(acc) | MR(acc) | CR(acc) | MPQA(acc) | Subj(acc) | TREC(acc) | CoLA(mcc) | Avg |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Llama-3-8B | 62.84 | 26.06 | 59.65 | 72.90 | 51.75 | 52.80 | 8.40 | 0.00 | 41.80 |
| RAEE (Llama-wiki) | 55.50 | 21.40 | 54.30 | 61.60 | 57.15 | 51.55 | 13.00 | 0.00 | 39.31 |

Rebuttal-Table 2. Exit layers of the Llama-3-8b and RAEE (Llama-wiki) across eight classification tasks.

| Layers | SST-2 | SST-5 | MR | CR | MPQA | Subj | TREC | CoLA | Avg |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Llama-3-8B | 32.00 | 32.00 | 32.00 | 32.00 | 32.00 | 32.00 | 32.00 | 32.00 | 32.00 |
| RAEE (Llama-wiki) | 29.41 | 28.83 | 29.30 | 28.75 | 30.13 | 29.31 | 30.34 | 30.56 | 29.58 |

As shown in Rebuttal-Table 1-2, we evaluate RAEE (Llama) with the wikitext-based retrieval database, termed RAEE (Llama-wiki). Unsurprisingly, RAEE (Llama-wiki) performs poorly, although it can exit earlier. However, this performance drop cannot simply be attributed to the out-of-domain dataset. There are three key points that jointly impact the model performance as well as the model inference efficiency. 

1. **The task type in the process of building the retrieval database and the inference with early exit should be aligned**. To verify this point, we have conducted an experiment on the summarization tasks such as CNN/DailyMail and XSum while using the wikitext-based retrieval database, as shown in Rebuttal-Table 3. Experimental results demonstrate that RAEE can improve performance and accelerate inference even though the domain of the retrieval database is out of the distribution. 
    
    Rebuttal-Table 3. Performance of the Llama-3-8b and RAEE (Llama-wiki) on generation tasks.
    
    |  | ROUGE-L | Layers |
    | --- | --- | --- |
    | CNN/DailyMail Llama-3-8B  | 8.95 | 32.00 |
    | CNN/DailyMail RAEE (Llama-wiki) | 14.01 | 29.60 |
    | XSum Llama-3-8B  | 5.22 | 32.00 |
    | XSum RAEE (Llama-wiki) | 7.15 | 28.82 |
2. **Data quality determines the generalization of the RAEE’s retrieval database of exit information**, impacting the quality of the exit distribution approximation through neighbors’ exit information. We have also conducted experiments to show why RAEE achieves poor performance when using the wikitext-based retrieval database. As shown in the Rebuttal-Table 4, we evaluate the accuracy of the next-token prediction of the backbone model and RAEE on the wikitext-2-v1, where the token with the maximal probability is chosen as the next token (accuracy can better explain the claim than that of perplexity). Experimental results show that the backbone can only achieve an accuracy of 53.90 for predicting the next token. With RAEE which corrects some predictions with early exit, the accuracy still maintains a low level of 57.20. These results demonstrate that there is a considerable volume of sentences in the wikitext dataset that the backbone model cannot make correct predictions even if RAEE corrects the predictions. Those cases also result in no exit of RAEE.
    
    Rebuttal-Table 4. Performance and exit layers of Llama-3-8b and RAEE (Llama-wiki) on wikitext training data.
    
    |  | acc | Layers |
    | --- | --- | --- |
    | Llama-3-8B  | 53.90 | 32.00 |
    | RAEE (Llama-wiki)  | 57.20 | 30.00 |
</content>
</thread_object>

</thread>


Metareview:
metareview:
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.