Title: Beyond Greedy Exits: Improved Early Exit Decisions for Risk Control and Reliability
Venue: NeurIPS 2025
Accepted: yes
Keywords: Early exits

Risk-controlled inference
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

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 × ) with a minimal performance drop ( < 2%) as compared to full model performance. Our source code is available at https://github.com/Div290/UAT .

1 Introduction

As Deep Neural Networks (DNNs) continue to scale in size and complexity, the cost of inference has become a major bottleneck-manifesting as increased latency, higher energy consumption, and a 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) [47] offer a compelling solution for adaptive inference under dynamic computational constraints. EEDNNs can produce intermediate predictions at various depths of the DNN, allowing the network to 'exit' the sample early without passing it 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, 29], natural language processing [12, 9], and vision-language systems [46]. 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 trade-off-balancing efficiency gains against potential degradation in predictive accuracy [30]. This trade-off is governed by the exit decisions that are crucial in EEDNNs.

In existing methods, the decision to exit depends on two variables: confidence and threshold. Confidence is usually defined as the entropy of the prediction [50], prediction consistency [52], or ensemble methods [8, 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 is allowed to exit without further processing (early exit). However, there are some crucial concerns in existing methods: 1) The confidence is not a 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. 1 discussed in the Appendix 5.1 shows this phenomenon, where a 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 an 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 a novel framework for dynamic threshold adaptation in EEDNNs, aimed at optimizing the trade-off between prediction performance and computational costs. In contrast to existing methods that rely on confidence at an intermediary layer exceeding a 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) [3] framework, where each arm represents a candidate exit threshold. The bandit agent aims to learn a 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-language understanding (e.g., image captioning [34], VQA [20]), text generation (e.g., summarization [25], QA [38]), and classification (e.g., NLI, sentiment analysis [48]). In 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 a dynamic early exit framework that adapts to distributional changes at test time, enhancing both efficiency and robustness in EEDNNs.

We apply the bandit framework to learn the optimal strategy. We develop an upper confidence-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 as 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: BranchyNet [47] first proposed the early exit method for image classification tasks. It uses classification entropy as a confidence metric. Shallow-deep [32] and MSDNet [29] improve upon BranchyNet by effectively choosing the thresholds based on the confidence threshold. ZTW [49], 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 [24], MuE [46], CapEEN[5] and FREE [10] implement early exits to vision language tasks.

For the Language tasks, and specifically in Large Language Models (LLMs), the EE methods have been popular [11, 15, 35, 13, 45, 21, 27, 14, 44, 31]. DeeBERT [50] first applied it to the BERT model. Later, PABEE [52] developed a patience-based method to decide exiting. BERxiT [51] proposes an

efficient fine-tuning strategy for EE-based models. LeeBERT [53] performs self-distillation between exits to effectively share the knowledge across layers with exits. ETFEE [31] 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 to 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-aware decisionmaking 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 to 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 × Y . A sample, denoted ( x, y ) , is drawn from an underlying probability distribution P . For classification tasks, the label space is given by Y = { 1 , 2 , . . . , K } , whereas for regression, the output set is a subset of R d .

3.1 The Notion of Risk

In statistical modeling, risk control mechanisms [2, 16] enhance the reliability of predictive models by imposing constraints on threshold-based decision processes. Let Ω denote a finite set of thresholds. Consider a model f τ parameterized by a threshold τ ∈ Ω , which outputs a set of predictions based on confidence levels. For a classification task, the set-valued predictor f τ : X → 2 Y is defined as f τ ( x ) = { ˆ y ∈ Y | p (ˆ y | x ) ≥ τ } , where p (ˆ y | x ) is the probability that model assigns label ˆ y to x . If no class probability exceeds τ , the predictor abstains by returning an empty set ( f τ ( x ) = ∅ ).

To assess reliability, we define a miscoverage loss that penalizes the exclusion of the true label as ℓ ( f τ ( x ) , y ) = ✶ [ y / ∈ f τ ( x )] , where ✶ [ · ] is the indicator function. The associated risk is the expected probability of miscoverage:

