Beyond Greedy Exits: Improved Early Exit Decisions for Risk Control and Reliability
Authors: Divya Jyoti Bajpai
Manjesh Kumar Hanawal
Venue: NeurIPS 2025 poster ()
Keywords: Early exits
Risk-controlled inference
Primary area: deep_learning
TL;DR: Risk awareness in Early Exit DNNs

Abstract: Early-Exit Deep Neural Networks enable adaptive inference by allowing prediction at intermediary layers, significantly reducing computational costs and latency. Most of the early exit strategies greedily exit a sample at an intermediary layer if the confidence in class prediction exceeds a predefined threshold that is set using a static validation set. This is problematic as the model might be overconfident in a wrong class. Also, they are not robust to distribution shifts encountered in deployment, which can undermine model trustworthiness and accuracy. To address these challenges, we propose UAT that adapts the threshold for exit decisions using a Multi-Armed Bandit framework, enabling online, unsupervised adjustment of exit decisions. UAT makes decisions based on a new reward function that assesses predictive certainty and its reliability to balance computational efficiency and prediction quality while penalizing unnecessary late exits. We provide guarantees on risk achieved by UAT and validate its performance on diverse tasks spanning vision-language understanding, text generation, and classification. Our framework demonstrates consistent improvements in speedup $(1.70-2.10\times)$ with a minimal performance drop $(<2)$\% as compared to full model performance.

