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

arXiv: https://arxiv.org/abs/2405.15198 · HTML: https://mlsanity.com/paper/2405.15198/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Placed in Generation and context integration, one of 15 directions in Retrieval-augmented generation.
- Builds on 47 of its 64 references that are in the map.
- First posted on arXiv 24 May 2024.

## Placed in

- Retrieval-augmented generation / [Generation and context integration](https://mlsanity.com/rag/generation-and-context-integration/)
- Efficient LLM inference and serving / [Early exit and edge inference](https://mlsanity.com/inference/early-exit-and-edge-inference/)

## Builds on (most cited first, 30 of 47)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Billion-scale similarity search with GPUs](https://mlsanity.com/paper/1702.08734/) (2017)
- [REALM: Retrieval-Augmented Language Model Pre-Training](https://mlsanity.com/paper/2002.08909/) (2020)
- [BLOOM: A 176B-Parameter Open-Access Multilingual Language Model](https://mlsanity.com/paper/2211.05100/) (2022)
- [LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding](https://mlsanity.com/paper/2308.14508/) (2023)
- [Don't Give Me the Details, Just the Summary! Topic-Aware Convolutional Neural Networks for Extreme Summarization](https://mlsanity.com/paper/1808.08745/) (2018)
- [Improving language models by retrieving from trillions of tokens](https://mlsanity.com/paper/2112.04426/) (2021)
- [Get To The Point: Summarization with Pointer-Generator Networks](https://mlsanity.com/paper/1704.04368/) (2017)
- [Making Pre-trained Language Models Better Few-shot Learners](https://mlsanity.com/paper/2012.15723/) (2020)
- [LaMDA: Language Models for Dialog Applications](https://mlsanity.com/paper/2201.08239/) (2022)
- [Leveraging Passage Retrieval with Generative Models for Open Domain Question Answering](https://mlsanity.com/paper/2007.01282/) (2020)
- [Generalization through Memorization: Nearest Neighbor Language Models](https://mlsanity.com/paper/1911.00172/) (2019)
- [LLM-Pruner: On the Structural Pruning of Large Language Models](https://mlsanity.com/paper/2305.11627/) (2023)
- [DeeBERT: Dynamic Early Exiting for Accelerating BERT Inference](https://mlsanity.com/paper/2004.12993/) (2020)
- [Confident Adaptive Language Modeling](https://mlsanity.com/paper/2207.07061/) (2022)
- [BERT Loses Patience: Fast and Robust Inference with Early Exit](https://mlsanity.com/paper/2006.04152/) (2020)
- [Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time](https://mlsanity.com/paper/2310.17157/) (2023)
- [FastBERT: a Self-distilling BERT with Adaptive Inference Time](https://mlsanity.com/paper/2004.02178/) (2020)
- [Nearest Neighbor Machine Translation](https://mlsanity.com/paper/2010.00710/) (2020)
- [Augmenting Language Models with Long-Term Memory](https://mlsanity.com/paper/2306.07174/) (2023)
- [Lift Yourself Up: Retrieval-augmented Text Generation with Self Memory](https://mlsanity.com/paper/2305.02437/) (2023)
- [SLEB: Streamlining LLMs through Redundancy Verification and Elimination of Transformer Blocks](https://mlsanity.com/paper/2402.09025/) (2024)
- [Not All Layers of LLMs Are Necessary During Inference](https://mlsanity.com/paper/2403.02181/) (2024)

## Built on by (most cited first, 4 of 4)

- [GM-Skip: Metric-Guided Transformer Block Skipping for Efficient Vision-Language Models](https://mlsanity.com/paper/2508.18227/) (2025)
- [River-LLM: Large Language Model Seamless Exit Based on KV Share](https://mlsanity.com/paper/2604.18396/) (2026)
- [GaiaFlow: Semantic-Guided Diffusion Tuning for Carbon-Frugal Search](https://mlsanity.com/paper/2602.15423/) (2026)
- [Efficient Block-Layer Parallel Inference for Vision-Language-Action on Hybrid Architectures](https://mlsanity.com/paper/2608.14586/) (2026)

## Abstract (arXiv)

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