For a given tolerance, ϵ ∈ (0 , 1) and δ ∈ (0 , 1) , we say that the threshold set ˆ Ω controls the risk with within tolerance ϵ with probability at least 1 -δ if

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.

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 a probability distribution over Y denoted by p i ( ·| x ) = p ( ·| x, θ i ) = EC i ( h i ) where h i = ϕ i ( x ) is the hidden representation output by the i th layer of the model with i ∈ { 1 , 2 , . . . , L } and θ i is the set of parameters consisting of the EEDNN parameters till layer i comprising of both EC i and main model parameters. The final prediction layer, indexed by L , corresponds to the standard DNN where all layers are utilized for inference. Each EC i estimates the class label as ˆ y i = arg max y ∈Y ˆ p i ( y | x )

with a confidence score C i := C i ( x ) = max y ∈Y p i ( y | x ) . C i measures the certainty associated with the prediction. However, this confidence might not be very reliable as it 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 τ i ∈ [0 , 1] for the first time. Following previous works [50, 30], a single global threshold is considered across all the ECs , i.e., τ i = τ for every i ∈ { 1 , 2 , . . . , L } . Then, for a given τ , the probability that x exits with a label ˆ y is given by

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 a set ˆ Ω that satisfies Eqn. 1, we aim to identify threshold τ ∗ = min { τ ∈ ˆ Ω } . The proportion of samples that exit earlier is higher for τ ∗ than any other τ in ˆ Ω .

̸

For a given threshold τ , let p τ ( y = ˆ y | x ) denote the probability that the EEDNNs makes errors on sample ( x, y ) . Then the risk of the EEDNNs is defined as

̸

Our objective is to find a threshold τ for which R ( τ ) lies within a given tolerance value with high probability. The value of τ that meets the requirement is a 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 τ 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 in two stages.

In the first stage, we perform offline training of the EEDNN and a linear layer network that will allow us to 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 τ 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 a threshold τ , let ˆ y be the label predicted by EEDNNs on sample ( x, y ) . Define p τ ( y = ˆ y | ˆ y, x ) be the probability that the predicted label is the true label. Then p τ ( y = ˆ y | x ) = p τ (ˆ y | x ) · p τ ( y = ˆ y | x, ˆ y )

The proof is straightforward and is given in the Appendix A.1. p τ ( y = ˆ y | x, ˆ y ) is the probability that the model's prediction is correct. Note that p τ (ˆ y | x ) can be estimated during inference as it does not depend on the label. Then, the lemma suggests that we can compute risk by estimating p τ ( y = ˆ y | x, ˆ y ) . The next subsection discusses training a neural network to estimate it in an offline fashion.

3.3 Exit and risk function training

We train a reliability function (neural network) g that assigns a reliability score to the class probabilities output by the exit classifiers. We use a modified loss function that is 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 a sample ( x, y ) ∼ P , we define a modified loss is as:

The L CE ( p i ( ·| x ) , y ) in the first term is the standard cross-entropy loss that encourages the model to correctly classify the input. It is multiplied by a factor 1 + g ( · ) , to link the sample-wise loss with the reliability function g , which estimates the confidence of the model's prediction. The minimization of L i 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 g is penalized for activating on high-loss (uncertain) samples and is incentivized to concentrate its high values on low-loss samples.

When g ( · ) ≥ 0 . 5 , it indicates that the model output is reliable. ϕ ( g ) in the second term denotes the fraction of samples for which the model's output is reliable. Φ( c -ϕ ( g )) acts as a regularizer to enforce that the model is reliable for at least c fraction of samples. The function Φ( a ) = max(0 , a ) 2 penalizes the model when the coverage ϕ ( g ) falls short of the desired threshold c i . This regularizer restricts g from giving a trivial solution g ≡ 0 , 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 at least c -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 correctness probability p ( y = ˆ y | x, ˆ y ) . To account for the cost associated with deeper exits [32], the final loss across all exits is computed as a weighted average L = ∑ L i =1 i ·L i ∑ L i =1 i , which assigns higher weights to exits at deeper layers where inference cost is higher.

3.4 Dynamic thresholding during inference