Beyond Greedy Exits: Improved Early Exit Decisions for Risk Control and Reliability
Divya Jyoti Bajpai & Manjesh Kumar Hanawal
Department of Industrial Engineering and Operations Research Indian Institute of Technology Bombay Powai, Maharashtra, India {divyajyoti. bajpai, mhanawal/@iitb. ic
Abstract
Early-Exit Deep Neural Networks enable adaptive inference by allowing prediction at intermediary layers, significantly reducing computational costs and latency. Most of the early exit strategies greedily exit a sample at an intermediary layer if the confidence in class prediction exceeds a predefined threshold that is set using a static validation set. This is problematic as the model might be overconfident in wrong class. Also, they not robust to distribution shifts encountered in deployment, which can undermine model trustworthiness and accuracy. To address these challenges, we propose UAT that adapts the threshold for exit decisions using a Multi- Bandit framework, enabling online, unsupervised adjustment of exit decisions. UAT makes decisions based on new reward function that assesses predictive certainty and its reliability to balance computational efficiency and prediction quality while penalizing unnecessary late exits. We provide guarantees on risk achieved by UAT and validate its performance on diverse tasks spanning vision-language understanding, text generation, and classification. Our framework demonstrates consistent improvements in speedup (1.70 2.10x) with minimal performance drop (< 2%) as compared full model performance. Our source code is available at https: / thub. Som/Div290//u
Introduction
As Deep Neural Networks (DNNs) continue to scale in size and complexity. the cost of inference has become major bottleneck manifesting as increased latency, higher energy consumption, and growing carbon footprint 37 39.. These challenges highlight the need for efficient and sustainable deployment strategies for large models. Moreover, real-world environments often impose dynamic computational constraints, driven by resource scarcity or energy availability [40]. To remain practical and responsive, modern DNNs must adapt computational load in real time and consume resources as much as necessary, as per the complexity of the samples.
Early-Exit DNNs (EEDNNs) [7 offer compelling solution for adaptive inference under dynamic computational constraints. EEDNNs can produce intermediate predictions at various depths of the DNN, allowing the network "exit" the sample early without passing through all the layers. This mechanism enables dynamic resource utilization and significantly reduces computation and inference latency. This makes EEDNNS well-suited for deployment under resource constraints across diverse domains, including computer vision [47] 4 (2). natural language processing [12] 9. and vasion-language systems [6). However, this flexibility comes at a cost: initial layers are more prone to making riskier predictions, which degrade the trust in the model [36). Consequently, EEDNNs face a fundamental wrde-off-balanciin efficiency gains against potential degradation predictive accuracy (3). This trade-off is governed by the exit decisions that are crucial in BEDNNs.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

In existing methods. the decision to exit depends on two variables: confidence and threshold. Confidence usually defined as the entropy of the prediction [50], prediction consistency [[2]. or ensemble methods g 49 Threshold is a value against which confidence is compared. If the confidence value exceeds the threshold, the model is considered confident enough, and the sample allowed to exit without further processing (early exit). However, there are some crucial concerns in existing methods: 1) The confidence is not reliable metric as the model may get overly confident at the initial layers on a wrong class, making the sample exit with an incorrect label [36]. Fig. [] discussed in the Appendix [.7 shows this phenomenon, where large number of samples can exit from the intermediary layer with high confidence on the wrong class. 2) The threshold to make an exit decision is precomputed using a small validation set, assuming it generalizes well to the inference dataset distribution. But this assumption does not hold in real-world conditions, where distribution shifts are common. This defeats the goal of low-risk and reliable inference and necessitates the need for improved strategy that incorporates the reliability of model output in the confidence metric and dynamically adjusts thresholds as per the input data distributions.
During inference, the data arrives in an online fashion. Hence, the exit decisions need to be adapted in an online and unsupervised manner (a quality also critical for zero-shot tasks). To address these challenges, we introduce novel framework for dynamic threshold adaptation in EEDNNs, aimed at optimizing the trade-off off between prediction performance and computational costs. In contrast to existing methods that rely confidence an intermediary layer exceeding fixed threshold, our early exit strategy uses a confidence metric that accounts not only for the model's predictive certainty but also the reliability of that certainty.
We derive an exit strategy using the Multi-Armed Bandit (MAB) framework, where each arm represents candidate exit threshold. The bandit agent aims to learn threshold that optimizes a custom reward that involves the confidence metric and a penalty for late exits to discourage unnecessary computation. We theoretically prove that optimizing this metric aligns with maximizing the true class probability, providing a principled foundation for early-exit decisions.
Our framework promotes efficient inference without compromising performance, and we demonstrate its effectiveness across diverse tasks, including vision-lang Languge understanding (e.g.. image captioning 34]. VQA (20). text generation (e.g., summarization [25], QA (38), and classification (c.g.. NLI, sentiment analysis
summary, our contributions are as follows:
We introduce the notion of reliability scores on the confidence of the early exit output to improve the trustworthiness of EEDNNS.
We propose dynamic early exit framework that adapts to distributional changes at test time, enhancing both efficiency and robustness EEDNNS.
We apply the bandit framework to learn the optimal strategy. We develop an upper confidence-bar based algorithm and establish a theoretical bound on the risk achieved by it.
We validate our approach on a broad set of tasks, demonstrating consistent gains in both computational efficiency and predictive accuracy.
2
Related works
Early Exit methods have been applied to various tasks such image classification, text classification, text generation and vision language tasks to reduce the computational latency and inference latency. Below, we discuss some of the works on risk control in EEDNNs.
Early Exits: Nra [47] first proposed the carly exit method for image classification tasks. uses classification entropy
confidence metric. Shallow- deep (32) and MSDNet [29] improve upon NranchyNet by effectively choosing the thresholds based on the confidence threshold. ZTW [4]. JEI-DNN [18] and BEEM [8] propose different methods for assessing the confidence. ZTW uses multiple classifiers to decide exiting. JEI-DNN learns a gating function during training to decide exiting, while BEEM builds a confidence measure using ensemble methods. DeeCAP 224, MuE
[46. CapHEN[S] and FREE [10] implement early exits to vision language tasks.
For the I Language tasks, and specifically in Large Language Models (LLMs), the EE methods have
been popular [1] 15 35
13
45
21 27 14 44 RT. DeeBERT (50) first applied it to the BERT model.
Later, PABEE
552 developed patience-base baaie method to decide exiting. BERxiT [51) proposes an

---

efficient tinettunigg strategy for EE-based models. LeeBERT [33 performs self-distillation between exits effectively share the knowledge across layers with exits. ETFEE [3] adds an adapter on top of the transformer layer and an entangled frame classifier to make exits learn better. CeeBERT [6] and DAdEE 7 provide methods to adapt the EEDNN various domains. CALM 41 extends the idea of EEs to language generation tasks.
Risk in EEDNNS: As EEDNNS consist of multiple exit classifiers, the initial layers of the EEDNNS with low-level features might leave the sample at risk. EERC (30] shows that the threshold, if chosen
properly, can also minimise the risk of EEDNNS. MIE-LAP 36 addresses the overconfidence in EEDNNs by uncertainty quantification. It proposes an approach for uncertainty-awalt decision- making by leveraging the last layer Laplace approximation implementation.
While existing early-exit approaches determine a fixed confidence threshold using a validation set and apply it uniformly during inference, this static strategy often permits overconfident yet incorrect predictions to exit prematurely. In contrast, our method introduces three key innovations: 1) Dynamic adaptation to test-time distribution: We tailor the exit decisions based on the distributional characteristics of the test data. 2) Confidence calibrated by reliability: Instead of maximizing the raw confidence scores, we explicitly account for their reliability, leading more robust exit decisions. 3) System-level risk modeling: Rather than comparing early exits against the final layer's predictions, we model the overall risk of the complete system, enabling our method to potentially surpass the performance of the base DNN itself, an important limitation overlooked by prior work.
3 Preliminaries
Consider the input-output space denoted as X x y. A sample, denoted (z,, y), is drawn from an underlying probability distribution P. For classification tasks, the label space is given by y {1,2.... KJ. whereas for regression, the output set is a subset of Rd.
3.1 The Notion of Risk
In statistical modeling, risk control mechanisms 16] enhance the reliability of predictive models by imposing constraints on threshold based decision processes. Let Ω denote a finite set of thresholds. Consider model fr parameterized by threshold on confidence levels. For a classification task, the set-valued predictor f- : X + 2 is defined as (x) (ύ y p(ū æ) 2 7). where p(ū 2:) is the probability that model assigns label ý to no class probability exceeds т, the predictor abstains by returning an empty set (f-(z) = 0).
Ω, which outputs set of predictions based
To assess reliability, we define a miscoverage loss that penalizes the exclusion of the true label as ((f((), 3) = 1 ly & f=(],,, where [[:] is the indicator function. The associated risk is the expected probability of miscoverage:
R(T) = (P [e[ (z), 8)] = (G Pr(y & fx)))
For a given tolerance, E (0. 1) and 8 E (0, 1), we say that the threshold set Ω controls the risk with within tolerance with probability at least 1 8 if
(I) The advantages of risk control frameworks are significant: (i) they require no assumptions on the data distribution P, making them distribution free and (ii) they can be applied post-hoc to any model. These features make them particularly valuable for practical deployment in uncertain environments.
(()(T) s <) 2 8, Ó.
3.2 Early Exit Mechanisms in Deep Neural Networks
Early Exit Deep Neural Networks (EEDNNs) extend the architecture of the DNNs by attaching Exit Classifiers (ECs) at the intermediate layers. The EC at layer intermediate layer i produces
probability distribution over y denoted by (=(r)r) = p(-|z: 0:) = EC, (h+) where h; = (x(x)) is the hidden representation output by the ith layer of the model with i (1, 2, ...,, L) and 0, is the set of parameters consisting of the EEDNN parameters till layer i comprising of both EC; and main model parameters. The final prediction layer, indexed by I,, corresponds to the standard DNN where all layers are utilized for inference. Each EC; estimates the class label ÿ arg max (v)2))
3

---

with confidence score C, := C₂(r) yy (r(1)z). C, measures the certainty associated with the prediction. However, this confidence might not be very reliable as provides confidence in the predicted class, which might be wrong due to overconfident EEDNNS.
During inference, the EEDNN exits from a layer where the confidence of EC exceeds the predefined threshold E [o, 1] for the first time. Following previous works [50] 30), a single global threshold considered across all the ECs, i.e., T2 for every {1, L}. Then, for given т, the probability that r exits with a label ÿ is given by
pr (0(110) b) = |PL00 if < 4 k E [1,2 .................. - (p a) where i min { E {I....,L ,, 1} Cx 2 r}, otherwise
As the class label assigned by the EEDNNs relies on the confidence score, their reliability is important for trustworthy predictions. The threshold directly influences the balance between computational cost and model accuracy. Lower values encourage premature exits, reducing computational overhead but potentially sacrificing performance.
The risk metric defined in the previous subsection can be naturally applied to EEDNNs. As samples exit if the confidence score exceeds a given threshold, we can associate risk for each threshold and aim to identify the threshold that yields the least risk. However, the task of the threshold in early exits is not only limited to minimizing the risk but also reducing the inference latency. Hence, a threshold that is the smallest from the set ÿ is desired. That is, instead of looking for set 52 that satisfies Eqn. we aim to identify threshold T" = min(T E Ω2}. The proportion of samples that exit earlier higher T than any other in f2.
For given threshold r, let p-(y A r)zz) denote the probability that the EEDNNs makes errors on sample (z, v). Then the risk of the EEDNNs is defined as
R(T) ((z,))>P (pr + 0|>)].
Our objective is find a threshold for which R(T) lies within given tolerance value with high probability. The value of 7 that meets the requirement is function of the sample distribution P, and the distribution seen in real deployment could drift from that of the training samples used to train the EEDNNS. The best value of r depends on the distribution of the input sample, and hence needs to adapt to any drift in the sample distributions. This adaptation requires knowledge of risk. However, this information cannot be estimated during the inference time as ground truth labels are not available. To address this challenge, we perform the learning two stages.
In the first stage, we perform offline training of the EEDNN and a linear layer network that will allow us express the risk without knowing the ground truth labels. This training involves estimating the reliability of the confidence scores output by the exit classifiers. In the second phase, we use this proxy for risk and learn the best threshold 7 by maximizing a reward function using the Multi-Armed Bandit (MAB) framework. The following Lemma gives a simple relation that we use to obtain a proxy for a risk function.
Lemma 3.1. Given threshold r, let ÿ be the label predicted by EEDNNs sample (a, v). Define pr = 0 [0 x) be the probability that the predicted label the true label. Then Pr = (1x)
Pr (0 x) Pr (y (1 ()
The proof straightforward and is given in the Appendix A. Pr - -0. (0r 9) is the probability that the model's prediction is correct. Note that pr (ū]=r) can be estimated during inference as it does not depend on the label. Then, the lemma suggests that we can compute risk by estimating pr (y 0|., 9). The next subsection discusses training neural network to estimate it in offline fashion.
3.3 Exit and risk function training
We train a reliability function (neural network) that assigns a reliability score to the class proba- bilities output by the exit classifiers. We use a modified loss function that used to train the exit classifiers in EEDNNs to train g. In the conventional EEDNNs, the objective at each exit is to maximize the probability of the correct label using cross-entropy loss. We augment this objective with the reliability function g. For sample (x, y) P. we define modified loss is as:

---

C; = (cn)(p((())) + s(px(+1o),
+ ((g))
(3)
The (CE((:(+|a)) y) in the first term is the standard cross-entropy loss that encourages the model to correctly classify the imput. It is multiplied by a factor 1 + 9-), to link the sample-wise loss with the reliability function 9, which estimates the confidence of the model's prediction. The minimization of L, causes g to assign high values to samples with low cross-entropy (i.e., correctly and confidently classified), and low values to those with high loss (i.e., ambiguous or misclassified samples). Intuitively, this creates a competitive dynamic where is penalized for activating on high-loss (uncertain) samples and is incentivized to concentrate its high values on low-loss samples.
When g+) 2 0.5, it indicates that the model output reliable. <(g) in the second term denotes the fraction of samples for which the model's ss output is reliable. $bc - <(g)) acts as a regularizer to enforce that the model is reliable for at least e fraction of samples. The function ((a) max(0, a)² penalizes the model when the coverage ((g) falls short of the desired threshold c.. This regularizer restricts g from giving a trivial solution y = O, as such a solution would provide no meaningful discrimination between easy and hard samples. The role of &(c ((g)) is to prevent this collapse by ensuring that the model is reliable on least e-fraction of the training samples. c can be a function of the exit layer that can be chosen based on validation dataset (see Appendix B.2 We empirically observe that our loss does not impact the model's predictive performance (see Appendix B.9).
Once trained, g can be used as a proxy for the model's s correctness probability p(y = 9 2, 9). To account for the cost associated with deeper exits (32), the final loss across all exits is computed as weighted average L = Ee i.C, inference cost is higher. ا
which assigns higher weights to exits at deeper layers where
3.4 Dynamic thresholding during inference
We first describe the MAB setup to dynamically adapt to the threshold during inference. In MAB setup, decision-maker iteratively selects actions, adapting to unknown environment. The goal is to identify an action that gives the highest reward. In our setup, the actions are the set of thresholds.
We next define the reward for each threshold. For given threshold/action т, let C+ denote the confidence model has in the predicted class when the sample exits from F.C,, We define C 1 (p.((lz)) as a score of underlying risk (unreliable) in the confidence score of the model at EC; We set 1 - Cý as an approximation of (r = 5| û). Then Lemma 3.1 indicates that the P- r|æ) can be well approximated as C² (111))) when sample exits from EC₁. However, threshold that maximizes C; (1 C%) may result in suboptimal efficiency gains as it will penalize the thresholds that are forcing the sample to move to deeper layers. To account for this fact, we define reward for threshold
on a sample as follows:
C (1-C) (1-C)- -(), if - (1-11110))) C 2i< (c
r(7) = (1- -1- C1-))--))) =fi=
(4)
where (b)) is a penalty to disincentivize the sample from moving to deeper layer unnecessarily. In our setup, we define (i) = X. where À could be interpreted as the processing cost per layer.
The expected reward for arm E Ω is E[r(r)] = E
where P(t) is the probability that the sample exits from ith layer and expectation is with respect to the randomness of confidence and reliability scores. Let 7" = arg max, en [[((T)) denote the optimal threshold for a given Now consider a policy π that selects the threshold T E 52 based on past observations. For a given number of rounds, the performance of policy is evaluated using the expected cumulative regret defined as R(T,, T') = ΣT= E(r(r*) r(r)), where the expectation is with respect to randomness in the selection of thresholds induced by the past sample. A policy π that satisfies is said to be sub-linear and plays most of the time the optimal threshold. Our goal is to develop an algorithm that is sub-linear and has small linear component. As we will show later, this translates to achieving a small empirical risk defined as R(T))
A- A-
TT/T + 0
-T Pr, (y) = (()))) where (æx, ve) the input sample, ÿ is the predicted label, and 7; is the threshold selection by л in round t.
5

---

