Jointly-Learned Exit and Inference for a Dynamic Neural Network
Authors: florence regol
Joud Chataoui
Mark Coates
Venue: ICLR 2024 poster (2024)
Keywords: early-exit dynamic networks
efficient inference
Primary area: general machine learning (i.e., none of the above)

Abstract: Large pretrained models, coupled with fine-tuning, are slowly becoming established as the dominant architecture in machine learning. Even though these models offer impressive performance, their practical application is often limited by the prohibitive amount of resources required for $\textit{every}$ inference. Early-exiting dynamic neural networks (EDNN) circumvent this issue by allowing a model to make some of its predictions from intermediate layers (i.e., early-exit). Training an EDNN architecture is challenging as it consists of two intertwined components: the gating mechanism (GM) that controls early-exiting decisions and the intermediate inference modules (IMs) that perform inference from intermediate representations. As a result, most existing approaches rely on thresholding confidence metrics for the gating mechanism and strive to improve the underlying backbone network and the inference modules. Although successful, this approach has two fundamental shortcomings: 1) the GMs and the IMs are decoupled during training, leading to a train-test mismatch; and 2) the thresholding gating mechanism introduces a positive bias into the predictive probabilities, making it difficult to readily extract uncertainty information. We propose a novel architecture that connects these two modules. This leads to significant performance improvements on classification datasets and enables better uncertainty characterization capabilities.

Published as a conference paper at ICLR 2024
JOINTLY-LEARNEI EXIT AND INFERENCE FOR A DY- NAMIC NEURAL NETWORK : JEI-DNN
Florence ,egol", Joud Chataoui" Mark Coates
McGill University, International Laboratory on Learning Systems (ILLS), Mila Québec Al Institute As {florence. robert-regee joud. chataoui @maal mmg. 11 ca
ma Park. csates emcll ca
ABSTRACT
Large pretrained models, coupled with fine-tuning, are slowly becoming estab- lished as the dominant architecture in machine learning. Even though these mod- els offer impressive performance, their practical application is often limited by the prohibitive amount of resources required for every inference. Early-exiting dynamic neural networks (EDNN) circumvent this issue by allowing a model to make some of its predictions from intermediate layers (i.e., early-exit). Training an EDNN architecture is challenging as it consists of two intertwined compo- nents: the gating mechanism (GM) that controls early-exiting decisions and the intermediate inference modules (IMs) that perform inference from intermediate representations. As a result, most existing approaches rely on thresholding con- fidence metrics for the gating mechanism and strive to improve the underlying backbone network and the inference modules. Although successful, this approach has two fundamental shortcomings: 1) the GMs and the IMs are decoupled during training, leading to a train-test mismatch; and 2) the thresholding gating mecha- nism introduces a positive bias into the predictive probabilities, making it difficult to readily extract uncertainty information. This leads to significant performance improvements classification datasets and enables better uncertainty characteri- zation capabilities.
INTRODUCTION
The dominant approach to improve machine learning models is to develop larger networks that can handle every potential sample. As result, despite very impressive performance, the resource overhead is huge (Scao et al., 2023). The push for larger model size is often driven by the need to handle a small percentage of samples that are particularly challenging to infer (Bolukbasi et al.) 2017); most inferences do not need the full power of large network to be successfully executed. Nonetheless, most traditional neural network (NN) models have a fixed processing pipeline. This means that every sample, simple complex, is processed the same way.
To tackle this inefficiency, dynamic networks have been introduced (see (Han et al. 2022a) for a review). These models adapt the computational processing to the specific sample being processed. Early-exit dynamic networks (EEDNs) tailor their depth to the sample, allowing easy-to-infer sam- ples exit at shallower layers. Compared conventional neural networks, BEDNs incorporate two additional components: 1) Intermediate Inference Modules (IMs) receive sample's representation their respective network depth and generate predictions; Gating Mechanisms (GMs) decide which intermediate inference module should be employed to derive the final prediction.
Most EEDNs employ a simple threshold-baad based gating mechanism. A threshold GM applies thresh- olds to confidence metrics obtained from inference modules. To integrate more sophisticated GMs, the gating mechanism can be treated post-training add-on component. Given resource budget and a set of pre-trained IMs, the optimal gating decisions are learned by taking into account both performance and inference cost. Unfortunately both strategies, the IMs and GMs are decoupled during training: the IMs are trained the full dataset even though they will only infer subset of the data at inference time. This inevitably creates a mismatch between the train and test distributions.
Equal contribution

---