We first describe the MAB setup to dynamically adapt to the threshold during inference. In a MAB setup, a decision-maker iteratively selects actions, adapting to an 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 a given threshold/action τ , let C i τ denote the confidence model has in the predicted class when the sample exits from EC i . We define C i g = 1 -g ( p i ( ·| x )) as a score of underlying risk (unreliable) in the confidence score of the model at EC i . We set 1 -C i g as an approximation of p τ ( y = ˆ y | x, ˆ y ) . Then Lemma 3.1 indicates that the p τ ( y = ˆ y | x ) can be well approximated as C i τ (1 -C i g ) when sample x exits from EC i . However, a threshold that maximizes C i τ · (1 -C i g ) 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 a threshold τ ∈ Ω on a sample as follows:

where ψ ( · ) is a penalty to disincentivize the sample from moving to a deeper layer unnecessarily. In our setup, we define ψ ( i ) = λ · i where λ could be interpreted as the processing cost per layer.

The expected reward for arm τ ∈ Ω is E [ r ( τ )] = ∑ L i =1 E [ C i τ (1 -C i g ) -λ · i ] P ( i ) where P ( i ) is the probability that the sample exits from i th layer and expectation is with respect to the randomness of confidence and reliability scores. Let τ ∗ = arg max τ ∈ Ω E [ r ( τ )] denote the optimal threshold for a given λ . Now consider a policy π that selects the threshold τ t ∈ Ω 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 =1 E ( r ( τ ∗ ) -r ( τ t )) , where the expectation is with respect to randomness in the selection of thresholds induced by the past sample. A policy π that satisfies R ( π, T ) /T → 0 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 a small linear component. As we will show later, this translates to achieving a small empirical risk defined as ˆ R ( π ) = 1 -∑ T t =1 p τ t ( y t =ˆ y | x t , ˆ y t ) T where ( x t , y t ) is the input sample, ˆ y t is the predicted label, and τ t is the threshold selection by π in round t .

4 Algorithm

We develop an algorithm named UAT. Its pseudo-code is given in 1. The inputs to the algorithm are the exploration constant γ and the penalizing term ψ ( i ) . 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 τ t . UCB indices are obtained by taking the weighted sum of the empirical average rewards and the confidence bonuses with γ as the weight factor. If C i τ (1 -C i g ) at the i th layer is greater than τ t , 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 ( τ t )) and the empirical mean ( Q ( τ t )) of the played arm. Note that the algorithm is applied in the inference phase. In Table 6, we empirically show how different components contribute to the reward function and discuss in Appendix B.1. 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 -C i g ) approximates the value of p τ (ˆ y = y | x, ˆ y ) with probability (1 -δ 1 ) . Let the risk associated with τ ∗ be bounded by ϵ ∗ . Then, for sufficiently large T and given tolerance ϵ , UAT achieves P ( ˆ R ( π ) ≤ ϵ d ) ≥ (1 -δ 1 )(1 -δ ′ ) . where ϵ d = ϵ + ϵ ∗ , λ ≤ ϵ L and δ ′ is a constant.

The detailed proof is given in the appendix. Note from the theorem that ϵ ∗ is 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 ϵ d = ϵ + ϵ ∗ . This theorem also suggests that a lower value of λ can lead to lower risks, which is true based on our reward function. If λ is chosen small, then more emphasis is given to the C i τ (1 -C i g ) value which maximizes p ( y = ˆ y | x, ˆ y ) . The proof also provides us a way to fix the λ value, as we want the maximum λ value such that the risk levels are below ϵ . Hence we choose λ = ϵ L in this work. This theorem bounds the risk of the EEDNN where the threshold to decide exiting is chosen based on the policy π given by algorithm 1. Also, observe from the proof outline that it suggests that the risk of the EEDNN will be greater than ϵ ∗ , which denotes the optimal risk. The proof steps also provide the link of risk with the regret.

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.

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), 2) Text entailment (RTE, QQP) and 3) Natural Language Inference (MNLI, QNLI). For Question-Answering tasks, we use the SQuAD 2.0 [38] dataset. For text summarization tasks, we use a diverse set of datasets covering different domains and styles: 1) Dialogue summarization (SAMSum) [25], 2) Long-form technical summarization (BIGPATENT) [43], 3) Multi-document summarization (MultiNews) [23], and 4) News summarization (CNN/DailyMail) [28]. For machine translation tasks, we use the IWSLT De-En dataset [17], which consists of translated transcripts from German to English. It serves as a standard benchmark for evaluating translation quality on spoken language content. For vision language tasks, we use the COCO [34] and NoCaps [1] datasets for Image captioning, VQAv2 [26] dataset for visual question answering and VisDial [20] for visual dialogue.

