BEEM: Boosting Performance of Early Exit DNNs using Multi-Exit Classifiers as Experts
Authors: Divya Jyoti Bajpai
Manjesh Kumar Hanawal
Venue: ICLR 2025 Poster (2025)
Keywords: Early Exits; Expert-based exiting
Primary area: other topics in machine learning (i.e., none of the above)
TL;DR: An exiting criteria for early exits with theoretical analysis

Abstract: Early Exit (EE) techniques have emerged as a means to reduce inference latency in Deep Neural Networks (DNNs). The latency improvement and accuracy in these techniques crucially depend on the criteria used to make exit decisions. We propose a new decision criterion BEEM where exit classifiers are treated as experts and aggregate their confidence scores. The confidence scores are aggregated only if neighbouring experts are consistent in prediction as the samples pass through them, thus capturing their ensemble effect. A sample exits when the aggregated confidence value exceeds a threshold. The threshold is set using the error rates of the intermediate exits aiming to surpass the performance of conventional DNN inference. Experimental results on the COCO dataset for Image captioning and GLUE datasets for various language tasks demonstrate that our method enhances the performance of state-of-the-art EE methods, achieving improvements in speed-up by a factor $1.5\times$ to $2.1\times$. When compared to the final layer, its accuracy is comparable in harder Image Captioning and improves in the easier language tasks. The source code is available at https://github.com/Div290/BEEM1/tree/main.

Published as a conference paper at ICLR 2025
BEEM: BOOSTING PERFORMANCE OF EARLY EXIT DNNs USING MULTI-EXIT CLASSIFIERS AS EXPERTS
Divya Jyoti Bajpai & Manjesh Kumar Hanawal
Department of Industrial Engineering and Operations Research
Indian Institute of Technology Bombay Powai, Maharashtra, India idivyajyot.l eayyot..ajppa.... bajiai, mhanawalJ@litbb.acc
ABSTRACT
Early Exit (EE) techniques have emerged as a means to reduce inference lateney in Deep Neural Networks (DNNs). The latency improvement and accuracy in these techniques crucially depend on the criteria used to make exit decisions. We propose a new decision criterion BEEM where exit classifiers are treated as perts and aggregate their confidence scores. The confidence scores are aggre- gated only if neighbouring experts consistent in prediction the samples pass through them, thus capturing their ensemble effect. A sample exits when the ag- gregated confidence value exceeds threshold. The threshold set using the error rates of the intermediate exits aiming to surpass the performance of con- ventional DNN inference. Experimental results on the COCO dataset for Image captioning and GLUE datasets for various language tasks demonstrate that our method enhances the performance of sfte-off-thearr EE methods, achieving im- provements in speed-up by factor 1.5x to 2.1x. When compared to the final layer, its accuracy comparable in harder Image Captioning and improves the easier language tasks. The source code for this work is publicly available at https::/gittieii
INTRODUCTION
Transformer bassdd models (Devlin al., 2018: Radford et al. 2019 Cornia al., 2020 Luo al. 2021 Li et al. 2022 2023) have set new benchmarks in performance across diverse tasks and domains through their prowess in capturing semantic information and dependencies using attention mechanisms (Vaswani et al. 2017). However, the scale and complex structure of these models poses a challenge, particularly in terms inference speed, limiting their practicalities in resource- constrained scenarios. Also, these models susceptible overthinking issues (Zhou et al. 2020
Zhu 2021
which degrades their performance in terms of accuracy and inference speed.
To address these challenges, various techniques have been proposed, including direct network prun- ing (Zhu & Guptal 2017 Fan et al.] 2019) Michel et al., 2019). knowledge distillation (Sun et al,
2019| et al. 2019 Jiao al.) 2019), quantization methods (Zhang al. 2020 Bai al. 2020 and adaptive inference (Zhou et al.) 2020; Xin et al. [020bb Geng et al.
Kim et al. 2022
2021
Liu et al. 2020p. Early Exit (EE) methods (Terrapitaayano al. 2016 Zhou et al. 2020 Fei al. 2022) is one of the adaptive inference methods where intermediate classifiers (exits) are added af- ter every layer. The difficulty of the sample is determined using confidence in the prediction, and the sample is inferred early based on the confidence score exceeding pre-defined threshold. The confidence score becomes crucial part of the inference process and decides the sample hardness.
EE strategies either perform confidence-based exiting (Xin et al. 20206 Bajpai & Hanawal 2024a) or a patience-based exiting (Zhou et al. 2020) depending on the prediction consistencies treating each classifier equally. Recently EEIC (Sun et al. 1 decided on exiting based majority voting between the exits. This method also treats each classifier equally. These methods either consider the confidence of individual exits or utilize the predictions made by exits to define the confidence scores. For instance, in [ an input sample is currently processed till the third exit; for confidence- based exiting, it checks the confidence at the third exit, ignoring all the information gathered from previous

---

Published as a conference paper at ICLR 2025
exits. The patience-ber basee exiting requires predictions of all exits to be consistent if it wants make an exit. Also, prediction from the first classifier is treated as important as the third, which should not be the case deeper layers have more information. Similar is the case with majority voting, where all the classifiers are treated equally. They also do not utilize the confidence available from previous layers, thus discarding available information.
This necessitates the requirement of an EE strategy that can utilize the information available at the exits effectively mitigate the overthinking issue while speeding up the inference. Also, the existing methods do not offer any theoretical viewpoint or make strong assumptions, e.g., all the layers have the same error rate, which makes them less desirable.
We present an EE mechanism named BEEM: Boosting Performance of Early Exit DNNs using Multi-Exit Classifiers as Experts motivated by ensemble learning (Dong al. 2020). improve performance of EEDNNS. We treat each intermediate exit classifier as an expert that outputs con- fidence values on the labels for each input. This confidence score is then weighted based on the expert's accuracy in predictions or the associated prediction cost, i.e., higher weights to deeper ex- its and vice versa. By treating each exit as an expert, BEEM ensures that the model leverages the strengths of each exit and does not discard the scores of the previous layers if their predictions are in agreement. To determine if a sample can exit at the ith classifier, we accumulate weighted confidences of the immediate previous layers whose predictions are in agreement with the ith clas- sifier. In case of disagreement with the immediate previous layer, the confidence score resets to the score of the current exit, ignoring past aggregated values. This score is subsequently compared to predefined threshold for EE decisions.
The exit decision each layer is based on cumulative confidence score exceeding threshold value. The thresholds play a pivotal role as they offer means to model the trade-off between accu- racy and latency. In Section 3.4 we introduce a novel approach to determine the threshold values for different exits by converting the problem of choosing thresholds to simple linear program. We utilize the error rate of exits to set the threshold values, forcing the exit classifier perform better than the final classifier of DNNs. In Section 3.5 we also perform theoretical analysis to derive condition based on the error rate the intermediary layer under which BEEM performs better than the vanilla DNN inference.
We experiment with widely adopted Pre-Trained Language Models (PLMs) and encoder-decoder models to perform experiments on GLUE (Wang et al. (019a) and COCO dataset (Lin et al. 2014). We show that BEEM outperforms all the previous EE methods in terms of speed as well as accuracy. BEEM increases the inference speed by 1.5x 2. x with accuracy close to the final layer. For casier NLP tasks such as sentiment analysis, BEEM even outperforms the final layer in terms of accuracy.
In summary, contributions
follows:
We propose new criteria make EE decisions in DNNs. It combines the confidence score of the intermediary exit classifier to produce an ensemble effect to make PLMs more effi- cient, robust, and adaptive.
We provide method to set threshold values in BEEM by analyzing the error rates of exit classifiers (Section 3.4). This not only helps BEEM achieve better speed-up but also improves accuracy compared to inference at the last layer. We also derive a condition under which this performance is guaranteed.
Extensive evaluation showed speed-up improvement in both GLUE and COCO datasets. For the GLUE dataset, accuracy also improves with speedup due to reduction in the overthinking issues of the DNNs.
2 RELATED WORK
Early exit methods applied for various tasks such as image classification, image captioning and NLP tasks to reduce the computational resources and inference latency.
Early exits in Image tasks: 2016) uses classification entropy at each attached exit to decide whether to infer the sample at the side branch based on the entropy of prediction. Shallow-degy deep (Kaya et al. 2019) and MSDNet
For image classification tasks, BranchyNet (Teerapitayaanon et al.
2

---

Published as a conference paper at ICLR 2025
C
Borr ow
Layer L
Layer
Layer
Addilic oparstion
1e Rolooa ccarsiee
Pos
0.91
Pos
2
Layer
Layer
Layer
0.88
Pos
0.176
0.241
Layer
Layer 2
0.65
0.66
+
0.065
0.065
Prediction score
Layer
Patience Counter
Prediction score
Embedding
Embedding
Embedding
Superb product!
Superb product!
Superb product!
DeeBERT
(b) PABEE
(c) Ours
Figure 1: Comparison between (a) DecBERT, which uses the confidence available at each exit the
metric or deciding early inference (set to 0.9). (b) PABEE, which uses the consistency in prediction the confidence metric (set to 2) and (c) BEEM that uses the weighted confidence S, (weights
= =0011, 0.2, 1.2)) and threshold a = 0.2. In BEEM, by appropriately considering information from previous classifiers, correct prediction made early which was not the case with others.
(Huang et al. 2017) improve upon BranchyNet by effectively choosing the thresholds based on the confidence distribution. Similar architectures (Pacheco et al. 2021 Dai et al.] (020) split the NN to be deployed on edge and cloud. SEE (Wang al. (019cc work in service outage scenarios.
FlexDNN (Fang et al.) 2020) and Edgent (Li et al. 2019) focus mainly on the most appropriate Neural Network (NN) depth. Other works such as Dynexit (Wang et al. 2019b) focus on deploying the multi-exit NN in hardware. It trains and deploys the NN on Field Programmable Gate Array (FPGA) hardware while Paul al. (Kim & Park 2020) explains that implementing a multi-exi NN on FPGA board reduces inference time and energy consumption. ZTW (Sun et al. 2021 uses a combination of probability distribution to decide to exit and the combination is learned during training reducing its generalization capabilities. JEI-DNN (Regol et al.) 2023) on the other hand uses a gating mechanism for inference where which gate will be opened for a sample is learned during training. In parallel vein, the MuE, DeeCap and CapEEN (Tang et al. 2023 Fei et al.) 2022 Bajpaí & Hanawal 2024b) model employ a distinctive approach to apply early exits to image captioning. DeeCap only applies to the decoder, while MuE applies to the encoder and the decoder. CapEEN makes the exiting process more robust to noisy images.
Early exit in PLMs: Multiple approaches have been proposed to effectively apply early exits to PLMS and solve multiple NL. tasks (Liu et al. 2021a Xin et al., 2020aa Zhou et al.) 2020; Banino al. 2021 Balagansky & Gavrilov 2022] Sun al. 2022 Ji et al. (2023). DeeBERT (Xin al.)
2020a), ElasticBERT (Liu et al. 2021a) and BERxiT (Xin et al. (021). BERxiT proposes an effi-
cient fine-tuning strategy for the BERT model with attached exits. DeeBERT is obtained by training
the exit points attached before the last module to the BERT backbone separately. In contrast, Elas- ticBERT obtained by training all the exit points attached to the BERT backbone jointly. PABEE (hhou et al. 2020) is another multi-exit model that makes the exit decision based on the stability of the predictions after different exits. LeeBERT (Zhu] 2021) proposed a self-distillation framework
that has similar exiting criteria PABEE. ETFEE Ji et al. 2023) adds an adapter top of the transformer layers and an (Entangled Frame) ETF classifier to make intermediate exits learn better.
CeeBERT (Bajpai & Hanawal 2024a) and DAdEE (Bajpai & Hanawal 2024c) propose multiple methods to adapt early exits to various domains in an unsupervised setup. |Bajpai et al. (2024b a);
Bajpai & Hanawal (2024d) utilize early exits for distributed inference setup.
Our approach differs from past works as 1) Unlike previous studies, BEEM utilizes the ensemble learning principles by treating each exit an expert. 2) Our work proposes early exiting method
3

