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

W1: Thanks for your comments. Unfortunately, the online date of this paper on ArXiv is October 6, which is later than our submission to ICLR 2025 (October 1, AOE). The techniques proposed in RAEE are novel, while this DIMEE can only be treated as a follow-up work of RAEE.

W2: Thanks for your comments. We carefully reviewed all the listed works and will discuss them in the related works. Specifically, ZTW [3], MSDNet [5] are specifically designed for CNN-based models. JEI-DNN[2], PALBERT [4], CEEBERT[6], and ETFEE [7] all require fine-tuning the whole backbone models. In our paper, we target the scenarios of no update on backbone models’ parameters.

[2] Florence Regol, Joud Chataoui, Mark Coates. Jointly-Learned Exit and Inference for a Dynamic Neural Network : JEI-DNN. CoRR abs/2310.09163 (2023)

[3] Maciej Wolczyk, Bartosz Wójcik, Klaudia Balazy, Igor T. Podolak, Jacek Tabor, Marek Smieja, Tomasz Trzcinski. Zero Time Waste: Recycling Predictions in Early Exit Neural Networks. NeurIPS 2021: 2516-2528

[4] Nikita Balagansky, Daniil Gavrilov. PALBERT: Teaching ALBERT to Ponder. NeurIPS 2022

[5] Gao Huang, Danlu Chen, Tianhong Li, Felix Wu, Laurens van der Maaten, Kilian Q. Weinberger. Multi-Scale Dense Networks for Resource Efficient Image Classification. ICLR 2018

[6] Divya Jyoti Bajpai, Manjesh K. Hanawal. CeeBERT: Cross-Domain Inference in Early Exit BERT. ACL (Findings) 2024: 1736-1748

[7] Yixin Ji, Jikai Wang, Juntao Li, Qiang Chen, Wenliang Chen, Min Zhang. Early Exit with Disentangled Representation and Equiangular Tight Frame. ACL (Findings) 2023: 14128-14142

W3: Thanks for your comments. Please see the General Response to Out-of-Domain Issues.

W4: Thanks for your comments. It seems there are a lot of misunderstandings. First, RAEE doesn’t have any classifier and doesn’t tune any parameters of backbone models, which is totally different from [5] as you mentioned. The main idea of RAEE is to retrieve the exit information from a pre-built retrieval database and then compute the exit layer according to the top-k nearest neighbors’ exit information. There is no “final layer classifier” in RAEE. In contrast, RAEE doesn’t add any extra components to the backbone model while only passing a parameter of the exit layer to stop the inner loop early. For more novelties claims, please refer to the response to Q1 of the reviewer Jjgv.

The comparisons to baselines are also fair. For RAEE, we only use the training dataset to create the retrieval database. For baselines containing classifiers (HashEE, DeeBERT, AdaInfer), we only tune the classifiers on the training dataset. For CALM requiring training the backbone models, we follow its default settings on the threshold and only perform its inference. For SLEB, we follow its settings to prune the backbone models.

W5: Thanks for your suggestions. However, we believe it is unnecessary to draw such a clustering figure. The reasons lie in the following,

1. Data with similar representations would exit at similar layers, but data with different representations may also exit at similar layers.
2. In RAEE, there are multiple possible exit layers for each input, which makes it difficult to draw all of them in one figure.

Besides, Figure 2 in the paper can better describe the distribution approximation with top-k nearest neighbors’ exit information.

W1: Thanks for your great suggestions. The main focus of PagedAttention is designing a new attention algorithm that optimizes memory usage and inference efficiency. This type of work is orthogonal to our work. For example, with RAEE’s early exit, PageAttention can dynamically save the computations of exit layers’ KV cache.

For comparable inference latency of T5-L, RoBERTa/ElasticBERT, the reasons lie that extra encoding and retrieving time offsets the benefits of early exit. For AdaInfer, it has a poor exit classifier and always exits very early with wrong predictions. For SLEB, there is a hyperparameter that decides how many layers would be removed from the backbone model. For a fair comparison, we chose similar exit layers for SLEB. Although SLEB can achieve a faster inference, it performs poorly compared to RAEE.

W3: Thanks for your suggestions. We evaluate the proposed RAEE on some representative generation tasks, such as CNN/DailyMail and XSum. Rebuttal-Table 11 shows the results of applying RAEE on two summarization tasks. Experimental results also demonstrate the efficacy of the proposed RAEE.

Rebuttal-Table 11. Performance of the Llama-3-8b and RAEE (Llama) on generation tasks.

W4: Thanks for your comments. Since the early exit frameworks compared in this work are not suitable for all backbone models. Specifically, HashEE and DeeBERT are specifically designed for BERT-liked models; CALM is specifically designed for T5-based models; SLEB is specifically designed for decoder-only models. We have tried our best to implement AdaInfer over all kinds of backbone models whose design is less related to the architecture of backbone models.

W2 & W6: Thanks for your insightful suggestions. We collect the statistics of the most-used exit layers in Rebuttal-Table 12. Table 8 in the appendix also shows the correlation between the type of task and exit layers. These results show the diversity of exit layers across different tasks and different inputs.

Rebuttal-Table 12. Most-used exit layers of the RAEE (Llama) and RAEE (Llama) Corr. over different distance metrics with different backbone models. RAEE (Llama) Corr. means just contains the data that is correctly predicted.

W7: Thanks for your suggestions. We have improved all figure captions.