Metrics: We report key metrics such as accuracy for text classification, F1 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 are preferred for better and fair comparison.

Setup: For the training phase, we augment the pre-trained models with a linear output layer to serve as an exit. For the architecture of the g , we use a single linear layer with shared parameters across the exits that serves as a risk predictor. Following the existing methods [8], the inference is performed on a per-instance basis, setting the batch size to 1 . The value of λ is fixed to ϵ L , while the set of candidate thresholds is chosen as ten equally spaced values between 0.5 and 1.0 (including 1.0). An ablation study over different λ values is also performed in Appendix B.3. It aligns with scenarios where low-latency is critical, such as processing individual requests from different users [42]. Similar to other baselines [30], we also choose ϵ = 0 . 01 for all the experiments. An ablation study over ϵ

can be found in Appendix B.5. More hyperparameter details can be found in Table 11 and 12 in the Appendix.

Baselines: 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 are the baselines where we consider various early exit methods, such as PABEE [52], patience-based 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 a gating function to performs an exit. Finally, BEEM [8] is a 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 to minimize risk and fix the overconfidence issue in the image classification tasks; we extend this to other tasks as well. Finally, we have EERC [30] that theoretically provides a simple extension of EE models to provide a method to choose the value of the exit threshold such that the risk is minimized. We use the same code and hyperparameters without any changes to get the results of the baselines.

Results on Text classification: In Table 1, 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 better than the full model performance due to the impact of overthinking [32, 52]. The EE-based methods have more focus on the speedup part while having a 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 an inference. However, this is only observed in easier tasks such as text classification. On the other hand, the risk-based methods have a lower loss in accuracy but observe a hit in speedup as their focus is solely on reducing the risk, while achieving minimal efficiency gains.

Our method balances risk and efficiency using the reward function and minimizes the risk while getting improved efficiency.

Results on Language Modeling: In Table 2, we report the results on tasks such as text summarization, Question-answering 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 a very high static threshold was also getting a 2 points ROUGE-L score drop with a higher loss in efficiency. Due to this fact, there is a performance drop in the risk-based as well as EE-based methods shown in Table 2. Risk-based 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 to switch from a hard criterion to exit to a soft one, helping it balance accuracy and efficiency.

Results on the vision-language tasks: In Table 3 and 4, we report the results on various visionlanguage tasks using the BLIP-2 backbone with ViT-g as the encoder and FlanT5-xl 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 as 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 5, we provide the BLEU-4 scores for image captioning tasks, where additional noise is added to the test split of the dataset. The model, BLIP-2-FlanT5xl, was trained on pristine (undistorted) images, and during inference, we add different levels of distortions in terms of Gaussian noise with mean zero and standard deviation σ ; 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 5 suggest that our method'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 a fixed threshold as done by other methods.

5.1 Ablation Study

Analysis of the overconfidence issue In Figure 1, 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 4 th 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.

0

1

2

3

4

5

6

7

8

9

10

11

Layer

0.0

0.2

0.4

0.6

0.8

1.0

True Class Confidence

0.0%

1.8%

2.0%

12.4%

10.2%

7.9%

8.2%

7.6%

7.4%

7.5%

7.4%

7.0%

6 Conclusion

In this work, we introduced a risk-tolerant approach for Early-Exit Deep Neural Networks (EEDNNs) that balances accuracy and efficiency under dynamic test-time conditions. We first designed a reliability-driven confidence metric that quantifies the trustworthiness of exit predictions, and seamlessly integrated it into a 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 a 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. It is 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 all the existing methods use a global threshold across all the exits to assess the confidence, there is scope to have a local threshold for individual exits, then the model can be more efficient.

References

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 8948-8957, 2019.

Anastasios N Angelopoulos, Stephen Bates, Adam Fisch, Lihua Lei, and Tal Schuster. Conformal risk control. arXiv preprint arXiv:2208.02814 , 2022.

Peter Auer, Nicolo Cesa-Bianchi, and Paul Fischer. Finite-time analysis of the multiarmed bandit problem. Machine learning , 47(2):235-256, 2002.