4 Algorithm
We develop algorithm named UAT. Its pseudo-code is given in [ The inputs to the algorithm are exploration constant 7 and the penalizing term (b()). For the first [Ω]] samples, the algorithm plays each arm once. In the subsequent rounds, it plays arm with the highest Upper Confidence Bound (UCB) index denoted as тІ- UCB indices are obtained by taking the weighted sum of the empirical
average rewards and the confidence bonuses with Y the weight factor. If C²(1 C2) at the ith layer is greater than тr. then the sample exits; otherwise, the sample is passed to the next layer in the backbone. If the sample does not exit at any intermediate classifier, then it is inferred at the final layer. Finally, the algorithm updates the number of pulls (N)( and the empirical mean (Q(Te) of the played arm. Note that the algorithm is applied in the inference phase. In Table 6 empirically show how different components contribute to the reward function and discuss in Appendix [B.]] We also show in Appendix B.8 that UAT has negligible computational cost. The following result establishes the average risk achieved by UAT.
Theorem 4.1. Let (1 C2) approximates the value of p+(ū y/ar,u) with probability (1 - 51). Let the risk associated with T be bounded by e". Then, for sufficiently large T and given tolerance 6, UAT achieves (P(((T)) I <d) 2 (1 -.) - 8'). where ea = & and is a constant.
The detailed proof given the appendix. Note from the theorem that E equivalent to the full model risk (that is always there), and the additional risk of EEDNN due to early exits is Then the risk in EEDNN performance is = c c e". This theorem also suggests that lower value can lead to lower risks, which is true based emphasis is given to the C²(1- - C2) value which maximizes p(y = 9 Jr, 9). The proof also provides way fix the À value, as want the maximum À value such that the risk levels below c. Hence we choose = I in this work. This theorem bounds the risk of the EEDNN where the threshold to decide exiting is chosen based the policy æ given by algorithm Also, observe from the proof outline that it suggests that the risk of the EEDNN will be greater than e", which denotes the optimal risk. The proof steps also provide the link of risk with the regret.
reward function. If À is chosen small, then more
5 Experiments
In this section, we provide empirical results of our work on various tasks such as text classification, language modelling and vision language tasks.
Algorithm 1 UCB-based Adaptive Thresholds (UAT)
I: Input: vii), 1
2: Initialize: Play each threshold once. Observe r(T) and set Q(T) + 0,N(T) + 1
3: for [S2] + Ω2 + ..... do
4: Observe an instance Zt
In(t) + 1 N(T)
5: + arg max 7ER Q(T)
6:
for i to L do
7:
Pass till layer i and apply threshold 7; and observe S,, C -1 (1-C%) if St 2 Ta and < L then
8:
9:
Infer at layer and exit
(T(TT) + St v()), (.(T))) NN-))) < tet
10:
N((T()
11:
break
12:
else i then
13: Process and infer at the last layer.
rT(T.) St. ((L)), Ni(TT) N. Q((T.) Ej (6))))))))
Ne(
end
16:
end for
17:
end for

---

Model/Datasets
SST-2 95.2
MNLI
RTE
QNLI
QQP 72.5
SQUAD
Speedup 1.00x
Full model
86.0
67.7
92.0
83.0
EE models
PABEE
94.8
85.0
67.2
91.1
72.3
82.0
1.33x
ZTW
94.5
85.1
90.6
72.0
81.7
1.70x
MuE
94.9
85.5
67.3
91.3
72.3
82.2
1.59x
JEI-DNN
95.1
85.6
67.5
91.5
72.2
82,3
1.49x
BEEM
95.4
86.1
67.6
91.9
72,4
82.4
1.79x
Risk- bisk-based methods on EE models
CALM
95.5
85.9
67.7
91.6
72.3
82,2
1.39x
MIE LAP
95.3
85,8
67.5
91.7
72.1
81.8
1.45
EERC
95.0
85.6
67.3
91.5
71.9
81.6
1.50x
Ours
86.6
68.3
92.4
73.0
83.2
1.85
Table 1: Results on the BERT-large model showcasing accuracy and average speedup (Speedup).
Summarization
QA
MT IWSLT
Speedup
Model
SamSum 48.82
CNN
Multi-News
BIGPAT.
SQuAD
Full model
41.15
37.62
49.68
90.63
39.19
1.00x
EE- baseb methods
PABEE
42.25
37.54
32.97
43.71
87.39
35,80
1.90x
ZTW
43.96
39.07
33.62
45.38
88.02
36.75
1.83 x
MuB
43.48
38.11
33.73
45.02
87.53
36.32
1.82
JEI-DNN
44.45
37.78
34.52
45.39
88. 14
36.94
2.05
BEEM
45.19
39.82
35.07
47.43
88.28
37.06
1.97
Risk based EE methods
CALM
46.73
40,29
36.14
48.02
88.79
37.84
1.70x
MIP LAP
44.87
39.58
35.92
47.16
88.47
37.91
1.76x
EERS
46.79
40.52
36.83
47.96
89.02
37.68
1.63x
Ours
48.09
40.97
37.10
48.74
90.59
38.62
2.10
Table 2: Results on the T5-large model showcasing task metrics with average speedup (Speedup).
Datasets: For text classification tasks, we use the standard GLUE [48] datasets that consist of three types of classification tasks: 1) Sentiment analysis (SST-2). Text entailment (RTE, QQP) and 3) Natural Language Inference (MNLI, QNLI). For Question-Answen Answering tasks, we use the SQuAD 2.0 38 dataset. For text summarization tasks, we use diverse set of datasets covering different domains and styles: 1) Dialogue summarization (SAMSum) 25). 2) Long-form form technical summarization (BIGPATENT) 43 3) Multi-document summarization (MultiNews) (23] and 4) News summarization (CNN/DailyMaii) 28 For machine translation tasks, we use the IWSLT De-En dataset [17,, which consists of translated transcripts from German to English. serves as a standard benchmark for evaluating translation quality spoken language content. For vision language tasks, we use the COCO (34) and NoCaps [T] datasets for Image captioning, VQAv2 (26) dataset for visual question answering and VisDial [20] for visual dialogue.
Metrics: We report key metrics such accuracy for text classification, score for Question answering tasks, Rouge-L score for text summarization, BLEU-4 (B4), CiDER (C), Spice (S) and Meteor (M) for image captioning. VQA accuracy for VQA tasks and Mean Reciprocal Rank (MRR) for visual dialogue. Following existing baselines [22] 46). speedup is calculated as the acceleration of average inference time per token compared to the full model. We note that the metrics reported in (30], such as performance gap risk, which is EE model performance subtracted from full model performance, can easily be calculated using existing performance metrics, hence, we do not specifically report. Also, performance metrics preferred for better and fair comparison.
Setup: For the training phase, we augment the pre-trained models with a linear output layer serve as an exit. For the architecture of the 9. we use a single linear layer with shared parameters across the exits that serves as risk predictor. Following the existing methods R the inference is performed on per-instance basis, setting the batch size to 1. The value of is fixed to i. while the set of candidate thresholds chosen as ten equally spaced values between 0.5 and 1.0 (including 1.0). ablation study over different À values is also performed in Appendix B.3] It aligns with scenarios
where low-latency critical, such as processing individual requests from different users 42). Similar
to other baselines [30], also choose 0.01 for all the experiments. An ablation study over
7

---