W8: Thanks for your suggestions. Tip-Adapter is a training-free method for learning better representation for CLIP, there is no clear relationship with early exit frameworks, except for the training-free setting.

There are several vector databases that use clustering and product quantization techniques, such as FAISS.

Q1: Thanks for your comments. Theoretically, there would be the case that two or more exit layers have the same probability. However, we collect all exit layers and corresponding probability in all experiments, and no exit layer exhibits the same probability for exiting, which shows that it is a very low probability case.

Q2: Thanks for your comments. It only requires less than 6MB. For more details, please refer to the response to the Q2 of the reviewer Jjgv.

W1. Thanks for your comments. We have revised the related work of retrieval-based augmentations.

W2. Thanks for your suggestions. We have revised Figure 3 with new subscripts for better readability and aesthetic appeal.

Q1: Thanks for your comments. There are clearly some differences between the proposed RAEE and existing mainstream early exit frameworks, including the referred Predictive Exit.

1. The referred Predictive Exit is now only suitable for CNN-based neural networks, such as VGG-19 and ResNet-34, and applying it to the existing general transformer-based large language models still remains unexplored. The implicit patterns between CNN-based models and transformer-based models are significantly different.
2. Although the referred Predictive Exit and the proposed RAEE both model the early exit problem as a distribution prediction problem, it is worth noticing that the proposed RAEE requires no parameters update of the backbone models, while most existing early exit frameworks [1, 2, 3] require jointly fine-tuning the backbone model and the early exit classifiers. Fine-tuning large language models is quite costly, even though using LoRA techniques.
3. Although the proposed RAEE and the referred Predictive Exit both introduce hyperparameters, the ablation studies show that RAEE is not sensitive to the hyperparameters, which makes the deployment much easier. However, the Predictive Exit requires setting the starting layer to determine the next exit layer, which significantly impacts the model performance, according to their papers.
4. More importantly, the proposed RAEE is more interpretable than those early exit frameworks using learned classifiers. As shown in Figure 2, we can draw the exit distribution with the retrieved examples’ exit information and make the exit predictions.

In summary, the proposed RAEE is suitable for existing state-of-the-art large language models, requires no parameters update of backbone models, is less sensitive to hyperparameters, and is more interpretable.

[1] Xiangjie Li, Chenfei Lou, Yuchi Chen, Zhengping Zhu, Yingtao Shen, Yehan Ma, An Zou, Predictive Exit: Prediction of Fine-Grained Early Exits for Computation- and Energy-Efficient Inference. AAAI 2023: 8657-8665.

[2] Florence Regol, Joud Chataoui, Mark Coates. Jointly-Learned Exit and Inference for a Dynamic Neural Network : JEI-DNN. ICLR 2024.

[3] Divya Jyoti Bajpai, Manjesh K. Hanawal. CeeBERT: Cross-Domain Inference in Early Exit BERT. ACL (Findings) 2024: 1736-1748.

Q2: Thanks for your comments. As shown in Table 5 in the manuscript, the index size and database size are relatively small compared to the backbone model, which is less than 6 MB in total. And Table 4 also shows that increasing the amount of data in the retrieval database would enhance the RAEE’s performance. So, when deploying such LLMs in resource-constrained scenarios, the resource requirements for the retrieval database should be the last consideration. For example, the memory should be preferentially allocated to the model weights and KV cache, then the retrieval database.

Q3: Thanks for your comments. Since the early exit frameworks compared in this work are not suitable for all backbone models. Specifically, HashEE and DeeBERT are specifically designed for BERT-liked models; CALM is specifically designed for T5-based models; SLEB is specifically designed for decoder-only models. We have tried our best to implement AdaInfer over all kinds of backbone models whose design is less related to the architecture of backbone models.

Q4: Thanks for your comments. We have conducted an ablation study on different thresholds in Rebuttal-Table 10. Experimental results show that the performance drop of the proposed RAEE is quite small on average.

Rebuttal-Table 10. Performance of RAEE using different thresholds across eight classification tasks with RoBERTa-Large.

W1: Thanks for your insightful suggestions. We rebuild the database with different distance metrics, such as inner product, and retrieve the top-k nearest neighbors based on the corresponding distance metrics. As shown in Rebuttal-Table 7-8, although the RAEE with different backbones based on the inner product achieves a bit poorer performance, the performance difference is quite small (less than one on average), demonstrating the robustness of the proposed RAEE over distance metrics.

Rebuttal-Table 7. Performance of RAEE over different distance metrics.

Rebuttal-Table 8. Exit layers of RAEE over different distance metrics.

W2: Thanks for your comments. Please refer to the General Response to Out-of-Domain Issues.

W3: Thanks for your suggestions. We collect the statistics of RAEE’s and AdaInfer’s correct exit layers to show that the reduction of inference time benefits from correct exit layers rather than exiting with wrong predictions very early. As shown in Rebuttal-Table 9, RAEE has similar exit layers for all the data and the correctly predicted data. This demonstrates that RAEE does effectively help accelerate the inference without sacrificing performance.

We also collect the correct exit layers of AdaInfer, which show similar conclusions. However, due to the poor exit layer classification and poor model performance, AdaInfer makes many early exits with wrong predictions so that it can achieve quite a small inference latency but a poor performance.

Rebuttal-Table 9. Exit layers of RAEE and AdaInfer on all predictions and correct predictions.

W4. Thanks for your comments. We have improved all figure captions with detailed introductions.