Divya J. Bajpai, Aastha Jaiswal, and Manjesh K. Hanawal. I-splitee: Image classification in split computing dnns with early exits. In ICC 2024 - IEEE International Conference on Communications , pages 2658-2663, 2024.

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 Linguistics: EMNLP 2024 , pages 6458-6472, Miami, Florida, USA, November 2024. Association for Computational Linguistics.

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: ACL 2024 , pages 1736-1748, Bangkok, Thailand, August 2024. Association for Computational Linguistics.

Divya Jyoti 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.

Divya Jyoti Bajpai and Manjesh Kumar Hanawal. Beem: Boosting performance of early exit dnns using multi-exit classifiers as experts. arXiv preprint arXiv:2502.00745 , 2025.

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.

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 Computational Linguistics: ACL 2025 , pages 23599-23615, Vienna, Austria, July 2025. Association for Computational Linguistics.

Divya Jyoti Bajpai and Manjesh Kumar Hanawal. A survey of early exit deep neural networks in nlp. arXiv preprint arXiv:2501.07670 , 2025.

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.

Nikita Balagansky and Daniil Gavrilov. Palbert: Teaching albert to ponder. Advances in Neural Information Processing Systems , 35:14002-14012, 2022.

Andrea Banino, Jan Balaguer, and Charles Blundell. Pondernet: Learning to ponder. arXiv preprint arXiv:2107.05407 , 2021.

Ankur Bapna, Naveen Arivazhagan, and Orhan Firat. Controlling computation versus quality for neural sequence models. arXiv preprint arXiv:2002.07106 , 2020.

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.

Mauro Cettolo, Marcello Federico, Luisa Bentivogli, Jan Niehues, Sebastian Stüker, 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-14, 2017.

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.

Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Yunxuan Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, et al. Scaling instruction-finetuned language models. Journal of Machine Learning Research , 25(70):1-53, 2024.

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.

Maha Elbayad, Jiatao Gu, Edouard Grave, and Michael Auli. Depth-adaptive transformer. In In Proc. of ICLR , 2020.

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.

Alexander R Fabbri, Irene Li, Tianwei She, Suyi Li, and Dragomir R Radev. Multi-news: A large-scale multi-document summarization dataset and abstractive hierarchical model. arXiv preprint arXiv:1906.01749 , 2019.

Zhengcong Fei, Xu Yan, Shuhui Wang, and Qi Tian. Deecap: Dynamic early exiting for efficient image captioning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 12216-12226, 2022.

Bogdan Gliwa, Iwona Mochol, Maciej Biesek, and Aleksander Wawer. Samsum corpus: A human-annotated dialogue dataset for abstractive summarization. arXiv preprint arXiv:1911.12237 , 2019.

Yash Goyal, Tejas Khot, Douglas Summers-Stay, Dhruv Batra, and Devi Parikh. Making the v 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.

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.

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.

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:1703.09844 , 2017.

Metod Jazbec, Alexander Timans, Tin Hadži Veljkovi´ c, Kaspar Sakmann, Dan Zhang, Christian A Naesseth, and Eric Nalisnick. Fast yet safe: Early-exiting with risk control. arXiv preprint arXiv:2405.20915 , 2024.

Yixin Ji, Jikai Wang, Juntao Li, Qiang Chen, Wenliang Chen, and Min Zhang. Early exit with disentangled representation and equiangular tight frame. In Findings of the Association for Computational Linguistics: ACL 2023 , pages 14128-14142, 2023.

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.

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 19730-19742. PMLR, 2023.

Tsung-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 Vision-ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13 , pages 740-755. Springer, 2014.

Xiangyang Liu, Tianxiang Sun, Junliang He, Lingling Wu, Xinyu Zhang, Hao Jiang, Zhao Cao, Xuanjing Huang, and Xipeng Qiu. Towards efficient NLP: A standard evaluation and A strong baseline. 2021.

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.

David Patterson, Joseph Gonzalez, Quoc Le, Chen Liang, Lluis-Miquel Munguia, Daniel Rothchild, David So, Maud Texier, and Jeff Dean. Carbon emissions and large neural network training. arXiv preprint arXiv:2104.10350 , 2021.

Pranav Rajpurkar, Robin Jia, and Percy Liang. Know what you don't know: Unanswerable questions for squad. arXiv preprint arXiv:1806.03822 , 2018.

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.