Models
in the COCO Karpathy test
VQAv2
VisDial
Spd.
B4
S
M
Acc
MRR
B-IP-2-V-F
44.0
145.8
25.3
31.9
81.8
45.9
1.00x
EE-based models
PABEE-BLA
35.8
122.6
21.9
26.5
75.2
36.7
1.55x
ZTW
36.3
121.1
21.7
26.0
76.9
37.8
1.59x
MuE
37.5
126.8
22.4
27.7
78.0
38.4
x,49
JEL-DNN
38.2
128.4
22.6
28.0
78.4
38.1
1.63x
BEEM
40.7
135.9
23.8
28.9
79.6
39.6
1.68x
Risk-based model
CALM
41.8
139.3
24.4
29.1
79.8
41.9
1,46X
MIP LAP
41.2
137.4
24.1
28.9
78.6
40.2
1.55x
EERS
42.5
140.3
24.7
29.5
80.0
42.6
1.38x
Ours
43.3
142.9
25.0
31.1
81.1
42.9
1.72x
Table 3: Results vision-language tasks (Visual Question Answering (VQA), test split of COCO and VisDial for visual dialogue) on BL.I--2VVTTFlla model with average speedup (Speedup).
in-domain
near-domain
out-domain
full-dataset
Speedup
Models
C
S
C
S
C
S
C
S
BLIP-2 2 VIT FTS
123.7
16.3
120.2
13.9
124.8
13.1
121.6
15.8
1.00x
EE-based models
LIBEE-BLIP
117.6
15.2
114.2
14.6
117.4
14.2
113.2
14.5
1.48 x
ZTW
119.1
15.5
115.4
14.8
119
14.4
115.9
14.8
1.35x
MuE
118.3
15.3
114.8
14.7
119.5
14.4
115.6
14.7
1.62
JELDNN
119.8
15.5
116.2
15.0
119.9
14.6
116.1
14.9
1.58
BEEM
120.2
15.7
117.5
15,3
119.6
14.5
116.8
15.1
Risk-based EE methods
CALM
121.0
15.9
117.8
15.5
119.8
14.7
117.1
15.2
1.47x
MIP LAP
120.6
15.8
117.2
15.2
119.3
14.3
116.9
15.1
1,53x
EERS
121.4
16.0
117.6
15.4
120.2
14.7
117.3
15.3
1.38x
Ours
122.5
16.2
118.9
15.7
122.7
15.0
119.5
15.6
1.75x
Table 4: Results of DLLIP-2-ViT-Flan'55 model on Nocaps dataset.
can be found in Appendix B.5 More hyperparameter details can be found in Table 11 and 12 Appendix.
We consider three types of baselines to compare our model: 1) Full model: This is the baseline showing conventional DNN performance. 2) Early Exit models: These the baselines where we consider various early exit methods, such as PABEE [52), patience-
Baselines:
exiting, and
ZTW [49], on the other hand, uses aggregation of confidence scores across the exits, MuE (46) uses the similarity score of the hidden representations to decide exiting. JEI-DNN [18] learns gating function to performs an exit. Finally, BEEM (8) is method that utilizes ensemble methods to decide exiting. 3) Risk-based methods for EEs: In this, we consider methods that consider the risk factor in the EE models. CALM 41 considers the risk of the T5-large model with EEs for text generation tasks, we extend this to other tasks as well. MIE LAP (36) considers aggregating the scores across exits minimize risk and fix the overconfidence issue the image classification tasks: we extend this to other tasks as well. Finally, we have EERC [30 that theoretically provides simple extension of EE models to provide a method to choose the value of the exit threshold such that the risk minimized. We use the same code and hyperparameters without any changes to get the results of the baselines.
Results on Text classification: In Table I we provide results on the text classification tasks (GLUE datasets) and the Question answering task over the BERT-large model. We observe that our method performs better than all the existing methods, where sometimes the performance of our method is The EE-based
better than the full model performance due to the impact of overthinking [32 52]. methods have more focus on the speedup part while having higher loss in accuracy. An important observation is that BEEM, the EE-based method, outperforms risk-based methods due to its ensemble methods for performing classification. On the other hand, the risk-hased methods have a lower loss accuracy but observe hit in speedup their focus is solely on reducing the risk, while achieving minimal efficiency gains.
inference. However, this is only observed in easier tasks such as text
8

---

Model
Noise 0.1
Noise 0.5
Noise 1.0
B4
C
Speed
B4
Speed
C
Speed
Full model
43.1
142.5
1.00x
39.8
133.2
1.00x
36.4
128.6
1.00x
EFF -rabbsed methods
JEI-DNN
36.7
128.9
1.58x
31.2
107.5
1.68x
26.1
93.9
1.74x
BEEM
38.6
132.0
1.60x
34.7
122.3
1.71x
30.6
104.5
1.22x
Risk based methods
CALM
39.1
131.5
1.55x
34.8
122.7
1.63
31.5
110.3
1.72x
EERS
39.6
134.3
x.47
36.0
126.4
1.54x
32.7
116.9
1.68x
Ours
41.3
139.7
1.70x
38.2
130.5
1.77x
35.0
125,5
1881x
Table 5: Results on the COCO test split by adding different levels of distortions the test images.
Our method balances risk and efficiency using the reward function and minimizes the risk while getting improved efficiency.
Results on Language Modeling: In 2 we report the results on tasks such as text summarization, Questionnansserte and Machine translation. Our method consistently outperforms existing baselines both in terms of performance and speedup, due to its dynamic adjustments of thresholds based on the test dataset. One key observation from the existing tasks is that very high static threshold was also getting a points ROUGE-L score drop with a higher loss in efficiency. Due to this fact, there is performance drop in the risk-based as well as EE-based methods shown in Table 2 methods have a lower drop in performance as compared to EE-based methods. Our method, on the other hand, dynamically adjusts the threshold and has an option switch from hard criterion exit to soft one, helping it balance accuracy and efficiency.
Results on the vision-language language tasks: In Table 3 and 4 we report the results on various vision- language tasks using the BLIP-2 backbone with ViT-g as the encoder and F5- as the decoder. We observe that the efficiency gains of our method are the highest with minimal performance drop as compared to other risk-based as well as EE-based methods. The results on the Nocaps dataset in Table 4 suggest that our method is robust to distribution changes compared to other methods as its performance dip is minimal when the images are from near-domain or out-of-domain This is an advantage of the online adaptation of the threshold, where it learns the threshold based on the incoming dataset distribution.
Robustness to distribution changes: In Table we provide the BLEU-4 scores for image captioning tasks, where additional noise added to the test split of the dataset. The model, Fl. xl, was trained pristine (undistorted) images, and during inference, we add different levels of distortions in terms of Gaussian noise with mean zero and standard deviation o;; a higher standard deviation will have more noise in the image. Noise in images is a common real-world scenario where the model suffers with a loss in performance as shown in the Table 5 when it is trained on pristine images. The results from Table S suggest that method's 's performance loss is small compared to other methods. The risk- based methods, specifically the EERS method, show some robustness to the noise, but due to their assumption that the test dataset is representative of the validation dataset, their performance is also significantly low. The reason for our method performing better is its dynamic adaptation of the threshold based on the incoming dataset distribution instead of fixed threshold as done by other methods.
5.1 Ablation Study
Analysis of the overconfidence issue In Figure [ we provide the box plots for the true class confidence values across the layers of the BERT-large model on QNLI dataset; we only plot the initial 12 layers out of 24 layers. Observe from the figure that the initial 3 layers have not learnt much information, hence mostly distributed around 0.5. After that, we see that the box plots are raised to a higher value. However, observe that there are around 12% samples at the 4th layer that are highly confident (more than probability 0.7 assigned to the wrong class as it is a binary classification task) on the wrong class, i.e., the model assigns a higher weight in output probability to the wrong class. This necessitates checking the reliability of the confidence. If not properly monitored these many samples with exit the backbone due to this high fake confidence and cause overall performance drop to be significant.

---

1.0
7.8%
7.5%
7.hh
True Class Confidence
0.6
0.2
10
11
Figure I: True class confidence on QNLI dataset across number of layers. The number top of the plots show the percentage of samples below 0.3 class confidence. As it binary classification task, any sample below 0.3 has a high confidence on the wrong class,
Conclusion
In this work, we introduced risk-tolerant approach for Early-Exit Deep Neural Networks (EEDNNs) that balances accuracy and efficiency under dynamic test-time conditions. We first designed reliability-dhit confidence metric that quantifies the trustworthiness of exit predictions, and seam- lessly integrated it into reward-based formulation. To optimize this reward, we proposed an upper confidence bound based algorithm UAT, which dynamically selects the optimal exit threshold during inference, enabling the model to adapt to distributional shifts without requiring labels. We established theoretical bound on the risk achieved by UAT. Experimental results on various tasks such as text classification, language modeling, and vision-language tasks further prove the effectiveness of our method. In this work, we considered a uniform exit threshold across all the exit layers. interesting to quantify the gains achievable by allowing the threshold to adapt to each layer.
7 Limitations
While our method effectively adapts exit thresholds based on prediction reliability, it relies on the quality of the learned confidence function g. Although the existing methods use global threshold across all the exits to assess the confidence, there is scope to have local threshold for individual exits, then the model can be more efficient.
Acknowledgements
Divya Jyoti Bajpai is supported by the Prime Minister's Research Fellowship (PMRF), Govt. of India. Manjesh K. Hanawal thanks funding support from SERB, Govt. of India, through the Core Research Grant (CRG/20220008807) and MATRICS grant Targeted Programme. We also thank funding support from Amazon IIT- (IT-Bombay AI-ML Initiative (AIAIMLI).
and DST-Innria
References
[1] Harsh Agrawal, Karan Desai, Yufei Wang. Xinlei Chen, Rishabh Jain, Mark Johnson, Dhruv Batra, Devi Parikh, Stefan Lee, and Peter Anderson. Nocaps: Novel object captioning at scale. In Proceedings of the IEEE/CVF international conference on computer vision, pages 5948-8957, 2019.
[2] Anastasios N Angelopoulos, Stephen Bates, Adam Fisch, Lihua Lei, and Tal Schuster. Confor- mal risk control. arXiv preprint arXiv:2208.0288 2022.
10

---

[3] Peter Auer, Nicolo Cesa-Biis -isa-Bianchi, and Paul Fischer. Finite-time analysis of the multiarmed bandit problem. Machine learning, 47(2):235-256, 2002.
[4] Divya J. Bajpai, Aastha Jaiswal, and Manjesh K. Hanawal. 1-splitee: Image classification in split computing dnns with early exits. In ICC 2024 IEEE International Conference on Communications, pages 2658-2666 2024.
[5] Divya Jyoti Bajpai and Manjesh Kumar Hanawal. CapEEN: Image captioning with early exits and knowledge distillation. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen, editors, Findings of the Association for Computational Linguisties: EMNIP 2024, pages 2258-6472, Miami, Florida, USA, November 2024, Association for Computational Linguistics.
[6] Divya Jyoti Bajpai and Manjesh Kumar Hanawal. CeeBERT: Cross-domain inference in early exit BERT. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar, editors, Findings of the Association for Computational Linguistics: ACI. 2024, pages 1736-1748, Bangkok, Thailand, August 2024. Association for Computational Linguistics.
[7] Divya Iyoti Bajpai and Manjesh Kumar Hanawal. DAdEE: Unsupervised domain adaptation in early exit PLMs. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen, editors, Findings of the Association for Computational Linguistics: EMNLP 2024, pages 6389-6400, Miami, Florida, USA, November 2024. Association for Computational Linguistics.
[8] Divya Jyoti Bajpai and Manjesh Kumar Hanawal. Beem: Boosting performance of early exit dnns using multi-exit classifiers experts, arXiv preprint arXiv:2502.00745, 2025.
[9] Divya Jyoti Bajpai and Manjesh Kumar Hanawal. Distributed inference on mobile edge and cloud: An early exit based clustering approach. In ICC 2025 IEEE International Conference on Communications, pages 5425-5430, 2025.
[10] 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, editors, Findings of the Association for Compulational Linguistics: ACL 2025, pages 23599-23615, Vienna, Austria, July 2025. Association for Computational Linguistics.
[11] Divya Jyoti Bajpai and Manjesh Kumar Hanawal. A survey of early exit deep neural networks in nip. arXiv preprint arXiv:2501.07670, 2025.
[12] Divya Jyoti Bajpai, Vivek Kumar Trivedi, Sohan L Yadav, and Manjesh Kumar Hanawal. Splitee: Early exit in deep neural networks with split computing. In Proceedings of the Third International Conference on AI-ML Systems, AIMLSystems '23, New York, NY, USA, 2024. Association for Computing Machinery.
[13] Nikita Balagansky and Daniil Gavrilov. Palbert: Teaching albert to ponder. Advances in Neural Information Processing Systems, 35: 1402-14401, 2022.
[14] Andrea Banino, Jan Balaguer, and Charles Blundell. Pondernet: Learning to ponder. arXiv preprint arXiv:2107.05447, 2021.
[15] Ankur Bapna, Naveen Arivazhagan, and Orhan Firat. Controlling computation versus quality for neural sequence models. arXiv preprint urXiv:2002.07106 2020.
[16] Stephen Bates, Anastasios Angelopoulos, Lihua Lei, Jitendra Malik, and Michael Jordan. Distribution-free, risk-controlling prediction sets. Journal of the ACM (JACM), 68(6):1-34, 2021.
[17] Mauro Cettolo, Marcello Federico, Luisa Bentivogli, Jan Niehues, Sebastian Stükker, Katsuitho Sudoh, Koichiro Yoshino, and Christian Federmann. Overview of the iwslt 2017 evaluation campaign. In Proceedings of the 14th International Workshop on Spoken Language Translation, pages 2-1 2017.
[18] Joud Chataoui, Mark Coates, et al. Jointly-learned exit and inference for a dynamic neural network. In The Twelfth International Conference on Learning Representations, 2023.
11

---

[19] Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Yunxuan Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, al. Scaling insruction-finetun language models. Journal of Machine Learning Research, 25(70):1-33, 2024.
[20] Abhishek Das, Satwik Kottur, Khushi Gupta, Avi Singh, Deshraj Yadav, José MF Moura, Devi Parikh, and Dhruv Batra. Visual dialog. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 326-335, 2017.
[21]
Maha Elbayad, Jiatao Gu, Edouard Grave, and Michael Auli, Depth-adaptive transformer. In In Proc. of ICLR, 2020.
[22] Mostafa Elhoushi, Akshat Shrivastava, Diana Liskovich, Basil Hosmer, Bram Wasti, Liangzhen
Lai, Anas Mahmoud, Bilge Acun, Saurabh Agarwal, Ahmed Roman, et al. Layer skip: Enabling early exit inference and self-speculative decoding. arXiv preprint arXiv:2404.16710, 2024.
[23] Alexander R Fabbri, Irene Li, Tianwei She, Suyi Li, and Dragomir R Radev. Multi-news: A large-scale multi-documeett summarization dataset and abstractive hierarchical model. arXiv preprint arXiv:1906.01749 2019.
[24] Zhengcong Fei, Xu Yan, Shuhui Wang, and Qi Tian. Deecap: Dynamic early exiting for efficient image captioning. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 12216-12226 2022.
[25] Bogdan Gliwa, Iwona Mochol, Maciej Biesek, and Aleksander Wawer. Samsum cor- pus: A human-annotated dialogue dataset for abstractive summarization. arXiv preprint arXiv: 12237, 2019.
[26] Yash Goyal, Tejas Khot, Douglas Summers-Stay Dhruv Batra, and Devi Parikh. Making the in vqa matter: Elevating the role of image understanding in visual question answering. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 6904 6913, 2017.
(27] 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. arXiv preprint arXiv:2111.00230., 2021.
(28) Karl Moritz Hermann, Tomas Kocisky, Edward Grefenstette, Lasse Espeholt, Will Kay, Mustafa Suleyman, and Phil Blunsom. Teaching machines to read and comprehend. Advances in neural information processing systems, 28, 2015.
[29] 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. arXiv preprint arXiv: 10::::1700.0000 09844, 2017.
[30] Metod Jazbec, Alexander Timans, Tin Hadži Veljković, Kaspar Sakmann, Dan Zhang, Chris- tian A Naesseth, and Eric Nalisnick. Fast yet safe: Early-exiting with risk control. arXiv preprint arXiv:2405.20911. 2024.
[31] 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 2023, pages 14128 12142, 2023.
[32] Yigitcan Kaya, Sanghyun Hong, and Tudor Dumitras. Shallow-deep networks: Understanding and mitigating network overthinking. In International conference on machine learning, pages 3301-3310 PMLR, 2019.
[33] Junnan Li. Dongxu Li, Silvio Savarese, and Steven Hoi. Blip-2: Bootstrapping language-image pre-training with frozen image encoders and large language models. In International conference on machine learning, pages 17730-19742. PMLR, 2023.
[34] Tsung-Yi Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In Computer VCCV 2014: 13th European Conference, Zurich, Switzerland, September 6- 12, 2014, Proceedings, Part V 13, pages 740-755. Springer, 2014.
12

---

[35] Xiangyang Liu, Tianxiang Sun, Junliang He, Lingling Wu, Xinyu Zhang, Hao Jiang, Zhao Cao, Xuanjing Huang, and Xipeng Qiu. Towards efficient NLP: standard evaluation and strong baseline. 2021.
[36] Lassi Meronen, Martin Trapp, Andrea Pilzer, Le Yang, and Arno Solin. Fixing overconfidence in dynamic neural networks. In Proceedings of the IEEE/CVF winter conference on applications of computer vision, pages 2680-2690, 2024.
[37] David Patterson, Joseph Gonzalez, Quoc Lc, Chen Liang, Lluis- Lluis-Miquel Munguia, Daniel Rothchild, David So, Maud Texier, and Jeff Dean. Carbon emissions and large neural network training. arXiv preprint arXiv: 104. 10350, 2021.
[38] Pranav Rajpurkar, Robin Jia, and Percy Liang. Know what you don't know: Unanswerable questions for squad. arXiv preprint arXiv: /80.00220 2018.
[39] Siddharth Samsi, Dan Zhao, Joseph McDonald, Baolin Li, Adam Michaleas, Michael Jones, William Bergeron, Jeremy Kepner, Devesh Tiwari, and Vijay Gadepally. From words to watts: Benchmarking the energy costs of large language model inference. In 2023 IEEE High Performance Extreme Computing Conference (HPEC), pages 1-9. IEEE, 2023.
[40] Simone Scardapane, Michele Scarpiniti, Enzo Baccarelli, and Aurelio Uncini Why should add early exits to neural networks? Cognitive Computation, 12(5):954 966, 2020.
[41] Tal Schuster, Adam Fisch, Jai Gupta, Mostafa Dehghani, Dara Bahri, Vinh Tran, Yi Tay, and Donald Metzler. Confident adaptive language modeling. Advances in Neural Information Processing Systems, 35:17456-17411
2022.
[42] Roy Schwartz, Gabriel Stanovsky, Swabha Swayamdipta, Jesse Dodge, and Noah A Smith. The right tool for the job: Matching model and instance complexities. arXiv preprint arXiv:2004.00 07453, 2020.
[43] Eva Sharma, Chen Li, and Lu Wang. Bigpatent: A large-scale dataset for abstractive and coherent summarization. arXiv preprint arXiv: 2019.
[44] Tianxiang Sun, Xiangyang Liu, Wei Zhu, Zhichao Geng, Lingling Wu, Yilong He, Yuan Ni, Guotong Xie, Xuanjing Huang, and Xipeng Qiu. A simple hash-based early exiting approach for language understanding and generation. arXiv preprint arXiv:2203.01670 2022.
[45] Tianxiang Sun, Yunhua Zhou, Xiangyang Liu, Xinyu Zhang, Hao Jiang, Zhao Cao, Xuanjing Huang, and Xipeng Qiu. Early exiting with ensemble internal classifiers. arXiv preprint arXiv:2105. 13 13792, 2021.
[46] Shengkun Tang. Yaqing Wang. Zhenglun Kong, Tianchi Zhang, Yao Li, Caiwen Ding, Yanzhi Wang, Yi Liang, and Dongkuan Xu. You need multiple exiting: Dynamic early exiting for accelerating unified vision language model. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10781-10791, 2023.
[47] Surat Teerapittayanon, Bradley McDanel, and Hsiang-Tsung Kung. Branchynet: Fast inference via early exiting from deep neural networks. In 2016 23rd International Conference on Pattern Recognition (ICPR), pages 2464-2469. IEEE, 2016.
[48] 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. 2019. In the Proceedings of ICLR.
[49] Maciej Wolezyk, Bartosz Wójcik, Klaudia Bałazy, Igor T Podolak, Jacek Tabor, Marek Šmieja, and Tomasz Trzcinski. Zero time waste: Recycling predictions in early exit neural networks. Advances in Neural Information Processing Systems, 34:2516-2555 2021.
[50] Ji Xin, Raphael Tang, Jaejun Lee, Yaoliang Yu, and Jimmy Lin. Deebert: Dynamic early exiting for accelerating bert inference. arXiv preprint arXiv:2004. 12993, 2020.
[51] Ji Xin, Raphael Tang, Yaoliang Yu, and Jimmy Lin. Berxit: Early exiting for bert with better fine-tuning and extension to regression. In Proceedings of the 16th conference of the European chapter of the association for computational linguistics: Main Volume, pages 91-104, 2021.
13

---

[52] Wangchumshu Zhou, Canwen Xu, Tao Ge, Julian McAuley, Ke Xu, and Furu Wei. Bert loses patience: Fast and robust inference with early exit. Advances in Neural Information Processing Systems, 33:183300-1833 2020.
[53] Wei Zhu. Leebert: Learned early exit for bert with cross-level optimization. In Proceedings of the 59th Anmual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference Natural Language Processing (Volume 1: Long Papers), pages 2968, 2021.
NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA].
[NA] mcans either that the question is Not Applicable for that particular paper or the relevant information Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers
integral part of your paper submission. They visible to the reviewers, area chairs, senior chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors their evaluation. While
"[Yes] is generally preferable to "[No] ". it is perfectly acceetable to answer "[No] provided a proper justification is given (e.g.. "error bars are not reported because it would be too computationally expensive" or used"). In general, answering "[No] or "[NA] " is
"we unable to find the license for the dataset grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write a justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to question, in the justification please point the section(s) where related material for the question be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: All the claims made are well justified through theoretical and empirical means.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can he expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
14

---

Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss limitation in Section 7
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model seel-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated practice what the implications would
The authors should reflect the scope of the claims made, e-g., if the approach was only tested on a few datasets or with a few runs. general, empirical results often depend on implicit assumptions, which should he articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures hecause fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might that reviewers discover limitations that aren't It acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each thcoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: The Lemma 3.1 and theorem 4.1 are proven in Appendix A. and ?? respectively.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by
formal proofs provided in appendix or supplemental material.
Thcorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification:
We provide all hyperparameter details and other details in Table 11 12 and Section S
Guidelines:
The answer NA means that the paper does not include experiments.
the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible verifiable.
15

---

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, the contribution is specific model empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same datasct, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to hosted model (e.g., in the of large language model), releasing of model checkpoint, or other means that are appropriate to research performed.
While NeurIPS does not require relcasing code, the conference does require submissions provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should deseribe the architecture clearly and fully.
(c) If the contribution is new model (c.g., large language model), then there should either be a way to access this model for reproducing the results a way to reproduce the model (c.g.. with an open-source dataset or instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome to describe the particular way they provide for reproducibility. In the of closed-source models, it may be that access to the model is limited some way (e.g.. to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification:
The code to implement is given in anonymized repository with the link in the supplementary.
Guidelines:
The answer NA means that paper does include experiments requiring code.
Please see the NeurlPS code and data submission guidelines httpp:://nips.ccpppiii
Muides/CodeSuuee saiooooicy for more details.
While encourage the release of code and data, understand that this might not be possible,
so "No" an acceptable answer. Papers cannot rejected simply for including code, unless this central to the contribution (e.g.. for a new open-source benchmark).
The instructions should contain the exact command and environment needed
reproduce
the results. See the NeurlPS code and data submission guidelines Settpe:/nnipas cc cc/public/ for more details.
guides/CodeSubaat Stte//CoooSuubisiiieeee heabbbmissiooPooiooococ
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental sefting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification:
All the hyperparameter details given in Table T1 12 and Section 5
Guidelincs:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
16

---

7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate informa- tion about the statistical significance of the experiments?
Answer: [Yes]
Justification:
The stability of our proposed algorithm can assessed from the figure 2a However, similar
to existing methods [3 19. we have not provided the error bars for the other tasks that require heavy computation.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, statistical significance tests, at least for the experiments that support the main claims
the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given
experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to library function, bootstrap, etc.)
The assumptions made should given (c.g., Normally distributed errors).
should be clear whether error the standard deviation or standard error of mean.
is OK to report I-sigma error bars, one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CL, if the hypothesis of Normality of errors not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative enror rates). If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and referenee the corresponding figures tables in the text.
Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification: Given Appendix
B.7
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code of ethies
Does the research conducted the paper conform, every respect, with the NeurIPS Code of Ethics htttpe////errrpss.cccpbbleebhib wwtppp/ppppp.pcccccace..aaaaaaaa ctt.a////eecccccuuueeeiiiees
Question:
Answer:
[Yes]
Justification:
We ahide by NeurIPS Code of Fthics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
17

---

Answer: [NA]
Justification: Our method is mostly about minimizing risk and improving efficiency, does not add to any societal impact.
Guidelines:
The answer means that there is no societal impact of work performed.
the authors answer NA or No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative socictal impacts include potential malicious unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g. deploy- ment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate point out that an improvement in the quality generative models could be used to generate deepfakes for disinformation. the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the tcchnology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release models, providing defenses in addition attacks, mechanisms for monitor- ing misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g.. pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: We do not release any new model data, and use only models and datasets that are open-source for research purposes.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines restrictions to access the model implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We have cited all the owners of assets that we have used and properly respected the license and terms of use.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website), the copyright and terms of service of that source should he provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperssi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide help determine the license of a dataset.
18

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's
13. New assets
Question: Are introduced in the paper well documented and is the documentation provided alongside assets?
Answer: [NA]
Justification: introduce no new assets.
Guidelines:
The answer means that paper not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their sub- missions structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an amonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: We have not done any crowdsourcing.
Guidelines:
The answer NA that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should included the main paper.
According to the NeurIPS Code of Ethics, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country institution) were obtained?
Question:
Answer: [NA]
Justification: We do not involve crowdsourcing.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing rescarch with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. you obtained IRB approval. you should clearly state this in the paper.
recognize that procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applica- ble), such the institution conducting the review.
16. Declaration LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, non-standard component of the editing, or formatting purposes does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
methods in this research? Note that if the LLM is used only for writing,
Answer: [NA]
19