---

Published as a conference paper at ICLR 2025
that utilizes each expert based its strengths. 3) We also provide method to set the thresholds using the error rates of the exit classifiers to perform better than the final classifiers.
3 PROBLEM SETUP
We start with a pre-trained DNN and attach exit classifiers after each layer. We provide details on training the PLMs for language tasks and the encoder-decoder backbone for image captioning,
3.1 PLMS
Training: BEEM requires the training of cxit classifiers that provide predictions based their respective layer outputs. Let D denote the dataset distribution with the label class C employed for backbone training. Given input sample (2r, v) DD the loss for ith exit is calculated as:
(,) (ce(f((z 0). + LLsps.PL)
(1)
where fi 0) represents the output of the exit attached at the i-th layer, 0 denotes the set of all learnable parameters, LCE denotes the cross-entropy loss and is the KL-divergence loss used to additionally train the exits with soft labels from the final layer. KL divergence (Kullback-Leibler divergence) used in knowledge distillation because it measures how well one probability distri- bution (the student model's 's predictions) approximates another (the teacher model's predictions). In the context of knowledge distillation, KI. divergence serves as a key component to transfer "soft knowledge" from the teacher the student. Let p: P.(c( denote the probability distribution over the set of output classes at the ith layer, where (.(e(z) denotes the estimated probability that belongs class c. We simultaneously optimize parameters for all exit classifiers, following the
methodology proposed by Kaya al. (2019). The loss function is defined as L Eig iC.
sidering the weighted average to account for the relative inference cost of each exit classifier where I. denotes the number of layers in the model. Note that the importance of KI--divergence loss is well-explained in Zhu (2021). Following this training, the model is ready for inference.
Inference: We illustrate the inference process of BEEM Fig. I As the input instance goes through layers 1, 2, L sequentially, the classifier attached to that layer predicts class label dis- tribution. For ith exit classifier, let C, denote the confidence in the estimate at the ith exit. We define C as the maximum of the estimated probability class, i.e., C; := maXeec (c(c))). We denote D: = arg max,ec (e(e()), the prediction of ith exit. Based on the confidence scores we define weighted confidence score, denoted S; as:
{
S + W.C, W.C.
J:-1 ys if 9:-1 + d::
S. 5
The inference process halts when S; 2 a, where a represents a predefined threshold, and exits with label di. Otherwise, the sample is processed in the next layer and the process completes. this condition is never met at any exit classifiers, label is assigned by the classifier at the final layer. This allows a sample to exit the backbone early if the condition satisfied, avoiding traversal through all layers.
3.2 ANCODER-DECODER MODELS
For the image captioning task where the objective is to generate caption for input image, we start with a pre-trained encoder and decoder model. We use the Swin Trans- former (Liu et al.) (021c) as encoder and GPT-2 (Radford et al.) 2019) decoder. We attach exits to the decoder of the backbone. The backbone is trained using cross entropy and KL-divergence loss where loss for ith exit could be written as:
Encoder-decoder model:
T
- = I
v1a-1). + KL(pi.
t=1
where 0 is the collection of all the parameters, x is the input image, T is the caption length, Virr the gruuh caption. P is the probability vector the vocabulary V for ith exit. Its eth
4

---

Published as a conference paper at ICLR 2025
0) where P, is the probability distribution output over V by the ith exit. Note that L here is the number of layers in the decoder. Similarly, we define pf as the probability vector for the final decoder layer. The overall loss across all the exits is the same as for the PI.M training.
component could be written pi()) = (:
Caption inference: We predict the caption in autoregressive manner. This entails making token-by-token prediction for a given image. In this case, the confidence could be formulated as C. ey y P.(v ( 0) where P; is same as defined above and # is the predicted caption till (1 1)th word. A token will be predicted at which exit is decided by equation 2 For an input image z, we start the caption with begin of the sentence token and then the inference process stops after the end of the sentence token is predicted.
early exit methods like DeeBERT and ElasticBERT, relying on soft- max scores, tend to be overly confident toward single class and classifier. Such methods also face the consequenees of ignoring information obtained from previous classifiers as they progress to the next one. This limitation is addressed by patience-based approaches like PABEE, which de- cide to exit when predictions show consistency across multiple classifiers. However, patience-based methods treat each classifier equally and underutilize valuable information available in terms of prediction scores, which affects the adaptability of the model.
In Figure ! confidence-beid
Contrarily, BEEM captures the confidence available at each exit and assigns weights to each clas- sifier based on its accuracy or cost. It takes into account the consistency in predictions by reducing the score to the current classifier's weighted confidence when predictions are inconsistent. This unique approach in BEEM incorporates both patience and confidence to make predictions. Note that the confidence score S, given in equation [ can predict hard samples early as if the predictions of initial classifiers are consistent but with low confidence, the summed-up S; score makes them exit early as shown Figure ! Also, it effectively mitigates errors arising from a single classifier while considering the confidence in predictions and weighing them based on their performance.
3.3 ASSIGNING WEIGHTS
In this section, we provide methods set the weights for the exits.
1) Cost vector: First, we consider weights as the cost of getting inference from the exit classifier where the cost could be in the form of w; Ai where is the processing cost of the one exit and since the layers are identical, the cost is a multiple of À for deeper layers.
2) Accuracy: We also consider the weights the accuracy of each classifier. The accuracy could be calculated on a validation dataset. This will provide weights to exits depending on how much accurate particular expert (exit) is.
Note that the major difference between the existing methods is the cost-based weights have task to reduce the overall cost while sacrificing some accuracy while the accuracy-bass based methods will focus more on accuracy. Note that using accuracy- baased weights can also improve the efficiency that comes because of overthinking issues. As in the accuracy-basedd we know the true capability of each exit.
3.4 CHOICE OF THRESHOLDS
We can choose the threshold values in two ways, one way to choose the best-performing threshold on the validation set, and the other is based on forcing error rates to be smaller than the error rate of the final classifier.
1) Classical method: We choose the search space for threshold a E S = {0.3, 0.6, 0.9, 1.2, 1.5). The values of w, [0, C, E [0, 1] imply that S, < ie. the score at any exit layer cannot be greater than the number of layers I, as S: is multiple of two values between 0 and 1, it becomes very small and is added almost T. times. We choose the best-performing threshold on the validation set in terms of accuracy.
2) Using error rates:
Let us consider that chise represents the number of samples that exit at ith classifier with misclassification, Cttop represents the number of samples that exit at the tth classifier. Note the Cstop Σラ= 1 (C,203 can also be considered the coverage of the ith classifier and
5

---

Published as a conference paper at ICLR 2025
- (0,20 where is the total number of samples in the dataset. p is the error rate of the final classifier, then we observe that our algorithm will perform better than final layer if of /mise/Cstop < p for every exit classifier t. The above condition tells that the fraction of the samples that have exited and misclassified (i.e., error rate) at tth classifier should be less than the error rate of the final classifier. If the above condition is satisfied for all the exits then we are guaranteed that BEEM outperform the final classifier of the PLM. The objective is to maximize speedup while satisfying the above condition.
Emise
Catop
Note that the error rate depends on the threshold a, higher value of the threshold will lower the error rate as then samples with higher confidence will exit reducing the chance of misclassifica- tion. Observe that, we can find the threshold value a for the rth classifier such that the condition emisc /Cttop < is satisfied on the validation set. We define 9a, as the eπrrr rate associated with the threshold 04 for the 1th classifier. We can set the threshold by solving the optimization problem.
minimize a:ES
subject to Oo. P.
where the set S {0.5, 1, 1,..., 5, L3 is the search space for the thresholds. Note that have added L in the search space so that the problem always remains feasible. By solving Eq. 3 our method finds the optimal threshold that has maximum speedup while performing better than the final layer. Also, observe that the above problem has small computational complexity as the minimization is over a very small finite set.
3.5 THEORETICAL ANALYSIS
Theorem 3.1. Consider an early exit PLM with L layers. Let p denote the error rate of the final classifier and the error probability of ith exit classifiers be qi such that q; < (1=((1/p-110;- holds
as
for all exit layers' 1, ...... L- where o, and b, are constants for given exit i. Then, the error probability of BEEM is better than i... it performs better than the final layer.
The proof of the theorem is given in the Appendix A. Note that the above theorem proves the gen- eral condition for better performance of BEEM and does not depend on the threshold values a. a; denotes the ratio of the probability of exiting with one change in prediction to the probability of exit-
ing with zero changes prediction till ith classifier and b,
where q max 19ax(41.....000 .4
while qeain mi.............................. ) Observe that as we move deeper into the backbone the bound be- comes tighter which makes sense as deeper layers are more likely to be accurate. Also, the bound is inversely proportional to the error rate of the final layer, if the error rate of the final layer is smaller, then the bound gets tighter. Previous method Zhou al. (2020) had a very strong assumption while providing a similar condition for their method, it assumed that all the classifiers have the same error rate which is not true. If we impose the same condition the bound simplifies to q. a.+((1/--1))
which is a more simplified and stronger bound than PABEE (Zhou et 2020).
EXPERIMENTS
In this section, we provide the details of the experiments performed in this work.
Datasets: We evaluate our approach using the GLUE benchmark datasets (Wang et al. 2019a) Our assessments encompass diverse tasks, such as sentiment classification using the Stanford Sen- timent Treebank (SST-2), Natural Language Inference (NLI) tasks with Multi-Genre Natural Lan- guage Inference (MNLI), Question Natural Language Inference (QNLI), and Recognizing Textual Entailment (RTE). For Paraphrase Similarity Matching, we include Microsoft Research Paraphrase Matching (MRPC) and Quora Question Pairs (QQP), while Linguistic Acceptability is measured us- ing The Corpus of Linguistic Acceptability (CoLA). In instances where datasets comprise multiple units. report the arithmetic mean. We exclude the WNLI task, following previous works (Devlin
et al. 2018 Zhu 2021
Zhou et al. 2020). For captioning, we use the COCO (Lin et al. 2014)
dataset.
Baselines: We compare against the vanilla DNN exiting and other techniques that speed up DNN inference. The baselines follows:

---

Published as a conference paper at ICLR 2025
Model/Data
SST-2
MNLI
RTE
QNLI
QQP
Acc
Speed
Acc
Speed
Acc
Speed
Acc
Speed
Ace
Speed
Dev set
ALBERT
92.4
1.00x
84.5
1.00x
77.9
L.0xx
91.3
1.00x
1.00x
ALBERT-91
1.33x
1.33x
1 1.33x
-2.7
1.33x
-1.5
1.33x
DeeBERT
-2.3
1.72x
-2.9
1.65x
-3.1
1.78x
-1.9
1.57x
-2.5
1.81x
ElasticBERT
-2.1
1.75x
-2.3
1.71x
-2.7
1.81x
-1.7
1.66x
-2.1
1.78x
FastBERT
-1.1
1.85x
1.61x
1.79x
-0.8
1.71x
1.88x
PABEE
-0.1
1.87x
-0.5
1.85x
-0.7
1.64x
-0.6
1.81x
-0.2
1.68x
ZTW
-0.2
1.64x
-0.3
1.67x
+0.2
1.63x
-0.3
1.75x
-0.1
1.71x
PCEEBERT
+0.1
1.24x
0.0
1.31x
+0.3
1.27x
-0.1
1.21x
+0.1
1.37x
LeeBERT
0.0
1.78x
-0.2
1.74x
-0.1
1.59x
+0.1
1.79x
-0.2
1.97x
PALBERT
1.54x
1.61x
+0.3
1.45x
1.59x
1.63x
JEI-DNN
-0.1
1.77x
40.1
1.67x
0.0
1.35x
-0.1
1.43x
+0.2
1.57x
BEEM-C
0.0
1.71x
+0.1
2.03x
+0.4
1.79x
0.0
1.90x
0.0
1.93x
BEEM-A
+0.4
1.98x
+0.3
1.96x
+0.7
1.89x
+0.2
1.92x
+0.5
2.09x
Test set
ALBERT
92.3
1.00x
84,2
1.00x
72.1
1.00x
90.9
1.00x
80.1
1.00x
ZTW
-0.4
1.61x
-0.5
1.52x
+0.1
1.64x
-0.1
1.59x
-0.5
1.81x
LeeBERT
-0.5
1.79x
-0.9
1.88x
1.68x
1.72x
-0.3
1.86x
PALBERT
-0.3
1.49x
-1.1
1.72x
+0.2
1.27x
-0.4
1.51x
-0.3
1,50x
JEI-DNN
-0.1
1.35x
-0.7
1.59x
0.0
1.36x
-0.2
1.39x
0.0
1.47x
BEEM-C
-0.2
1.98x
-0.4
1.95x
+0.1
1.74x
+0.1
1,81x
+0.1
1.97x
BEEM-A
+0.4
1.91x
2.06x
+0.6
1.77x
+0.5
1.88x
1.95x
Table 1: Main results: This table compares BEEM against all the state-of-the-arr early exiting base- lines. We report the accuracy (Acc in %) and Speed-up (Speed).
1) Final layer: The final layer of the DNN model, referred to as the "final layer" in Table
2) Reducing layers: We use only the first 9 layers of the DNN model with single output layer, denoted as DNN-9L. This serves performance lower bound since it employs EE techniques.
3) Early-exit models: DeeBERT (Xin al. 2020b) and ElasticBERT (Liu et al.) 2021b): Use fixed confidence thresholds for early exits. FastBERT (Liu et al. 2020): Uses a self-distillation framework to train intermediate exits. PABEE (Zhou et al.) 2020) and LecBERT (Zhu) 202I) uses prediction stability, with LeeBERT incorporating knowledge distillation. ZTW (Sun et al. 2021): :
combines the output probability outputs across all the layers and trains the weights provided as additional parameters for training. PCEEBERT (Zhang al. 2022): Combines confidence and patience metrics, similar to PABEE. MuE (Tang et al. 2023 Uses hidden representation similarity for early exits, applied the BERT-base model. PALBERT model (Balagansky & Gavrilov 2022):
State-of- satte-of-the-are fe-the-art art methods that face adaptation challenges due to training dataset bias. PALBERT uses Lambda layers (Banino et al., 202T). JEI-DNN (Regol et al., (023) performs exiting using a gating mechanism where it learns a probability distribution over all exits and decidesto exitg based on that.
DeeCAP (Fei al., 2022) is specifically for image captioning that uses an imitation network to mimic the behaviour of the decoder model.
We utilized the codebases of existing methods get the results, all the results were obtained using the hyperparameters given in their available codes. Note that for the encoder-decoder model, we extend the ideas of DeeBERT, FastBERT, PABEE, and LeeBERT to the decoder of the backbone.
4.1 EXPERIMENTAL. SETUP
Our experiments conducted single NVIDIA RTX 2070 GPU, The runtimes given below.
Training. For the training phase, we augment the pre-trained BERT/ALBERT model with linear output layer after each intermediate layer serve exit point. We conduct grid search over batch sizes of {8, 16, 32} and learning rates (2e-5, 3e-5, 4e-5, 5e-5) using the Adam (Kingma & Ba 2014) optimizer.
Incorporating an early-stoppiin mechanism, we select the best model based on the validation set. These parameters fixed 16 batch size and 3e-5 learning rate for the encoder-decoder backbone.

---

Published as a conference paper at ICLR 2025
Model/Data
RTE
CoLA
Acc
Speed
Acc
Speed
BERT
69.3
1.00x
57.8
1.00x
BERT-9L
-1.8
1.33x
-2.1
1.33x
DeeBERT
-2.5
1.47x
1.21x
BlasticBERT
-2.2
1.52x
-1.2
1.18x
FastBERT
-0.8
1.44x
-0.2
1.24x
PABEE
-1.1
L.62x
-0.1
1.16x
ZTW
1.52x
1.48x
LeeBERT
-0.6
1.60x
-0.1
1.28x
PALBERT
-0.5
1.32x
-0.6
1.19x
JEL-DNN
-0.2
1.30x
-0.3
1.18x
BEEM-C
+0.0
1.30x
BEEM-A
+0.2
1.70x
+0.3
1.49x
Table 2: Results on the BERT backbone on the GLUE datasets.
We report accuracy (in %).
Models/Metric BLEU-1
BLEU-
METEOR
CIDEr
SPICE
ROUGE-L
Speedup
Final-Exit 82.5
42.3
32.2
147.1
26.7
61.3
1.00x
Decoder-9L
76.5
37.1
29.3
134.8
23.2
57.9
1.33x
DeeBERT
70.1
32.3
26.9
110.2
20.9
50.7
1.35x
ElasticBERT
71.4
32.8
27.6
114.6
21.4
51.6
1.37x
PABEE
72.7
33,9
27.9
115.6
21.9
52.3
1.30x
FastBERT
75.0
35.6
28.2
119.5
22.1
53.7
1.42x
LeeBERT
77.3
38.7
29.4
129.2
23.0
55.9
1.39x
DeeCap
77.5
39.2
29.9
132.8
23.2
56.9
1.60x
MuF
79.3
40.5
30.9
139.4
24.9
59.7
1.64x
BEEM-C
81.8
41.5
31.7
145.1
25.9
60.1
1.71x
BEEM-A
82.4
42.1
32.0
146.5
26.3
60.9
1.67x
Table
Results showing that BEEM outperforms the other baselines on test split of COCO dataset.
The training time has an average GPU runtime of around 10 hours on a dataset, with the COCO dataset exhibiting the highest runtime (~ 26 hours).
Inference: Following the previous methodology on input-adaptive inference (Teerapitaayanon et al,, 2016 Kaya al.) 2019), the inference is performed on per-instance basis, setting the batch size to 1. This aligns with scenarios where low latency is critical, such as processing individual requests from different users (Schwartz et 2020). The reported values represent the median results from 5 runs with different seeds as small datasets such as CoLA and RTE have high variance in perfor- mance. For performing inference, the average runtime was < 20 minutes for NLP datasets. For COCO dataset the runtime was 5 hours on the Karpathy test split.
Metric. We report the speed-up ratio metric for measuring time reduction remain consistent with the previous methods. Speed-up could be defined Σ I.xn; where n, the number of Σ sxn;
samples exiting from the ith layer. For the image captioning task n4 is the number of words exiting from the ith layer. This metric could be interpreted as the increase in speed of the model as compared to the naive (AL) model. This metric can be converted to expected time reduction rate.
In Table 2 we present results wherein classifiers are assigned weights based the cost of each classifier denoted as BEEM-C and where the weights are set using the accuracy on the validation set, we denote by BEEM-A.
RESULTS
this section, we highlight and discuss the key findings our work. Tables I and present the results when ALBERT and BERT serve the backbone models, respectively. BEEM consistently outperforms all previous baselines by a significant margin. A major observation is a notable en- hancement in BEEM compared to the performance of (AL)BERT models, except for minor setback on the MNLI dataset. The improvement in accuracy by BEEM may be attributed to the thresholds being chosen after solving the constraint optimization exclusively the validation
8

---

Published as a conference paper at ICLR 2025
dataset. Table results the COCO dataset and observes significant improvements by using BEEM.
The substantial accuracy drop observed in DeeBERT and ElasticBERT results from direct compar- ison with entropy, neglecting the information utilized by preceding classifiers. Conversely, PABEE, LeeBERT, FastBERT, and ETFEE employ patience-basee based early exit criteria, posing stringent cri- terion for exiting. ZTW is one of the works that weights the classifier and utilizes the ensemble techniques but suffers from poo generalization as weights learned restricting better generaliza- tion as well as adding complexity. JEI-DNN uses the gating mechanism to decide whether to exit and does utilize the information of multiple available classifiers. Similar is the case with DecCAP and MuE, and for image captioning. they do not perform any knowledge distillation, further re- ducing the performance. These methods do not account for the confidence available each exit, assigning them equal weight irrespective of their varying confidence level prediction. This lack of consideration impacts the adaptiveness of early exit models.
BEEM-C and BEEM-A are the two variants of our proposed method. In the results, can observe that BEEM-A consistently outperforms BEEM-C for all the datasets in terms of accuracy and most of the datasets in terms of speed-up. This gain for BEEM-A could be attributed to the assumption in BEEM-C that the cost of exits (experts) was set by assuming that it is directly proportional to the accuracy but this is not true due to the overthinking issue. Still BEEM-C performs better than previous baselines on the datasets in which the overthinking issue minimal. The main advantage still tune the cost À based
of BEEM-C is that since the thresholds
fixed in our setup, the speed-up (see section A.3 needed which is unavailable in BEEM-A.
ABLATION STUDY AND ANALYSIS
In this section, we provide the results of our method on (AL)BERT large models. In the Appendix, perform analysis of the behaviour of parameters and À (see Appendix A.2 A.3 This analysis shows that our methods not only have better performance but also better models for the accuracy-efficiency trade-off i.e., the drop in accuracy of BEEM was lower when speedup increases as compared to others.
6.1
PLM SIZE
In Table we analyze BEEM's per- formance on ALBERT-Large models, each with 24 layers.
Data
RTE
CoLA
QQP
Acc
Spd
Acc
Spd
Acc
Spd
AB-L
80.5
1.00x
60.9
1.00x
1.00x
Our results show significant acceler- ation in processing speed, especially for larger models, due their inher- This effi-
Our-A
+1.8
2.04x
+1.3
2.85x
+0.1
3.33x
B-I.
70.9
1.00x
64.3
1.00x
91.2
1.00x
Our-A
+0.5
1.81x
+0.9
1.71x
+0.3
2.51x
ent overparameterization. ciency gain underscores BEEM's po- tential for optimizing large architec- tures.
Table 4: This table provides results on the large variants of (AL)BERT models compared with BEEM-A. AB-L ALBERT-Large and B-L is BERT-Large.
Furthermore, BEEM notably im- proves accuracy by mitigating overthinking, where models focus on irrelevant features. This issue is more pronounced in larger models, making BEEM particularly effective. Our findings demonstrate that BEEM enhances performance and speedup for large-scale transformer-base PLMs, becoming increasingly effective with larger model sizes.
6.2 CHOICE OF THRESHOLDS
In table 5 we compare results when the thresholds chosen based the equation 3 and when the thresholds are set using the vanilla method i.e. best-performing on the validation set. We can observe that, there is significant increase in the performance in both ALBERT models attributed to the choice thresholds made by equation Observe that setting thresholds by solving the equation improve both speedup well accuracy, this is as the equation finds the smallest threshold

---

Published as a conference paper at ICLR 2025
Method
SST-2
QQP
Our method
Acc
Spd
Acc
Spd
w/o fix
92.4
1.81x
90.4
2.05x
Base
fix
92.6
1.98x
91.1
2.09x
w/o 93.1
2.19x
91.1
2.95x
Large
W fix
93.4
2.31x
91.3
3.33x
Table This table compares the setting of thresholds based on the best-performing threshold on the validation dataset (w/o fix) and fixing the threshold after solving equation 3 (w fix) ALBERT- Base/Large models.
that improve the accuracy from the final layer. The thresholds are set such that each of them perform equivalent to the final layer.
7 CONCLUSION
In conclusion, our study introduces novel framework BEEM designed enhance the efficiency, robustness, and adaptability of early exiting strategies in DNNs. By leveraging multiple exit classi- fiers, where each exit treated as an 'cxpert', and their outputs are combined to create an ensemble effect. Our approach considers both prediction confidence and patience, leading improved perfor- mance and reduced latency, particularly advantageous in scenarios with strict latency requirements. Additionally, propose method for threshold selection, further enhancing the effectiveness of our approach. We also perform theoretical analysis to provide deep insights into our method. We experimentally validate that the speed-up observed was 1.5x 2.1x for various NLP and image captioning tasks.
8 LIMITATIONS
While the performance of our method is better than the final layer for NLP tasks, takes a hit for difficult tasks such image captioning. It happens the thresholds being dataset that might not generalize well on the test dataset i.c., the solution to the optimization problem 3 might not work for the test dataset. However, loss, BEEM effectively does that and performs better than all the existing early exit models and is comparable the final layer DNNs with a large improvement inference speed.
the validation
our objective is minimize the performance
ACKNOWLEDGEMEENTS
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/2022/088807) and MATRICS grant (MTR/2021/000645)) and DST-Inria Tar- geted Programme. We also thank funding support from Amazon IIT-Bombay AI-ML Initiative (AIAIMLI).
REFERENCES
Haoli Bai, Wei Zhang, Lu Hou, Lifeng Shang, Jing Jin, Xin Jiang, Qun Liu, Michael Lyu, and Irwin King. Binarybert: Pushing the limit of bert quantization. arXiv preprint arXiv:2012.15701 2020.
Divya J. Bajpai, Aastha Jaiswal, and Manjesh K. Hanawal. I-splitee: Image classification in split computing dnns with early exits. ICC 2024 IEEE International Conference on Communica- lions, pp.
2024a. doi: 10.1 109901111 11T./C/511166.02222 166.2024. 166.16024.00622995...
Divya Jyoti Bajpai and Manjesh Hanawal. CeeBERT: Cross-domain inference in early exit BERT. In Lum Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Findings of the Association for Com- putational Linguistics: ACL 2024, pp. 1736- 1748, Bangkok, Thailand, August 2024a. Associ- www.2//2..innnns-eeee
ation for Computational Linguistics. doi: 10. 101118655111/222 indd ngs .101 1017
URL hetps:
10

---

Published as a conference paper at ICLR 2025
Divya Jyoti Bajpai and Manjesh Kumar Hanawal. CapBEN: Image captioning with early its and knowledge distillation. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen (eds.), Findings of the Association for Computational Linguistics: EMNLP 2024, pp. 6458- 6472, Miami, Florida, USA, November 2024b. Association for Computational Linguistics. doi: 10. 06653///2024.fnninnnseeee
URL
Findings wwccings-enn.p..333 www...ppp.7..77..
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 (eds.), Findings of the Association for Computational Linguistics: EMNLP 2024, pp. 6389-6400, Miami, Florida, USA, November 2024c. Association for Computational Linguistics. doi: 10.28443/ finding--emnl.3311. URL httpp:://aclarthooooy.oooo http://22222..nnnnnneeeeeee findings-emnip
371/
Divya Jyoti Bajpai and Manjesh Kumar Hanawal. Distributed inference on mobile edge and cloud: An early exit based clustering approach, arXiv preprint arXiv:24 2024d.
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 A/-MI. Systems, AIML Systems '23, New York, NY, USA, 2024b. Association for
Computing Machinery. ISBN 9798400716492. doi: 10.1 10.114/766698563399977 URL. https:
//doi. 00......1111333311 011198885666668778 3639873
Nikita Balagansky and Daniil Gavrilov. Palbert: Teaching albert to ponder. Advances in Neural Information Processing Systems, 15:14002-14 2022.
Andrea Banino, Jan Balaguer, and Charles Blundell. Pondernet: Learning to ponder. arXiv preprint arXiv:2107.05407, 2021.
Marcella Cornia, Matteo Stefanini, Lorenzo Baraldi, and Rita Cucchiara. Meshed-memory trans- former for image captioning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10578-10587, 2020.
Xin Dai, Xiangnan Kong, and Tian Guo. Epnet: Learning to exit with flexible multi-branch net- work. In Proceedings of the 29th ACM International Conference on Information & Knowledge Management, pp. 235-244, 2020.
Jacob Devlin, Ming- bidirectional transformers for language understanding. arXiv preprint arXiv:1810.0480 2018.
Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep
Xibin Dong, Zhiwen Yu, Wenming Cao, Yifan Shi, and Qianli Ma. A survey on ensemble learning. Frontiers of Computer Science, 14:241-258, 2020.
Angela Fan, Edouard Grave, and Armand Joulin. Reducing transformer depth demand with structured dropout. arXiv preprint 1xXiv:190.111530 2019.
Biyi Fang, Xiao Zeng, Faen Zhang, Hui Xu, and Mi Zhang. Flexdnn: Input-adaptive on-device deep learning for efficient mobile vision. In 2020 IEEE/ACM Symposiun Edge Computing (SEC), pp. 84-95. IEEE, 2020.
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, pp. 12216-12226, 12226, 2022.
Shijie Geng, Peng Gao, Zuohui Fu, and Yongfeng Zhang. Romebert: Robust training of multi-exit bert. arXiv preprint arXiv:210100975,, 2021.
Gao Huang. Danlu Chen, Tianhong Li, Felix Wu, Laurens Van Der Maaten, and Kilian Q Wein- berger. Multi-scale dense networks for resource efficient image classification. arXiv preprint 2017.
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 Com- putational Linguistics: ACI. 2023, pp.
2023.

---

Published as a conference paper at ICLR 2025
Xiaoqi Jiao, Yichun Yin, Lifeng Shang. Xin Jiang, Xiao Chen, Linlin Li, Fang Wang. and Qun Liu. Tinybert: Distilling bert for natural language understanding. arXiv preprint arXiv: 10351, 2019.
Yigitcan Kaya, Sanghyun Hong, and Tudor Dumitras. Shallow deep networks: Understanding and mitigating network overthinking. In International conference machine learning, pp. 3301- 3310. PMLR, 2019.
Geonho Kim and Jongsun Park. Low cost early exit decision unit design for enn accelerator. In 2020 International SoC Design Conference (ISOCC), pp. 127-128. IEEE, 2020.
Schoon Kim, Amir Gholami, Zhewei Yao, Michael W Mahoney, and Kurt Keutzer. I-bert: Integer- only bert quantization. In International conference machine learning, pp. 5506-5518, PMLR, 2021.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1412. 69800, 2014.
En Li, Liekang Zeng, Zhi Zhou, and Xu Chen. Edge ai: On-demand accelerating deep neural network inference via edge computing. IEEE Transactions Wireless Communications, 19(1): 447-457, 2019.
Junnan Li, Dongxu Li, Caiming Xiong, and Steven Hoi. Blip: Bootstrapping Language-image pre- training for unified languue understanding and generation. In International conference on machine learning, pp. 12888-12900. PMLR, 2022.
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. pp. 19730-19743 PMLR, 2023.
Tsung-V Yi Lin, Michael Maire, Serge Belongie, James Hays, Pictro 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, pp. 740-755. Springer, 2014.
Weijie Liu, Peng Zhou, Zhe Zhao, Zhiruo Wang, Haotang Deng, and Qi Ju. Fastbert: a self self-distilling bert with adaptive inference time. arXiv preprint arXiv:2004.02178 2020.
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. 2021a. URL (nttts::: /tt.xiiivv...///////.a///aaaalv. org/ abs/ 22122111100031
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 strong baseline. arXiv preprint arXiv:21 2021b.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 10012-10022, 2021c
Yunpeng Luo, Jiayi Ji, Xiaoshuai Sun, Liujuan Cao, Yongjian Wu, Feiyue Huang, Chia-Wen Lin, and Rongrong Ji. Dual-level collaborative transformer for image captioning. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pp. 2286-2293, 2021.
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than onc? Advances in neural information processing systems, 32, 2019.
Roberto G Pacheco, Rodrigo S Couto, and Osvaldo Simeone. Calibration-aided edge inference offloading via adaptive model partitioning of deep neural networks. In ICC 2021-IEEE Interna- tional Conference on Communications, pp. 1-6. IEEE, 2021.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9 2019.
12

---

Published as a conference paper at ICLR 2025
Florence Regol, Joud Chataoui, and Mark Coates. lointly-leameed exit and inference for dynamic neural network: Jei-dnn. arXiv preprint arXiv:2310.0113,, 2023.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv: 2019.
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.
Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. Patient knowledge distillation for bert model compression. arXiv preprint arXiv: 100X:::9080000555,, 2019.
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.
Tianxiang Sun, Xiangyang Liu, Wei Zhu, Zhichao Geng, Lingling Wu, Yilong He, Yuan Ni, Guo- tong Xie, Xuanjing Huang, and Xipeng Qiu. simple hash-based early exiting approach for language understanding and generation. arXiv preprini arXiv:2203. 01670, 2022.
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 Computer Vision and Puttern Recognition, pp. 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), pp. 2464-2469, IEEE, 2016.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- tion processing systems, 30, 2017.
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. 2019a. In the Proceedings of ICLR.
Meiqi Wang, Jianqiao Mo, Jun Lin, Zhongfeng Wang, and Du. Dynexit: A dynamic early-exit strategy for deep residual networks. In 2019 IEEE International Workshop on Signal Processing Systems (SiPS), pp. 178-183. IEEE, 2019b.
Zizhao Wang, Wei Bao, Dong Yuan, Liming Ge, Nguyen H Tran, and Albert Y Zomaya. See: Scheduling early exit for mobile dnn inference during service outage. In Proceedings of the 22nd International ACM Conference on Modeling, Analysis and Simulation of Wireless and Mobile Systems, pp. 279-288, 2019c.
Ji Xin, Raphael Tang, Jaejun Lee, Yaoliang Yu, and Jimmy Lin. DeeBERT: Dynamic early exiting for accelerating BERT inference. In Proceedings the 58th Annual Meeting of Association for Computational Linguistics, pp.
Association for Computational Linguistics. 2020a.
Ji Xin, Raphael Tang. Jaejun Lee, Yaoliang Yu, and Jimmy Lin. Deebert: Dynamic early exiting for accelerating bert inference. arXiv preprint arXiv:2004. 20W.:9904.2993. 2020b.
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, pp. 91-104, 2021.
Wei Zhang, Lu Hou, Yichun Yin, Lifeng Shang. Xiao Chen, Xin Jiang, and Qun Liu. Ternarybert: Distillation-wware ultra-low bit bert. arXiv preprint arXiv: 12812, 2020.
Zhen Zhang, Wei Zhu, Jinfan Zhang, Peng Wang, Rize Jin, and Tae-Sun Chung. Peee-bert: Accel- erating bert inference via patient and confident early exiting. In Findings of the Association for Computational Linguistics: NAACL 2022, pp. 327-338, 2022.
13

---

Published as a conference paper at ICLR 2025
Wangchunshu Zhou, Canwen Xu, Tao Ge, Julian McAuley, Ke Xu, and Furu Wei. Bert loses pa- tience: Fast and robust inference with early exit. Advances in Neural Information Processing Systems, 33:1833001883335 18341, 2020.
Michael Zhu and Suyog Gupta. To prune, or not to prune: exploring the efficacy of pruning for model compression. arXiv preprint arXiv: 1710::17700875 2017.
Wei Zhu. Leebert: Learned early exit for bert with cross- optimization. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguisties and the 11th International Joint Conference Natural Language Processing (Volume 1: Long Papers), pp. 2968-2980, 2021.
A APPENDIX
A.1 PROOF OF THEOREM 3.1
Proof. For simplicity, we prove it for the binary classification case. For the samples that are not inferred intermediate exits, the misclassification probability will remain the same with without BEEM. Therefore, we only need to consider the case when the sample exits early and misclassified. We denote the probability that the sample exits early from exit t as p. and the probability that the exiting sample is misclassified as p, We next look for conditions under which
PE punie
Pt ptop < p. i.e.,
the fraction of samples early exiting being misclassified is less the error rate of the final classifier.
Let us consider two random variables X, and Y, where X, if the sample exits at /th classifier else 0 and Y, = 1 if the prediction /th classifier is correct else O. Now the probability that sample exits at the fth classifier could be written as:
P(X, = = P(Si < <a... a..... S < a. S₁ 2 a)
(4)
Let g.tii = min 92..... I ---1 } and qiaaz = 99,,92......4 .2..... We denote the prediction of classifier as St. Note that P(y" 11t) 1 qt. We define a counter Y, as:
1 if J U1
C2
if 1t-1 + 9t
(5)
We define C, = E/= c, Note that C, monitors the number of times the prediction has been changed till the /th classifier.
stop = P(X, = Pi
P(X, 1/C, = =c) =) (6)
==0
We denote P(G, 1|C, c) as Aç. We have
Pt stop
P(X,= Σ ACP(C, =).
By Total Probability law, we can say that
P(X, = = P(X, = 1|mc) P(mc)
P(X, = P(cc)
(7)
where me is misclassified and is correctly classified. The ratio is now:
14

---

Published as a conference paper at ICLR 2025
Pt
gaisc
stop
Pt
P(X, = 1 pmc) P(mc)
<p (8)
P(X, 1/mc)) + P(X: P(ce)
Simplifying this, we have,
P(X; P(ce) P(X, = 1 1mmcm.
1 p
For P(X; 1mm),, from the definition of the confidence score 5 we observe that the highest prob- ability of exiting at ith layer with miselassification if all the previous predictions are consistent, ie., St will be highest when all the previous exits misclassified. Hence we have
P(X; 1mmc t4 ... (
For P we observe that again by definition of the confidence score, we can lower bound it, since the lowest probability of exiting with correct classification /th classifier will be when all the previous classifiers had misclassification and then at /th classifier the prediction reversed. Hence, we have that
P = =l(c) >
14 1..e..... (1-- --4)) > - 9c) (9)
Now we have inequality as
tA,99 11 (1
at
t.A% 14pmoaze
P
We get the desired result by simplifying the above inequality for ge. We denote the constant term A
q max
Hence, we finally show with b,
A
@t
(10)
at + (11/p - 1)61
where and be are constant for each t.
This concludes the proof.
A.2 ANALYSIS OF THRESHOLDS &
In table we compare BEEM-A where thresholds are set by optimizing Eq against the case where threshold values are chosen using the validation dataset and are constant for all exits. We provide the comparison on (AL))ERTTbbase as well large models. We can observe a significant difference in accuracy and speedup. The improvement in accuracy and speedup is explained by the formulation of the optimization problem Sec. 3.4
Also, in Fig 2a plot the accuracy-speedup trade-off curves. We can observe that the change rate of decrease of the accuracy of BEEM by increasing speed-up smaller as compared to previous baselines. This extra stability by BEEM could be attributed to its characteristic of confirming the predictions from multiple exits (experts). Note that Fig [2 not assume the thresholds are fixed and varies them to show the stability of the approach.
15

---

Published as a conference paper at ICLR 2025
Sones-p
2.73
92.0
2.39
91.5
Accuraty
90.5
90.0
LeaBENT
89.0
Ours
1.
1.0
2.0
2.4
2.6
Speed-up
Jeaalass
(a) Accuracy-speedup trade-off for
(b) Trade-off by varying X.
A.3 ANALYSIS OF COST À AND # PARAMETERS
In BEEM-C, we introduce weights representing costs associated with utilizing exits. Figure 2b il- lustrates how altering these costs influences accuracy and speed-up trade-offss As the cost attributed to each exit classifier increases, we observe slight decline in accuracy accompanied by more sig- nificant enhancement in speed-up. This hyperparameter thus offers mechanism for modeling the trade-off between accuracy and speed-up, particularly as the thresholds a remain fixed. It's worth noting that adjusting the value of À impacts the quantity dot defined in Section 3.4 Consequently, modifications to À may induce changes in the values of a'.
Note that BERT-basel Lase/ has 10/340 Million Parameters with exits and ALBERT-base/LL. has 13/19 Million parameters with exits.
16


Reviews:
[{"review_id": "0YMwzjAEHy", "reviewer": "Reviewer_MFiX", "summary": "This paper studies the early exit (EE) technologies for language tasks and image captioning tasks. More specifically, instead of making decisions based on individual intermediate layers, this paper proposed to use weighted sum of multiple intermediate layers to achieve stronger results. Thresholds for each layer are set with error rate restrictions on the validation set, to make sure each exit classifier (on early exited examples) performs no worse than the final layer. Speed-up improvements are achieved on both language and image captioning tasks, while on language tasks there’s also a tiny quality boost.", "questions": "* In line 412, to calculate speed-up factors for image captioning tasks, how do you distinguish the cost from the vision encoder and the language encoder / decoder? \n* Are the vision / language backbones fine-tuned or frozen (where only the new linear layers are trainable)?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "* Extensive experiments are conducted on both language tasks and the image captioning task, with source code attached and plenty of baseline results reported.\n* The proposed early exit method not only speeds up the inference, but also improves model quality on language tasks.", "weaknesses": "* Speed-up factor is only computed based on the theoretical analysis (line 412), which is very misleading, especially for image captioning tasks where exit is only attached to the decoder but not encoder (line 201). In this case, do you count speed up only based on decoder layers? This is problematic, for example the baseline method MuE in Table3 applies early exit to both the encoder and the decoder (Figure2 of https://arxiv.org/pdf/2211.11152) and how is that calculated in line 386? Or do you also count encoder layers as a constant in this case as the inference cost is never reduced for encoders? If that’s the case, then it’s also not accurate due to the shape mismatch between the vision encoder and the language decoder. It would be more useful to report the actual end-to-end inference time speed up with and without early exit. \n* Per-layer threshold tuning depends on the validation set of each downstream task, which makes it not practical for broad use cases (e.g. zero-shot or few-shot tasks). Section 6.2 compares vallina threshold selection with the proposed threshold selection method, do you also have study showing the impact of validation set size? For example, if reduce the size of the validation set, would that significantly affect the proposed threshold setting method quality?\n* The proposed method is incremental, where it differs from the baseline methods of whether or not to use a weighted sum of the exit classifiers. And there’s also no ablation test to verify this point of using a weighted sum versus using simply individual exit classifiers. It's unclear whether we should treat DeeBERT and PABEE as ablation tests because the results are produced with the existing code (line 376) instead of a comparable setup (e.g. exactly the same code from this paper).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AhVLMZ9MKB", "reviewer": "Reviewer_9UE5", "summary": "The paper introduces BEEM, a novel framework for enhancing the performance of Early Exit Deep Neural Networks (DNNs) by treating exit classifiers as experts. The key idea is to aggregate confidence scores from these exit classifiers, making exit decisions based on the ensemble effect and consistency of predictions. A sample exits the network when the aggregated confidence exceeds a predefined threshold, which is determined by the error rates of intermediate exits. The paper presents experimental results on the COCO dataset for image captioning and GLUE datasets for language tasks, showing improvements in speed-up by a factor of 1.5× to 2.1× and comparable or improved accuracy compared to the final layer of DNNs. The significance lies in its ability to reduce inference latency while maintaining or enhancing accuracy, particularly useful for resource-constrained environments.", "questions": "1. While the paper demonstrates strong performance on NLP and image captioning tasks, it is unclear how well BEEM generalizes to other types of tasks or datasets, which could be a limitation in its applicability.\n\n2. The paper could provide more details on the computational complexity of implementing BEEM, especially regarding the aggregation of confidence scores and the potential overhead it introduces.\n\n3. The performance of BEEM is highly dependent on the choice of thresholds, and the paper could benefit from a deeper exploration of how sensitive the model is to these choices.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper introduces BEEM, a novel framework for enhancing the performance of Early Exit Deep Neural Networks (DNNs) by treating exit classifiers as experts. The key idea is to aggregate confidence scores from these exit classifiers, making exit decisions based on the ensemble effect and consistency of predictions. A sample exits the network when the aggregated confidence exceeds a predefined threshold, which is determined by the error rates of intermediate exits. The paper presents experimental results on the COCO dataset for image captioning and GLUE datasets for language tasks, showing improvements in speed-up by a factor of 1.5× to 2.1× and comparable or improved accuracy compared to the final layer of DNNs. The significance lies in its ability to reduce inference latency while maintaining or enhancing accuracy, particularly useful for resource-constrained environments.", "weaknesses": "1. The paper provides extensive experimental results demonstrating the effectiveness of BEEM in improving inference speed and accuracy across various NLP tasks and image captioning, which strengthens the credibility of the proposed method.\n\n2. The theoretical analysis providing conditions under which BEEM outperforms standard DNN inference adds depth to the paper and offers insights into its underlying mechanisms.\n\n3. The paper is well-organized, with clear explanations of the methodology, experiments, and results, making it accessible to readers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4pdLo3iqKP", "reviewer": "Reviewer_RQs8", "summary": "The authors propose a new decision criterion, BEEM, where exit classifiers are considered experts whose confidence scores are aggregated. A sample exits when the aggregated confidence surpasses a set threshold. The paper evaluates this approach on the COCO dataset for image captioning and GLUE datasets for various transformer-based language tasks.", "questions": "The paper seems somewhat rushed, with presentation issues that would benefit from an additional round of revision (resubmission basically). From my past experiences, these are unlikely to be fully addressed during the rebuttal period. Given the positive empirical results, I believe the paper would be stronger with more attention to clarity in writing, figure design, and structural organisation.\n\nHow much are you prepared to revise the presentation of the paper? If you are aiming for acceptance, please outline a specific plan to improve its readability and structure.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 1, "contribution": 3, "strengths": "The results look promising, as shown in Tables 1, 2, and 3.\n\nAs this area is outside my expertise, I would defer to my fellow reviewers on:\n- The relevance of the benchmark tasks and datasets used in the paper\n- The significance of the results\n- Any potential biases or issues within the experimental setup", "weaknesses": "Presentation issues:\n- The introduction and abstract delve too deeply into technical details from the outset. Consider starting with a high-level overview before moving into specifics, which are available in detail later.\n- Rework Figure 1: It should be consistent with other figures. For instance, S1 should connect to S2 with an arrow.\n- The title of Section 3 (\"Problem Setup\") feels off.\n- Section 3.1: The setup is challenging to follow. How are parameters denoted for exit classifiers? Are they trained independently or jointly with the model? Are the exit classifiers linear? Additionally, a brief (one-sentence) explanation of why the KL term is relevant would - improve clarity. Is the inference algorithm novel here, or adapted from previous work?\nSection 3.2: Describe the task (e.g., image captioning) first before discussing specific architectures.\n- Line 219: Statements like \"confidence-based early exit methods like DeeBERT and ElasticBERT...\" lack supporting evidence. Many claims are subjective, without empirical or theoretical backing.\n- Section 3.3: The cost vector isn't \"learned\" in the typical sense; it's set by the model developer. Adjust the wording accordingly.\n- Lines 256–259 are confusing; rephrase them. For instance, what is L? Is it the number of layers? Why is S_i < L important? What does i signify?\n- In Section 3, contributions are not clearly separated from known results in the literature. Improved subsectioning could help delineate the original contributions.\n\nNit:\n- Use \\citep instead of \\citet when appropriate (e.g., Line 46).\n- Line 215: There's no Equation 5.\n- Theorem 3.1: Avoid using t for the layer index, as it was previously used for tokens (Line 205). Clarify terms a_t and b_t, which are hard to interpret.\n- Table 5 should appear after Table 4.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eEvCEc7dzS", "reviewer": "Reviewer_mL6L", "summary": "BEEM presents an approach for reducing the latency of model inference by allowing classification decisions to be made at multiple points along the model, with early termination occurring if classification confidence exceeds a threshold. Instead of treating each early exit point independently, BEEM also considers a consensus along a span of exit points to be sufficient for early exiting, even if the confidence of the current exit point is below the threshold that would allow exit by itself.", "questions": "Equation 2:\nWhy do we only track the max prediction chain? Alternatively, could we just continue to aggregate weighted confidences until one of them reaches threshold? For example, if the early parts of the model are trying to decide between two classes, and thus have comparable (but oscillating max) confidence, in BEEM, these oscillations cause the S_i score to stay small, and thus requiring a deeper traversal once the two classes are resolved into one (as we need to build confidence for the single resolved class).\n\nSection 4:\nWhy was DNN-9L chosen? Probably a minor question, but unless that comes from a related work, I'm not sure where the 9 came from.\n\nGeneral:\nIt seems intuitive that perhaps some layers are only specifically good at EE-classifying a subset of classes. Have you observed any behavior like this? It would tie into my question about equation 2, which is that because you're relying on a chain of consistent max labels, that uneven error distributions could be harming your overall approach.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "Treating early exiting as an ensemble based consensus mechanism is a good idea, and weighting those exit classifiers based on their ability to classify effectively is also a good idea. As best I can tell, the approach does not introduce hyperparameters without providing a rigorous way to arrive at them, which I think helps with the general utility of the approach. Given that classifiers and next-token generation are major parts of modern AI, and this approach works on both, I think this work has significance.", "weaknesses": "I don't feel like there are any showstopping weaknesses that call into question the acceptability of the paper in general. So I will limit it to some constructive feedback, and then hopefully the authors can speak to my questions below.\n\nFigure 1:\n- I found this figure a little bit hard to parse, even after reading the rest of the paper. I think that if you had an arrow from the 0.065 box to the addition operator above it, it would be much more obvious what’s happening. \n\nSection 3.4.2:\n- Line 262.5: Should it be c_stop or c_stop^t?\n- Should c_misc^t be divided by c_stop on line 264, as it’s meant to be the count (“represents the number of samples …”), and thus `c_misc^t / c_stop^t` would be the `p^t` error rate? Otherwise we’re looking at `c_misc^t / c_stop^{2 (t)}` which doesn’t seem to be right. \n\nSection 4:\n- It perhaps would be nice to see the early exit behaviors of some of the techniques with a histogram over exit points, for example. Or even simply the average exit layer over each dataset. \n\nSection A.2\n- Broken table link on line 782\n\nNit:\n- Wrong parenthesis form on line 319\n- If you have space, it might be good to introduce what BEEM-A and BEEM-C mean in Table 1, as it doesn’t show up until the end of section 4.1.Metric. \n- Line 445 “poo [sic]” typo", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)