RAEE: A Robust Retrieval-Augmented Early Exit Framework for Efficient Inference
Authors: LIANMING HUANG
Shangyu Wu
Yufei Cui
Ying Xiong
Haibo Hu
Xue Liu
Tei-Wei Kuo
Nan Guan
Chun Jason Xue
Venue: ICLR 2026 Poster (2026)
Keywords: Early Exit; Retrieval Augmentation; Large Language Model
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: This paper introduces RAEE, a retrieval-augmented early exit framework that accelerates inference while enhancing performance.

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

Published as a conference paper at ICLR 2026
RAEE: A ROBUST RETRIEVAL-AUGMENTT EARLY EXIT FRAMEWORK FOR EFFICIENT INFERENCE
Lianming Huang' Shangyu 2 Yufei Cui Ying Xiong", Haibo Hu',
Xue Lin Tei- Wei Kuo", Nan Guan', Chun Jason Xue²
City University of Hong Kong
2 Mohamed bin Zayed University of Artificial Intelligence (MBZUAI)
3 MILA, McGill University National Taiwan University
ABSTRACT
Deploying large language model inference remains challenging due to their high computational overhead. Early exit optimizes model inference by adaptively re- ducing 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, robust Retrieval-Augmented Early Exit framework that not only enables early exit but also enhances model performance through corrective exit information at inter- mediate layers. This paper first demonstrates that the early exit problem can be effectively modeled distribution prediction problem, in which the distribution can be further approximated through the exit information of similar data. Subse- quently, 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 performance across eight downstream tasks.
INTRODUCTION
Large language models (LLMs) have been widely used various application scenarios due their
excellent performance (Thoppilan et al.) 2022 Touvron et al.) 2023 Scao et al. 2022). However, the model inference efficiency still poses great challenge due to high computational overhead and memory requirements (Dao et al.) 2022 Liu et al., (023). Early exit has emerged as an advanced model pruning method, which dynamically terminates the LLMs inference when a certain confidence
criterion is met, thereby reducing latency and memory consumption |Valicenti et al. 2023 2023
Ma et al.
Existing early exit frameworks (Liu et al. 2020 Zhu 202 Xin al. 2020 Fan al. 2024 can
be categorized into three types: training-baan based, semi- crani-traiing based, and training ffre approaches.
Training- based methods (Zhul 2021 Zhou et al. 2020 Zhu et al. 2023 Bae et al. 2023 Schuster
et al. 2022) jointly optimize internal c classifiers and the backbone model, incurring substantial training overhead. Semi-training-baat
methods (Fan et al. 2024) freeze the backbone model and only train
the lightweight classifiers, but they often heavily rely on manual feature engineering and may not
generalize well. Truining-free methods (Sun al. 2022) typically
heuristic exit criteria that lack adaptability and often lead to performance degradation compared to the full model. Notably, most
early exit frameworks trade off accuracy for speed (Fan al. 2024 Sun et al. 2022 Schuster et al.
2022
Bae et al. 2023)
AAthhrr contributed equally to this research.
1 Corresponding author.

---

Published as a conference paper at ICLR 2026
1.0
$7.60
97.70
23,15
Probabiliiy
100%
0.6
Acuuray
31.95
10%
0.4
60
32.63
20%
1011 1213223334 2017 2920 Layer index (a)
5972
5573 MPGA
SUB)
(b)
Figure 1: (a) Correct prediction probabilities for test sample and its top-8 nearest neighbors from the SST-2 dataset. (b) The accuracy the full model and the model after correcting from the intermediate layers, and the correct ratio. All results are collected with the backbone model RoBERTa--Lage.
Different from existing early exit frameworks, we observe that early exit can not only accelerate the inference, but also act as corrective mechanism when the full model makes the wrong predictions. Furthermore, we also observe that the exit behaviors are highly consistent across semantically similar data. These insights challenge the conventional trade-off scheme early exit frameworks and also suggest that the key to ideal early exit strategy lies in adaptively selecting the exit layer based on the behavior similar instances.
Based those observations, propose RAEE, Retrieval-Augmentee Early Exit framework that improves both efficiency and accuracy without training classifiers. RAEE leverages the insights where similar inputs have similar exit behaviors to determine the exit layer. Specifically, retrieval database from external data by recording exit layers where correct predictions occur. During the inference, RAEE retrieves the exit information of the top-k similar examples and aggregates the results to determine the final exit layer.
construct
We conduct comprehensive experiments evaluate the proposed RAEE and various comparison methods on eight downstream tasks. Experimental results demonstrate that RAEE can accelerate the model inference while achieving robust model performance. Codes are available at
The main contributions of this paper are:
We model the early exit problem as a distribution prediction problem and demonstrate that the exit information of similar data can approximate the exit distribution;
We propose RAFE, a robust retrieval-augmentt early exit framework, which leverages an external database guide the early exit;
Experimental results show that the proposed RAEE can not only accelerate the model inference but also significantly improve the model accuracy, even sometimes surpassing the full model, which outperforms traditional early exit frameworks.
2 MOTIVATIONS
The primary goal of this work is develop a robust and efficient early exit framework without training any classifiers or model parameters. To this end, we explore a retrieval-based diagram, which offers simple yet effective way to augment the early exit framework during the inference stage. In this section, we begin with the definition of the early exit problem, then present key observations that motivate our approaches.
2

---

Published as a conference paper at ICLR 2026
2.1 PROBLEM STATEMENT
Formally, the early exit problem can be defined as follows: Given a backbone model M with m layers and an input x, early exit framework aims to design an exit function classifier / f(a) to determine whether to exit at the layer 1. The final prediction y is then transformed from the intermediate output states by of the I-th layer. And the final prediction probability can be formulated as P Z) = P(y (()))), where f(x) is trained or built on the external data D.
2.2 OBSERVATION I: EARLY EXIT AS A CORRECTIVE MECHANISM
Conventional early exit frameworks are primarily designed accelerate the inference the expense of model accuracy. However, they may overlook the great potential of carly exit corrective mechanism. This perspective indicates that intermediate layers can sometimes make better predictions than the output of the final layer. Therefore, argue that early exit should be considered not only as an acceleration technique, but more as a dynamic mechanism to correct the model's ' output for each individual input.
To validate the above claims, we conducted analysis experiments using RoBERTa-large on the SST-2 datasets in Figure Figure (a) shows an example of the probabilities of predicting the correct answers of a testing sample and their top-8 nearest neighbor samples over different layers. As shown in Figure (a). the model can already make correct predictions within the intermediate layers (layers 10-15), rather than performing all computations to obtain the final predictions. This indicates that if the model exits the inference from the intermediate layers, it may not degrade the overall performance.
Furthermore, we collect the correct prediction ratios of exiting from intermediate layers when the final predictions are wrong in Figure (b). When the model makes the correct predictions, the way of exiting from the intermediate layers can achieve the same accuracy. This is because, in the worst case, the model can still obtain the correct predictions at the final layer even if the model cannot predict correctly in any intermediate layers. Figure | (b) shows that when the full model makes wrong predictions, 90.66% on average of those wrong predictions can be corrected by the outputs of intermediate layers if can exit properly.
2.3
OBSERVATION 2: CONSISTENT EXIT BEHAVIORS SIMILAR DATA
Building on the corrective potential of early exit, a critical question arises: how can we reliably identify the optimal exit layer for a given input? We hypothesize that inputs with similar semantic content should exhibit similar optimal exit behaviors.
To validate this hypothesis, we build a retrieval database based on the exit behaviors of training data and collect the exit information of the top-8 nearest neighbors. As shown in Figure (a), we observe that not only does the test sample show higher confidence intermediate layers, but top-8 nearest neighbors display remarkably consistent pattern. Their probabilities of correct predictions are almost the same high at these intermediate layers (layers 10-15) compared to the final layer. This consistency suggests great opportunity to leverage the top-k nearest neighbors' exit behaviors to approximate the given input's 's exit behavior.
The above observations motivate to propose a retrieval-augmentte early exit framework, which can not only accelerate the model inference by exiting from the intermediate layers but also further correct the model's wrong predictions, thus improving model accuracy.
3 METHODOLOGY
This section presents the serrivall-uugment early exit framework in detail. First, this paper outlines the process of collecting exit features and constructing the retrieval database to facilitate early exit. Then, this paper introduces the aitreval-auugmenter early exit framework, denoted as RAEE.
3.1 COLLECTING THE EXIT FEATURES AND BUILDING THE RETRIEVAL DATABASE
This paper uses the collected exit features as the keys and values within the retrieval database. To avoid introducing too much retrieving overheads, this paper only retrieves once at the beginning of
3

---