---

Justification:
We have used LLM only for writing and phrasing.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer our LLM policy htttps://neurips.cco http://www...p..ooooooottr////// for what should should not be described.
20

---

A Appendix
A.1 Proof of Lemma 3.1
Let p(y]z) denote the probability of the true label y given an input æ. Consider an carly-exit model that outputs a probability distribution p(-) over the set of classes for every threshold т. Let pr (x)(z) be the true class probability given by the model, (i]r) be the probability of prediction given by the model and P-(y 9) be the probability that predicted class is the true class.
For fixed r, write the probability that the prediction made by the model correct
Pr(y = b)) = P- = D 0) pr p-(0 x))
p. (ū]x)) p. (v 012, 5)
max (y / = arg max (ū)xr) p. (y = g) TED
Fstt
Hence, we prove that the threshold r that maximizes p- (W]lz) pr = a/ 9) value maximizes the chances of getting the prediction correct.
Proof of Theorem 41
Before moving prove the theorem, will first bound the single-run regret of the UAT algorithm.
At each round 1, .......22 T, the algorithm selects and receives reward r(7)).
Regret Definition: We define the realized (actual) regret for single
R(T)
MA.) Σ A,N
where N.(T) is the number of times arm i was pulled up to time T.
The UAT Algorithm: At cach time t, the UAT algorithm selects the arm
2 T N:(1)
arg max (i((1) V
where: (t()) is the empirical mean reward of arm i up to time N:) is number of times arm has been pulled until time t.
High-Probability Bound Using Hoeffding's Inequality: By Hoeffding's inequality, for any 2 1 and any e >> 0:
P [lAt pel $ 2exp( 2nc?).
Define the confidence interval:
2log T
2 log N.(1)
CL.(1 = (1(()) - -V N:(1) A:(1) +
In one particular round, true reward mean for ith arm lies in this interval with probability 5 where 5 = 1/7
Using the union hound over all arms and times, with probability at least 1 1/" that the true reward mean lies the interval over all the rounds.
Bounding the Number of Suboptimal Pulls: Let bound N:(T") for any suboptimal arm with A > 0.
Suppose arm is selected time f. For this happen, the UAT index for arm must be least large that for the optimal arm i":
2 log T Ni(l) 2 A: (1) +
2 T
(a(() +
No (()
all confidence intervals are valid (which occurs with high probability). then:
log T
H4 +
(t p
21

---

Rearranging,
2og T (1)
A 2
Ni(1) log T A
Thus, the number of times pulled is most:
8 log
N:(T) -
Total Regret Bound: Using R(T) = Eia,s (TN,(T), we obtain:
R(T) &
Σ A
log +1
E
log T
A. + 4.).
PA;20
Finally: With probability at least O(/TT) the realized regret of UAT satisfies:
R(T) Σ
8log T
A: + ).
B(T)
62,50
Let ()) denote model's confidence and P- empirical risk could be written as:
7 v, () calibrated correctness probability. The
T
R(TT) 21
L
p-(y Ú|., 9)
The UAT algorithm that provides threshold = ग (2s) using the policy v. for simplicity, drop the index ł. With probability atleast 8' = 11 7. the single run regret of UAT algorithm can be bounded as:
R(T) B(T)
where the reward r(7) C2 C.) A (ū) pr(y (]= i) Ai consider (1 - C)) approximates the pr 1 a. 1) with probability 51. Choose = [1-d2](1 5') then with probability 8:
Σ [p. (0]=)) (y 0 =. 1) p-. (01=)) 01.. 0) A(i - x0 s (T))
(5)
Rearranging terms and bounding exit differences +11 L:
Σ (+)()) - P-)()) s B(T) + XLT
(6)
where p-(-) 2 p- ()|==)) 322000
Using pr (:) S
Σ Pr ETソ Σ P (()) ALT
Dividing by T and substituting into the risk definition:
R( 1 - (-)+
3(T)
+ +
T
As we assume that R' =1 $ R Pr. Pr. (1) 50
3(T)
R((T) S +
XL
T
For large T. A(T) -.77
0 (since 3(T) Ollog T)). To ensure R(TT) S E:
- _ E"
XL 5 á
L
Therefore, if we select À < f, then with probability at least = (1 - -8)(1- - 5'), the empirical risk R(T)) is bounded This completes the proof.
L
22

---

Model
SamSum
COCO
SQUAD
ROUGE-L
Speedup
B4
Speedup
F1
Speedup
C
46.49
x
40.72
1.62x
81.9
x.57
C A
44.81
1.70x
39.24
x.79x
80.6
1.62x
1 1- C
48.16
1.29x
43.38
1,35x
82.3
1,48
C Cg) X-
47.92
1.68x
42.27
1.56x
81.7
1.68x
48.34
1.35x
43.51
1,48x
83.3
1.60x
C.)
48.09
1.98
43.32
1.79x
83.2
1.82x
Table 6: Performance using different components in the reward function.
e-values
c=0.01
c0.05
c0..1
Model
Risk
Speed
Risk
Speed
Risk
Speed
CALM
04.3
1.37x
04.9
1.54x
05.3
1.92x
MIPLAP
08.1
1.76x
08.6
1.87x
09.7
2.04x
EERS
04.2
1.68x
04.4
1.79x
05.8
2.13x
Ours
01.5
2.08x
02.3
2,25x
03.1
2.52x
Table 7: The impact on performance gap risk (in %) and speedup when the values of e are varied.
B
More Ablation study
B.1 Importance of components of reward function
In T we show the performance of model different tasks and datasets when different components of the reward function given in equation 4 are used. The tasks used are summarization on the SamSum dataset using T5-large, image captioning on the COCO test split using BLIP-2-FlanTS-. xl and Question Answering task using SQuAD on the BERT-large model. We observe that simply maximizing the prediction confidence term C; might lower the performance due to riskier exits while maximizing only the correctness probability -C²) has better improvements, but causes underconfident exits in such cases the model might be less confident, increasing the chances of wrong predictions even with good correctness probability. However, when we maximize the product, it becomes more robust as model confidence and reliability on confidence are both maximized, giving us better performance, As we add the penalty term, there is a slight drop in performance but a boost in efficiency, due to which we select it as the final objective in our sctup.
B.2
The value of c
The choice of coverage threshold depends on the prior knowledge of how complicated the dataset is based on the given distribution and model, as its task make g produce higher values for fraction of samples that easy for the model. We do not priori know this information, hence we rely on validation dataset to assess the task complexity. Also, as the exits might find the sample easy or hard based on its information that might be different for different exits. Hence we choose dependent on i, denoted as e, During training of the first epoch, we do not have any prior knowledge of the task complexity hence set to be 0 for all é, but after the first epoch we define e x Σm= (A,==)) where Okj, i.c., the fraction of samples for which the exit is correct. This assesses the task complexity, making the g function produce higher scores hased complexity.
B.3 Risk-E fisk-Efffciency trade-off
In Figure 2b we show the risk-efficiency trade-off by plotting the percentage of performance drop compared final layer vs Speedup for the SQuAD dataset on the BERT-large model by changing the parameter A. For other baselines, we change their trade-off parameter and plot the results. The results suggest that our method consistently has lower risk while having a higher speedup. While we keep the value of = 0.01, i.e the percentage of R. 5 c. Our method keeps the risk lower than 1% till speedup of 2.3x after which the risk crosses the desired hound, while other risk-based methods very less tolerant to keep risk below 1% when we increase in speedup even when we restrict their permissible risk limit to 1% similar to ours.
B.4 Regret Analysis
we plot the average cumulative regret while choosing the thresholds during inference using different policies on the QNLI dataset. The results are average across runs where in every run the dataset was randomly reshuffled and fed to the algorithm. We plot the cumulative regret of final layer when all the samples always exit from the final layer, randomly selected, where random threshold was assigned to each sample and fixed
Figure 2a
23