Simone Scardapane, Michele Scarpiniti, Enzo Baccarelli, and Aurelio Uncini. Why should we add early exits to neural networks? Cognitive Computation , 12(5):954-966, 2020.

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-17472, 2022.

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.07453 , 2020.

Eva Sharma, Chen Li, and Lu Wang. Bigpatent: A large-scale dataset for abstractive and coherent summarization. arXiv preprint arXiv:1906.03741 , 2019.

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.

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.13792 , 2021.

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.

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.

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.

Maciej Wołczyk, Bartosz Wójcik, Klaudia Bałazy, Igor T Podolak, Jacek Tabor, Marek ´ Smieja, and Tomasz Trzcinski. Zero time waste: Recycling predictions in early exit neural networks. Advances in Neural Information Processing Systems , 34:2516-2528, 2021.

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.

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.

Wangchunshu 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:18330-18341, 2020.

Wei Zhu. Leebert: Learned early exit for bert with cross-level optimization. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers) , pages 2968-2980, 2021.

NeurIPS Paper Checklist

The checklist is designed to 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 on how to answer these questions. For each question in the checklist:

You should answer [Yes] , [No] , or [NA] .

[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.

Please provide a short (1-2 sentence) justification right after your answer (even for NA).

The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.

The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation. While "[Yes] " is generally preferable to "[No] ", it is perfectly acceptable to answer "[No] " provided a proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] " or "[NA] " is not grounds for rejection. While the questions are phrased in a binary way, we acknowledge that the true answer is 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 a question, in the justification please point to the section(s) where related material for the question can 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 be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

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 are not discussed in the paper.

The authors are encouraged to create a 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 well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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, a worse outcome might be that reviewers discover limitations that aren't 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 not penalize honesty concerning limitations.

3. Theory assumptions and proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: The Lemma 3.1 and theorem 4.1 are proven in Appendix A.1 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 are encouraged to provide a 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.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental result reproducibility

Question: 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)?

Answer: [Yes]

Justification: We provide all hyperparameter details and other details in Table 11, 12 and Section 5.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a 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 a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, 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 a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in 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.

5. 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, as described in supplemental material?

Answer: [Yes]

Justification: The code to implement is given in an anonymized repository with the link in the supplementary.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/public/ guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https://nips.cc/public/ guides/CodeSubmissionPolicy ) for more details.

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 as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental setting/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 to understand the results?

Answer: [Yes]

Justification: All the hyperparameter details are given in Table 11, 12 and Section 5.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment statistical significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information 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 [33, 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, or statistical significance tests, at least for the experiments that support the main claims of 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 a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is 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 error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. 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 to reproduce the experiments?

Answer: [Yes]

Justification: Given in 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 runs as 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 make it into the paper).

9. Code of ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We abide by NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If 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 or 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?

Answer: [NA]

Justification: Our method is mostly about minimizing risk and improving efficiency, it does not add to any societal impact.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment 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 to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On 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 as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: We do not release any new model or data, and use only models and datasets that are open-source for research purposes.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 or restrictions to access the model or 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 a best faith effort.

12. Licenses for existing assets

Question: Are the creators or 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 a URL.

The name of the license (e.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 be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

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 are encouraged to reach out to the asset's creators.

13. New assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: We introduce no new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized 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 means that the paper does not involve crowdsourcing nor 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 be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in 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 or equivalent for research with human subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: We do not involve crowdsourcing.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the 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 applicable), such as the institution conducting the review.

16. Declaration of LLM usage

Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.

Answer: [NA]

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 as any important, original, or non-standard components.