Published as a conference paper at ICLR 2026
Build Train Example
EMB
Layer 1
Layer
Layer
Layer
Logits
LM Head
(layer .p).
(0 O''')x
(L.p') x
(2,p')
(.P)
(m.
Embeddiing as keys
Encoder
Retriever Index
Database
Embeddiings 88 queries
Exit Features
(layer index
Exit at atyer
Test Example
Layer 1
Layer
Layer j
Layer m
Logits
Head
Inference
collect features
model forward
build retriever
retrieve exit layer
activated modules
Inactivated modules
wrong prediction
correct prediction
Figure 2: The overview of the retrieval-augmentee early exit framework. During the build phase, a retrieval database is constructed from the collected exit features, including layer indexes and their corresponding probabilities for correct predictions. During the inference phase, the framework retrieves similar data's exit information based data embeddings to guide the model in selecting the optimal exit layer.
the backbone model. Consider the training data D {xxtrain y/ vfrain),.... (xtrain y backbone model M with m layers (CI...., ... In this context, as shown in the top part of Figure the keys I. are input embeddings of the training data, which can be obtained from an extra encoder model &, such as BERT (Devlin et al. 2019), or the outputs of embedding layers in the backbone model Memb.
} and
k= -
(1)
For the values, this paper collects set of possible exit layers I, and corresponding probabilities for each embedding ei, i.e., {1/ ((2.v() 1)) where m, indicates the number of possible exit layers for the embedding é;. The layer chosen as the exit layer is determined by whether the outputs of this layer h can be used make the right predictions ý compared the training labels yrrain, Then, the values V are all sets of possible exit layers,
== (1 = (00.00)7
(2)
We follow the same dataset splitting used in the LM-BFF (Gao et al. 2021 the collecting process requires parameters the model to update, only model inference performed.
After collecting keys and values for the retrieval databases, this paper uses state-of-the-aus approximate nearest neighbor search indexing, such FAISS (Johnson et al. 2019), and efficient key-value stores to build the retrieval database. More details can be found in Algorithm C.T of Appendix
3.2 THE RETRIEVAL-AUGMENTT EARLY EXIT FRAMEWORK
In this section, this paper then presents a retrieval -uggeented early exit framework named RAEE to optimize the model inference. RAEE regards the exit layer as random variable taking values in the set of {1...., m) where is the total number of layers the backbone model M. The probability mass function P(= I)) represents the probability of the case that the backbone model exits at the layer With the gold label, we can observe that the random variable z follows an unknown discrete distribution F. Then, this paper shows how leverage the retrieval database to approximate the distribution F.
Given input æ. RAEE first retrieves top-k nearest neighbors {Ei...., .k)....,)).. ) where each neighbor v has m, possible exit layers. Naturally, we can approximate the distribution by estimating the probability function ((= I),
4

---

Published as a conference paper at ICLR 2026
Table 1: Model performance of different methods eight downstream tasks. 'RB-L:, 'EB-L, and
'TS-L' refer to RoBERTa-Large. ElastiBERRT-Laree and T5-Large, respectively.
Methods
SST-2
SST-5
MR
CR
MPQA
SUBJ
TREC
CoLA
Avg
RB-L
83.60
34.98
80.80
79.55
67.60
51.45
32.40
2.03
54.05
EB-L
51.15
22.35
49.25
48.65
48.05
48.85
17.60
0.11
35.75
T5-L.
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 7B
49.08
28.64
50.05
50.10
50.00
48.05
14.40
-0.79
36.19
Backbone: RoBERTa-Lon Large,
ElasticBERT-Larg
HashEE (EB-L.)
49.08
14.16
49.95
50.05
50.00
50.00
27.00)
0.00
36.28
DeeBERT (RB-L.)
52.29
18.05
50.60
50.00
75.95
80.85
16.20
0.00
42.99
AdaInfer (RB-L)
50.92
24.48
50.00
50.00
60.90
50.85
22.60
-1.62
38.52
RAFE (RB-L)
84.63
33.57
81.55
68.05
78.55
84.05
62.40
14,48
63.41
Backbone: T5-Large
CALM (T5-L)
51.72
23.17
49.25
50.55
49.80
49.90
18.00
0.00
36.55
AdaInfer (T5-L)
50.11
28.14
50.35
49.80
46.30
49.95
26.00
5.22
38.23
RAFE (T5-L)
52.98
26.56
50.80
51.60
55.65
49.90
39.80
12.20
42.44
Backbone: Llama-3-8B
SLEB (Llama)
54.01
21.09
51.10
49.45
55.65
49.95
14.00
0.92
37.02
AdaInfer (Llama)
53.21
18.05
53.50
50.00
49.95
47.55
16.20
0.00
36.06
RAEE (Llama)
73.05
35,25
66.45
57.95
75.05
90.05
51.80
9.55
57.39
Backbone: Gemma-77
SLEB (Gemma)
50.69
19.82
49.95
49.95
50.00
52.10
12.80
0.00
35.66
AdaInfer (Gemma)
50.92
12.62
50.00
50.00
50.00
50.60
22.60
0.00
35.84
RABE (Gemma)
73.17
32.40
66.75
56.75
75.60
90.15
40.00
10.46
55.66
& Σ
P(z=11a0 x)
P(us x) Si.
i=1
my
S.
Σ
(any(() = && v 2 2)) vi.
Where is the indicator function that returns if the condition is true and otherwise, any(-) (-) is the function that returns true one condition is true and false otherwise, 7 is the threshold for filtering the layers, the inner loop only count once since there is at most one possible exit layer of neighbor that is equalt to I. Since different neighbors should have different contributions to the probability function ( 1). RAEE uses the reciprocal of the scaled distance between each neighbor and the query to estimate the contribution,
P(u x) min ((distance (U,, x)}$
distance(v;,
Then, RAEE designs a function f(x) to determine the exit layer, which selects the layer that maximizes the probability function P(= = l),
f()) arg max P(2 x))x).
Notably, when multiple exit layers have the same maximal probability, RAEE selects the earliest one.
The bottom part of Figure Seguree]shrr shows the inference workflow of RAEE. Specifically, RAEE first feeds the inputs into the backbone model for the label predictions and the same encoder used 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 exit layers Finally, RAEE stops
nearest neighbors, RAEE computes the exit layers based the Equations the forwarding at the calculated exit layer and passes the intermediate outputs the exit layer to the final prediction layer, e.g., LM Head in language models, to obtain the final predictions. This
5

---

Published as a conference paper at ICLR 2026
RB-L
Adainfer DeeBERT
RAEE Forwarding
EB
Liama
HashÈE
SLEB
RAEE Retrieving
MasticMEnT
-40
latency
120- 200
120-
(a) SST-2
(b) SUBJ
180
140
100-
Ulea
Gemma
1:0-
Rolk ERTty
Ulaaa
Gemma
140-
Inference istency
120-
120
Gemerit
(c) TREC
(d) CoLA
Figure 3: Inference latency of RAEE compared with various methods on selected downstream tasks. The backbone models used in the comparisons include LoBERTa-LargeEElasicEETTTT Ta-Large, Llama-3-8B, and Gemma-7B.
implemented based on the Transformer library, passing the exit layer as a parameter into the 'forward('' function and stopping the inner iteration based on the exit layer. More details can be found in Algorithm C.2 Appendix C
EXPERIMENTS
In this section, we first introduce the dataset and the experimental setup. Then, we presented the main results of different methods eight downstream tasks. We also conducted analysis experiments and ablation studies to show the impact of these factors on RAEE performance.
4.1 DATASET AND EXPERIMENTAL SETUP
Datasets We conducted comprehensive experiments across eight downstream tasks from GLUE benchmark (Wang et al. 2019), covering sentiment analysis, opinion polarity analysis, grammatical judgment, natural language inference, paraphrasing, etc.
Experimental Setup The proposed RAEE was implemented using the PyTorch framework and Transformer. We evaluated methods based on the backbone models RoBERTa-Land Large (Lu et al. 2019) and L5-Large (Raffel al., 2020) on one NVIDIA GeForee RTX 4090 with 24GB GPU memory, while Llama-3-8B (Dubey et al. 2024) and Gemma-7B (Mesnard et al. 2024) on one NVIDIA A100 GPU with 40GB GPU memory. The experiments were conducted two settings: training-free settings for methods that require parameter updates, and fine-tuning settings for methods that require searching for the optimal parameter exit classifiers. The evaluation metric accuracy, except for the Matthew correlation coefficient for the CoLA task. The number of retrieved nearest neighbors of RAEE is set to 12 in the experiments. The threshold r of RAEE is set to 0.9.
To validate the effectiveness, 1) RoBERTa-Larr (Liu et al. 2019). a state-of-the- encoder model, where the prompt-basee
compared RAEE with three types of methods. Pretrained Models:
version (Gao
al. 2021 is used; 2) ElasticBERT (Liu al. 2022). pre-trained multi-exit

---

Published as a conference paper at ICLR 2026
Table 2: Model performance and inference latency of RAFE and RAEE with only correctly predicted examples in the retrieval database. The backbone model is Llama-3-8B.
Models
SST-2
SST-5
MR
CR
MPQA
SubJ
TREC
CoLA
Ayg
Performance f
Llama
62.84
26.06
59.65
72.90
51.75
52.80
8.40
41.80
RAEE w/o
60.55
24.52
57.30
53.55
56,65
81.70
20,80
0.00
44.38
RAFE 73.05
35.25
66.45
57.95
75.05
90.05
51.80
9.55
$7.39
Latency (ms)
Llama
122.27
122.13
122.03
121.78
121.82
121.70
122.52
122.06
122.04
RACE w/o
37.65
115.26
37.98
31.69
54.08
34.59
112.03
91.34
64.33
RAEE
52.33
65.47
53.83
34.65
55.02
33.25
81.74
84.34
57.58
transformer model, where the large version is used in this paper: 3) T5-Large (Raffel et al. 2020),
versatile wansformer-based model for various NLP tasks; 4) Llama-3-8B (Dubey et al. 2024 a pre-trained model with strength in specific language scenarios; 5) Gemma- 7B (Mesnard et al. 2024), model with the potential for outstanding performance in specific settings. Training-Free Methods: 1) HashEE (Sun et al. 2022), a hash-based early exit approach with ElasticBERT-Large as its backbone model; CALM (Schuster et al. 2022), classical entropy- ttreeholldingg..... based early exit method with Largar as its backbone model, where the training-free setting is applied; 3) SLEB (Song et al. 2024). method that eliminates redundant transformer blocks. Semi-Training Methods: 1) Adalnfer (Fan et al., 2024). an SVM SVM-based early exiting method with our reproduced version; 2) DeeBERT (Xin al., 2020p. a classical entropy-
early exiting method its backbone model; The templates are listed in the Appendix B More
with RoBERTa-Laad
details about the experimental setup can be found in Appendix
4.2 MAIN RESULTS
Table 1 presents the main results, comparing the performance of RAEE against different types of methods across eight downstream tasks. Experimental results show that RAEE achieves the highest average performance 63.41 with RoBERTa-Large RAEE with Gemma-7B achieves the maximal improvement over baseline models, while with RoBERTa-Large it boosts performance from 36.28 to 63.41 over comparison methods. Across eight downstream tasks, RAEE consistently improves the model performance compared to current state-of-the-ane early exit frameworks.
Figure shows the inference latency RAFE and comparisons on eight downstream tasks. We show the inference latency on the selected four tasks, which covers different task types. For million- level backbone models, such as RoBERTa-Land ElasticBERT-Large. T5-Large, RAEE can achieve comparable inference efficiency. Since the inference speeds of those backbone models are already fast enough, adding too many components for early exit would degrade the inference efficiency like HashEE and DecBERT. However, for billion-level backbone models, the acceleration of RAEE significant. This benefits from the effectively predicted early exit layers. Although Adainfer with the backbone Llama-3-8B and Gemma-7B much faster than RAEE, it only achieves comparable performance of backbone models. RAEE significantly improves the performance of those backbone models and also reduces the inference latency by nearly half.
4.3 REASONS FOR SIGNIFICANT PERFORMANCE IMPROVEMENT
The main results demonstrate that the proposed RAFE can significantly outperform backbone models, which is not an intuitive result compared previous early exit methods. The reasons for this lie in that the collected exit information guides the RAEE an error corrector. This means RAEE can leverage exit information from cases correctly predicted by intermediate layers but missed by the full-backbone models without early exit.
To further validate the above claims, we performed an analysis using the retrieval database, which contains exit information only from the data correctly predicted by the full Llama-3-8B without early exit. As shown Table RAEE w/o refers to the one built on only correctly predicted examples. can be that RAEE w/o achieves comparable performance baselines while accelerating the
7

---

Published as a conference paper at ICLR 2026
Table 3:
The impact of retrieval number k the distribution approximation.
SST-2
SST-5
MR
CR
MPQA
SUBJ
TREC
CoLA
Avg
2 79.82
32.99
76.40
68.40
77.55
81.75
61.00
7.89
60.73
4
82.22
33.17
79.20
68.05
78.65
83.60
63.60
14.74
62.90
8
84.52
34.12
80.70
68.05
78.90
84.20
63.20
12.19
63.24
12
84.63
33.57
81.55
68.05
78.55
84.05
62.40
14.48
63.41
16
84.98
32.17
82.05
69.00
77.90
84.00
62.60
13.00
63.21
20
85.44
32.26
81.80
68.95
78.05
83.50
62.40
12.40
63.10
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%
84.63
31.76
82.80
65,85
75.70
81.00
58.80
8.77
61.16
50%
83.37
32.85
80.65
66.95
77.90
83.60
61.20
11.73
62.28
100%
84.63
33.57
81.55
68.05
78.55
84.05
62.40
14.48
63.41
inference process. This is because test examples, for which the full backbone correctly predict, RAEE w/o almost achieves the same correct predictions while exiting early. However, due to a lack of exit information on examples where the full backbone fails to predict, RAEE w/o also fails predict on the test data where backbone models fail. Therefore, when providing the exit 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
shows Impact of Top-k: Table shows the effect of varying retrieval numbers distribution approximation. With RoBERTa-Large RAEE improves from 60.73 to 63.41 as k increases, indicating that more retrieved exit information enhances approximation. However, beyond & 12, performance drops to 60.10, likely due to noise from less relevant retrievals. This suggests that limited amount of relevant exit data suffices for accurate approximation.
Retrieval Database Size: Table the performance of RAEE with different sizes of retrieval databases. As the database size increases, the performance of RAEE with the backbone model RoBERTa-Lane increases significantly from 61.16 to 63.41 on average. This demonstrates that collecting more data can improve the generalization of RABE, thus approximating the exit distribution more accurately.
4.5 OUT-OF-DOMAIN PERFORMANCE
Table 5: Model performance (ROUGE-L) and
layer usage on generation tasks. The backbone model is Llama-3-8B
To investigate out-of-domain issues, we con- ducted experiments on summarization tasks, in- cluding MNll (See et al. 2017) and XSum (Narayan et al., 2018), using retrieval database based on Wixi (Merity et al. al. 2017). as shown in Table 5 The results demonstrate that RAEE enhances performance and exits car- lier despite the retrieval database being out of distribution. Detailed experimental settings are provided in Appendix F
Models
CNN/DailyMail
XSum
Performance (ROUGE-L) Llama
8.95
5.22
RAEE Wiki
14.01
7.15
Layers
Llama
32.00
32.00
RAEE Wiki
29.60
28.82
8

---

Published as a conference paper at ICLR 2026
Table 6: The building of RAFE and different comparison methods. The results collected with the backbone model RoBERTa-Large (ElasticBERT for HashEE).
Model
SST-2
SST-5
MR
CR
MPQA
Subj
TREC
CoLA
Avg
Time Cost (seconds)
RAEE
91.75
111.80
112.09
27.43
110.26
103.99
71.19
108.59
92. 14
HashEF
5.85
14.13
14.68
2.82
10.40
13.84
7.50
5.81
9.38
AdaInfer
91.85
124.78
129.11
41.77
122.96
120.52
66.43
114.19
101.45
Storage Overheads (MB)
RAEE (Index)
3.4
3.7
3.8
2.0
3.8
3.6
3.1
3.7
3.4
RAFE (DB)
2.5
2.0
3.0
0.8
2.5
2.7
1.0
2.6
2.1
4.6 BUILDING OVERHEADS
We present the building overheads in Table 6 including database size, index size, and building time costs for different methods. Specifically, the building time cost for RAEE refers to the time required build the retriever, while the building time cost for AdaInfer pertains the time needed for training the classifier. In the case of HashEE, the building time cost corresponds to the time taken to build the hashing buckets. For RAEE, the average time required to build the retrieval database under minutes on single NVIDIA GeForce RTX 4090, which is considered an acceptable overhead in comparison the time involved fine-tuning. The index and database sizes are relatively small and can be considered negligible in comparison to the size of the backbone model.
RELATED WORK
5.1 EARLY EXIT FRAMEWORK
Early exit inference is popular pruning method to reduce computation overhead in text tasks. Most
current works [Bae et al. 2023 Kong et al. 2022 Ji et al. [023] Wolczyk et al. 2021 Hooper et al.)
2023) introduce clarsifiers in each layer determine whether the inference should continue. (Xin
et al.
2020: Liu et al.) 2020 Zhou et al. 2020 He et al.) 2021 Bajpai & Hanawal 2025b) train
classifiers by aligning intermediate and final outputs, then apply early exit based threshold. (Liao
et al. 202T) incorporates the past and future information to predict the early exit layer. Instead of
training a neural network as classifiers, (Fan et al. 2024 only fit the machine learning classifiers on
the extracted features for the early exit. (Zhu) [021 Zhu et al. 2021 2023 Zhang et al. (023a) focus
on designing novel loss functions for training more robust classifier. (Li
2021) incorporates
sentence-level features as well as token-level features to predict the early exit. Different from those
works, our method does not require training the classifier. (Sun et al. (022) proposes hash-based early exit method that uses the hashing functions to map tokens exit layers. (Bajpai & Hanawal 2024) proposes an online learning algorithm for BERT, dynamically setting exit points based on confidence thresholds. (Regol et al. 2023) proposes jointly-leamed framework for early exiting and
inference in dynamic neural networks, integrating gating mechanisms and intermediate inference
modules. (Balagansky & Gavrilov 2022)
introduces deterministic Q-exit criterion and revising the
model architecture. (Jabbec et al. 2024) further develops theoretically grounded, risk-controlled
exit rule that can wrap around existing dynamic models. Our method uses pre-built databases to predict exit layers, achieving better generalization and even outperforming the full backbone model.
Other works (Jazbec et al., 2023] Li et al., 2023 Huang al., 2018 Bajpai & Hanawal] 2025a] Meronen et al. (2024) focus on designing early exit frameworks for image classification tasks or multimodal tasks. They are optimizations in the different domains compared to our works.
5.2 RETRIEVAL-BASED AUGMENTATIONS
Retrieval- based augmentations et al. 2022 Wang et al. 2023 Xiong et al. 2023 Cui et al. 2023
Wu et al. 2024a have been widely used various natural language processing (NLP) tasks and achieved remarkable performance. Current works mostly leverage external knowledge databases to augment generator models various text-generation tasks, such language modeling (Khandelwal

---

Published as a conference paper at ICLR 2026
et al. 2020 Lewis al. 2020 Borgeaud et al.) 2022).
question-answee answering (Guu al. 2020
Izacard & Grave 2021 D. machine translation (Khandelwal et al., 2021; Wang et al. 2022). dialogue system (Cheng al. (023). Those works focus only improving generation quality, while approach simultaneously accelerates inference and enhances model performance. (Zhang et al. 2023b) stores the generation trajectories of diffusion models that differ from target 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.
CONCLUSION
This paper models the early exit problem a distribution prediction problem and observes that similar data's 's exit information can be used to approximate the distribution. Based on the observations, this paper proposes a retrieval- -aggeeetd early exit framework named RAEE. Experimental results show that RAFE can accelerate the model inference while significantly improving the model performance.
7 ACKNOWLEDGEMENT
This work was partially supported by Guangdong Provincial 1+1+1 Grant.
REFERENCES
Sangmin Bae, Jongwoo Ko, Hwanjun Song, and Se- Young Yun. Fast and robust early-exiting framework for autoregressive language models with synchronized parallel decoding. In Proceedings of the 2023 Conference Empirical Methods in Natural Language Processing (EMNLP), pp. 5910-5924, 2023.
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. Longbench: A bilingual, multitask benchmark for long context understanding. In Lun-Wei Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2024, Bangkok, Thailand, August 11-16, 2024, pp. Association for Computational Linguistics, 2024.
Divya Jyoti Bajpai and Manjesh K. Hanawal. Ceebert:
inference in early exit BERT. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Findings of the Association for Computational Linguistics, ACL 2024, Bangkok, Thailand and virtual meeting, August 11-16, 2024, pp. 1736-1748. Association for Computational Linguistics, 2024.
Divya Jyoti Bajpai and Manjesh Kumar Hanawal. FREE: fast and robust vision language models with early exits. In Wanxiang Che, Joyce Nabende, Ekaterina Shutova, and Mohammad Taher Pilehvar (eds.), Findings of the Association for Computational Linguistics, ACL 2025, Vienna, Austria, July 27 August 1, 2025, pp.
Association for Computational Linguistics, 2025a.
Divya Jyoti Bajpai and Manjesh Kumar Hanawal. BEEM: boosting performance of early exit dnns using multi-exit classifiers as experts. The Thirteenth International Conference on Learning Representations, ICLR 2025, Singapore, April 24-28, 2025. OpenReview.net, 2025b.
Nikita Balagansky and Daniil Gavrilov. PALBERT: teaching ALBERT ponder. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NewrlPS 2022, New Orleans, LA, USA, November 28 December 9, 2022, 2022.
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Mil- lican, George van den Driessche, rean-Baptise Lespiau, Bogdan Damoc, Aidan Clark, Diego de Las Casas, Aurelia Guy. Jacob Menick, Roman Ring, Tom Hennigan, Saffron Huang, Loren Maggiore, Chris Jones, Albin Cassirer, Andy Brock, Michela Paganini, Geoffrey Irving, Oriol Vinyals, Simon Osindero, Karen Simonyan, Jack W. Rae, Erich Elsen, and Laurent Sifre. Improv- ing language models by retrieving from trillions of tokens. In Proceedings of the 2022 International Conference Machine Learning (ICML), pp. 2206-2240, 2022.
10

---

Published as a conference paper at ICLR 2026
Xin Cheng, Di Luo, Xiuying Chen, Lemao Liu, Dongyan Zhao, and Rui Yan. Lift yourself up: Retrieval-augmente text generation with melf-memory. In Proceedings of the Advances in Neural Information Processing Systems 36: Anmual Conference on Neural Information Processing Systems (NeurIPS), 2023.
Yufei Cui, Ziquan Liu, Yixin Chen, Yuchen Lu, Xinyue Yu, Xue (Steve) Liu, Tei-Wei Kuo, Miguel Rodrigues, Chun Jason Xue, and Antoni B. Chan. Retrieval-augment multiple instance learning. In Proceedings of the Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023 (NeurIPS). 2023.
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Proceedings of the Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems (NeurIPS), 2022.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre- training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (NAACL-HILT), pp. 4171-4188 Association for Computational Linguistics, 2019.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, and et al. The llama 3 herd of models, 2024. URL
https://arxiv org/ abs
www./og/////////eoo/oooooe 2407 21783
Siqi Fan, Xin Jiang, Xiang Li, Xuying Meng, Peng Han, Shuo Shang, Aixin Sun, Yequan Wang. and Zhongyuan Wang, Not all layers of Ilms are necessary during inference. CoRR, abs/2403.02181, 2024.
Tianyu Gao, Adam Fisch, and Danqi Chen. Making pre-trained language models better few-shot learners. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguis- tics and the 1 11th International Joint Conference on Natural Language Processing (ACL/IJCNP), pp. 381 3816-3830. Association for Computational Linguistics, 2021.
Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Ming- Wei Chang. Retrieval augmented language model pre-training. In Proceedings of the 37th International Conference on Machine Learning (ICML). pp. 3929-3938, 2020.
Xuanli He, Iman Keivanloo, Yi Xu, Xiang He, Belinda Zeng, Santosh Rajagopalan, and Trishul Chilimbi. Magic pyramid: Accelerating inference with early exiting and token pruning. CoRR, abs/2111.00255 2021.
Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Hasan Gene, Kurt Keutzer, Amir Gholami, and Yakun Sophia Shao. SPEED: speculative pipelined execution for efficient decoding, CoRR, abs/2310. 12072, 2023.
Gao Huang, Danlu Chen, Tianhong Li, Felix Wu, Laurens van der Maaten, and Kilian Q. Weinberger. Multi-scale dense networks for resource efficient image classification. In 6th International Confer- ence on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30- - May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018
Gautier Izacard and Edouard Grave. Leveraging passage retrieval with generative models for open domain question answering. In Proceedings of the 16th Conference of the European Chapter of the Association for Computational Linguistics (EACL): Main Volume, pp. 880, 2021.
Metod Jazbec, James Urquhart Allingham, Dan Zhang, and Eric T. Nalisnick. Towards anytime classification in early-exit architectures by enforcing conditional monotonicity. In Proceedings of the Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems (NeurIPS), 2023.
Metod Jazbec, Alexander Timans, Tin Hadzi Veljkovic, Kaspar Sakmann. Dan Zhang, Christian An- dersson Naesseth, and Eric T. Nalisnick. Fast yet safe: Early-exiting with risk control. In Amir Globersons, Lester Mackey, Danielle Belgrave, Angela Fan, Ulrich Paquet, Jakub M. Tomezak,
11

---

Published as a conference paper at ICLR 2026
and Cheng Zhang (eds.), Advances in Neural Information Processing Systems 38: Annual Confer- ence on Neural Information Processing Systems 2024, NeurlPS 2024, Vancouver, BC, Canada, December 10- 15, 2024, 2024.
Yixin Ji, Jikai Wang, Juntao Li, Qiang Chen, Wenliang Chen, and Min Zhang. Early exit with disentangled representation and equiangular tight frame. In Findings of the Association for Computational Linguistics (ACL), pp. 14128-14142, 2023.
Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scale similarity search with GPUs. IEEE Transactions on Big Data, 7(3):535-547, 2019.
Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generalization through memorization: Nearest neighbor language models. In Proceedings of the 8th International Conference on Learning Representations (ICLR), 2020.
Urvashi Khandelwal, Angela Fan, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Nearest neighbor machine translation. In Proceedings of the 9th International Conference on Learning Representations (ICLR), 2021.
Jun Kong. Jin Wang, Liang-Chie Chih Yu, and Xuejie Zhang. Accelerating inference for pretrained language models by unified multi-perspective early exiting. In Proceedings of the 29th International Conference on Computational Linguisties (COLING), pp. 4677 7686, 2022.
Patrick S. H. Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, Sebastian Riedel, and Douwe Kiela. Retrieval- -uggrente generation for smowledge-intensiee NLP tasks. In Proceedings of the Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems (NeurIPS), 2020.
Xiaonan Li, Yunfan Shao, Tianxiang Sun, Hang Yan, Xipeng Qiu, and Xuanjing Huang. Accelerating BERT inference for sequence labeling via early-exit. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (ACL/UJCNP),, pp. 189 199, 2021.
Yuhang Li, Tamar Geller, Youngeun Kim, and Priyadarshini Panda. SEENN: towards temporal spiking early exit neural networks. In Proceedings of the Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems (NeurIPS), 2023.
Zonglin Li, Ruiqi Guo, and Sanjiv Kumar. Decoupled context processing for context augmented language modeling. In Proceedings of the Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems (NeurIPS), 2022.
Kaiyuan Liao, Yi Zhang, Xuancheng Ren, Qi Su, Xu Sun, and Bin He. A global past-future early exit method for accelerating inference of pre-trained language models. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics (NAACL-HLLT): Human Language Technologies, pp. 2013-2023, 2021.
Weijie Liu, Peng Zhou, Zhiruo Wang, Zhe Zhao, Haotang Deng, and Qi Ju. Fastbert: a self-distilling BERT with adaptive inference time. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics (ACL), pp. 6035-6044, 2020.
Xiangyang Liu, Tianxiang Sun, Junliang He, Jiawen Wu, Lingling Wu, Xinyu Zhang, Hao Jiang, Zhao Cao, Xuanjing Huang, and Xipeng Qiu. Towards efficient NLP: A standard evaluation and A strong baseline. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics (NAACL): Human Language Technologies, pp. 3288-3303, 2022.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer. and Veselin Stoyanov. Roberta: A robustly optimized BERT pretraining approach. CoRR, abs/1907. 16622, 2019.
12

---

Published as a conference paper at ICLR 2026
Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivastava, Ce Zhang, Yuandong Tian, Christopher Ré, and Beidi Chen. Deja vu: Contextual sparsity for efficient Iims at inference time. In Proceedings of the 2023 International Conference on Machine Learning (ICML), pp. 22137-22176. 2023.
Xinyin Ma, Gongfan Fang, and Xinchao Wang. LIm-pruner: On the structural pruning of large language models. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Proceedings of the Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems (NeurlPS), 2023.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017.
Lassi Meronen, Martin Trapp. Andrea Pilzer, Le Yang, and Amo Solin. Fixing overconfidence in dynamic neural networks. In IEEE/CVF Winter Conference on Applications of Computer Vision, WACV 2024, Waikoloa, HI, USA, January 3-8, 2024, pp. 2668-2678. IEEE, 2024.
Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, Pouya Tafti, Léonard Hussenot, Aakanksha Chowdhery, Adam Roberts, Aditya Barua, Alex Botev, Alex Castro-Ros, Ambrose Slone, Amélie Héliou, Andrea Tacchetti, Anna Bulanova, Antonia Paterson, Beth Tsai, Bobak Shahriari, Char- line Le Lan, Christopher A. Choquette- Reid, Elena Buchatskaya, Eric Ni, Eric Noland, Geng Yan, George Tucker, George- -ringgen Mu- raru, Grigory Rozhdestvenskkiy. Henryk Michalewski, Ian Tenney, Ivan Grishchenko, Jacob James Keeling, Jane Labanowski, Jean-Baptiste Lespiau, Jeff Stanway, Jenny Brennan, Jeremy Chen, Johan Ferret, Justin Chiu, and et al. Gemma: Open models based on gemini research and technology. CoRR, abs/2403.0823 08295, 2024.
Clément Crepy, Daniel Cer, Daphne Ippolito, David
Austin,
Shashi Narayan, Shay B. Cohen, and Mirella Lapata. Don't give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization, In Ellen Riloff, David Chiang, Julia Hockenmaier, and Jun't Jun'ichi ichi Tsujii (eds.), Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, Brussels, Belgium. October 31 November 4, 2018. pp. 1797-1807. Association for Computational Linguistics, 2018
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res., 21:140:1-140:67, 2020.
Florence Regol, Joud Chataoui, and Mark Coates. Jointly-leareed exit and inference for a dynamic neural network JEI-DNN. CoRR, abs/2310.09163, 2023.
Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Ilie, Daniel Hesslow, Roman Castagné, Alexandra Sasha Luccioni, François Yvon, Matthias Gallé, Jonathan Tow, Alexander M. Rush, Stella Biderman, Albert Webson, Pawan Sasanka Ammanamanchi, Thomas Wang, Benoît Sagot, Niklas Muennighoff, Albert Villanova del Moral, Olatunji Ruwase, Rachel Bawden, Stas Bekman, Angelina McMillan-Major, Iz Beltagy, Huu Nguyen, Lucile Saulnier, Samson Tan, Pedro Ortiz Suarez, Victor Sanh, Hugo Laurençon, Yacine Jernite, Julien Launay, Margaret Mitchell, Colin Raffel, Aaron Gokaslan, Adi Simhi, Aitor Soroa, Alham Fikri Aji, Amit Alfassy, Anna Rogers, Ariel Kreisberg Nitzav, Canwen Xu, Chenghao Mou, Chris Emezue, Christopher Klamm, Colin Leong, Daniel van Strien, David Ifeoluwa Adelani, and et al. BLOOM: A 176b- parameter open-access multilingual language model. CoRR, abs/221 1.05100, 2022.
Tal Schuster, Adam Fisch, Jai Gupta, Mostafa Dehghani, Dara Bahri, Vinh Tran, Yi Tay, and Donald Metzler. Confident adaptive language modeling. In Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022.
Abigail See, Peter J. Liu, and Christopher D. Manning. Get to the point: Summarization with pointer- generator networks. In Regina Barzilay and Min- Yen Kan (eds.), Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics, ACL 2017, Vancouver, Canada, July 30 August 4, Volume 1: Long Papers, pp. 1073-1083, Association for Computational Linguistics. 2017.
13

---

Published as a conference paper at ICLR 2026
Jiwon Song, Kyungseok Oh, Taesu Kim, Hyungjun Kim, Yulhwa Kim, and Jae-Joon Kim. SLEB: streamlining Ilms through redundancy verification and elimination of transformer blocks. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27. 2024. OpenReview.net, 2024.
Tianxiang Sun, Xiangyang Liu, Wei Zhu, Zhichao Geng, Lingling Wu, Yilong Hc, Yuan Ni, Guotong Xie, Xuanjing Huang, and Xipeng Qiu. A simple hash-based early exiting approach for language understanding and generation. In Findings of the Association for Computational Linguistics (ACL), 2022.
pp.
Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv Kulshreshtha, Heng-Tzz Tze Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du, YaGuang Li, Hongrae Lee, Huaixiu Steven Zheng, Amin Ghafouri, Marcelo Menegali, Yanping Huang. Maxim Krikun. Dmitry Lepikhin, James Qin, Dehao Chen, Yuanzhong Xu, Zhifeng Chen, Adam Roberts, Maarten Bosma, Yanqi Zhou, Chung-Ching Chang, Igor Krivokon, Will Rusch, Mare Pickett, Kathleen S. Meier-Hellsteem. Meredith Ringel Morris, Tulsee Doshi, Renelito Delos Santos, Toju Duke, Johnny Soraker, Ben Zevenbergen, Vinodkumar Prabhakaran, Mark Diaz, Ben Hutchinson, Kristen Olson, Ale- jandra Molina, Erin Hoffiman-John, Josh Lee, Lora Aroyo, Ravi Rajakumar, Alena Butryna, Matthew Lamm, Viktoriya Kuzmina, Joe Fenton, Aaron Cohen, Rachel Bernstein, Ray Kurzweil, Blaise Agüera y Arcas, Claire Cui, Marian Croak, Ed H. Chi, and Quoc Le. Lamda: Language models for dialog applications. CoRR, 11s/2201.082399 2022.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurélien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models. CoRR, abs/2302. 13971, 2023.
Tim Valicenti, Justice Vidal, and Ritik Patnaik. Mini-gpts: Efficient large language models through contextual pruning. CORR, abs/2312.12688 2023.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In Proceedings of the 7th International Conference Learning Representations (ICLR), 2019.
Shuohang Wang, Yichong Xu, Yuwei Fang, Yang Liu, Siqi Sun, Ruochen Xu, Chenguang Zhu, and Michael Zeng. Training data is more valuable than you think: A simple and effective method by retrieving from training data. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (ACL), pp. 3170-3179, 2022.
Weizhi Wang, Li Dong, Hao Cheng, Xiaodong Liu, Xifeng Yan, Jianfeng Gao, and Furu Wei. Augmenting language models with long-term memory. In Proceedings of the Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems (NeurIPS), 2023.
Maciej Wolezyk, Bartosz Wójcik, Klaudia Balazy, Igor T. Podolak, Jacek Tabor, Marek Smieja, and Tomasz Trzeinski. Zero time waste: Recycling predictions early exit neural networks. In Proceedings of the Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems (NeurIPS), pp. 2516-2228, 2021.
Shangyu Wu, Ying Xiong, Yufei Cui, Xue Liu, Buzhou Tang, Tei-Wei Kuo, and Chun Jason Xue. Improving natural language understanding with computation-efficcent retrieval representation fusion. CORR, 00./2/401.02999 2024a.
Shangyu WU, Ying Xiong, Yufei Cui, Haolun Wu, Can Chen, Ye Yuan, Lianming Huang, Xue Liu, Tei-Wei Kuo, Nan Guan, and Chun Jason Xue. Retrieval-augmented generation for natural language processing: A survey. CoRR, A3s/2407.113193, 2024b.
Ji Xin, Raphael Tang. Jaejun Lee, Yaoliang Yu. and Jimmy Lin. Deebert: Dynamic early exiting for accelerating BERT inference. In Proceedings of the 58th Anmual Meeting of the Association for Computational Linguistics (ACI.), pp. 2246-225 2020.
14

---

Published as a conference paper at ICLR 2026
Ying Xiong, Xin Yang, Linjing Liu, Ka-Chun Wong, Qingcai Chen, Yang Xiang, and Buzhou Tang. EARA: improving biomedical semantic textual similarity with entity-aligned attention and retrieval augmentation. In Findings of the Association for Computational Linguistics (EMNLP), pp. 8760-8771, 2023.
Jingfan Zhang, Ming Tan, Pengyu Dai, and Wei Zhu. LECO: improving early exiting via learned exits and comparison- based exiting mechanism. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (ACL): Student Research Workshop, pp. 298-309,
2023a.
Kexun Zhang, Xianjun Yang. William Yang Wang, and Lei Li. Redi: Efficient learning-free diffusion inference via trajectory retrieval. In Proceedings of the 2023 International Conference on Machine Learning (ICML), pp. 41770-41788 2023b.
Wangchunshu Zhou, Canwen Xu, Tao Ge, Julian J. McAuley, Ke Xu, and Furu Wei. BERT loses patience: Fast and robust inference with early exit. In Proceedings of the Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems (NeurIPS), 2020.
Wei Zhu. Leebert: Learned early exit for BERT with cross-level optimization. Proceedings of the 59th Annual Meeling of the Association for Compulational Linguisticy and the 1 11th International Joint Conference on Natural Language Processing (ACL/IJCNP),, pp. 2968-29800 2980, 2021.
Wei Zhu, Xiaoling Wang, Yuan Ni, and Guotong Xie. GAML-BERT: improving BERT early exiting by gradient aligned mutual learning. In Proceedings the 2021 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 3033-3044, 2021.
Wei Zhu, Peng Wang, Yuan Ni, Guotong Xie, and Xiaoling Wang. BADGE: speeding up BERT inference after deployment via block-wise bypasses and divergence based early exiting. In Proceedings of the The 61 Annual Meeting of the Association for Computational Linguistics (ACL): Industry Track, pp. 500-509, 2023.
15

---

Published as a conference paper at ICLR 2026
A THE USE OF LARGE LANGUAGE MODELS
We used a large language model mainly for language polishing (grammar, phrasing, and clarity) in paper. All LLM suggestions were manually reviewed by the authors ensure accuracy and consistency with our research objectives and writing style.
B TEMPLATES ON ALL TASKS
Table 7 provides an overview of the manual templates and selected label words used for each dataset with the backbone model RoBERTa-Larg. Large (Liu et al. (209)) in this paper. These templates and label words were created following LM-BFF (Gao et al. 2021).
Table 7: Templates and label words with the backbone model RoBERTa-LLarge.
Task
Prompts
Label word
SST-2
[CLS] was [MASK]. [SEP]
"t"rrtter, "I":""reat" "0":""eeriblee.""""" "bad",
SST-5
[CLS] It was [MASK]. [SEP]
"2": "okay", "good","""""" "great"
MR
[CLS] It [MASK]. [SEP]
CR
[CLS] æ was [MASK]. [SEP]
"O":"terrible", "1":"great"
MPQA
[CLS] It [MASK]. [SEP]
"O":"terrible" "1"="great"
SUBJ
[CLS] This is [MASK]. [SEP]
"0": "ubjecttc,,,,"" "I"":"ojecctive"
TREC
[CLS] [MASK] æ: [SEP]
"0": "essreesriipion" "I"tt"Enitt" "xxpe'Exprsss"o""
"3"." ""mman"""""""LLcattiii" "4"::"oocttion" ".""."
CoLA
[CLS] It was [MASK]. [SEP]
"0": "Ono"ric"r,,""" "I":"correct"
Table 8 provides overview of the manual templates and selected label words used for each dataset with the backbone model T5-Large (Raffel et al.) 2020), Llama-3-8B (Dubey et al.) (2024) and
Gemma-7B (Mesnard al. 2024) in this paper.
Table 8: Templates and label words with the backbone model Ta-Large, Llama-3-8B and Gemma- 7B.
Task
Prompts
Label word
SST-2
What is the sentiment of the sentence Print negative or positive. The answer is What is the sentiment of the sentence "?
"0"; "O""atvaa,,""" "I":"positive"
SST-5
"teriilrill""""" "1" "bad",
Print terrible, bad, okay, good or great. The answer is What is the sentiment of the sentence x Print negative or positive. The answer is What is the sentiment of the sentence x: Print negative or positive. The answer is
"2": "okay", "3": "good",
"4": "great"
MR
"negatvev""" "I"." "p""iiivvee""
CR
"0": "negatvvv,""" "p"t"positvve"
MPQA
What is the sentiment of the sentence x Print negative or positive. The answer is What is the subjectivity of the sentence Print subjective or objective. The answer is Print the category for the sentence x: description, entity, expression, person, location or quantity. The answer is the sentence æ grammatically acceptable?
"("" "negatvvv,""" "1":"positive
SUBJ
"O":"""bjectti "T"o"obecc
TREC
"O": "de'drccppption"" """" "1""tettity"
"2" """"rreesiii"""""" """"iioon"""""""""""""""""
"4":"location". "5""hinntt""""""tttttt"""""" "O":"no", "1":"yes"
CoLA
Print yes. The answer is
16

---

Published as a conference paper at ICLR 2026
Algorithm C.1 Collect the exit features keys and values for building the retrieval database.
Input: Given the training data available for model inference D {<xfrrnn "f (2"p) Pi ")"))) } backbone model M with m layers (Ci,...: L.h). encoder (None value means no encoder provided).
Output: Keys and values V.
1: K. = [1 V =
2: for 1. do do
3: V; -]:
4:
ho Memm (afaaa)
5:
for j=l..... do
6:
h (hg-11))
1 Compute intermediate outputs at layer j Predict from the layer j
7:
logits Mimmhhea(hh)) g = arg max logits;
8:
9:
P = max (sottmax
is equal to y grasn then Add (j.r.) into end end for
11:
1 Store the possible exit layer */
12:
14:
Add U2 into V:
15:
if E is None then
16:
Add he into K::
1a Store backbone embeddings when encoder model
end if end for
18:
19:
if & is not None then
20:
Add all E(a)rinn) into IC:
f* Store encoder embeddings //
21:
end if
return KC, V:
C DETAILED ALGORITHMS OF RAEE
Algorithm C.1 collects keys and values for building the retrieval database. For each sample training data D, the backbone model M is traversed layer by layer to compute the hidden state hj and corresponding logits. If prediction exiting information, including the layer j and the probability pi. is added the sample's value list (Line 2-1 When the encoder E is unavailable (Line 16), RAEE utilizes the hidden states from the backbone model M as embeddings for indexing. The specific layer from which the hidden states are extracted is treated as hyperparameter that the user can define.
certain layer j matches the sample's true label y/ the
Algorithm performs model inference with retrieval-augmenned early exiting. When the encoder E is unavailable (Line 5), RAEE utilizes the hidden states from the backbone model M as embed- dings for querying. The specific layer from which the hidden states are extracted is treated hyperparameter that the user can define.
D EXIT LAYERS
Table compares the average exit layers of the RABE method against other method types across eight downstream tasks. Experimental results show that the RAEE method can exit earlier, thus reducing computational overhead during model inference. This result also aligns with the expectations in the motivation example. This suggests that the RAEE method can accurately approximate the gold exit layer distribution by using the retrieval database. Although Adalnfer exits earlier than the RAFE method, it exhibits quite poor performance, as shown in Table [ The reason may be that the collected features can only provide limited information for the SVM, thus resulting unstable prediction performance.
IMPLEMENTATION DETAILS
This section lists the implementation details.
17

---

Published as a conference paper at ICLR 2026
Algorithm C.2 Model inference with the synchronized retrieval early exit.
Input: Input 2, backbone model M with layers (f1,...: L. encoder &, indexing I, top-k, the exit layer determination function f()).
Output: Final prediction ÿ.
1: An = (z)s(z))
2: if not None then
3: Equery = {());
18 Encode the inputs when the encoder available
4: else
5: Equery = ho:
1 Use the embeddings of backbone model */
6: end if
7: [(N,, [= = I[Qquerr: k);
1 Retrieve the possible exit layers */
8: I= f [(((V), disi)...... (uk, ))):
1 Obtain exit layer
9: for lrri=,...,fu
10: 4; 1i(
1a Perform model inference with early exit */
11: end for
logits Mimdlead(hl):
Predict based on the layer hy outputs
i arg max logits;
14: return y:
Table 9: Exit layers of RAEE and different types of methods on 8 downstream tasks. The sum of the number of layers in the encoder and the decoder counts the number layers for T5-large (Raffel
et al. 2020
Model
SST-2
SST-5
MR
CR
MPQA
Subj
TREC
CoLA
Avg
RB-L
24
24
24
24
24
24
24
24
24
EB-L
24
24
24
24
24
24
24
24
24
T5-L
48
48
48
48
48
48
48
48
48
Llama-3-8B
32
32
32
32
32
32
32
32
32
Gemma-78
28
28
28
28
28
28
28
28
Backbone: RoBERTa- Large, ElasticBERT Large
DeeBERT
22.95
24.00
23.33
8.98
15.90
10.36
24.00
18.31
18.48
AdaInfer (RB-L)
1.00
0.00
1.46
1.00
18.00
1.10
0.00
4.00
3.32
RAEE (RB-L)
18.55
13.93
18.71
15.35
17.20
13.59
12.82
12.48
15.33
Backbone: T5-L
AdaInfer (T5-L)
6.34
0.00
7.72
0.00
1.00
1.00
0.00
1.00
2.13
RAEE (T5-L)
22.27
18.74
21.88
26.84
18.05
19.06
27.29
18.55
21.59
Backbone: Llama-3-8B
SLEB (Llama)
13.00
13.00
13.00
13.00
13.00
13.00
13.00
13.00
13.00
Adalnfer (Llama)
4.00
0.00
3.18
3.00
1.00
4.71
0.00
2.00
2.24
RAEE (Llama)
11.77
15.70
12.43
7.04
12.83
6.58
20.06
21.04
13.43
Backbone: Gemma-7B
SLEB (Gemma)
11.00
11.00
11.00
11.00
11.00
11.00
11.00
11.00
11.00
AdaInfer (Gemma)
1.00
0.00
1.04
1.00
3.00
1.00
0.00
2.00
1.13
RAEE (Gemma)
11.00
17.62
11.70
3.29
14.72
0.51
9.50
20.06
11.05
For DeeBERT (in al. 2020). we use RoBERTa-Large its backbone model. Since DeeBERT(Xin et al. (020) is a classical entropy-thredd sengooy-thresholdng--bbbeee requires first fine-tunin fine-tuning the backbone model on the downstream task and then updating all but the last off-ramp, for fair comparison. we only update the off-ramp in DecBERT 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
early-exit method,
experimental settings for DeeBERT (Xin et 2020) remain as default.
For CALM (Schuster et al. 2022).
T5- T (Raffel et al.) (020)) its backbone model. CALM (Schuster et al.] 2022) is also a classical thresooditrs based early- exit method, and evaluate it under the training-free setting.
18

---

Published as a conference paper at ICLR 2026
For (LLEB(Song al. 2024), use Llama-3-8h Dubey et al. 2024) and Gemma- 7B (Mes nard et al. 2024) as its backbone model. (Snn et al. of early exit methods by eliminating redundant transformer blocks. Since the proposed RAFE exits at about 40% layers, for a fair comparison, we also set the hyper-parameter
2024) tackles the limitation
num_ www.rremoeeeeleeee blocks of SLEB et al. 2024) int(60% num_layers) for comparable efficiency.
F EXPERIMENTAL SETUP FOR OUT-OF-DOMAIN PERFORMANCE
This section details the experimental setup for the out-of-domain performance analysis.
To investigate out-of-domain issues, constructed a retrieval database based WikiText-2- vI (Merity et al. 2017). Since the text dataset lacks gold labels, we adopted the next-token prediction task setting, where the next token of each input sentence serves as the gold label.
Specifically, we first segmented the entire text dataset into individual sentences to preserve semantic integrity. Then, based the maximum input length of the backbone model, for sentences shorter than this limit, we designated the last meaningful token the gold label. For longer sentences, we applied sliding window of the model's maximum input length, selecting the last meaningful token within each window as the gold label. Finally, we collected exit information using the approach described in this paper.
Using the retrieval database constructed from WikiText (Merity et al. 2017). we conducted summa- rization experiments on CNN/Daily Mail (See et al.) (017) and XSum (Narayan et al. (018). The evaluation metrics for these tasks align with those used for summarization tasks in LongBench (Bai ct al. 20t4 the zero-shot setting.
G RETRIEVED EXAMPLES OF RAEE
We show two examples from the SST-2 task and their retrieved top-k data samples. As shown in Table TO and Table [I] the retrieved samples are semantically similar the query sentence, demonstrating the proposed RAEE's efficacy.
H LAYER-WISE EARLY-EXIT DISTRIBUTIONS
To show a visualization of carly-exit behavior, we report the proportion of samples that exit at each layer for eight downstream tasks. Figures and 5 show the layer-wise early-exit distributions of RAEE when applied to RoBERTa-large and LLaMA-3-8B, respectively. For each task, the r-axis denotes the exit layer index and the y-axis denotes the fraction of test samples that terminate that layer. These figures present results consistent with the latency trends in Figure 3 providing more fine-grained evidence of how RAEE redistributes computation across layers. For example, RAEE exits mostly at the 1st and 1 1th layers in Figure 5(g), leading to a substantial speedup as shown in Figure3jbb In contrast, RAEE exits predominantly around the 27th layer in Figure(5ih), backbone. For the
which results in smaller speedup in Figure () when using the LLaMA-3-8 RoBERTa-large backbone, the exit distributions in Figure and Figure 4 g)) are similar, which
consistent with the comparable inference latencies reported in Figures b) and Figure (d).
19

---

Published as a conference paper at ICLR 2026
Layer Distribution on
Layer Distribution on
Estimated Exit Distribution
0.40
0.35
0.30
0.3
0.25
Estimated Exit Distribution
Propertion
0.2
0.20
0.15
0.1
0.10
0.05
0.0
0.00
Layer
Layer
CoLA
(b) CR
Exit Layer Distribution on MPQA
Exit Layer Distribution on MR
0.30
0.25
0.30
Estimated Exit Distribution
0.25
0.20
0.20
Proportion
0.15
Exit Distribution
0.15
0.30
0.10
0.05
0.05
0.00
0.00
10 11 13 Exit Layer
20
30 13
Exit Layer
(c) MPQA
(d) MR
Exit Layer Distribution on SST-2
Exit Layer Distribution SST-5
0.30
Samples
Estimated Exit Distribution
0.25
0.20
0.15
0.15
Estimated Exit Distribution
0.10
0.05
0.05
0.00
0.00
Layer
Layer
(e) SST-2
(f) SST-5
Exit Layer Distribution on SUB)
Exit Layer Distribution on TREC
0.40
Samples les
0.35
Estimated Exit Distribution
0.35
Estimated Exit Distribution
0.30
0.30
0.25
0.20
0.20
0.15
0.15
0.10
0.05
0.05
0.00
0.00
10 Exit Layer
21
Exil Layer
SUBJ
(h) TREC
Figure 4: Layer-wise early-exit distributions of RAEE with RoBERTa-larg backbone on eight downstream tasks. For each task, the r-axis is the exit layer and the y-axis is the proportion of samples that terminate at that layer.
20

---

Published as a conference paper at ICLR 2026
Exit Layer Distribution on CoLA
Exit Layer Distribution on CR
Estimated Exit Distribution
Samples
Estimated Exit Distribution
0.5-
Proportion of
03
0.2
0.0
0.0
tass to Te miehinn i 1408 2019 nor 2r7m0s
Exit Layer
Exit Layer
(a) CoLA
(b) CR
Exit Layer Distribution MPQA
Exit Layer Distribution MR
0.4
Samples
Estimated Exit Distribution
M
0.25
0.3
0.20
Progortion
0.2
Estimated Exit Distribution
0.15
0.10
0.05
Layer
Layer
(c) MPQA
(d) MR
Exit Layer Distribution on SST-2
Exit Layer Distribution on SST-5
0:35
Estimeted Exit Distribution
Proportion Samples
0.4-
Estimated Exit Distribution
0.25
0.3
0.20
Proportion
0.15
02
0.10
0.1
0.05
0.00
0.0
191144 58 2013
12 16 29 20 27
Exit Layer
Exit Layer
(e) SST-2
(1) SST-5
Exit Layer Distribution SUB]
Exit Layer Distribution TREC
Proportion Samples
0.5
Estimated Exit Distribution
Propertion of Samples
0.30
Estimated Exit Distribution
0.25
0.20
0.15
0.2
0.10
0.1
0.05
Layer
Layer
(g) SUBJ
(h) TREC
Figure 5: Layer-wise early-exit distributions of RAEE with LLaMA-3-8B backbone eight downstream tasks. For each task, the --axis is the exit layer and the y-axis the proportion of samples that terminate that layer.
21

---

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

---

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


Reviews:
[{"review_id": "dY88sM9rpu", "reviewer": "Reviewer_niT3", "summary": "The paper presents a retrieval-augmented early exit framework that enhances the overall model performance by classifying the input at one of the intermediate exit points. The proposed RAEE model can outperform the full model performance while improving the inference speed of the model. The method does not require any training or additional classifiers to which makes the method efficient. The method is tested upon eight downstream tasks which further validates the method importance.", "questions": "See weaknesses.", "limitations": "", "rating": 2, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The method improves the efficiency of a model by early exiting where the decision to make an exit is based on the analysis of similar previously seen samples helping in training free-exiting.\n\n2. The method utilises the data characteristics like word embedding to map the incoming sample to the existing database and then directly assigning the layer to exit based on the samples in database to which the incoming sample closely resembles.", "weaknesses": "1. While authors claim that the method is zero-shot while they use the training samples to create the database making the method supervised, this claim makes the proposed method confusing.\n\n2. The provided method cannot be easily generalised to other tasks except classification which is a major issue with this work.\n\n3. Why there are different baselines with different backbone, this makes it unfair for a fair comparison all methods should be tested on the same underlying backbone model. For instance, DeeBERT can be applied to T5 and Gemma but authors do not report.\n\n4. The methodology of the work closely aligns with an existing work [1] still no comparison with it is found.\n\n5. There are multiple methods which prove and show that if monitored properly early exits can outperform the full models which are not compared and not even cited which shows that the method has a less hold of the existing literature [2], [3], [4], [5] and many more.\n\n[1] https://ieeexplore.ieee.org/document/11161749\n\n[2] https://arxiv.org/abs/2502.00745 [ICLR 2025]\n\n[3] https://proceedings.neurips.cc/paper_files/paper/2024/file/ea5a63f7ddb82e58623693fd1f4933f7-Paper-Conference.pdf\n\n[4] https://openaccess.thecvf.com/content/WACV2024/papers/Meronen_Fixing_Overconfidence_in_Dynamic_Neural_Networks_WACV_2024_paper.pdf\n\n[5] https://aclanthology.org/2025.findings-acl.1209.pdf", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GjMIW9B0N3", "reviewer": "Reviewer_qNzk", "summary": "This paper proposes RAEE (Retrieval-Augmented Early Exit), a training-free early exit framework that leverages a retrieval database built from the training set to dynamically select an optimal exit layer during inference. The key insight is that early exit can serve not only as an acceleration mechanism but also as a **corrective** one—intermediate layers sometimes outperform the final layer. RAEE models early exit as a distribution prediction problem and approximates this distribution by aggregating exit behaviors of semantically similar training examples retrieved via FAISS. Experiments across eight GLUE tasks and multiple backbone models (RoBERTa-Large, T5-Large, Llama-3-8B, Gemma-7B) show that RAEE consistently reduces latency (by nearly half for large models) while **improving accuracy**, often surpassing the full model baseline.", "questions": "1. The performance gain is dramatic on CoLA but marginal on SST-2. Can the authors provide an analysis of **which types of tasks or data distributions benefit most** from RAEE? Is it related to task difficulty or the full model’s calibration?\n\n2. RAEE requires ground-truth labels to build the database. Have the authors considered a **self-supervised variant** that uses high-confidence predictions (e.g., $p > 0.95$) as pseudo-labels? This would extend applicability to unlabeled data.\n\n3. In Algorithm C.1, the key is the embedding from the backbone’s first layer ($h_0$). How sensitive is performance to the choice of embedding layer or pooling strategy? A brief ablation would strengthen the method’s robustness claim.\n\n4. The OOD experiment uses a next-token prediction setup on WikiText. Could RAEE be adapted to **classification tasks without labels** in the OOD setting? If not, what are the fundamental limitations?\n\n5. Given the observed variance in baseline performance, could the authors provide more details about their baseline implementation and evaluation protocol? Specifically, were the reported baseline scores obtained through standard evaluation procedures or customized setups?\n\n6. As the retrieval database scales up, how does the method maintain inference efficiency? Have the authors tested the framework with larger database sizes (e.g., 100K+ examples) and measured the corresponding latency impact?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1. Originality: Proposes a novel perspective—early exit as a corrective mechanism—and implements it via retrieval augmentation, avoiding the need for trainable classifiers.\n2. Quality: Comprehensive experiments across 8 tasks and 4 model families; includes ablation on $k$, database size, and OOD generalization.\n3. Clarity: Figures 1–2 and Algorithms C.1–C.2 make the method transparent. The “correct ratio” analysis (Figure 1b) is particularly compelling.\n4. Significance: Offers a practical, zero-shot solution to LLM inference efficiency that simultaneously boosts accuracy—a rare win-win in model compression.", "weaknesses": "1. Dependency on in-distribution labeled data: RAEE requires access to the **training set with ground-truth labels** to build the database. This limits applicability in zero-shot or unsupervised settings. The paper acknowledges this but does not explore alternatives (e.g., using pseudo-labels or confidence-based proxies).\n2. Inconsistent gains across tasks: Performance improvements vary widely (e.g., CoLA: +12.45%; SST-2: +1.03% with RoBERTa). The paper does not analyze **why**—e.g., whether task difficulty, label distribution, or model calibration plays a role.\n3. Embedding quality sensitivity: The method assumes that the query embedding (from backbone or BERT) captures semantic similarity well. No ablation compares different embedding sources (e.g., [CLS] vs. mean pooling vs. Sentence-BERT).\n4. Limited OOD validation: Only tested on summarization with WikiText→CNN/XSum. More diverse OOD shifts (e.g., domain, style, language) would better assess robustness.\n5. Scalability concerns: The retrieval component's latency with larger databases remains unexplored, raising practical deployment questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Lmv9uMV8TS", "reviewer": "Reviewer_jtM8", "summary": "This paper proposes RAEE, a robust Retrieval-Augmented Early Exiting framework for efficient inference. RAEE first builds the retrieval dataset that records the correct exit layer, then retrieves the nearest neighbor to get the exit layer. Experimental results show its effectiveness.", "questions": "See weakness.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Using $k$NN-based retrieval method for early exiting is novel. \n2. This paper is simple, effective, and intuitive, even without additional technical contributions.\n3. Experimental results are good.", "weaknesses": "1. The paper lacks a visualization of early-exit layers. What proportion of data exits early at each layer for each dataset?\n\n2. Why does your method introduce additional time overhead on T5 in Figure 3?\n\n3. Is the proposed method effective for supervised methods? For example, for a T5 model fine-tuned on SST-2.\n\n4. Section 4.5 may not be very solid. Could you provide the performance when using a fixed layer, for example, the six values from layers 27 to 32?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)