---

600
Finel Jayer
3.5
OUTS CAUM
Randomly_seiien selected
3.0
500
Fixed, Uhreshold (Existing methods)
Cumulative Regret
Our method
400
300
Risk
200
100
1000
2000
3000
4000
5000
Samples
Speedup
(a) Cumulative regret plots for different policies to choose the thresholds.
(b) Risk vs Speedup curve obtained by varying the trade-off parameter A.
Figure 2: Comparative plots showing (a) cumulative regret for different threshold selection policies and (b) the risk-speedup trade-off.
threshold where static threshold was utilized similar existing EE and risk based methods. Our method observed smallest regret as compared to others. This indicates importance of learning a policy for threshold selection instead of fixing it to constant value.
B.5 Changing values
Table we provide results on changing the that guides R < c on the SamSum dataset with Ta-large model, note that, changing value changes the value of A, allowing riskier samples to make an early exit. We report the percentage of increase in R as compared to the final layer. Observe from the table that even when risky samples upgraded confidence metric that allows reliable. While imposing a hard criterion the closest to the required threshold. After relaxing the criteria, all methods except MIPLAP do not violate it, still, our method had the lowest performance gap risk as compared to other methods.
allowed to exit early, method's 's drop in performance is very low. This is duc carly exit only when risk methods
confident and the confidence is also
0.01,
violated that, while our being
B.6 Some important examples
Table 10 we provide some examples on the classification datasets such as, SST-2 and QNLI, where we report the True label, Fake confidence (average model confidence over the true class), and risk score (average Cy)). high that the model outputs
Observe from the Table that the sample gets wrong prediction. In such cases, becomes important flag those samples that done by the C value in case. For instance, consider sample "the under-7 crowd", which means that review negative, but model prohahly understands the numbers as ratings and develops high confidence over the wrong class.
ambiguous, the chances
B.7 Computational setup
our experiments, we have used a setup NVIDIA A6000 GPUs. The highest run time was observed during training the BLIP-2 model for im- age captioning, which took 15 hours to train for 15 epochs. The average GPU runtime for the BLIP-2 model across tasks was 9 hours. The inference time for BLIP-2 was less than 20 minutes across all the tasks. For the T5-large model, the average runtime observed was 6 hours, with the highest the Ques- tion Answering dataset due to large number epochs. For inference, required less than 15 minutes across all datasets and tasks. Finally for the BERT-large model, the maximum runtime was on MNLI dataset with 1 hour of GPU runtime across epochs. The average runtime across GLUE tasks was 25 minutes. The average inference runtime was less than 2 minutes for all the tasks.
Model
SamSum
QA
MT
TS-large
1035s
664s
846s
TS-large w UAT
1048s
672s
855s
Table 8: Comparison of inference times (in sec- onds) across tasks on T5-large model.