Please refer to our LLM policy ( https://neurips.cc/Conferences/2025/LLM ) for what should or should not be described.

A Appendix

A.1 Proof of Lemma 3.1

Let p ( y | x ) denote the probability of the true label y given an input x . Consider an early-exit model that outputs a probability distribution p ( · ) over the set of classes for every threshold τ . Let p τ ( y | x ) be the true class probability given by the model, p τ (ˆ y | x ) be the probability of prediction given by the model and p τ ( y = ˆ y ) be the probability that predicted class is the true class.

For a fixed τ , we can write the probability that the prediction made by the model is correct as:

Hence, we prove that the threshold τ that maximizes p τ (ˆ y | x ) · p τ ( y = ˆ y | x, ˆ y ) value maximizes the chances of getting the prediction correct.

Proof of Theorem 4.1:

Before moving to prove the theorem, we will first bound the single-run regret of the UAT algorithm.

At each round t = 1 , 2 , . . . , T , the algorithm selects an arm τ t and receives a reward r ( τ t ) .

Regret Definition: We define the realized (actual) regret for a single run as:

where N i ( T ) is the number of times arm i was pulled up to time T .

The UAT Algorithm: At each time t , the UAT algorithm selects the arm

where: ˆ µ i ( t ) is the empirical mean reward of arm i up to time t . N i ( t ) is the number of times arm i has been pulled until time t .

High-Probability Bound Using Hoeffding's Inequality: By Hoeffding's inequality, for any n ≥ 1 and any ϵ > 0 :

Define the confidence interval:

In one particular round, true reward mean for i th arm lies in this interval with probability 1 -δ where δ = 1 /T 2

Using the union bound over all arms and times, with probability at least 1 -1 /T that the true reward mean lies in the interval over all the rounds.

Bounding the Number of Suboptimal Pulls: Let us bound N i ( T ) for any suboptimal arm i with ∆ i > 0 .

Suppose arm i is selected at time t . For this to happen, the UAT index for arm i must be at least as large as that for the optimal arm i ∗ :

If all confidence intervals are valid (which occurs with high probability), then:

Rearranging,

Thus, the number of times arm i is pulled is at most:

Total Regret Bound: Using R ( T ) = ∑ i :∆ i > 0 ∆ i N i ( T ) , we obtain:

Finally: With probability at least 1 -O (1 /T ) , the realized regret of UAT satisfies:

Let p τ (ˆ y | x ) denote the model's confidence and p τ ( y = ˆ y | x, ˆ y ) its calibrated correctness probability. The empirical risk could be written as:

The UAT algorithm that provides a threshold τ t = π ( x t ) using the policy π , for simplicity, we drop the index t . With probability atleast δ ′ = 1 -1 T , the single run regret of UAT algorithm can be bounded as:

where the reward r ( τ ) = C i τ · (1 -C i g ) -λ · i = p τ (ˆ y | x ) · p τ ( y = ˆ y | x, ˆ y ) -λi as we consider (1 -C i g ) approximates the p τ ( y = ˆ y | x, ˆ y ) with probability δ 1 . Choose δ = (1 -δ 1 )(1 -δ ′ ) then with probability δ :

Rearranging terms and bounding exit differences | i -i ∗ | ≤ L :

where p τ ( · ) ≜ p τ (ˆ y | x ) p τ ( y = ˆ y | x, ˆ y ) .

Using p τ (ˆ y | x ) ≤ 1 :

Dividing by T and substituting into the risk definition:

As we assume that R ∗ = 1 -1 T ∑ T t =1 p τ ∗ ( · ) ≤ ϵ 0

For large T , β ( T ) T → 0 (since β ( T ) = O (log T ) ). To ensure ˆ R ( π ) ≤ ϵ :

Therefore, if we select λ ≤ ϵ L , then with probability at least δ = (1 -δ )(1 -δ ′ ) , the empirical risk ˆ R ( π ) is bounded by ϵ . This completes the proof.

B More Ablation study

B.1 Importance of components of reward function

In Table 6, we show the performance of our model on 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-FlanT5-xl and Question Answering task using SQuAD on the BERT-large model. We observe that simply maximizing the prediction confidence term C i τ might lower the performance due to riskier exits while maximizing only the correctness probability (1 -C i g ) has better improvements, but causes underconfident exits as 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 setup.

B.2 The value of c

The choice of coverage threshold c depends on the prior knowledge of how complicated the dataset is based on the given distribution and model, as its task is to make g produce higher values for a fraction of samples that are easy for the model. We do not a priori know this information, hence we rely on a 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 a c dependent on i , denoted as c i . During training of the first epoch, we do not have any a prior knowledge of the task complexity hence we set c i to be 0 for all i , but after the first epoch we define c i = 1 N ∑ N j =1 ✶ { ˆ y ij = y } where ˆ y ij , i.e., the fraction of samples for which the exit is correct. This assesses the task complexity, making the g function produce higher scores based on complexity.

B.3 Risk-Efficiency trade-off

In Figure 2b, we show the risk-efficiency trade-off by plotting the percentage of performance drop as compared to the final layer vs Speedup for the SQuAD dataset on the BERT-large model by changing the parameter λ . 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 G ≤ ϵ . Our method keeps the risk lower than 1% till a speedup of 2 . 3 × after which the risk crosses the desired bound, while other risk-based methods are 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

In Figure 2a, we plot the average cumulative regret while choosing the thresholds during inference using different policies on the QNLI dataset. The results are average across 5 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 a random threshold was assigned to each sample and fixed

0

1000

2000

3000

4000

5000

Samples

0

100

200

300

400

500

600

Cumulative Regret

Final_layer

Randomly_selected

Fixed_threshold (Existing methods)

Our method

(a) Cumulative regret plots for different policies to choose the thresholds.

1.4

1.6

1.8

2.0

2.2

2.4

2.6

Speedup (

)

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

Risk % (

)

Ours

CALM

MIE-LAP

EERC

(b) Risk vs Speedup curve obtained by varying the trade-off parameter λ .

threshold where a static threshold was utilized similar to existing EE and risk based methods. Our method observed the smallest regret as compared to others. This indicates the importance of learning a policy for the threshold selection instead of fixing it to a constant value.

B.5 Changing ϵ values

In Table 7, we provide results on changing the ϵ that guides R G ≤ ϵ on the SamSum dataset with T5-large model, note that, changing ϵ value changes the value of λ , allowing riskier samples to make an early exit. We report the percentage of increase in R G as compared to the final layer. Observe from the table that even when more risky samples were allowed to exit early, our method's drop in performance is very low. This is due to an upgraded confidence metric that allows for an early exit only when it is confident and the confidence is also reliable. While imposing a hard criterion of ϵ = 0 . 01 , all the risk methods have violated that, while our being 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.

B.6 Some important examples

In 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 C i g ). Observe from the Table that as the sample gets more ambiguous, the chances are high that the model outputs a wrong prediction. In such cases, it becomes important to flag those samples that is done by the C i g value in our case. For instance, consider the sample 'the under-7 crowd', which means that the review is negative, but the model probably understands the numbers as ratings and develops high confidence over the wrong class.