Published as a conference paper at ICLR 2024
The reliance on threshold GMs also makes it more difficult to obtain confidence levels for predic- tions. The information required for this determination has already been "consumed" in the gating decision; moreover, applying threshold introduces an overconfidence bias. In setting where model can produce different outputs using varying levels of computational resources, confidence measures are essential. They enable end-users to make informed decisions, whether to accept cost-efficient output or to request additional computation for a more reliable answer. Currently, state-of-the-art EEDNs do not offer confidence information.
Contributions: We propose a novel learning procedure for the GMs and IMs given a fixed backbone network. Our approach involves joint training so it directly avoids train-test mismatch and provides good uncertainty characterization. The method introduces a new approach for modeling the proba- bility of exiting at particular inference module. We optimize a loss that jointly assesses accuracy and inference cost, and formulate the minimization as a bi-level optimization task. Each level of the bi-level optimization is simpler than the overall problem, leading to more stable learning. We show empirically that the approach leads to better overall inferencee/eerfommacce trade-off. The benefits are threefold: 1) we close the training gap between IMs and GMs, which leads to better perfor- mance; 2) the architecture produces reliable uncertainty characterization in the form of conformal intervals and well-calibrated predicted probabilities; and 3) our approach sigmificantly outperforms other rechttecture-annstfi early-exit training procedures, while its generality facilitates integration into most state-of-the-art architectures, including those tailored for dynamic-exit
2 RELATED WORK
We briefly summarize the most relevant literature. Appendix contains further discussion. EEDN-tailored Architectures with threshold GMs: Much of the EEDN literature focuses on designing architectures that better lend themselves to early-exiting (EEDN-tailored architectures). BranchyNet (Teerapittayanon et al.) (016) augments the backbone network with early-exit branches and is trained using a weighted loss. Shallow-deep networks (SDNs) quantify the network confusion by comparing the outputs of different IMs (Kaya et al. 2019). Efficient end-to-end training is dif- ficult because earlier branches influence the performance of later ones. Subsequent EaDN-tailored architectures, such as MSDNet (Huang et al.) 2018) and RANet ( (Yang et al.] 2020). alleviate this via dense connections or increased resolution. Training can be improved by reducing conflictual gradient updates using position-based based rescaling (Li et al. 2019) meta-learning (Sun et al. 2022).
conoiiiinn
Reintegrating computations from earlier layers can boost performance via ensembling (Wolezyk) et al., 2021 Liu et al., 2022, Passalis al. 2020). All of these works rely on confidence score thresholds to decide whether to early-exit. The geometric ensembling of |Woczyk et al. (2021),
where learnable weights adjust the contribution of each IM to the final classification decision, has connections to our exit probability modeling. Many recent advances in the EEDN field also rely on simple threshold GMs (Wang et al., [220 Han al., 2023b Wang et al. 2021 Chen et al. 2022b). In general, the thresholds are treated as hyperparameters and assigned values using a validation set.
Learnable GMs, fixed IMs: A second class of EEDNS considers the task of learning GMs for frozen pre-trained backbones and IMs. Bolukbasi et al. (2017) train the gates sequentially using a top-down approach. In EPNet (Dai et al. 2020) the gate-training is formulated as Markov de- cision process. PTEENet (Lahiany & Aperstein 2022) employs much simpler gate architecture. Karpikova et al. (2023) design learnable gates with a focus on confidence metrics for images. All of these models learn gating functions after having trained and frozen all of the IMs. While this two-phase procedure simplifies training, it ignores the coupling that inherent to these two tasks.
Addressing the train-test mismatch: Ignoring the influence of the GMs on the IMs by training learnable gates after freezing the IMs introduces a train-test mismatch (Yu et al., 2023 Han et al.) 2022b). BoostNet YYu et al. 2023) proposes an architecture-agnsti training procedure inspired by gradient boosting to close the train-test mismatch. Deeper IMs are trained with an emphasis on samples that are incorrectly classified by earlier IMs. [Han et al. (2022b) cast the problem as a meta-learning task where a weight prediction network (WPN) predicts the gate that will lead to the lowest classification loss and weighs the inference module losses accordingly. Although both methods are good steps towards aligning the training of the IMs with the gating mechanism, they do not fully close the train-test gap in practice. Both methods still employ threshold GMs, and the proposed specialized training of the IMS is not directly tied to which gate is ultimately selected by
Coode reproduce our experiments is available at our Github repository
2

---

Published as a conference paper at ICLR 2024
the threshold GMs. Scardapane et al. (2020) propose differentiable branching, approach where gates and IMs are jointly learned. A straightiorward modeling of the gate exit probabilities and joint optimization of all parameters results in unstable training that is slow to converge.
3 PROBLEM SETTING
We consider classification problem with training dataset D = (x,, (1 W where Rp denotes the input and y/ IC its corresponding classification target (K = [I....,K] KJ). We are given fixed network NN: RDD pretrained on the same task that can be decomposed into composition of L layers: (x) o h -- o...o h (x), where denotes the softmax operator. Let z' be the 1-th intermediate representation of NN such that = h)zz")) with z"=
Our goal to augment this fixed network (the backbone) with additional trainable components such that we can obtain final prediction Y; with its associated predicted probability vector P. at reduced inference cost. This setting categorized as intermediate classifiers-only training (IC-
only training) in the review by Laskaridis et al. (2021).
To measure the performance of our model, we consider three types of metrics: (i) performance-
based (the accuracy of 1); (ii) uncertainty-baase (the calibration error and the inefficiency score of the conformal intervals obtained from p:): and (iii) efficiency-based (the inference cost to obtain ÿ)
Inference cost: In line with BoostedNet (Yu et al. (023) and L2W-DEN (Han et al. 2022b) the inference cost ICs, for a single prediction ii represents the number of multiply-add (Mul-Add) operations to obtain is We approximate the expected inference cost E [[C as:
1 N Σ IC
E[IC] v IC
Uncertainty metries: We assess calibration of the model via the expected calibration error (ECE) of the probability of the predicted class (the maximum probability of p, p.max = max(p:)). To es- timate the ECE, we follow the approach in (Guo et al. (017). More details concerning computation of the inference cost and ECE are provided in Appendix 9.2
To obtain conformal intervals, we use the typical conformal score 8; -D; and compute a con- formal threshold 7 vV,a from a validation set V so that we can guarantee a coverage of 1 (See Appendix 9.3 for details). A conformal interval C₁ for sample x; is then constructed by threshold- ing the conformal score s; C, (k: s
(k) < vV..), where (A)) denotes the k-th element of
We measure the performance of our predicted conformal intervals C, with two metrics: empirical coverage a, which computes the empirical probability that the ground truth is contained in the interval, and inefficiency (C), which is the average cardinality of the set intervals:
1 N
C
CI.
1 1-
Σ
[[ e C.).
A well-behaved conformal prediction model has small inefficiency |C| (intervals are small aver- age) while maintaining an empirical coverage close to the desired one: 1 a 1-1
4 METHODOLOGYY JEI-DNN
We introduce our method named: Jointly-Learned Exit and Inference for a Dynamic Neural Network (JEI-DNN). Our model consists of augmenting pretrained and fixed network by attaching L clas- sifiers and L gates to the intermediate layers of the network architecture. We refer to each these classifiers
inference module (IM). The (th IM, fè, (-), uses the h representation z! of the input
x, model class probabilities (x.) = f(()) and obtain prediction ÿf arg max, Po I.()) (xx)). We set the last IM to be the original classifier of the pretrained network: fd = h The output of the tth gate, e' e {0, 1), is modelled by random variable E that specifies whether the I-th exit gate is open. If a sample reaches the [-th gate and it open, then exits via the th IM. Figure depicts the JEI-DNN learning procedure and highlights the bi-level optimization used for joint training.
3

---

Published as a conference paper at ICLR 2024
Exily 3
Joint training via bi-level
$ 5² 1000) 8
2 0.85
P(D-8x
-(8"-1s)
optimization GMM training
(xxx) = s
2 0.1
P(R
X
%
fd (x) = 8
=
1xx) 0s
X
0
IM fg uses
Each gate gf independently ouiputs
valid distribution
Using all previous P(G]x). the prob. P( = 1/x) of gate of being open computed
P(G [xx exiting
g
produced by layer to make prodiction
prob. of exiting at given
layer is derived from independent gate outputs
Pg
uncertuinty stats. ma(0,
hackhone laver I
Figure 1: Left: Modelling of GM and how early exiting achieved. Right Illustration of the mutual influence of GM and IM.
The assumption is that some samples are "easy enough" be classified by the early classifiers (z'). Since the last classifier should be able to classify the entire dataset, the last gate is always set to one, i.e., e² = 1. Denoting the indicator function by [[-], the output of our model is:
((x) = g (xx where I arg min I[e' (x) (x)== = 1]. e'(x) (s))>((E(X=x)
(3)
k[4
We can view each IM fg as solving subtask over subset of data D' D, where the partitioning is determined by the gates, and D' D* = Ø for j k and Uie[z =D. Since the E' are random, we can also define the distribution P(G) of exiting at gate G' E [L]. This equals the probability that the sample does not exit at any previous gate multiplied by the probability that it exits gate l:
[-]
P(G= = 1X= P(E
(LEP((X))= 5-1
(4)
The first network layers and the first / IMs and gates must be evaluated in order exit at the th layer. Assuming fixed-size input samples, the inference cost IC'; associated with obtaining prediction from the 1th IM can be approximated as a constant IC that only depends on l. For all predictions j obtained at level / we then have IC2g IC Appendix provides more detail about
how IC' is computed. During training we use a normalized version IC murm relative cost incurred by exiting at / compared to going through the full network (i.e., all L layers).
IC' that represents the IC
We can then formulate our loss as the expectation over the exited gates G|X and the data distribution X, Y of the traditional cross entropy loss CCE (Y, G (X)) added to the associated inference cost I0G scaled by a parameter that controls the importance of the inference cost:
Ex.XEx X
(Y,D, GIX (X)) A.C.GIX J. norm
(5)
N E
L
C
Pa ACCl norm) P(G ((G=llx;))
Our goal is to simultancously learn the parameters A of the IMs and choose P(G = l/xx) ) in order to minimize the expected loss as shown in Figure We target P(G = l[x:) rather than P(F xx) for j 1....,I because directly learning the P( (P) leads to optimization issues. The product of probabilities across multiple layers can quickly vanish or saturate to 1.
Modeling P(G[x;). In practice, only need to evaluate P(G l/x,) if the dynamic evaluation has reached that point, i.e., everything up I (xt) has been evaluated and is accessible. There- fore, our parameterization of P(G = 7|x;) can include any intermediate values calculated by the base architecture, the inference modules, and the gates up to and including layer 1. We denote this aggregation of information as es
This however also implies that we cannot directly model the multiclass distribution P(G[x:) with the traditional softmax approach since it would require (z)) to be evaluated for all 1, This would
4

---

Published as a conference paper at ICLR 2024
defeat the purpose of exily-exiting, Instead, model each P(G = 1|x,) sequentially, starting from the first layer 1 = 1, with learnable variable o (c(²) )e()ee [001]. We set P(G = 1/x;) to the minimum of the learnable variable (of( e [0, 1] and 1 - Σÿ- (c,)). Since the model must exit at the
final gate if it has not already done so, we have g² gc -1. Thus:
Pa 1|xx)= (c
H Pa(G= = [[x)) min (c.c) 9, Σ (co)) for 1= = j=1
(8)
The min operator Eq. ( ensures that the distribution. The loss is then proper approximation of the expectation in Eq. ( It also guarantees that P = x). which is calculated using 44 and governs the decision to exit, lies in [0. 1].
L} form valid probability
This formulation leads to much more stable model for the gate probabilities compared to param- eterization through P(E [x). A deviation in one the g'ss has an additive rather than multiplicative impact on subsequent ( = l|x,)s. The left diagram in Figure depicts the modeling of P(G Our approximated loss from equation 6 is then:
1
E(CEE 1=111
N L Σ
(yi, (x::)) ATC min (d)(c,)) Eo(( j=1
L
Any parameterization of the IMs fa and gates function I' (c can be adopted. However, the modules must be small compared to the size of one layer of the backbone model to prevent them from substantially increasing the computational cost of the entire layer during inference. In our experiments, f/ and g's (c(")) are parameterized as simple one-layer neural networks and uncertainty metries are extracted from c.e) to serve as input to d (.) A detailed description of our modules is included in Section Appendix (eapddxxX.2jinceee innludss the inference cost of the gates and IMs demonstrate that the added computation is minimal when compared to the inference cost of a single layer. (The total added computation amounts to less than 0.01% % of the Mul-Adds of the backbone.)
Uncertainty prediction: : Next we describe how we use the predicted probabilities pl to form conformal intervals. The conformal intervals are derived from a conformal threshold vV,a that is computed using a validation set V. However, since we can view the IMS solving different subtasks over different subsets of the data DD' we should compute a different conformal threshold for each subtask. Hence, we form one validation set per gate, V", where V
{i:l (G i V),
and compute conformal threshold from that subset: T ,a If there are too few samples in the validation set associated with gate I, we r :a general threshold vV.a computed using the entire validation set. The conformal intervals are then constructed as follows:
C, (k:; 1 p/ (A) (x:) < 1 ,1~ P(G[x:))
(10)
There are many other ways to construct the sets to compute the threshold in the EDNN setting. In Appendix 9.3 we show that a variety of methods produce similar results.
We present candidate designs for the gates and IMs for conereteness, but note that other choices possible within the JEI-DNN framework.
Gate design (c²5): To ensure lightweight design, we construct a small number of fea- tures by computing uncertainty statisties from the output i of the 1-th IM, [x;).. Hence (m), We choose m(f, x;)
the 1-th gate can be represented as (c(") (x,,)x(xx)), he (xx;), ()ooo mx," where p!.max(xx) is maximum predicted probability, hxx), is the entropy, h (x,) is the entropy scaled by temperature, and mar (x,) is the differ- ence between the two most confident predictions. For more detail see Appendix 9.9
IM design (2:"): The IMs (z+a) all have the same architecture. We reduce cost by limiting their input size and number of layers. The input to f is only z and the IM is a single layer NN. The lightweight design allows to insert exit branches at every layer. This gives model greater exit flexibility, in contrast with many existing approaches (Chiang et al.) [021]] Lahiany & Aperstein
2022
Li et al. 2023 IIhan et al. 2023). which typically only few exit branches.
5

---

Published as a conference paper at ICLR 2024
4.1 OPTIMIZATION
Directly optimizing the loss in equation is challenging because of the min operator:
1 N
0. 0" arg min (E =11-1
(V.. (x.))) XIC min (9's (c ) -
(II)
0,9 N
However, there is distinction between the parameters 0 and ó which makes this loss good can- didate for bi-level optimization formulation. (See (Chen et al. 2023a) for a survey.) Hence, using Co(a) _ (vi, (x:))
XIC norm:
rewrite equation as follows:
8² arg min Cout arg min
1
N
Cg
(o) (xx)),
(12)
N
N 1 (c (x.)) + PC (13)
s.t. (0) arg min L arg min mig
N 1-1
Since the gate probabilities Pa(G ([x;) take as input the aggregation of any intermediate values that were previously caleulated c,', it could be dependent the 0. To make that dependence explicit in the calculation of the gradient of 0, we denote Po(G = !(x,) Go (m)(0, x.)). The derivative of
the loss with respect to the two different set of parameters is given by:
Oh
1
N OC (vi, (v..P((x.) (x.)) Po = 1|x:) @Go(m ).)) Cocas:
3 N
(14)
00
00
00
1 N dmin ( 1- Ca-cor- (15)
Ocout
N. (P
ΣΣ
Crr
ΣΣ
06
06
Eq. (15) is undefined when the min operator terms are equal. We address this issue below when describing how to optimize the gates.
Optimizing the IMs: By inspecting equation 14 we can recognize that the left term is the same gradient that is encountered for a straightforward weighted cross-entropy loss. Interestingly, our principled approach leads to objective with term that is similar to the one proposed by Hann ct al. (2022b) to address the train-test mismatch issue. In our case, the weights emerge directly from our proposed gating mechanism: (& lJx,). The right term is driven by the impact of the
0 parameter on the gates: (m((..)) In practice, we observe that ignoring the right term of the 80
gradient does not impact performance and leads to faster convergence, as we show in Appendix 9.7
Optimizing the gates The second derivative in equation 15 is more challenging, making direct optimization undesirable. Instead, we construct surrogate binary classification tasks to train the g For given sample x; construct binary targets for (e )..,),... - (c;L-1) by evaluating the relative cost of each gate (a) and determining which of the L gates has the lowest cost, denoted as l" arg min, C/- Ch~(())- This is the gate at which x; should exit. Hence, we set the binary target of gate 1" to 1, and the binary targets for all of the preceding gates to zero. As for the subsequent gates, since the sample is supposed to exit at 1*, assume that it should also exit any later gate, so we set the binary targets of all subsequent gates to I as well. Hence, the targets of our binary tasks are given by t = 0 for j < l' and t! 1 for j 2 l"
+f for our gates g
The connection between those surrogate tasks and the initial objective from equation can be estab- lished by showing that they can share the same solution under some assumptions. We demonstrated this is Appendix 9.4] Hence, instead of following the gradient from equation 15 we approximate
by the gradient of the surrogate loss:
ott d
1
ΣΣ 11
(16)
N
This amounts to summing the gradients of the losses of independent binary classification tasks. If the tasks are highly imbalanced, we compute the class imbalanced ratio on validation set at each epoch and use weighted class training. Following the bi-level optimization procedure, training carried out by alternating between optimizing using the gradients in equation [6]and equation 14 In practice, we start by training all IMS on the full dataset in a warmup stage as described Section Algorithm I in Appendix 9.5 provides detailed exposition of the entire algorithm.

---

Published as a conference paper at ICLR 2024
5 EXPERIMENTS
We use the vision transformers T2T-VTT-7 and T2T-ViT-14 (Yuan et al. 2021 pretrained on the ImageNet dataset (Deng et al. 2009) which we then transfer-learn FARI00, CIFAR 100-LT (Krizhevsky 2009) and SVHN (Netzer et al. 2011). For dataset details, see Appendix P.1.T For transfer-learning, we use the procedure from (Yuan et al. 2021). We pro- vide checkpoints for all these models in our code. Our backbone NN for each dataset the frozen transfer-leamde model. We augment the backbone with gates and intermediate IMs at every layer. We generate results by varying the inference cost parameter À over the range 0.01 10. Rather than sample E' ~ p( E. ()) to decide whether to exit, we use the deterministic decision P (E'E) 0.5. Appendix P.I.TJcontaing further details concerning hyperparameters and experimental procedure.
datasets: CIFARIO, CI-
Training procedure: Our training procedure consists two phases: (i) warm-up; and (ii) bi-level optimization. In the warm-up phase the first [-1 IMs are trained in parallel all samples for fixed number of warm-up epochs. This ensures that intermediate IMs are performing reasonably well when we start training the gates. During the bi-level optimization, we alternate between opti- mizing the gate parameters & and the IM parameters 0. See Appendix 9.1.4
9.1.4 for more detail.
0.90
0.95
0.9
80% end acc
0.85
0.90
0.8
W
Accuracy
Boostnet L2W-DEN
0.75
80% end
0.70
0.75
80% end acc
0.65
0.70
0.4
0.5
0.6
0.7
0.8
0.9
0.3
0.4
0.5
0.7
0.8
0.3
0.4
0.5
0.6
0.7
0.8
(48.06x10*)
Mul-Add (0.24x10*)
(48.06x10*)
Figure 2: Accuracy vs Mul-Add of: Left CIFAR10M 100 (12t-14); Middle SVHN (12t-7); and Right CIFAR10OLT (12t-14). The x-axes are scaled by the full model inference cost, Mul-Add (TC4). Baselines: We compare our proposal with the following baselines:
BoostedNet (Yu et al. 2023 and L2W-DEN (Han et al. 2022b are state-of of-the-art benchmarks with architecture agnostic training procedures for general- purpose networks.
MSDNet (Huang et al. 2018) () and RANet (Ynng et al. 2020) are DEDN-tailored architectures. We also compare with the improved training procedures in et al. 2023 and (Han et al.) (022b).
To obtain the uncertainty metrics, we calibrate the predicted p by first applying the post-calibration algorithm of temperature scaling from (Guo al.) (017) on the output logits using a validation set. Since all baselines also require validation set to select gating thresholds, we split the validation set into two V y² y2; y is used to compute the gating thresholds and 2 used for hyperparameter tuning, early stopping, setting the conformal threshold, and calibrating the output.
5.1 RESULTS
Observation 1:
Joint Training of the GMs and the IMs closes the train-test mismatch which leads significant performance improvement.
presents results on ad- ditional datasets and architectures. Focusing on the region exceeding 80% of the end accuracy of the total network, our suggested approach significantly outperforms the baselines for all datasets. Outside of that range, is on par or better. Moreover, the 95% confidence intervals around the mean of the accuracy of our proposed approach (shaded areas) are significantly tighter. The CIs of all methods are wider for CIFARI00-LT. The imbalance leads to noisier prediction, which makes it particularly challenging for the baselines. To provide further insight into why model outper- forms the baselines, we analyze the results in the subsequent sections. Appendix 9.6 provides an ablation study which demonstrates the value of learnable gates and joint training.
Figure depicts performance versus inference cost curves. Appendix 9.8

---

Published as a conference paper at ICLR 2024
1.0
1.0 b)
Accuracy exited points (D)
0.8
Accuracy all points (D)
0.8
method
0.6
e00 oopp
0.6
Boostnet
L2W-DEN
JEI-DNN
0.4
0.4
| (IMs usage)
5
0.2
Boostnet accuracy
0.2
40%
L2W-DEN accuracy
JEI-DNN accuracy
20
0.0
0.0
2.5
8 10
12
14
10
12
14
IMs (I)
IMs (1)
Figure 3: Decomposition of the contributions of the IMs to the final accuracies (depicted by dotted lines) for CIFARI00. The operational point is marked by a star in the left panel of Figure 2 a) Accuracy of each IM, fg. evaluated only on their exited samples (D'). b) Accuracy of each IM, fg- on the full dataset D. The size circle is proportional to the number of samples exited for a trial. Observation 2: Training the GMs leads to a better gate selection.
Because GMs are jointly trained with the IMs and take into account the of inference, the GMs can avoid poorly performing gates or gates that provide only marginal performance improve- ment. Figure shows which gates selected for exit by our proposed method (JEI-DNN) and the baselines. We see that JEI-DNN avoids using early IMs altogether and focuses on IMs to 10. The joint training approach thus effectively addresses the EEDN placement task. Rather than employing fewer more capable, complex GMs and IMs selecting where to place them, as in (Li et al. 2023). our approach introduces very simple, lightweight GMs and IMs, and learns to focus the IMs that can perform satisfactorily for given inference budget. In contrast, the threshold GMs employed by the baselines lead more evenly distributed exits. Since the threshold GMs make decisions by thresholding confidence metrics derived from predicted probabilities, it is more difficult to avoid exiting at early poorly performing IMs (these IMs are overly-confidee for some samples).
Moreover, the thresholding approach relies on the as- sumption that IMs are well calibrated, which not guar- anteed (Guo et al. 2017 Minderer et al.) 202T). Exper- imentally, we verify that it is not the case and that cal-
ECE on all points (D)
0.4
ibration is positively correlated with IM depth (see Fig- ure 4. Consequently, calibration is also correlated with accuracy, which is in line with the findings of Minderer et al. (202T). poor calibration leads to inaccurate pre- dicted probabilities, which introduces noise in the inputs to the threshold GM. As a result, the exit selection is also more variable, leading to less predictable accuracy-IC trade-off (as exhibited by the much wider Cls in Figure
0.2
0.0
10
IMs (I)
Figure ECE of the IMs averaged over all baselines all datasets with 12t-7.
Observation Better gate selection concentrates training on better IMs.
Figure b)) highlights that the improvement of our model does not come from IMs that are superior over all of the training data. Baseline and JEI-DNN IMs perform similarly. The improvement in IC-accuracy trade-off is explained by the trend of IM accuracies in conjunction with IM usage. JEI- DNN exits very few samples IMs 1-4 and 11-14. The early IMs are too inaccurate. The accuracy gain achieved by postponing exit to late IMs (11-14) is not justified by the increased inference cost. The GMs must be able to accurately direct easier samples to earlier IMs. JEI-DNN clearly achieves this. For IMs 6-8, where most samples exit, the accuracy is higher over exiting samples (Fig ((a) than all samples (Fig. (F)..) By contrast, for gates 9-10, where only a few challenging samples exit, the exit accuracy is lower. For JEI-DNN, the most challenging samples exit at gates 9-11; I; for the baselines, they exit at gate 13. Despite having access to less informative features, JEI-DNN achieves better average accuracy for these challenging samples at considerably lower inference cost.
8

---

Published as a conference paper at ICLR 2024
0.75
60
0.10
Boostnet L2W-DEN
0.50
10 40
0.25
JEI-DNN
20
0.05
0.00
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.01
0.02 0.03 Requested a
0.04
0.05
Mul-Add (48.06x10f)
Mul-Add (48.06x109)
Figure 5: Uncertainty metrics on CIFAR-100 (12t-14). a) ECE vs Mul-Add. b) Inefficiency vs Mul-Add for an empirical coverage bounded by 1 â 2 95% on CIFAR-I00 (12t-14). e) Average empirical á vs requested a. Appendix 9.8 presents results on other datasets and architectures.
Observation 4: Learnable GMs leads to better uncertainty characterization capability
Fia- highlights that even with post-calibration, the baselines' predicted probabilities remain ex- tremely poorly calibrated. JEI-DNN's calibration is much better and the calibration error correlates negatively with aceuracy. This is desirable; a higher inference cost should lead to better accuracy and better calibration. When the baselines provide more accuracy, the calibration error is much worse. This disparity has two causes: (i) JEI-DNN refrains from using the early, worst- corst-calibrate gates (Fig- 4 (ii) the baselines must incraase the GM thresholds to achieve higher accuracy, but this leads to a bias - samples only exit if the predictions are very confident. We now consider conformal intervals. These should circumvent the threshold GM bias they are derived via validation set and have probabilistic guarantees. Fig. b) shows that JEI-DNN offers significantly tighter confor- mal intervals [C] for given constraint on empirical coverage å. The empirical coverage values for JEI-DNN also closer to the requested coverages (see Fig. (()).
Observation 5: JEI-DNN can enhance SOTA dynamic architectures
We examine whether JEI-DNN hance a state-of-the-art EEDN-taiiored architecture. We select the Dynamic Perceiver Han et al.] (2023b) with MobileNetV3 (Howard
82
Acc. Nill network
Goteed accurary N
76
Acc. full network
6a% cost of full network
Gated sccures
80
al. 2019)
full network
72-
as the backbone. The combination (DynPerc+JEI-DNNN includes learnable gates. We train for 310 epochs. During the last 10 epochs for DynPenc+JEII-DNN, JEI-DNN is used to jointly train gates and IMs. Fig. [6 shows that this leads to significant accuracy improvements for both CIFAR100 and ImageNet.
70
DynPerc:
DynPerc.
DynPerc+JBIDNN
66
DynPerc.+JIIDN
74 6.5 Figure ceiver CIFAR100 (left) and ImageNet (right) (Mo- bileNetV3).
7.0
7.5
8.0 8.5
9.0
6.5
7.0
7.5
8.0
8.5
Mul-Adds (x10f)
Mul-Adds (x10")
DynPereeiver:JJIIDNN vs.
DynPer-
LIMITATIONS
Our proposal outperforms the baselines in terms of the accuracy-IC trade-off and uncertainty char- acterization, but it has some drawbacks. Since we must specify the inference cost parameter X, we only obtain one accuracy/IC point per training. To obtain different values, need to retrain the model with a changed À value. By contrast, the baseline methods that employ threshold GMs can modify the inference time after deployment by changing the threshold. Second, cannot directly target specific accuracy or inference cost. In practice, there is a constraint on one of the met- rics, we either must train multiple models, adjusting À until the constraint is met, or use an ad-hoc approach by applying different thresholds to the S/ values. In summary, if the desired trade-off be- tween inference cost and accuracy is likely to vary over time, and it is acceptable to sacrifice some accuracy, then the approach taken by the baselines would be preferable.
7 CONCLUSION
We have introduced a novel early-exit dynamic neural network learning procedure, JEI-DNN, that is compatible with off-the-shelf backbone architectures. By augmenting backbone with lightweight, trainable gates and inference modules that are jointly optimized, we close the train-test gap, leading to better performance of inference modules. The approach can improve ofate-of-the-ar EEDN- tailored architectures and provides significantly better uncertainty characterization. The inference modules are better calibrated and tighter conformal intervals are obtained for desired coverage.

---

Published as a conference paper at ICLR 2024
8 ACKNOWLEDGMENTS
Cette recherche a été financée par le Conseil de recherches en sciences naturelles et en génie du Canada (CRSNG), [numéro de référence 260250].
We acknowledge the support of the Natural Sciences and Engineering Research Council of Canada (NSERC), [funding reference number 260250).
We also thank Boris Oreshkin for kindly providing additional computational resources.
REFERENCES
Tolga Bolukbasi, Joseph Wang, Ofer Dekel, and Venkatesh Saligrama. Adaptive neural networks for efficient inference. In Proc. Int. Conf. Machine Learning. (ICML), 2017.
Can Chen, Xi Chen, Chen Ma, Zixuan Liu, and Xue Liu. Gradient-based bi-level optimization for deep learning: A survey, arXiv preprint: arXiv 2207.11719, 2023a.
Mengzhao Chen, Mingbao Lin, Ke Li, Yunhang Shen, Yongjaan Wu, Fei Chao. and Rongrong Ji. CF-ViT: A general toorsc-to-fine method for vision transformer. In Proc. of the AAAI Conf. on Artificial Intelligence, 2023b.
Xinshi Chen, Hanjun Dai, Yu Li, Xin Gao, and Le Song. Learning to stop while learning to predict. In Proc. Int. Conf. Machine Learning. (ICML), 2020.
Zhourong Chen, Yang Li, Samy Bengio, and Si Si. You look twice: Gaternet for dynamic filter selection in cnns. In Proc. IEEE/CVF Conf. on Computer Vision and Pattern Recognit. (CVPR), 2018.
Chang-Han Chiang, Pangfeng Liu, Da-Wei Wang, Ding- Yong Hong, and Jan-Jan Wu. Optimal branch location for cost-effective inference on branchynet. In Proc. IEEE Int. Conf. on Big Data (Big Data), 2021.
Xin Dai, Xiangnan Kong, and Tian Guo. EPNet: Learning to exit with flexible multi-branc net- work. Proc. Int. Conf. on Inf. & Knowledge Management (CIKM), 2020.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In Proc. JEEE/CVF Conf. on Computer Vision and Pattern Recognit. (CVPR), 2009.
Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, Barret Zoph, Liam Fedus, Maarten P Bosma, Zongwei Zhou, Tao Wang, Emma Wang, Kellie Webster, Marie Pellat, Kevin Robinson, Kathleen Meilr-Hellster, Toju Duke, Lucas Dixon, Kun Zhang, Quoc Le, Yonghui Wu, Zhifeng Chen, and Claire Cui. GLaM: Efficient scaling of language models with offture-of-exppets. In Proc. Int. Conf. Machine Learning. (ICML), 2022.
Vijay Prakash Dwivedi, Chaitanya K. Joshi, Anh Tuan Luu, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Benchmarking graph neural networks. Journal of Machine Learning Research, 2022.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Jour. of Machine Learning Research, 23(120):1-33 2022.
Yoav Freund and Robert E. Schapire. A decision-theoretic generalization on-line learning and an application to boosting. J. Comp. System Sciences, 55( 1591)11991,3, 1997.
Alex Graves. Adaptive computation time for recurrent neural networks. ArXiv, abs/ has/1603308983. 2016. htt.httts: hstp:///.pp.......t...mmmmm
semant www.chala...oooooooo ccc.cholaa..org/oorrrttttt:::00.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q. Weinberger. On calibration of modern neural networks. In Proc. Int. Conf. Machine Learning, (ICML), 2017.
10

---

Published as a conference paper at ICLR 2024
Shashank Gupta, Subhabrata Mukherjee, Krishan Subudhi, Eduardo Gonzalez, Damien Jose, Ahmed H. Awadallah, and Jianfeng Gao. Sparsely activated mixture-of-experts are robust multi- task learners. arXiv preprint: arXiv 2204.07689. 2022.
Dong-Jun Han, Jungwuk Park, Seokil Ham, Namjin Lee, and Jaekyun Moon. Improving low- latency predictions in multi-exit neural networks via block-dependent losses. IEEE Transactions on Neural Networks and Learning Systems, pp. 1-9, 2023a.
Y. Han, G. Huang, Song, L. Yang, H. Wang, and Y. Wang. Dynamic neural networks: A survey. IEEE Trans. on Pattern Analysis.; Mach. Intell., 44(1 14(11):4436-7556 2022a.
Yizeng Han, Yifan Pu, Zihang Lai, Chaofei Wang, Shiji Song, Junfen Cao, Wenhui Huang, Chao Deng, and Gao Huang. Learning to weight samples for dynamic early-exiting networks. In Proc. European Conf. on Computer Vision (ECCV), 2022b.
Yizeng Han, Dongchen Han, Zeyu Liu, Yulin Wang, Xuran Pan, Yifan Pu, Chao Deng, Junlan Feng. Shiji Song, and Gao Huang. Dynamic perceiver for efficient visual recognition. In Proc, IEEE Int. Conf. Computer Vision (ICCV), 2023b.
Andrew Howard, Mark Sandler, Bo Chen, Weijun Wang. Liang-Chieh Chen, Mingxing Tan, Grace Chu, Vijay Vasudevan, Yukun Zhu, Ruoming Pang, Hartwig Adam, and Quoc Le. Searching for mobilenetv3. In Proc, IEEE Int. Conf. Computer Vision (ICCV), 2019.
Gao Huang, Danlu Chen, Tianhong Li, Felix Wu, Laurens van der Maaten, and Kilian Weinberger. Multi-scale dense networks for resource efficient image classification. In Proc. Int. Conf. on Learning Reprexentations (ICLR), 2018.
Fatih Ilhan, Ling Liu, Ka-Ho Chow, Wenqi Wei, Yanzhao Wu, Myungjin Lee, Ramana Kompella, Hugo Latapie, and Gaowen Liu. EEnet: Learning to early exit for adaptive inference. arXiv preprint: arXiv 2301.07099, 2023.
Borui Jiang and Yadong Mu. Russian doll network: Learning nested networks for sample-adaptive dynamic inference. In Proc. Int. Conf. on Computer Vision Workshops (ICCVW), 2021.
Polina Karpikova, Ekaterina Radionova, Anastasia Yaschenko, Andrei Spiridonov, Leonid Kostyushko, Riccardo Fabbricatore, and Aleksei Ivakhnenko. FIANCEE: Faster inference of adversarial networks via conditional early exits. In Proc. IEEE/CVF Conf. on Computer Vision and Pattern Recognit. (CVPR), 2023.
Yigitcan Kaya, Sanghyun Hong, and Tudor Dumitras. Shallow-deep networks: Understanding and mitigating network overthinking. In Proc. Int. Conf. Machine Learning. (ICML), 2019.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
Assaf Lahiany and Yehudit Aperstein. PTEENet: Post-trained early-exit neural networks augmen- tation for inference cost optimization. IEEE Access, 10:69680-69687, 2022.
Stefanos Laskaridis, Alexandros Kouris, and Nicholas D. Lane. Adaptive inference through early- exit networks: Design, challenges and directions. In Proc. Int. Workshop on Embedded and Mobile Deep Learning, 2021.
Hao Li. Hong Zhang, Xiaojuan Qi, Ruigang Yang, and Gao Huang. Improved techniques for train- ing adaptive deep networks. In Proc. IEEF/CVF International Conference on Computer Vision (ICCV), 2019.
Xiangjie Li, Chenfei Lou, Yuchi Chen, Zhengping Zhu, Yingtao Shen, Yehan Ma, and An Zou. Pre- dictive exit: Prediction of fine-grained early exits for computation- and energy-efficient inference. In Proc. AAAI Conf. on Artificial Intell., 2023.
Shiwei Liu, Tianlong Chen, Zahra Atashgahi, Xiaohan Chen, Ghada Sokar, Elena Mocanu, Mykola Pechenizkiy, Zhangyang Wang, and Decebal Constantin Mocanu. Deep ensembling with over- head for either training or testing: The all-round blessings of dynamic sparsity. In Proc. Int. Conf. Learning Representations (ICI.), 2022.
11

---

Published as a conference paper at ICLR 2024
Liheng Ma, Chen Lin, Derek Lim, Adriana Romero-Soriana Puneet K. Dokania, Mark Coates, Philip Torr, and Ser-Nam Lim. Graph inductive biases in transformers without message passing. In Proc. Int. Conf. Machine Learning. (ICML), 2023.
Matthias Minderer, Josip Diolonga, Rob Romijnders, Frances Hubis, Xiaohua Zhai, Neil Houlsby, Dustin Tran, and Mario Lucic. Revisiting the calibration of modern neural networks. In Proc. Adv. Neural Info. Process. Syst. (NeurIPS), 2021.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y. Ng. Reading digits in natural images with unsupervised feature learning. In Proc. NIPS Workshop on Deep Learning and Unsupervised Feature Learning 2011,2011.
Nikolaos Passalis, Jenni Raitoharju, Anastasios Tefas, and Moncef Gabbouj. Efficient adaptive infer- ence for deep convolutional neural networks using hierarchical early exits. Pattern Recognition, 105, 2020.
Svetlana Pavlitska, Christian Hubschneider, Lukas Struppek, and J. Marius Zollner. Sparsely-gen gated mixture-of-expert layers for CNN interpretability. In Proc. Int. Joint Conf. on Neural Networks (IJCNN) 2022.
Lorena Qendro, Alexander Camphell, Pietro Lio, and Cecilia Mascolo. Early exit ensembles for uncertainty quantification. In Proc. of Machine Learning for Health, 2021.
Xuwei Qian, Renlong Hang, and Qingshan Liu. Rex: An efficient approach to reducing memory cost in image classification. In Proc. AAAI Conf. on Artificial Intell., 2022.
Doyen Sahoo, Quang Pham, Jing Lu, and Steven C. H. Hoi. Online deep learning: Learning deep neural networks on the fly. In Proc. Int. Joint Conf. Artificial Intell. (IJCAI), 2018
Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Ilié, Daniel Hesslow, Roman Castagné. Alexandra Sasha Luccioni. François Yvon, Matthias Gallé, Jonathan Tow, Alexan- der M. Rush, Stella Biderman, Albert Webson, Pawan Sasanka Ammanamanchi, Thomas Wang, Benoît Sagot, Niklas Muennighoff, Albert Villanova del Moral, Olatunji Ruwase, Rachel Baw- den, Stas Bekman, Angelina McMillan Major,, Iz Beltagy, Huu Nguyen, Lucile Saulnier, Samson Tan, Pedro Ortiz Suarez, Victor Sanh, Hugo Laurençon, Yacine Jernite, Julien Launay, Margaret Mitchell, and Colin Raffel. Bloom: A por--pareettr open-access multilingual language model. arXiv preprint: arXiv 2211.05100, 2023.
Simone Scardapane, Danilo Comminiello, Michele Scarpiniti, Enzo Baccarelli, and Aurelio Uncini. Differentiable branching in deep networks for fast inference. In ICASSP 2020 2020 IEEE In- ternational Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 4167-4171, 2020.
Tal Schuster, Adam Fisch, Jai Gupta, Mostafa Dehghani, Dara Bahri, Vinh Q. Tran, Yi Tay, and Donald Metzler. Confident adaptive language modeling. In Proc. Adv. Neural Info. Process. Syst. (NeurIPS), 2022.
J. Seon, J. Hwang, Mun, and B. Han. Stop or forward: Dynamic layer skipping for efficient action recognition. In Proc. IEEE/CVF Winter Conf. on Applications of Computer Vision (WACV), 2023.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In Proc. Int. Conf. on Learning Representations (ICLR), 2017.
Yi Sun, Jian Li, and Xin Xu. Meta-GF: Training dynamic-depth neural networks harmoniously. In Proc. European Conf. on Computer Vision (ECCV), 2022.
Surat Teerapittayanon, Bradley McDanel, and H.T. Kung. Branchynet: Fast inference via early exiting from deep neural networks. In Proc. Int. Conf. on Pattern Recognition (ICPR), 2016.
Yulin Wang, Kangchen Lv, Rui Huang. Shiji Song, Le Yang, and Gao Huang. Glance and focus: dynamic approach to reducing spatial redundancy in image classification. In Proc. Ady. Neural Info. Process. Syst. (NeurIPS), volume 33, 2020.
12

---

Published as a conference paper at ICLR 2024
Yulin Wang. Rui Huang, Shiji Song, Zeyi Huang, and Gao Huang. Not all images worth 16x16 words: Dynamic transformers for efficient image recognition. In Proc. Adv. Neural Info. Process. Syst. (NeurIPS), 2021.
Maciej Wolezyk, Bartosz Wójeik, Klaudia Balazy, Igor T. Podolak, Jacek Tabor, Marek Smieja, and Tomasz Trzcinski. Zero time waste: Recycling predictions in early exit neural networks. In Proc. Adv. Neural Info. Process. Syst. (NeurIPS), 2021.
Zuxuan Wu, Tushar Nagarajan, Abhishek Kumar, Steven Rennie, Larry Davis, Kristen Grau- man, and Rogerio Feris. BlockDrop: Dynamic inference paths in residual networks. In Proc. IEEE/CVF Conf. on Computer Vision and Puttern Recognit. (CVPR), 2018.
Le Yang, Yizeng Han, Xi Chen, Shiji Song, Jifeng Dai, and Gao Huang. Resolution adaptive net- works for efficient inference. In Proc. IEEE/CVF Conf. on Computer Vision and Pattern Recognit. (CVPR), 2020.
Haichao Yu, Haoxiang Li, Gang Hua, Gao Huang, and Humphrey Shi. Boosted dynamic neural networks. In Proc. AAAI Conf. on Artificial Intell., volume 37, 2023.
Li Yuan, Yunpeng Chen, Tao Wang, Weihao Yu, Yujun Shi, Zihang Jiang, Francis EH Tay, Jiashi Feng, and Shuicheng Yan. Tokens-to-token ViT: Training vision transformers from scratch on imagenet. In Proc, IEEE Int. Conf. on Computer Vision (ICCV), 2021.
Wangchunshu Zhou, Canwen Xu, Tao Ge, Julian McAuley, Ke Xu, and Furu Wei. BERT loses patience: Fast and robust inference with early exit. Proc. Adv. Neural Info. Process. Syst. (NeurIPS), volume 33, 2020.
Jinguo Zhu, Xizhou Zhu, Wenhai Wang. Xiaohua Wang, Hongsheng Li, Xiaogang Wang, and Jifeng Dai. Uni-Perceiver-MoC: Learning sparse generalist models with conditional MoEs. In Proc. Adv. Neural Info. Process. Syst. (NeurIPS), 2022.
APPENDIX
EXPERIMENT DETAILS
9.1.1 DATASETS
The CIFARIO and CIFAR100 (Krizhevsky) 2009) datasets both consist of 60,000 32 32 coloured images. CIFARIO has 10 classes and CIFARTO0 has 100 classes. We follow 75%-8.3%-116.6
vrain-alliation-tess split. The validation set is used for hyperparameter tuning and carly-stopping. ImageNet (Deng al. 2009) consists of 1.2 million training images spanning 1,000 classes. We reserved 50,000 these images to be used as validation set and used another 50,000 images as a test set. Following the data transform of Han al. (2023b), the images cropped to be of size 224x224. We also use the cropped-digit SVHN (Netzer et al. (011) dataset which consists of 99,299 32 coloured images of house numbers where the target is the digit in the center. We use 68.6%- 5%-26.2% train-validation-tee. split. The CIFAR100-LT (long-tail) dataset consists of the regular
CIFAR 100 images but with imbalanced classes. We use imbalance factor if 100 conuac))
(lastoomaal We follow the pre-processing procedure from (Yuan et al., 202T). For datasets, we use a training batch size of 64. We present 95% confidence intervals on the means of our results. These are generated using a bootstrap procedure over 10 trials derived by partitioning the test set into 10 subsets.
9.1.2 EXPERIMENT PARAMETERS
We use the Adam optimizer with learning rate of 0.01 with weight decay of 5e-4. We use batch size of 64 for CIFARIO, CIFAR100 and CIFAR 10OLT, and a batch size of 256 for SVHN. We train until convergence using early stopping with maximum of 15 epochs (E) for CIFAR10 and SVHN, 20 for CIFARI0OLT and 30 for CIFARIOOLA For the number warmup epochs E, we perform hyperparameter search over the range (1, (1,..., F./2).
13

---

L2W-DEN Published as a conference paper at ICLR 2024
9.1.3 TRAINING TIME AND MEMORY USAGE
We report the average training time on a GPUx machine and number parameters used for both architectures in Table T Overall, our approach takes a little longer to train and has negligible number of additional parameters. Once trained, the baselines do have the advantage that heuristic adjustment of the gating threshold can achieve different IC/accuracy trade-offs.
Table 1:
Memory usage and training time comparison. *We note that training time is to obtain a single IC/accuracy trade-off model (one point in curve).
JEI-DNN
Boostnet
T2T-7 (CIFARIO)
time* (min)
13 15.42K
13 15.42K
num. parameters
15.45K
T2T-14 (CIFAR100)
time* (min)
22
24 500.50K
num. parameters
500.56K
500.50K
9.1.4
WARM-UP LOSS
We now provide more thorough exposition of the warm-up phase which precedes the bi-level optimization phase. As mentioned, during warm-up we train all IMs in parallel on all samples. This done to ensure that IMs have an acceptable accuracy and can potentially be chosen as exit points when we start training the gates. Earlier IMs have lower representational capacity and thus require more training. To that end, we increase the learning rate of earlier IMs using a simple position-base scaling. The warm-up loss for the 1th IM can thus be written as:
1 N N
C warmup
(1.. (x:)))
(17)
where Iro is set to 0.01 as a hyperparameter.
9.2 INFERENCE COST AND EXPECTED CALIBRATION ERROR (ECE)
Following BoostNet YYu et al. 2023) and L2W-DEN (Han et al. 2022b). we measure the inference cost in terms of multiply-ade operations (Mul-Add). We use Meta's FAIR fvcore library for computing Mul-Adds. We obtain results that match the Mul-Adds reported in Yuan et al. 2021 for the T2T-ViT architecture. Since our training procedure directly relies on the inference cost (see equation
Inference cost:
6) we measure the number of Mul-Adds at each exit before training, accounting for the layer of the backbone network, denoted h', the IM f/ and the gate 9 This value is treated ignore any possible decrease in Mul-Adds that might arise from sparsity in certain samples and thus use worst-case scenario value that depends only on the network architecture and the input size. We can thus define the inference cost incurred for exiting at layer I:
function of / that constant across all samples exiting at layer l. That is,
IC Mul-Add(h*)) + Mul-Add( )+ Mul-Add (g.
(18)
Tables and 3 show the number of Mul-Adds at each exit of the base T2T-ViT-7 architecture on CIFARIO and the base T2T-ViT-14 architecture CIFAR100 respectively well as the added Mul-Adds incurred from augmenting the backbone network with exit branches (IMs and gates) at every layer. We note that overall, this leads to almost insignificant inference cost increase of less than 0.003% for the former and 0.01 % for the latter. As a result of our very lightweight gate design, optimal branch placement is unnecessary.
14

---

Published as a conference paper at ICLR 2024
Table 2: Comparison of the number of Mul-Adds per layer between the base VTT-7 backbone
and JEI-DNN. The middle rows show the net increase due to augmenting with IMs and gates.
IC'
IC
IC
IC*
IC
IC
IC
T2T V
414.3M
538.2M
662.1M
786M
909.9M
1034M
1158M
Cum. I from IMs
+2570
+5140
+7710
+10280
+12850
+ +15420
+15420
Cum. I from GMs
+96
+192
+288
+384
+480
+576
+576
JEI-DNN
414.3M
538.2M
662.1M
786M
909.9M
1034M
1158M (t 0.003%
Comparison of the number of Mul-Adds per layer between the base T2T- SET-VIT-14 backbone Table 3:
and JEI-DNN. The middle rows show the net increase due to augmenting with IMs and gates.
IC'
IC²
IC
IC
IC
IC
IC
T2T VIT 14
626,2M
947.7M
1269M
1590M
1912M
2233M
2555M
Cum. from IMs
+38.5K
+77K
+115.5K
+154K
+192.5K
+231K
+269.5K
Cum. from GMs JEI-DNN
+906
+1812
+2718
+3624
+4530
+5436
+6342
626.3M
947.7M
1269M
1591M
1912M
2233M
2555M
IC
IC
IC
IC
IC
IC
IC
T2T ViT 14
2876M
3198M
3519M
3841M
4162M
4483M
480SM
Cum. from IMs
+308K
+346.5K
+385K
+423.5K
+462K
+500.5K
+500.5K
Cum. from GMs
+7248
+8154
+9060
+9.96K
+10.87K
+11.78K
+1 1.78K
JEI-DNN
2876M
3198M
3519M
3841M
4162M
4483M
4805M (1 0.01% )
Expected Calibration Error (ECE): To estimate the ECE, we approximate the probability that y is correct given a predicted pr We evaluate the accuracy over a set of samples ACC(Q)
[ Eie [ = yi] and compare it to the average predicted probability P(Q) O Eieg Pi The
predicted probability values are sorted (denoting vi)) the sorting index such that P(() <Pu()=1)
and split into B bins following (Guo al. 2017) to approximate the ECE:
B Qal N
ECE
where Qs = ()))
(19)
9.3 CONFORMAL INTERVALS
In this section, present other ways obtain conformal intervals in the EEDN setting and show that they are all approximately equivalent. One simple approach is to consider all samples as single group, ignoring where they exited. We then the entire validation set to select single thresh- old. An alternative is to derive a different threshold for each IM. To do this, we partition the valida- tion set into multiple smaller validation sets, one associated with each exit IM: V V2, U!..,... UVV Fach val' contains only the samples that exit at the 1-th IM. In determining conformal interval threshold for the 1-th IM, we then use only the samples in V and the predicted probabilities for that specific IM: (x,). For each sample, we calculate score 1 (* Given the set of scores { and a coverage value - a, the conformal interval threshold value is determined by estimating the 1 percentile of (s:).
We construct four methods by combining different choices:
General:
We ignore the gating and use the entire validation set to obtain single v value. We use this common threshold to generate the conformal sets for every exit IM. The score value for given sample is derived from the predicted probabilities of its selected exit gate.
IMs: : We compute one threshold value per gate, but we use the entire validation set V. The scores used compute the threshold for the 1-th IM are derived from P/ (x;). In this computation, we do not take any account of where sample exits.
Strict gating A strict version of what we used in our experiments. We compute one threshold value per IM, using only the samples in the validation set associated with that IM. We derive the
15

---

Published as a conference paper at ICLR 2024
score function using (x,)). If no samples in the validation set assigned to specific IM, ie; [D]| 0, we set 7/ V
0.
The version used to produce the experimental results in the puper. We use the strict gating approach, but if [V*| 20, we set 77 TV,e, where 7V.a is derived using the General gating approach above.
Gated:
strict gating
strict gating
strict gating
gated
gated
gated
IMs
IMs
IMs
general
general
general
3
3
1
0.8 Mul-Add (11.58x109)
Mul-Add (11.58x109)
Mul-Add (11.58x109)
Figure 7: Inefficiency Mul-Add on CIFAR10 (t2t-7) of the conformal intervals obtained from the different methods. (Lefft JEI-DNN, Middle Boostnet and Right I.2W-DEN. The x-axis are scaled by the inference cost of the full model Mul-Add (IC4).
Boostnet
L2W-DEN
JEI-DNN
5
2
1
0.4
0.5
0.6
0.7
0.8
Mul-Add (11.58x10*)
Figure 8: The conformal intervals obtained from the different methods, for CIFARIO 12t-vit-7.
Figure shows that all methods are approximately equivalent for each baseline. Figure shows that JEI-DNN achieves the smallest inefficiency for CIFAR10 12t-vit-7, irrespective of the method chosen to construct confidence intervals. Similar results are obtained for the other datasets. To ensure that the comparison is meaningful, we report the smallest average inefficiency with an em- pirical coverage that exceeds the targeted coverage 1 - a. To do this, for a targeted a = 0.05, we compute the conformal intervals from multiple thresholds, obtained by slowly decreasing the target a: (0050. 70=0.045 ....) until the empirical coverage is above our target: 1- > 1 a. We then report the inefficiency at that empirical coverage.
9.4 SURROGATE LOSS
this section, we provide the connection between the surrogate problem we introduced for training the ós and the initial problem defined in equation T2
16

---

Published as a conference paper at ICLR 2024
Here, we don't 't assume that the solutions are unique will denote the set of solutions by क. The initial problem is then given by:
N I C
Φ arg min
Po(G= (0)
(20)
N
We replace by the surrogate problem:
N L Σ
T -h arg min
Σ
L CF
(t.
(c))))
(21)
@
N
1-1--1
[ffj
where 4
(22)
1ifj21"
arg C%-(
(23)
To motivate this surrogate problem, we show that if there exists a solution @' $5 s.t. the associated (G lsx;)s can be equal to each P that minimizes:
L Σ E1
L. P!, with 1-1
2P!-1, Vi [N].
P. arg min
(24)
then prove that o' will also be an optimal solution for our surrogate problem (ø' (b"r))
That is, of also minimizes:
1 N ΣΣ Σ c
c
(25)
For given sample i, in order to find the probability vector P that minimizes Ele Co P. need allocate all the probability mass of P to the gate 1 with the lowest cost. That is, the solution to Eqn 24 given by:
1 if 1=1" arg miny C Cor(o) 0 o.w.
Pd
Vie [N]
(26)
From our assumption, we assume that solution of satisfies (G l[x;) pa We then show that such a o' would then also be solution to our surrogate problem.
In our surrogate problem, the associated gates d'gr have to satisfy:
if = arg min, Cz-(o), 0i < arg aig Cz-(o)
(27)
Proof:
For the trivial case where 1" 1, g/ 1, then Py = 1[x;) = by definition of (c')
and Par (G l|x;) (o
0 for
Suppose 1* 1, then 9's- = 0 for all 1 < l-. This results in (G l(x,) min Σj_1 0) 0. Once reach // 1, this results in Par ParGG= ('):)) min Σ 0)
and Pay (G l[x:) min (c 1) This concludes the proof.
Hence, since the optimized gates g'y given by equation 27 which corresponds to the solution of surrogate tasks, have shown that those d' would be shared as the solution of both problems.
9.5 ALGORITHM
17

---

Published as a conference paper at ICLR 2024
Algorithm Training
1: Input: num. of gates I., cost importance À parameter, num of epochs E, num of warm-up
epochs W bilevel switch count bisswitch
2: STATE-WARMUP
3: for we in {1...., WE} do
4: for X4, in batch do
5: for in (1....,... do
6:
Cb (x) = CE{(p()))))) (optimize each IM)
7:
Update the 0 parameters
8: end for
9: end for
10: end for
11: STATE=GATE
12: for in {WE, E) do
13: for X4, M/ in batch do
14:
STATE = CLASSIFIER then
15:
Get probability of exiting each gate 1: Pa(G 1|x;)
Optimize the weighted loss L3nh El- Pa = (y,,
else if STATE GATE then
18:
Compute the inference cost at each gate I: IC!
narm
19:
Compute the accuracy cost at each gate I: CC (y:,
20:
Store the total cost of that gate: Co-(o)) (C AIC morm
Compute the best gate [* arg min, Ch (o)
Optimize Cout Σt BC (1" end
(equation 16
23:
24: if batch idx bi switch then
25:
switch from STATE GATE to CLASSIFIER or from STATE CLASSIFIER to GATE end if
26:
27:
end for
28:
end for
18

---

Published as a conference paper at ICLR 2024
9.6 ABLATION STUDY
To verify that our joint training procedure is needed, we conduct an ablation study on the CIFARIO dataset where we adapt our algorithm to the two decoupled settings:
1. Learnable GMs, fixed IMs:
warm-up procedure, but avoid bi-level optimiza- tion, i.e., after warm-up we only optimize the gates' parameters, Φ. We follow Algorithm exactly, but never enter the CTASSTTFT F.R state. We train for the same total number of epochs.
We the
2. Threshold GMs: We also adapt algorithm the other setting of threshold GMs. We remove the learnable gates from the architecture. We first train the IMs on all the data and then use the threshold gating algorithm used by Yu al. (2023) and Han al. (022b).
The importance of joint training is illustrated Figure 9 The Learnable GMs, fixed IMs approach
is less stable but slightly outperforms the Threshold GMs approach.
Learnable GMs, fixed IMs
Threshold
95
JEI-DNN
90
Accuracy
80
75
70
0.50
0.55
0.60
0.65
0.70 0.75
0.80
0.85
Mul-Add (11.58x109)
Figure 9: Ablation study on CIFAR10 (12t-7). Jointly training the GMs and IMs is crucial; both decoupled approaches result in a performance decline.
9.7 OPTIMIZATION OF Lin
The complete gradient to optimize the inner gradient of our bi-level formulation given in equa-
tion
14
Ohin
1
N L OCCE (yi, (a,,)) Po(G=1[x)) (f), Ca(e)
(28)
08
N
08
00
our methodology, we mention that in practice, we ignore the right term of the gradient, and approximate the gradient using the term that corresponds to the gradient of weighted cross entropy
loss:
1 N OC (y). (xx:)) (G(x))) + (m,(m((, Coro
ΣΣ
(29)
00
N
1-1
00
00
acin
N
OC (yt. (xx))) (G=1|x,).
(30)
00
N
00
We made that decision following our observation that both approaches yield comparable outcomes, but opting for the complete gradient leads to slower convergence. We denote the JEI-DNN version Figure shows the test accuracy of both
using the full gradient equation 29 JEI-DNN-FULD models during training; we can see that both reach the same test accuracy, but that JEI-DNN reaches it considerably faster.
19

---

Published as a conference paper at ICLR 2024
100
test accuracy
method
JEI-DNN
JEI-DNN-FUI
JEI-DNN
JEI-DNN-Full
0.0
2.5
5.0
7.5
10.0
12.5
training epoch
Figure 10: Test accuracy during training both our proposed JEI-DNN and FULL..
9.8 ADDITIONAL RESULTS
We include additional performance versus inference cost curves in Figure I1] additional calibration versus inference cost curves in Figure and inefficiency versus inference cost curves in Figure 13 The same trends as presented in the main paper can be observed.
0.95
0.9
0.90
0.8
0.8
Accuracy
0.7
0.7
0.80
0.75
0.6
Boostnet L2W-DEN
0.5
0.5
JEI-DNN
0.4
0.65
0.4
80% end acc
0.4
0.7 Mul-Add (11.58x10*)
0.2
0.4
0.6
0.8
0.9
Mul-Add (48.06x10*)
Mul-Add (11.58x10*)
Figure 11: Accuracy Mul-Add of (Left) CIFARIO (t2t-7), Middle CIFAR10 (t2t 14) and Right CIFAR10OLT (121-7). The x-axes are scaled by the inference cost of the full model Mul-Add (IC²).
9.9 UNCERTAINTY STATISTICS
In this section we provide details concerning the uncertainty statistics that are used as input to our gate functions Io- We represent the gate input as m(0.x;) Lomaz (x))) (xoo mox)),)r(x))]) where:
p.aar (x:) max(p:) is the maximum predicted probability; .I(k)
h(x:)= Σx Po. log P
((K)) is the entropy the predictions;
.1((k) (
(x,)((x;) = Σn Po. log Po.: where Pa.i (22,01)
(.())
are scaled predictions (temp. = 0.5);
..(Kk) mar(x:) = pl. for kj arg max; Po. P Po.i and maxx P0. This the difference between the most confident predictions.
((k)
L((K))
20

---

Published as a conference paper at ICLR 2024
0.8
0.6
0.6
Boostnet
ECE
L2W-DEN
0.4
JEI-DNN
0.2
0.2
0.0
0.4
0.5
0.6
0.7
0.8
0.2
0.4
0,6
0.8
0.4
0.5
0,6
0.7
0.8
Mul-Add (11.58x109)
Mul-Add (48.06x109
Mul-Add (11.58x109)
Figure 12: Calibration Mul-Add of (Leff)) CIFARIO (12t-7), Middle CIFAR10 (12t-14) and Right CIFARI0OLT (t2t-7). The x-axes are scaled by the inference cost of the full model Mul-Add (1C%)
6
Boostnet L2W-DEN JEI-DNN
60
40
20
2
0.4
0.5
0.6
0.7
0.8
0.2
0.4
0.6
0.8
0.4
0.5
0.6
0.7
0.8
Mul-Add (11.58x10*)
Mul-Add (48.06x10")
Mul-Add (11.58x10")
Figure 13: Inefficiency vs Mul-Add of (Left) CIFARIO (12t-7), Middle CIFARIO (12t-14) and Right CIFAR100LT (12t-7). The x-axes are scaled by the inference cost of the full model Mul-Add (IC*).
9.10 RELATED WORK
In this appendix, we provide a more extensive and detailed discussion of related work, touching on some topies that are more tangentially related.
Architectures with threshold GMs The majority of the literature on early-exit neural networks has primarily focused on designing architectures that are amenable to efficient and effective early- exit. Generally, these approaches, the decision to exit at an earlier branch taken by comparing confidence scores extracted from the IMs to predetermined thresholds.
The idea is first introduced by BranchyNet (Terrapittayanon et al. 2016), which augments network with carly-exit branches and trains the augmented model end-to-end, by summing a weighted loss of all IMs. Training the whole network end-to-end can bring complications as earlier IMs nega- tively impact the performance of later IMs. With this in mind, MSDNet (Huang al. 2018) dense connections to reduce the influence of early IMs on later layers. The performance of earlier IMs is also improved by using representations of varying scales train with coarser features. In RANNet (Yang et al. 2020), spatial redundancy of the input is exploited, with higher resolutions used only for the more difficult samples. SDN (Kaya al.) 2019). off-the-shelf networks are augmented with intermediate IMs. Two training schemes are explored: freezing the backbone while training only the early exit branches or training the full network. The experimental results suggest that training the full network works better since the frozen backbone is optimized for the last layer. The interference of the early-exit branches on later layers is addressed via weighted loss scheme with weights that depend on the position of the exit. More recently, Han et al.) (2023b) proposed the Dynamic Perceiver, demonstrating that any CNN-based state-of fa-te-of-thean of-tthearr vision backbone can aug- mented with parallel computation path that is optimized for early classification. The introduction of such parallel branch can improve carly-exiting while reducing the interference of carly IMs on the performance of later IMs.
21

---

Published as a conference paper at ICLR 2024
Table 4: Classification of EEDN problem settings and approaches. A tick mark indicates that method is capable of providing this feature. Fixed backbone: The method can be applied to off-the- shelf backbones without re-training, i.c., does not require tailored EEDN backbone. Learnable The gates can be trained. Most methods use threshold GMs. Train IMs: The IMs are trained. Adaptive IM training: The training adapted to different IMs, either by introducing weights or training with different subsets of the training data. Joint Training: Gate and IM parameters are jointly learned. * Indicates the baselines we include in our work.
GMs:
Model
Fixed backbone
Learnable GMs
Train IMs
Adaptive IM Training
Joint Training
BranchyNot Tocrapittaynnon et 2016
MSDNet Hhang et
2018
SDN (Kaya ml.
2019
Wang al.
2020
Wang 2021
Chen et 2023g
DynPerceiver Han
2023b
ReX: Qian
2022
IMTA (1.
2019
Meta-GFF (unn et al. 2022
Bolukbasi ct al
2017
Predictive Exit (Li et ml
202.
EPNet Dai al.
2020
PTEENet (Lahiuny Aperstein
2022
EEN (Ilhan et al. 2093
Bleck-Dependent (Han cl.
2023n
"BoostedNet Yo al. 2023
HHn
2022b
DiffBranch (carddaaaee et
2020
JEI-DNN
Another line of work has investigated ways of improving the training and performance of EE ar- chitectures by tackling the issue of conflicting gradient updates on the shared backbone parameters. These arise because the training procedure is attempting to simultaneously optimize multiple dif- ferent exits. In IMTA (Li et al., 2019). the gradients computed when optimizing the intermediate IMs are rescaled based their position in the network to ensure they have bounded scale, ducing their variance. Additionally, knowledge distillation is employed to help earlier IMs learn from deeper ones. Meta-GF [Sun al. (2022) also addresses conflicting gradient updates, but rather than using position-based rescaling, the scaling factors are learnt via a meta-lcarning formulation. The motivation is that in large over-parameterized models, certain parameters of the backbone are more relevant to some exits than others. Therefore, Meta-GF takes into account each parameter's influence on the various exits before scaling the respective gradients.
Other research efforts have focused on customizing EEDNs to specific type of backbone archi- tecture, Wu al. (2018) introduced an EEDN architecture customized for Residual networks, while Chen et al. (2018) did the same for CNNs.
Fundamentally, all these studies rely on confidence score thresholds that are set as hyperparameters to decide whether to early-exit. Most recent EEDN advances fall into this category of a EEDN-
tailored architecture paired with either threshold GMs (Wang et al. 2220 Han et al. 2023b Wang
et al. 2021 Chen et al.] 2023b) or other types of fixed gating mechanisms (Qian et al. 2022
While most of these studies concentrate on image classification tasks, FEDNS have also seen ex- pansion into other domains, including language models (Zhou et al. 2020 Schuster et al.) (002),, as well as video frame classification (Seon et al. 2023). Grraves] (2016) introduces an RNN architec- ture where the model is trained to predict and to learn the number of computational steps required, which can be framed as an end-to-end EEDN model.
EEDN research efforts have also explored how fully leverage the additional computations per- formed by the previous gates and inference modules (Jiang & Mu (022 Wolezyk et al. 2021]. For example, Wolezyk al. 2002): Liu al. (2022) and Passalis al. (0020) all combine the outputs
22

---

Published as a conference paper at ICLR 2024
of the IMs using ensemble methods to enhance performance. Similarly, Qendro et al. (2021) also employ ensembling of IMs but target the uncertainty characterization of the overall network rather than inference cost. Of particular interest among the ensembling techniques is the geometric ensem- bling procedure of Wolczyk et al. (02T). When combining the predictions of the IMs, the output probability at IM is raised to the power wt, where wi is learnable parameter. This allows the model to learn to place less emphasis on the predictions from the weaker IMs at lower layers during the fusion process. In our proposal, we do not form ensembles, but we do learn the exit probabilities. If we interpret each IM softmax output as a probability distribution across labels, then the overall la- bel distribution for any given sample a mixture of the individual IM label distributions, where the mixture weights are the leamned exit probabilities. Thus, although our learning procedure does not perform ensembling for any specific classification decision, its probabilistic output can viewed as being drawn from an ensemble of classifiers.
Learnable GM; pre-trained and frozen backbone and IMs Differing from the dedicated archi- tecture and threshold GMs methods described above, another class of EEDN treats the IMs and backbone architecture a fixed, frozen component, and foeuses training the GMs. In some cases, the pretraining of the IMs is included in the methodology. In (Bolukbasi et al.) 20177) the learning is performed in two separate phases. First, the IMs are trained with the full training set and then the IMs are frozen and the gates are trained. The gates are trained sequentially starting from the deepest layer to the shallowest, allowing the gate training objective to be defined recursively. In EP- Net (Dai et al. 2020) the task of training the gates after freezing the IMs is formulated as a Markov decision process. The methodology involves designing a substantial, trainable neural network of multiple layers for the gate mechanism. This was critized by following works as being wasteful, im- posing an unacceptable inference overhead. PTEENet (Lahiany & Aperstein 2022) parameterizes the gates and directly learns the exit decisions using sigmoid neural networks. EENet (Ilhan et al.) 2023) adopts a different strategy by formulating the carly-cxit policy as a multi-objective optimiza- tion problem. Karpikova et al. (2023) design learnable gates, and specifically focus on confidence metrics tailored to images. Finally, there are tangential works such as (Li et al. 2023) that consider the related problem of optimizing gate placement for energy efficiency.
All of these models learn gating functions after having trained and frozen the IMs at every branch. While this two-phase procedure simplifies training. it ignores the coupling between gates and IMs the gating decisions change the distributions of samples presented to the IMs, implying that parameters should be re-learned.
Addressing the train-test mismatch Two recent works have proposed architecture-agnosti train- ing procedures to address the aforementioned train-test mismatch. In BoostNet (Yu et al., (023)). the training procedure is inspired by gradient boosting. Prediction reweighting is employed so that deeper IMs are trained with prediction reweighting to induce an emphasis on samples that are in- correctly classified by the earlier IMs. Han et al. (2022b) cast the problem as a meta-learning task where a weight prediction network ( (WPN) trained jointly with the early-exit network to predict the most likely exit of sample. The predictions of the WPN are used to assign weights to the losses of the different IMs. Consequently, easy samples should have a reduced influence deeper IMs while harder samples contribute more to their training. This architecture is related to ours but the training procedures differ in multiple important ways. Most importantly, we do not rely directly on thresh- olding uncertainty metrics for early exiting. We show that the uncertainty metrics are not a good proxy for accuracy for earlier IMs due to miscalibration. In practice, the inclusion of the weight prediction network does not close the train-test gap because it is trained using poorly calibrated la- bels. By jointly training the IMs and the GMs via bi-level optimization, our approach provides more reliable uncertainty characterization.
Scardapane et al. (2200) propose differentiable branching, an approach where gates and IMs are jointly learned. The differentiable branching uses straightforward modeling approach for the gate exit probabilities and optimizes all parameters jointly; as a result, training is highly unstable.
Chen et al. (2020) propose fully trainable architecture featuring a learnable policy component that is responsible for determining the number of layers used for inference. Although Chen et al. (020) do not target inference efficiency, their framework shares similarities with ours. The gating mechanism is modelled as a product of independent exiting events. However, the inference cost is not incorporated in the training objective. The gate exit distribution is controlled by an entropy
23

---

Published as a conference paper at ICLR 2024
regularization regularization; result, the procedure in (Chen al. 2020 cannot be used to construct an effective EEDN.
This concludes review of the most relevant EEDN contributions. We include categorization of the works in Table 4 As is evident from the table, our approach is the only one that (i) can be applied any layer-based off-the-shelf backbone network, without re-training the backbone; (ii) has learnable GMs and IMs; and (iii) jointly trains both the GMs and the IMs.
Sparsely-gated mixture-of-eepprts (MoE) A closely related field to our problem setting is the Mixture of Experts (MoE) framework (Shazeer et al. 2017 Fedus et al. 2022 Du et al. 2022). In this framework, the network also employs gating mechanism to make adaptive decisions based on the input it receives. However, in the MoE setting. the gating decisions are made across layers to incorporates multiple experts or weights (Shazeer et al. 2017: Du et al. 2022, Fedus et al. 2022). The goal is not to tailor inference time based the input complexity but rather enable the use of increasingly larger models by maintaining constant inference costs with sparse activation strategies. In this field, it is also considered important problem to bridge the gap between the gating decisions and the inference modules (Gupta et al. 2022 Zhu et al. 2022). The interpretability potential of such architectures also highlighted in (Pavlitska et al. 2022)
Online learning In online learning with deep neural networks the aim is to jointly learn represen- tations and a suitable architecture. However, given that the data is streamed serially during training, the optimal depth of the network may change more data is accumulated. To address this, Sahoo jet al. (2018) proposed to incorporate learnable GMs that produce convex combination of IMs to construct a prediction that is weighted average of all IM outputs and the final layer output. The GM training employs the hedge backpropagation algorithm proposed by Freund & Schapire 1997).
24

---

Published as a conference paper at ICLR 2024
9.11 JEI-DNN ON OTHER BACKBONE ARCHITECTURES
As previously mentioned, the JEI-DNN training procedure can be applied many different types of architecture. To demonstrate this, we use our procedure on the GRIT (Ma et al., (023) backbone. GRIT is transformer-basee model for graph data that achieves state-of of-teoof-the-rr performance on the most common graph benchmark datasets proposed by Dwivedi et al. (2022). As before, we augment the backbone with IMs and lightweight gates. Figure 14 shows that out procedure leads to better graph dataset. This demonstrates the flexibility of our training procedure, showing that it can achieve efficient inference variety of deep architectures addressing different learning domains.
accuracy-cost combination than Boostnet when trained on Super-Resolution-NNNSS
Accuracy full netwerk
Gated accuracy
90
85
GRIT JEIDNN
70
0.55
0.60
0.65
0.70
0.75
fraction of cost of full network
Figure 14: Comparison of the JEI-DNN training scheme and Boostnet on the GRIT backbone on the super-resolution graph MNIST dataset.
9.12 ADDITIONAL DATASET: IMAGENET
0.80
80% cost of full network
0.75
Accuracy
0.70
0.65
0.60
L2W-DEN
80% end acc
JEI-DNN
0.55
Mul-Add (45.7x10f)
Figure 15: Accuracy Mul-Add of ImageNet (t2t-14).
We provide the additional performance versus inference cost curves for the ImageNet dataset (Deng et al. (see) Appendix 9.1.] for details) using the vision transformer T2T-ViT-14. Since this model was trained on ImageNet, no transfer learning is needed. We follow the general procedure as described in paper; augment the backbone with gates and intermediate IMs at every layer. We generate results by varying the inference cost parameter over the range 0.01 to 10, Rather than sample F.
p( F. xx) to decide whether exit, we use the deterministic decision P(F') 0.5.
25

---

Published as a conference paper at ICLR 2024
In Figure 15 compare method with the most competitive baseline (L2W-DEN) for T2T- ViT- 14. In the main paper, we also provide results for ImageNet but using the specialized early-exit architecture Dynamic Perceiver a backbone in Figure (right).
We observe that by using the JEI-DNN training strategy we achieve an improved trade-off be- tween accuracy and inference cost for T2T-ViT-IE compared to L2W-DEN, the best-performing Sehhitectureagnnoste baseline training strategy. For the Dynamic Perceiver backbone, the incorpo- ration of JEI-DNN training lead reduction of inference cost of approximately 1-2 percent for operating points in the range from 50-65% of the inference cost of the original network. In this re- gion, the preserved aceuracy ranges from 85-95% of the accuracy the full network. The Dynamic Perceiver backbone has only four IMs, with three located relatively close to the final layer in the architecture. This imposes some limits on the flexibility of the exit decisions and means that more challenging for the JEI-DNN training to achieve an improvement for operating points where the accuracy close to the original accuracy.
The decrease in accuracy of JEI-DNN is steeper when the backbone is general architecture such as T2T-ViT-14 as opposed to an architecture tailored for early-exit (Dynamic Perceiver). JEI-DNN with T2T-VIT-14 sees its performance decrease to around 75% of its full performance with 80% of its total inference cost, whereas JEI-DNN with the DynPerceiver can maintain 92% of its full performance with only 60% of its total cost.
9.13 ADDITIONAL. BASELINE
We provide preliminary results for an additional baseline, DiffBranch (Scardapane et al., 2020)
in Figure 16 which significantly outperform. The performance of (Scardapane al.] (020)
abnormally ly low, and those results persisted even after our best effort at hyperparameter tuning and different possible implementations. Unfortunately, as the code is provided and implementation details were missing from the paper, we are unable to confirm that our implementation and therefore our results are correct. For these results, to ensure fair comparison, we employ the same warm-up procedure as used by our method. Without it, the IMs of DiffBranch (Scardapane et al. (2220) were not able reach reasonable accuracy. We also used the same parameterization of the IMs and GMs as JEI-DNN to adapt the architecture to the 12t-vit-7.
80
Accuracy
JEI-DNN
DiffBranching
60
40
0.2
0.4
0.6
Mul-Add (4.3x10%)
Figure 16: Accuracy vs Mul- SVNH (12t-7) of the DiftBranch baseline. The x-axes are scaled by the inference cost of the full model Mul-Add (IC4).
9.14 ADDITIONAL ACCURACY DECOMPOSITION ANALYSIS
In this section, we offer additional insights into how our algorithm behaves across a wide range of inference costs, illustrated through supplementary decomposition figures. In our main paper, Figure 3 corresponds to À 1 for the CIFAR 100 experiment with (2t-vit-14 (Figure 2). Here, we provide in Figure 17 the same decomposition for À = which imposes a greater penalty on computation and leads earlier exiting behaviour. Figure 18 shows the decomposition for À 0.5, which imposes a small penalty for additional computation, and leads to more samples exiting at later layers.
26

---

Published as a conference paper at ICLR 2024
We that for all À values, obtain similar trend that shifts from early gates to later gates based on the inference cost. For the larger value of À in Figure 17 JEI-DNN starts exiting samples at gate 3 and heavily uses the gates 4, 5, 6 and 7. For smaller penalty computation in Figure T8 JEI-DNN starts to exit samples much later, at gate 5, and mainly uses the gates 7, 8, 9, 10 and TT. These results highlight how the parameter can be used to control the trade-off between computation during inference and accuracy.
1.0
1.0
b)
Accuracy on exited points (D))
Accuracy all points (D)
0.8
0.8
method
05 Ptae
0.6
Boostnet
L2W-DEN
JEI-DNN
0.4
0.4
|D'| (IMs usage)
5
0.2
Boostnet accuracy L2W-DEN accuracy
0.2
10
15
JEI-DNN accuracy
20
0.0
0.0
6
8
10
12
14
6
8
10
12
14
IMs (/)
Figure 17: Decomposition of the contributions of the IMs to the final accuracies (depicted by dotted lines) for CIFAR100, with the inference cost À 3. Accuracy of each IM, f evaluated only
their exited samples (D'). b) Accuracy of each IM, fy. on the full dataset D. The size of circle is proportional to the number of samples exited for trial.
1.0
0000
1.0
b)
Accuracy exited points (D)
0.8
Accuracy all points (D)
0.8
method
0.6
0.6
Boostnet
L2W-DEN
JEI-DNN
0.4
0.4
| (IMs usage)
5
0.2
Boostnet accuracy
0.2
408
10
L2W-DEN accuracy
15
JEI-DNN accuracy
0.0
0.0
25
6
8
10
12
14
6
8
10
12
14
IMs (/)
IMs (/)
Figure 18: Decomposition of the contributions of the IMs to the final accuracies (depicted by dotted lines) for CIFAR100, with the inference cost 0.5.
9.15 OPTIMIZING THE GATES: BINARY CLASSIFICATION SURROGATE TASKS
In the section Optimizing the gates of our paper, we describe surrogate learning task that use to train our gating mechanism. In short, we train each of our g modules binary classifier, and we set the target e based on the the relative cost of each gate Cy,, We first find the index of the optimal gate = arg min, Cj. and set its target to tt = 1, then we set the targets of all the previous gates to 0, and the targets of all the subsequent gates 1.
27

---

Published as a conference paper at ICLR 2024
The intuition behind this is that, if sample misses the exit at the optimal IM the next best exit is the following one, at 1" +11. Underpinning this is an assumption that the accuracy of the IMs tends to be monotonically increasing for all samples the layer increases. The assumption would be wrong if, the best exit at I", but then the performance of the IM at the next gate is very poor and the IM should be avoided. (The desired target pattern 6 gates might then look like 0,0, 1, 1,0, 1, for example).
An alternative strategy is to only set gate's target there is later IM that reaches lower cost (considering both accuracy and inference overhead). So instead of our strategy, that we call ExitSubsequent:
== l
1 1> l*
(31)
0
<<l
we could use this other strategy, that we call ExitIfMin:
1=t²
e' 1> Cge C2.. Vjsl
(32)
0 O.W.
We verify that the ExitlfMin strategy yields similar results to our ExitSubsequent strateg by running experiment the SVHN dataset with 12t-vit-7. The results be viewed in Figure 19 where we can see that both approaches give the same results. This confirms that our assumption holds.
100
90.0
Accuracy
87.5
Accuracy
85.0
60
JEI-DNN
82.5
JEI-DNN with exit min
40
Samples previoulsy correct samples
80.0
0.5 0.6
Mul-Add (4.3x10')
Gate
Figure 19: JEI-DNN vs JEI-DNN with Exi- UtMin for SVHN dataset with 12t-vit-7 back- bone. The x-axis scaled by the by the full model inference cost, Mul-Add(IC2).
Figure 20: Accuracy each IM for all the samples vs the accuracy at each IM for the samples that correctly classified by the previous IM. SVHN dataset with 12t-vit-7.
We can also examine the progression of the accuracy at each IM for the samples that were correctly classified by previous IM. If see trend such that once a sample is correctly classified by an IM 1, it is also correctly classified by all subsequent IMs. then it suggests that it is indeed best to exit the sample as soon as possible. In Figure 20 we compare the conditional accuracy (conditioned on correct classification at the previous IM) with the overall accuracy. We see that the conditional accuracy is higher, and above 90% for all IMs, indicating that our strategy sensible.
We do note, however, that these results are for an architecture where the intermediate layers not trained with the goal of providing good features for classification. intermediate layers trained to do Kaya et al. (2019) observe that there can be phenomenon of "overthinking" for such networks, where an early IM can make correct classification, but later IMs make errors for the same sample. For such backbone, the ExitliMin strategy may be preferable.
shallow-deep network, the
9.16 END-TO-END TRAINING OF JEI-DNN WITH DEDICATED ARCHITECTURE
Our method is designed to be used on a fixed backbone architecture. This includes enhancing architectures that tailored for early-exit, have shown in main paper in Observation 5.
28

---

Published as a conference paper at ICLR 2024
We go step further and consider fully end-to-end trainable model that trains the backbone, the IMs, and the GMs. Instead of keeping the backbone fixed during training, we can alternate between our JEI-DNN procedure and training the backbone. To present results, again select as the backbone architecture the Dynamic Perceiver Han et al. (2023)) with MobileNetV3.
82
full network
81
80
Geted accuracy
79-
60% cost full netwwork
78-
77
76-
X
Dyn. Perceiver + JEI-DNN
Dyn. Perceiver
75
Dyn. Perceiver + JEI-DNN E2E
74
8.5
9.0
Mul-Adds
1e7
Figure 21: Comparison of Dynamic Perceiver, Dynamic Perceiver trained with the JEI-DNN pro- cedure (Dyn. Perceiver JEI-DNN) and Dynamic Perceiver trained end-to-end with the JEI-DNN procedure (Dyn. Perceiver JEI-DNN E2E) on CIFAR100.
The end-to-end combination (Dynamic Perceiver+JEI-DNNY includes learnable gates, learnable IMs, and learnable backbone. We compare it to the results we have presented with the Perceiver, Perceiver+JEI-DNN, which are both trained for 310 epochs. We train our Perc+JEI-DNN end-to- end by alternating between training JEI-DNN for 2 epochs, then the Perceiver for epoch. Figure shows that this produces similar results to Perceiver+JEI-DNN. but there is small improvement. For example, if require 80 percent accuracy (a reduction of only 1 percent from the accuracy of the full network), the end-to-end training means that we can reduce the required mul-adds by further 2.6%. These promising results indicate that further research into fully end-to-end training architecture has the potential to be a fruitful research direction.
29


Reviews:
[{"review_id": "0pY7Wq1l26", "reviewer": "Reviewer_KZAV", "summary": "This paper presents an early-exit dynamic neural network architecture, JEI-DNN, that augments a backbone with lightweight, trainable gates and intermediate classifier that are jointly optimized. The gates are trained through a surrogate binary classification task, focusing on the optimization for assigning the most cost-effective early-exit classifier for each input. The results are presented on the cifar and svhn datasets for the T2T-ViT backbone.", "questions": "- The assumption that a sample exited at a gate at layer $l$ should also exit from any late stage gate seems against the prevalent view of $overthinking$:\n\"Overthinking is computationally wasteful, and it can also be destructive when, by the final layer, a correct prediction changes into a misclassification.\" by Kaya et al. ICML 2019.\nHow did you make this assumption?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- Analysing the conformal intervals for the early exiting classifiers is interesting and it's promising to see the proposed method yields better uncertainty characterisation.\n- The ablation study in 8.6 clearly shows the merits of learnable GMs versus other design choices.\n- Figure 2 a) and b) very clearly show the early exiting patterns for the method and competing approaches.", "weaknesses": "- It is unclear what the major contribution of the paper is. The use of learned early exiting gates and intermediate classifiers, trained jointly with the backbone is common practice for many early exiting architectures. After all the approximations described in 4.1, the loss term for the gating modules turnes into the common practice of summing the losses of independent binary early-exiting classifiers. The main task loss is similarly aligned with prior work from Han et al. (2022b).\n\n- Constructing surrogate binary targets for the learned gates is also common, e.g. FrameExit by Ghodrati et al. CVPR 2021. \n\n- Figure 2a) The fact that samples do not exit at all from the first few early exiting layers in JEI-DNN is puzzling and I am wondering if it is due to the choice of specific hyperparameters that prevent early exiting from these layers. As can be seen, the accuracy of the IM at layer 5 & 6 is far higher that the overall performance. It is conceivable that a good accuracy higher than the green dashed line is still achievable by the earlier classifiers at least for a proportion of samples. Is there any hyperparamer that could potentially give more control into the exiting pattens of JEI-DNN?\n\n\n- All the results in the paper are limited to the T2T-ViT backbone. The authors should show the efficacy of their method for a larger variety of backbones, preferably to models that are established in the early exiting literature such as MSDNet, DenseNet, etc.\n\n- Comparison of the accuracies among T2T-Vit and MSDNet architectures in Fig 5 seem unfair. Most of the gain in accuracy comes because of the more powerful transformer-based backbone and not because of the efficacy of the early exiting approaches. In fact, MSDNet and RANet show more robust early exiting results compared to the results reported for T2T. E.g. MSDNet retains the original accuracy of the model after almost 50% compression. The performance of the proposed JEI-DNN approach in comparison drops very rapidly even with 25% compression.\n\n- The method is only evaluated on three small-scale datasets: CIFAR10 & 100, SVHN. The authors should consider expanding the evaluation to ImageNet.", "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": "6wULHk9xbn", "reviewer": "Reviewer_CAtU", "summary": "This submission introduces JEI-DNN, an early-exit model approach that can be applied on top of off-the-shelf backbones for image classification. The proposed methodology appends light-weight classifiers and trainable gates along the depth of a frozen backbone model and jointly optimises them with a custom and insightfully designed loss function. As a result, a better speed-accuracy trade-off is provided compared to several baselines considering both learnable and non-parametric exit policies, while offering improved prediction uncertainty.", "questions": "1. What is the role of the second (1-sum) term in Eq.8 ? Potentially, its use could be validated through an ablation experiment.\n\n2. How would the method perform on end-to-end trainable early-exit models? Would the training of the gates affect the convergence of the overall model? Should the method be considerably adjusted to be applicable in this setting? \n\n3. Is the proposed approach equally effective on the most commonly used ImageNet-1K benchmark?\n\nMinor Comments:\nSec3: The term \"IC-Only training\" is widely used in the community (including the cited survey paper) to denote \"Intermediate-classifier only training\", rather \"than inference cost-only training\" as stated in the manuscript\n\nPost-Rebuttal Edit: Following the clarification on Q1 and additional results included in the updated version of the manuscript for Q2 and Q3, I am increasing my score from 6(=BA) to 8(=A).", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "-The submission studies a very interesting problem, focusing on the emerging inference setting of input-dependent computation via early-exiting. \n\n-The manuscript offers a beautiful formulation of the examined task and provided solution, that can benefit the community as a principled definition of early-exit models. Additionally, the discussion of the manuscripts findings presents useful insights that can guide practitioners and inspire future research in the direction of EE models.\n\n-Most design choices in the adopted solution are well motivated and backed by practical insights.\n\n-The presented results indicate that the proposed method is effective and achieves a superior speed-accuracy trade-off to a wide set of baselines approaches.", "weaknesses": "-A few technical aspects of the paper remain unclear. Specifically, the use of the min operator between two terms in Eq.8 is not fully justified in the manuscript, nor experimentally verified by a relevant ablation. \n\n-The proposed methodology is only evaluated on frozen-backbone CNNs. Although this a practical sub-category of EE-models, further evaluation on end-to-end EE models (jointly training the backbone and exits), which are shown to achieve superior speed-accuracy trade-off would increase the contribution of the paper.\n\n-Although the proposed approach is compared with a wide set of baseline, ImageNet experiments are omitted, rendering the presented results less convincing. \n\nShould the above limitations be satisfactorily addressed, I am inclined to increase my score.", "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": "Lo47I5A08b", "reviewer": "Reviewer_1vfw", "summary": "This paper proposes to jointly learn the backbone parameters and the exiting strategies in a dynamic multi-exit model. Intermediate classifiers are constructed in a pre-trained backbone, and exiting gates are built as a binary classification head. The joint learning problem is formulated as a bi-level optimization problem. Experiments on several small datasets demonstrate that 1) a better trade-off between accuracy and efficiency is achieved; and 2) a better estimation of uncertainty (calibration of classification confidence) is achieved.", "questions": "See weaknesses.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The studied problem is of great interest;\n\n2. The motivation is clear;\n\n3. The proposed method is technically sound;\n\n4. The literature review is comprehensive;\n\n5. The experiments show the effectiveness in both accuracy-efficiency tradeoff and calibrated confidence.", "weaknesses": "1. **Lack of experiment on more advanced architectures**. Why did the authors select a T2T-ViT to construct the early-exiting model? Straightforwardly, the joint learning procedure can be directly applied in mature multi-exit models, such as the compared MSDNet, RANet, and the cited Dynamic Perceiver.  To my understanding, Dynamic Perceiver is the most recent work in this field, achieving SOTA performance in dynamic early exiting.\n    \n2. **Overclaiming contributions**. Based on the above point, the contribution of this paper should not include a new architecture. It is recommended to summarize the contribution as the joint learning procedure only. I believe this would already be significant enough if the learning method is shown effective on more SOTA architectures and on the ImageNet dataset.\n\n3. **Lack of experiments on ImageNet**. To my understanding, the learning approach does not need to be applied in a downstream task on the toy small-scale datasets. Experiments on ImageNet would be more convincing.\n\n4. **Presentation**. It is recommended to use figures to clearly show the motivation and the method pipeline.", "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": "jA5NQAqqo9", "reviewer": "Reviewer_smKZ", "summary": "The paper describes a novel mechanism to add trainable early exits (EEs) to a pre-trained neural network. It addresses a common concern of the majority of previous works, where the train procedure and the test procedure of the EEs were mismatched (e.g., jointly training all the early exits during training while thresholding their outputs during test). \n\nIn the proposed framework, each EE is associated to a probabilistic gate, and everything is trained jointly with a bilevel optimization problem, where the outer problem is defined over the gates and the inner problem over the EEs. The system is trained to also minimize execution time by training the gates to select the first exit having good accuracy for each pattern.\n\nThey also propose a novel conformal prediction (CP) strategy for EE networks, where different thresholds are selected for each exit (since each exit will observe a different subset of values at inference times).", "questions": "Apart from the topics discussed above, I have a general question about the bilevel task. Can you clarify why this is needed? There are techniques to train probabilistic blocks (e.g., reparameterization, score function estimators), which would allow to train all blocks simultaneously, or I am missing something? Also, what is the overall training time of your model compared to the alternatives?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The topic of the paper is important, since the method described can be used to reduce inference time (hence, power consumption) of any pre-trained model. The paper is generally well written and easy to read, although some additional visualization could be useful (see below). Results are relatively comprehensive, although I am concerned by the transfer learning from a large dataset (ImageNet) to a smaller, fundamentally similar dataset (CIFAR, also see below).", "weaknesses": "I find it hard to fully understand the method before reading through the entire paper. In particular: (a) the abstract has no mention of the paper's contributions; (b) the \"Contributions\" also does not explain how the method works; (c) there is no visual schema of the model; (d) some important details are described very late, for example \"Gate design\" is part of the experimental section but it is a crucial design decision. Overall it is understandable, but I think some minor reorganization and some additional visual descriptions (maybe a shortened pseudo-code?) can significantly help the reader.\n\nOn the novelty, I think a few methods could be added and discussed more. For example, the zero-time waste model (cited in the paper) has a geometric ensemble of the early exits which is very similar to their gating probability, if I understand correctly. As another example, differentiable branching (https://ieeexplore.ieee.org/document/9054209) considers a trainable exit strategy that also combines the EEs and that can be trained end-to-end with the rest of the network. Some of these methods can be also added to the experimental comparison.\n\nConcerning the last point, I am curious about transfer learning from a large, generalist dataset (ImageNet) to a small dataset which is basically in the same domain (CIFAR-10). Would it not make more sense to test directly on ImageNet? Improvements would be directly comparable with the original pre-trained model.\n\nEDIT AFTER REBUTTAL: the authors have added several baselines and improved the presentation of the method. I have raised my score to 8.", "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)