---

Model
SamSum (T3)
COCO (BLIP-2)
QNLI (BERT)
loss performance
43,9
Our loss performance
48.82
44.0
90.5
Table 9: Model performance trained using Conventional loss loss.
Example
True lbl.
Fake Conf.
C
SST-2
mysterious and brutal nature the under-7 crowd seems endless movie with two stars
positive
0.79
0.68
negative
0.91
0.95
negative
0.90
0.82
positive
0.82
0.86
QNLI
Where What university donated the land ...? A site chosen Houston, ...
Tought? The
lought primarily along the entailment 0.70
0.88
entailment
0.76
0.72
Table 10: Examples of samples achieving fake confidence, the table shows example, its true label (True lbl.), the average confidence the model on the wrong class (Fake Conf.) and the (1 - Cg) score that abstains the sample early.
B.8 Computational Complexity of UAT
Our method learns policy to decide the exit threshold that is learned over time based on the data distribution. However, the computational complexity of UAT is negligible it does involve any big operations. In every, step. requires maximizing over small finite set (set 10 elements) and maintain list of reward functions. Table we perform experiment to show the additional complexity our method. In this experiment exit was fixed to the final layer. i.c. there was no early exiting. We implement the UAT algorithm and do not change anything else, i.e., the exit is still from the final layer, but the only thing changed is that there is some additional computation happening due UAT. So the time wall-clock time added to T5-large with UAT will be Table computational complexity UAT negligible as the additional time is less than 15 seconds, further proving UAT's contribution during inference.
the computational complexity of UAT. Observe
B.9 Loss observations
used train function given our method in Equation 3 The Juggesss that adding the other components in the loss function does not impact the overall model performance. We check the final layer performance, that would have been the most affected by the additional components. However, the effect very insignificant, and safely use the loss function given Equation
Table we provide results of final layer when the conventional cross-entropy EEDNN i.e., just use the first term Lce from Equation Cantioonnannn you use the
Model
BLIP-2-FlanTS-Al
BERT-large GLUE 5
Tasks
Image captioning
VQA 5
Finetuning epochs
15
Finetuning dataset
Train Split
Train split
Train split 1000
Warmup steps
1000
1000
Learning rate
le-5
le-5
le-5
Batch size
16
16
32 (0.9, 0.999)
AdamW beta
(0.9, 0.999)
(0.9, 0.999)
Weight decay
0.05
0.05
0.05
Drop path
Image resolution
360
490
Prompt
"a photo of" 5
"Question:() () Answer"
Inference beam size
5
Table 11: Hyperparameter details the BLIP-2-Flam T5-xl and BERT-large model on various tasks.
25