B.7 Computational setup

In our experiments, we have used a setup of 5 NVIDIA A6000 GPUs. The highest run time was observed during training the BLIP-2 model for image 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 on the Ques-

tion Answering dataset due to a large number of epochs. For inference, it 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.

Table 9: Model performance trained using Conventional loss vs our loss.

Table 10: Examples of samples achieving fake confidence, the table shows an example, its true label (True lbl.), the average confidence of the model on the wrong class (Fake Conf.) and the (1 -C i g ) score that abstains the sample early.

B.8 Computational Complexity of UAT

Our method learns a policy to decide the exit threshold that is learned over time based on the data distribution. However, the computational complexity of UAT is negligible as it does not involve any big operations. In every, step, it requires maximizing over a small finite set (set of 10 elements) and maintain a list of reward functions. In Table 8, we perform experiment to show the additional complexity of our method. In this experiment exit was fixed to the final layer, i.e., 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 to UAT. So the time wall-clock time added to T5-large with UAT will be the computational complexity of UAT. Observe from Table 8, the computational complexity of UAT is negligible as the additional time is less than 15 seconds, further proving UAT's contribution during inference.

B.9 Loss observations

In Table 9, we provide results of the final layer when the conventional cross-entropy loss is used to train the EEDNN i.e., just use the first term L CE from Equation 3 and if you use the full loss function given in our method in Equation 3. The Table 9 suggests that adding the other components in the loss function does not impact the overall model performance. We check the final layer performance, as that would have been the most affected by the additional components. However, the effect was very insignificant, and we safely use the loss function given in Equation 3.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n* 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.\n\nQUESTIONS:\nThe 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.”\n\nLIMITATIONS:\n* 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.", "TITLE": "Official Review", "RECOMMENDATION": 2, "REVIEWER_CONFIDENCE": 5, "CLARITY": 2, "IMPACT": 1, "ORIGINALITY": 1, "DATE": "2025-06-19"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nYes, the authors adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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?\n\nQUESTIONS:\nSee weaknesses\n\nLIMITATIONS:\nIncluded but outside page limit", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-06"}