---

Model
T5-large MT
Tasks
Summarization
QA
Finetuning epochs
5
Finetuning dataset
Train Split
Train Split
Train Split
Warmup steps
1000
1000
1000
Learning rate
le-4
le-4
le-4
Batch size
8
8
16
AdamW beta
(0.9. 0.999)
(0.9, 0.999) 0.05 0
(0.9, 0.999)
Weight decay
0.05
0.05
Drop path
0
0
Image resolution
Prompt
summarize:
translate german english
Inference beam size
5
Table 12: Hypereparameters details on the Ta-lerge model for various tasks.
26


Reviews:
[{"review_id": "L7PEk0DWax", "reviewer": "Reviewer_hcfn", "summary": "The paper presents a framework for selecting adaptive thresholds for early-exit neural networks. The motivation is that existing methods rely on a fixed threshold chosen from a static validation set, which can lead to incorrect predictions when the model is over-confident in the wrong class under distribution shift at inference. This reduces accuracy and undermines model reliability. To avoid this, the paper proposes to adaptively update a single global threshold online during inference using an upper-confidence-bound (UCB) multi-armed-bandit framework, with a probabilistic guarantee that the overall miscoverage risk stays within a tolerance, while improving computational efficiency by exiting at early layers. Since ground-truth labels are unavailable during inference, a neural network is trained offline to estimate the reliability of the prediction. This estimated reliability is used to select the best threshold by maximizing a reward function, which combines the product of confidence and reliability with an exit penalty that increases for later layers. Experimental results show that the proposed dynamic-threshold selection achieves significant speed-ups with minimal drops in performance.", "questions": "See weaknesses", "limitations": "Included but outside page limit", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n* The motivation of the paper is clear and intuitive; the threshold for early exit decisions should be adaptive to account for distribution shift during inference.\n* The formulation is interesting, using confidence and reliability of prediction in addition to penalizing late exits to improve prediction quality and computational efficiency simultaneously.\n* Experimental results demonstrate the effectiveness of the approach on diverse tasks, where in some tasks the performance is better than the full model\n\n**Weaknesses**\n* The method still selects a global threshold which is adapted online based on previously observed samples. So if an OOD sample arrives, the exit decision is made based on the global threshold, and the model can still be overconfident and misclassify or exit at later depth. Instead, a per-sample based threshold would be more robust to abrupt distribution shifts. \n* Because the reliability function g is trained offline, it is possible that it drifts under distribution shifts and the estimates are unreliable at inference time due to sub-optimal threshold selection. How can we be sure of the robustness of g under distribution shifts?", "was_revised": true, "final_justification": "This is a well-written paper that addresses a critical problem in the early-exit literature. The proposed solution is reasonable and represents a step towards improving the real-world practicality of such models.", "rebuttal": ""}, {"review_id": "2cgx3zwX9K", "reviewer": "Reviewer_Cxuf", "summary": "The paper addresses the problem of selecting an appropriate exit threshold in early-exit neural networks. The core idea is to focus on confident but incorrect predictions. To tackle threshold selection, the authors introduce a formal notion of risk and apply a Multi-Armed Bandit approach guided by this notion. They support their method with theoretical justification and present empirical comparisons against state-of-the-art early-exit strategies.", "questions": "1. What is the definition of near-domain and out-of-domain in Section 5? Specifically, what techniques are used to generate samples that are not in the domain? Are there any limitations on how far the data can deviate from the original domain?\n2. From my experience, CeeBERT [7] and UCBEE use Multi-Armed Bandit (MAB) techniques for exit decision-making. However, the authors do not compare their method with these or any other MAB-based approaches in the main results (e.g., Tables 1–5). If I’m mistaken, please correct me. Otherwise, it would be helpful to clarify whether there is a specific reason such a comparison was omitted. If not, including it would be strongly encouraged.\n3. In my opinion, using a single threshold shared across all exits is a major limitation. Can the authors provide experiments with exit-wise thresholds? I understand that this may be challenging for networks with many intermediate classifiers. To make this more tractable, a simplified setup—such as a predefined threshold schedule (e.g., linear or exponential decay using a single additional parameter)—would already be valuable. An alternative approach would be to adopt dynamic thresholding, as proposed in Fixing Overconfidence in Dynamic Neural Networks [36], where a threshold vector is derived from a single “heaviness level” parameter. https://github.com/AaltoML/calibrated-dnn/blob/main/adaptive_inference.py#L251-L259\n\nI would consider raising my score if the authors address the above concerns and provide additional clarification or experiments where suggested.", "limitations": "Yes, the authors adequately addressed the limitations.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The paper is well-structured, easy to follow, and well-motivated.\n2. The definitions of the loss and reward functions are clear and intuitive.\n3. The authors perform distribution shift experiments, which are crucial for practical applications.\n4. The work includes multiple evaluations across various tasks.\n5. The included mathematical theorems enhance the understanding of the topic.\n\nWeaknesses:\n1. There is no evaluation on vision-only tasks, such as CIFAR-100 or ImageNet-1k. Many works in the early-exit literature also consider these benchmarks.\n2. Evaluation on CNNs would be beneficial as well, especially since all the models considered in the paper are transformers (which is understandable for NLP tasks)\n3. The authors do not provide performance scores (e.g., accuracy) for individual heads – that is, what the accuracy would be if all samples exited at the i-th intermediate classifier – nor do they report the exit rates (i.e., the proportion of samples exiting at each head). Including such information would offer valuable insights into the model's decision-making process.\n4. The authors consider a single-threshold strategy for exiting, which is known to be suboptimal. While it is commendable that the authors explicitly acknowledge this limitation, it may significantly affect the generalizability of the proposed improvements.", "was_revised": true, "final_justification": "The authors addressed my concerns. This is a good piece of work.", "rebuttal": ""}, {"review_id": "SnM4ouWodJ", "reviewer": "Reviewer_NGvD", "summary": "This paper proposes a reward function that assesses predictive certainty for early exit mechanism to balance computational efficiency and prediction quality. It utilizes a Multi-Armed Bandit framework to take into account it for an early exit decision. It evaluates the proposed method in terms of inference speed up and performance.", "questions": "The authors state “reliability of confidence” in multiple places, it is not clear what is the difference between “reliability of confidence” and “correctness/accuracy of confidence”. As far as I figure out, by reading the paper, it is ultimately the correctness of the confidence. Please clarify what the authors mean by “reliability.”", "limitations": "* One of the limitations is the novelty – the problem is not novel.\n\n* The another limitation is that the work contains conflicting assumptions and/or explanations with regard to the offline training mechanism (especially the pre-training mechanism) of the proposed Early exit mechanism. I do not believe it can be practically deployed in real-world applications.", "rating": 2, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 1, "originality": 1, "strengths_and_weaknesses": "* The problem is not novel.\n\n* “High confidence” depends on the threshold. If the model overall overaly confident, the threshold can be adjusted. Accordingly, depending on the threshold, the Fig. 2’s results can vary, and the authors’ claim on the over confidence can be not true.\n\n* The statement in line 59, “During inference, the data arrives in an online fashion.” is not always true. It depends on the applications/systems/ downstream tasks.\n\n* Thue authors state “reliability of confidence” in multiple places, it is not clear what is the difference between “reliability of confidence” and “correctness/accuracy of confidence”. As far as I figure out, by reading the paper, it is ultimately the correctness of the confidence. Please clarify what the authors mean by “reliability.” – Which situation would be reliable but incorrect and which situation would be not reliable but correct?\n\n* There is a conflicting or confusing point. The authors mentioned “During inference, the data arrives in an online fashion” in line 59. But, in line 160, for the first stage, they mention “we perform offline training of the EEDNN…”. If the model is already deployed to receive data for inference arriving an online fashion, how it can be offline-trained again? Then, the model cannot accommodate such online/realtime data, and also that is against the aim of the EEDNN for inference.", "was_revised": true, "final_justification": "This work does not have a significant novelty.", "rebuttal": ""}]

Author rebuttal:

Meta-review: This paper addresses the problem of risk control for early-exit neural networks (EENNs).  Whereas prior work has predominantly used confidence-based exiting (e.g. picking a confidence threshold on a validation set), this has the limitation of being non-robust to distribution shift.  This paper addresses this issue by formulating threshold-tuning as an online bandit problem, allowing the threshold to be calibrated for data that has undergone distribution shift.  Compelling results are shown on language and vision tasks.

The reviewers, in aggregate, had a favorable opinion of the work: 1 x reject, 2 x accept.  Reviewer NGvD criticized the work's novelty---specifically, the motivation for an online setting that necessitates calibrating an exit-threshold.  However, I concur with the other two reviewers, finding the work does indeed provide a novel approach to a perhaps niche but still practical setting of online calibration under data shift.  Yet, in subsequent drafts, the authors should work to better motivate the online setting.

While not mentioned by the reviewers, I see one issue with the draft that is in need of correction.  Section 2 claims "Rather than comparing early exits against the final layer’s predictions, we model the overall risk of the complete system, enabling our method to potentially surpass the performance of the base DNN itself, an important limitation overlooked by prior work".  This statement is incorrect: Jazbec et al. [30] do consider and support negative losses, which means that an early exit is performing better than the final exit.  See their Proposition 3 and surrounding discussion.  Please correct this claim in subsequent drafts.
Decision: Accept (poster)