Jointly-Learned Exit and Inference for a Dynamic Neural Network :
  JEI-DNN
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 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.
Florence Regol, Joud Chataoui, Mark Coates
# Jointly-Learned Exit and Inference for a Dynamic Neural Network : JEI-DNN

###### 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 _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.

## 1 Introduction

The dominant approach to improve machine learning models is to develop larger networks that can handle every potential sample. As a 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 a large network to be successfully executed. Nonetheless, most traditional neural network (NN) models have a fixed processing pipeline. This means that every sample, simple or complex, is processed the same way.

To tackle this inefficiency, dynamic networks have been introduced (see (Han et al., 2022) 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 samples to exit at shallower layers. Compared to conventional neural networks, EEDNs incorporate two additional components: 1) Intermediate Inference Modules (IMs) receive a sample's representation at their respective network depth and generate predictions; 2) Gating Mechanisms (GMs) decide which intermediate inference module should be employed to derive the final prediction.

The vast majority of EEDNs employ a simple threshold-based gating mechanism. The _threshold GM_ applies thresholds to confidence metrics obtained from the inference modules. To integrate more sophisticated GMs, the gating mechanism can be treated as a post-training add-on component. Given a 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 in both strategies, the IMs and GMs are decoupled during training; the IMs are trained on the full dataset even though they will
only infer a subset of the data at inference time. This inevitably creates a mismatch between the train and test distributions.

The reliance on threshold GMs also makes it more difficult to obtain confidence levels for predictions. The information required for this determination has already been "consumed" in the gating decision; moreover, applying a threshold introduces an overconfidence bias. In a setting where a 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 a 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 architecture and training approach 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. We show empirically that the approach leads to a better overall inference/performance trade-off. The benefits are threefold: 1) we close the training gap between IMs and GMs, which leads to better performance; 2) the architecture produces reliable uncertainty characterization in the form of conformal intervals and well-calibrated predicted probabilities; and 3) the generality of our approach facilitates integration into any state-of-the-art architecture, thus allowing it to significantly outperform dedicated EEDN backbone architectures.

## 2 Related work

In this section we briefly summarize the most relevant early-exit dynamic networks literature. Appendix 8.10 contains a more in-depth discussion as well as a larger coverage of peripheral fields, including online learning with adaptive architectures, and sparsely activated Mixtures of Experts.

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., 2016) augments the backbone network with early-exit branches and is trained end-to-end using a weighted loss. Shallow-deep networks (SDNs) additionally quantify the confusion of the network by comparing the outputs of different IMs (Kaya et al., 2019). Efficient end-to-end training is difficult because earlier branches influence the performance of later ones. Subsequent EEDN-tailored architectures, such as MSDNet (Huang et al., 2018) and RANet (Yang et al., 2020), alleviate this problem via dense connections or increased resolution. Training can be improved by reducing contil gradient updates on the backbone parameters using position-based rescaling (Li et al., 2019) or meta-learning (Sun et al., 2022). Reintegrating computations from earlier layers can boost performance via ensembling (Wolczyk et al., 2021; Liu et al., 2022; Passalis et al., 2020). In spite of their differences, all of these works rely on confidence score thresholds to decide whether to early-exit. Many recent advances in the EEDN field also rely on simple threshold GMs (Wang et al., 2020; Han et al., 2023; Wang et al., 2021; Chen et al., 2023b). 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 problem 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 a Markov decision process. PTEENet (Lahiany & Aperstein, 2022) employs a much simpler gate architecture. Karpikova et al. (2023) design learnable gates with a specific 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 is inherent to these two tasks.

Addressing the train-test mismatch:Ignoring the influence of the GM 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 (Yu et al., 2023) proposes an architecture-agnostic training procedure inspired by gradient boosting to close the train-test mismatch. Deeper IMs are trained with an emphasis on samples that are not correctly classified by earlier IMs. Han et al. (2022b) cast the problem as a meta-learning task where a _weight prediction network (WPN)_ is trained jointly with the inference modules. For a given sample, the VPN 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 the threshold GMs.

## 3 Problem Setting

We consider a classification problem with a training dataset \(\mathcal{D}=\{\mathbf{x}_{i},y_{i}\}_{i=1}^{N}\) where \(\mathbf{x}_{i}\in\mathbb{R}^{D}\) denotes the input and \(y_{i}\in\mathcal{K}\) its corresponding classification target (\(\mathcal{K}=\{1,\ldots,K\}\)). We are given a fixed network \(NN:\mathbb{R}^{D}\rightarrow\mathcal{K}\) pretrained on the same task that can be decomposed into a composition of \(L\) layers: \(NN(\mathbf{x})=\sigma\circ h^{L}\circ h^{L-1}\circ\cdots\circ h^{1}(\mathbf{x})\), where \(\sigma\) denotes the softmax operator. Let \(\mathbf{z}^{l}\) be the \(l\)-th intermediate representation of \(NN\) such that \(\mathbf{z}^{l}=h(\mathbf{z}^{l-1})\) with \(\mathbf{z}^{0}=\mathbf{x}\).

Our goal is to augment this fixed network (the backbone) with additional trainable components such that we can obtain a final prediction \(\hat{y}_{i}\) with its associated predicted probability vector \(\hat{\mathbf{p}}_{i}\) at a reduced inference cost. This setting is categorized as **inference cost-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 \(\hat{y}\)); (ii) uncertainty-based (the calibration error and the inefficiency score of the conformal intervals obtained from \(\hat{\mathbf{p}}_{i}\)); and (iii) efficiency-based (the inference cost to obtain \(\hat{y}\)).

**Inference cost:** In line with BoostedNet (Yu et al., 2023) and L2W-DEN (Han et al., 2022) the inference cost \(IC_{\hat{y}_{i}}\) for a single prediction \(\hat{y}_{i}\) represents the number of multiply-add (Mul-Add) operations to obtain \(\hat{y}_{i}\). We approximate the expected inference cost \(E[IC]\) as:

\[E[IC]\approx IC=\frac{1}{N}\sum_{i=1}^{N}IC_{\hat{y}_{i}}\;. \tag{1}\]

**Uncertainty metrics:** We assess calibration of the model via the expected calibration error (ECE) of the probability of the predicted class (the maximum probability of \(\hat{\mathbf{p}}_{i}\) : \(\hat{p}_{i}^{\max}=\max(\hat{\mathbf{p}}_{i})\)). To estimate the ECE, we follow the approach in (Guo et al., 2017). More details concerning computation of the inference cost and ECE are provided in Appendix 8.2.

To obtain conformal intervals, we use the typical conformal score \(\mathbf{s}_{i}\triangleq\mathbf{1}-\hat{\mathbf{p}}_{i}\) and compute a conformal threshold \(\tau^{V,\alpha}\) from a validation set \(\mathcal{V}\) so that we can guarantee a **coverage** of \(1-\alpha\). (See Appendix 8.3 for details). A conformal interval \(\hat{\mathcal{C}}_{i}\) for a sample \(\mathbf{x}_{i}\) is then constructed by thresholding the conformal score \(\mathbf{s}_{i}:\mathcal{C}_{i}=\{k;\mathbf{s}_{i}^{(k)}<\tau^{V,\alpha}\}\), where \(\mathbf{s}^{(k)}\) denotes the \(k\)-th element of \(\mathbf{s}\).

We measure the performance of our predicted conformal intervals \(\hat{\mathcal{C}}_{i}\) with two metrics: **empirical coverage**\(1-\hat{\alpha}\), which computes the empirical probability that the ground truth is contained in the interval, and **inefficiency**\(|\bar{\mathcal{C}}|\), which is the average cardinality of the set intervals:

\[|\bar{\mathcal{C}}|\triangleq\frac{1}{N}\sum_{i=1}^{N}|\hat{\mathcal{C}}_{i}| \,,\quad 1-\hat{\alpha}\triangleq\frac{1}{N}\sum_{i=1}^{N}\mathbb{1}[y_{i} \in\hat{\mathcal{C}}_{i}]\,. \tag{2}\]

A well-behaved conformal prediction model has small inefficiency \(|\bar{\mathcal{C}}|\) (intervals are small on average) while maintaining an empirical coverage close to the desired one: \(1-\alpha\approx 1-\hat{\alpha}\).

## 4 Methodology: JEI-DNN

We introduce our method named: Jointly-Learned Exit and Inference for a Dynamic Neural Network (JEI-DNN). Our model consists of augmenting a pretrained and fixed network by attaching \(L\) classifiers and \(L\) gates to the intermediate layers of the network architecture. We refer to each of these classifiers as an inference module (IM). The \(l^{th}\) IM, \(f^{l}_{\theta_{l}}(\cdot)\), uses the \(l^{th}\) representation \(\mathbf{z}^{l}_{i}\) of the input \(\mathbf{x}_{i}\) to model class probabilities \(\hat{\mathbf{p}}^{l}_{\theta}(\mathbf{x}_{i})=f^{l}_{\theta}(\mathbf{z}^{l}_{ i})\) and to obtain a prediction \(\hat{y}^{l}_{i}=\arg\max_{k}\hat{\mathbf{p}}^{l,(k)}_{\theta}(\mathbf{x}_{i})\). We set the last IM to be the original classifier of the pretrained network: \(f^{L}=h^{L}\). The output of the \(l^{th}\) gate, \(e^{l}\in\{0,1\}\), is modelled by a random variable \(E^{l}|X\) that specifies whether the \(l\)-th exit gate is open. If a sample reaches the \(l\)-th gate and it is open, then it exits via the \(l^{th}\) IM.

The assumption is that some samples are "easy enough" to be classified by the early classifiers \(f^{l}_{\theta}(\mathbf{z}^{l})\). Since the last classifier should be able to classify the entire dataset, the last gate is always
set to one, i.e., \(e^{L}=1\). Denoting the indicator function by \(\mathbb{I}[\cdot]\), the output of our model \(\hat{y}\) is:

\[\hat{y}(\mathbf{x})=\hat{y}^{l}(\mathbf{x})\text{ where }l=\operatorname*{arg\,min}_{l \in[L]}\mathbb{I}[e^{l}(\mathbf{x})=1],\quad e^{l}(\mathbf{x})\sim p(E^{l}|X =\mathbf{x}). \tag{3}\]

We can view each IM \(f^{l}_{\theta}\) as having to solve a subtask over a subset of the data \(\mathcal{D}^{l}\subset\mathcal{D}\), where the partitioning is determined by the gates, and \(\mathcal{D}^{j}\cap\mathcal{D}^{k}=\varnothing\) for \(j\neq k\) and \(\cup_{l\in[L]}\mathcal{D}^{l}=\mathcal{D}\). Since the \(E^{l}\) are random, we can also define the distribution \(P(G)\) of exiting at a gate \(G\in[L]\). This is equal to the probability that the sample does not exit at any previous gate multiplied by the probability that the sample exits at gate \(l\):

\[P(G=l|X=\mathbf{x})=P(E^{l}|X=\mathbf{x})\prod_{j=1}^{l-1}\left(1-P(E^{j}|X= \mathbf{x})\right). \tag{4}\]

The first \(l\) layers of the network as well as the first \(l\) IMs and gates must be evaluated in order to exit at the \(l^{th}\) layer. Assuming input samples of fixed size, the inference cost \(IC_{\hat{y}^{l}_{i}}\) associated with obtaining a prediction from the \(l^{th}\) IM can be approximated as a constant \(IC^{l}\) that only depends on \(l\). That is, for all predictions \(\hat{y}^{l}_{i}\) obtained at level \(l\) we have \(IC_{\hat{y}^{l}_{i}}=IC^{1}\). Appendix 8.2 provides more details concerning how \(IC^{l}\) is computed for our model. During training we use a normalized version \(\text{IC}^{l}_{\text{norm}}=\frac{\text{IC}^{l}}{\text{IC}^{l}}\) that represents the relative cost incurred by exiting at \(l\) compared to going through the full network (i.e., all \(L\) layers).

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 \(\mathcal{L}^{CE}(Y,\hat{\mathbf{p}}^{G|X}_{\theta}(X))\) added to the associated inference cost \(IC^{G|X}\) scaled by a parameter \(\lambda\) that controls the importance of the inference cost:

\[\mathcal{L} =\mathbb{E}_{Y,X}\mathbb{E}_{G|X}[\mathcal{L}^{CE}(Y,\hat{ \mathbf{p}}^{G|X}_{\theta}(X))+\lambda IC^{G|X}_{\text{norm}}], \tag{5}\] \[\mathcal{L} \approx\frac{1}{N}\sum_{i=1}^{N}\Big{(}\sum_{l=1}^{L}(\mathcal{ L}^{CE}(y_{i},\hat{\mathbf{p}}^{l}_{\theta}(\mathbf{x}_{i}))+\lambda IC^{l}_{ \text{norm}})P(G=l|\mathbf{x}_{i})\Big{)}\,. \tag{6}\]

Our goal is to simultaneously learn the parameters \(\theta\) of the IMs and choose \(P(G=l|\mathbf{x}_{i})\) in order to minimize the expected loss. We target \(P(G=l|\mathbf{x}_{i})\) rather than \(P(E^{j}|\mathbf{x}_{i})\) for \(j=1,\dots,l\) because directly learning the \(P(E^{j})\) leads to optimization issues. The product of probabilities across multiple layers can quickly vanish or saturate to 1.

Modeling \(P(G|\mathbf{x}_{i})\).In practice, we only need to evaluate \(P(G=l|\mathbf{x}_{i})\) if the dynamic evaluation has reached that point, i.e., everything up to \(f^{l}_{\theta}(\mathbf{x}^{l}_{i})\) has been evaluated and is accessible. Therefore, our parameterization of \(P(G=l|\mathbf{x}_{i})\) can include any intermediate values calculated by the base architecture, the inference modules, and the gates up to and including layer \(l\). We denote this aggregation of information as \(\mathbf{c}^{sl}_{i}\).

This however also implies that we cannot directly model the multiclass distribution \(P(G|\mathbf{x}_{i})\) with the traditional softmax approach since it would require \(f^{l}_{\theta}(\mathbf{x}^{l}_{i})\) to be evaluated for all \(l\). This would defeat the purpose of early-exiting. Instead, we model each \(P(G=l|\mathbf{x}_{i})\) sequentially, starting from the first layer \(l=1\), with a learnable variable \(g^{l}_{\phi}(\mathbf{c}^{sl}_{i})\in[0,1]\). We set \(P(G=l|\mathbf{x}_{i})\) to the minimum of the learnable variable \(g^{l}_{\phi}(\mathbf{c}^{sl}_{i})\in[0,1]\) and \(1-\sum_{j=1}^{l-1}g^{j}_{\phi}(\mathbf{c}^{\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\, \,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\, \,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\, \,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\, \,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\, \,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\, \,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\, \,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\, \,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\, \,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\, \,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\, \,
Any parameterization of the IMs \(f^{l}_{\theta}\) and gates function \(g^{l}_{\phi}(\mathbf{c}^{Sj}_{i})\) 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^{l}_{\theta}\) and \(g^{l}_{\phi}(\mathbf{c}^{\leq j}_{i})\) are parameterized as simple one-layer neural networks and uncertainty metrics are extracted from \(\mathbf{c}^{\leq j}_{i}\) to serve as input to \(g^{l}_{\phi}(\cdot)\). A detailed description of our modules is included in Section 5.1. Appendix 8.2 includes the inference cost of the gates and IMs to 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 1.12% of the Mul-Adds of the backbone.)

Uncertainty predictionNext we describe how we use the predicted probabilities \(\hat{\mathbf{p}}^{l}_{\theta}(\mathbf{x}_{i})\) to form conformal intervals. The conformal intervals are derived from a conformal threshold \(\tau^{\mathcal{V},\alpha}\) that is computed using a validation set \(\mathcal{V}\). However, since we can view the IMs as solving different subtasks over different subsets of the data \(\mathcal{D}^{l}\), we should compute a different conformal threshold for each subtask. Hence, we form one validation set per gate, \(\mathcal{V}^{l}\), where \(\mathcal{V}^{l}=\{i;l\sim P_{\phi}(G|\mathbf{x}_{i})\,,i\in\mathcal{V}\}\), and compute a conformal threshold from that subset: \(\tau^{\mathcal{V}^{l},\alpha}\). If there are too few samples in the validation set associated with gate \(l\), we set \(\tau^{\mathcal{V}^{l},\alpha}\) to a general threshold \(\tau^{\mathcal{V},\alpha}\) computed using the entire validation set. The conformal intervals are then constructed as follows:

\[\mathcal{C}_{i}=\{k;1-\hat{\mathbf{p}}^{l,(k)}_{\theta}(\mathbf{x}_{i})<\tau^{ \mathcal{V}^{l},\alpha}\,,l\sim P(G|\mathbf{x}_{i})\} \tag{10}\]

There are many other ways to construct the sets to compute the threshold in the EDNN setting. In Appendix 8.3, we show that a variety of methods produce similar results.

### Optimization

Directly optimizing the loss in equation 9 is challenging because of the \(\min\) operator:

\[\phi^{*},\theta^{*}= \operatorname*{arg\,min}_{\phi,\theta}\frac{1}{N}\sum_{i=1}^{N} \sum_{l=1}^{L}\bigl{(}\mathcal{L}^{CE}(y_{i},\hat{\mathbf{p}}^{l}_{\theta}( \mathbf{x}_{i}))+\lambda IC^{l}_{\text{norm}}\bigr{)}\min\left(g^{l}_{\phi}( \mathbf{c}^{\leq l}_{i}),1-\sum_{j=1}^{l-1}g^{j}_{\phi}(\mathbf{c}^{\leq j}_{i })\right)\!. \tag{11}\]

However, there is a distinction between the parameters \(\theta\) and \(\phi\) which makes this loss a good candidate for a bi-level optimization formulation. (See (Chen et al., 2023a) for a survey.) Hence, using \(C^{l}_{\theta(\phi)}=\mathcal{L}^{CE}(y_{i},\hat{\mathbf{p}}^{l}_{\theta}( \mathbf{x}_{i}))+\lambda IC^{l}_{\text{norm}}\), we can rewrite equation 11 as follows:

\[\phi^{*}=\operatorname*{arg\,min}_{\phi}\mathcal{L}^{out}\triangleq \operatorname*{arg\,min}_{\phi}\frac{1}{N}\sum_{i=1}^{N}\sum_{l=1}^{L}C^{l}_{ \theta^{*}(\phi)}P_{\phi}(G=l|\mathbf{x}_{i}), \tag{12}\] \[s.t.\quad\theta^{*}(\phi)=\operatorname*{arg\,min}_{\theta} \mathcal{L}^{in}\triangleq\operatorname*{arg\,min}_{\theta}\frac{1}{N}\sum_{ i=1}^{N}\sum_{l=1}^{L}\bigl{(}\mathcal{L}^{CE}(y_{i},\hat{\mathbf{p}}^{l}_{ \theta}(\mathbf{x}_{i})+\lambda IC^{l}_{\text{norm}})P_{\phi}(G=l|\mathbf{x}_{ i}). \tag{13}\]

Since the gate probabilities \(P_{\phi}(G=l|\mathbf{x}_{i})\) take as input the aggregation of any intermediate values that were previously calculated \(\mathbf{c}^{\leq l}_{i}\), it could be dependent on the \(\theta\). To make that dependence explicit in the calculation of the gradient of \(\theta\), we denote \(P_{\phi}(G=l|\mathbf{x}_{i})=G_{\phi}(m(\theta,\mathbf{x}_{i}))\).

The derivative of the loss with respect to the two different set of parameters is given by:

\[\frac{\partial\mathcal{L}^{in}}{\partial\theta}= \frac{1}{N}\sum_{i=1}^{N}\sum_{l=1}^{L}\frac{\partial\mathcal{L}^{CE }(y_{i},\hat{\mathbf{p}}^{l}_{\theta}(\mathbf{x}_{i}))}{\partial\theta}P_{ \phi}(G=l|\mathbf{x}_{i})+\frac{\partial G_{\phi}(m(\theta,\mathbf{x}_{i}))}{ \partial\theta}C^{l}_{\theta(\phi)}\,, \tag{14}\] \[\frac{\partial\mathcal{L}^{out}}{\partial\phi}= \frac{1}{N}\sum_{i=1}^{N}\sum_{l=1}^{L}\frac{\partial P_{\phi}(G=l| \mathbf{x}_{i})}{\partial\phi}C^{l}_{\theta^{*}(\phi)}=\frac{1}{N}\sum_{i=1}^{ N}\sum_{l=1}^{L}\frac{\partial\min\left(g^{l}_{\phi}(\mathbf{c}^{\leq l}_{i}),1- \sum_{j=1}^{l-1}g^{j}_{\phi}(\mathbf{c}^{\leq j}_{i})\right)}{\partial\phi}C^ {l}_{\theta^{*}(\phi)}\,. \tag{15}\]

We note that equation 15 becomes undefined when the two terms of the \(\min\) operator 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 an objective with a term that is similar to the one proposed by Han et al. (2022b) to address the train-test mismatch issue. In our case, the weights emerge directly from our proposed gating mechanism: \(P_{\phi}(G=l|\mathbf{x}_{i})\). The right term is driven by the impact of the \(\theta\) parameter on the gates: \(\frac{\partial G_{\phi}(m(\theta,\mathbf{x}_{i}))}{\partial\theta}\). In practice, we observe that ignoring the right term of the gradient does not impact performance and leads to faster convergence, as we show in Appendix 8.7.
Optimizing the gatesThe second derivative in equation 15 is more challenging, making direct optimization undesirable. Instead, we construct surrogate binary classification tasks to train the \(g_{\phi}^{l}\). For a given sample \(\mathbf{x}_{i}\) we construct binary targets for \(g_{\phi}^{1}(\mathbf{c}_{i}^{\leq 1}),...,g_{\phi}^{L-1}(\mathbf{c}_{i}^{\leq L -1})\) by evaluating the relative cost of each gate \(C_{\theta^{*}(\phi)}^{l}\) and determining which of the \(L\) gates has the lowest cost, denoted as \(l^{*}=\arg\min_{l}C_{\theta^{*}(\phi)}^{l}\). This is the gate at which \(\mathbf{x}_{i}\) should exit. Hence, we set the binary target of gate \(l^{*}\) 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 \(l^{*}\), we assume that it should also exit at any later gate, so we set the binary targets of all subsequent gates to 1 as well. Hence, the targets of our binary tasks \(t_{i}^{1},...,t_{i}^{L}\) for our gates \(g_{\phi}^{1}(\mathbf{c}_{i}^{\leq 1}),...,g_{\phi}^{L-1}(\mathbf{c}_{i}^{ \leq L-1})\) are given by \(t_{i}^{j}=0\) for \(j<l^{*}\) and \(t_{i}^{j}=1\) for \(j\geq l^{*}\).

The connection between those surrogate tasks and the initial objective from equation 12 can be established by showing that they can share the same solution under some assumptions. We demonstrated this is Appendix 8.4. Hence, instead of following the gradient from equation 15, we approximate it by the gradient of the surrogate loss:

\[\frac{\partial\mathcal{L}^{out}}{\partial\phi}\approx\frac{1}{N}\sum_{i=1}^{N }\sum_{l=1}^{L}\frac{\partial\mathcal{L}^{CE}(t_{i}^{l},g_{\phi}^{l}(\mathbf{ c}_{i}^{\leq l}))}{\partial\phi}. \tag{16}\]

This amounts to summing the gradients of the losses of \(L\) independent binary classification tasks. If the tasks are highly imbalanced, we compute the class imbalanced ratio on a validation set at each epoch and use weighted class training. Following the bi-level optimization procedure, training is carried out by alternating between optimizing using the gradients in equation 16 and equation 14. In practice, we start by training all IMs on the full dataset in a warmup stage as described in Section 5. Algorithm 1 in Appendix 8.5 provides a detailed exposition of the entire algorithm.

## 5 Experiments

We use the vision transformers T2T-ViT-7 and T2T-ViT-14 (Yuan et al., 2021) pretrained on the ImageNet dataset (Deng et al., 2009) which we then transfer-learn to the datasets: CIFAR10, CIFAR100-LT (Krizhevsky, 2009) and SVHN (Netzer et al., 2011). For dataset details, see Appendix 8.1.1. For transfer-learning, we use the procedure from (Yuan et al., 2021). We provide checkpoints for all these models in our code. Our backbone \(NN\) for each dataset is the frozen transfer-learned model. We augment the backbone with gates and intermediate IMs at every layer. We generate results by varying the inference cost parameter \(\lambda\) over the range 0.01 to 10. Appendix 8.1.1 contains further details concerning hyperparameters and experimental procedure.

### Our model

**Gate design \(g_{\phi}^{l}(\mathbf{c}_{i}^{\leq l})\):** To ensure a lightweight design, we construct a small number of features by computing uncertainty statistics from the output \(\hat{\mathbf{p}}^{l}\) of the \(l\)-th IM, \(f_{\theta}^{l}(\mathbf{x}_{i})\). Hence the \(l\)-th gate can be represented as \(g_{\phi}^{l}(\mathbf{c}_{i}^{\leq l})=g_{\phi}^{l}(m(\theta,\mathbf{x}_{i}))\). We choose \(m(\theta,\mathbf{x}_{i})=[\hat{p}_{i}^{l,\max}(\mathbf{x}_{i}),h^{l}(\mathbf{ x}_{i}),h_{pow}^{l}(\mathbf{x}_{i}),mar^{l}(\mathbf{x}_{i})]^{T}\) where \(\hat{p}_{i}^{l,\max}(\mathbf{x}_{i})\) is maximum predicted probability, \(h^{l}(\mathbf{x}_{i})\) is the entropy of predictions, \(h_{pow}^{l}(\mathbf{x}_{i})\) is the entropy of predictions scaled by a temperature, and \(mar^{l}(\mathbf{x}_{i})\) is the difference between the two most confident predictions. For more detail see Appendix 8.9. Rather than sample \(E^{l}\succ p(E^{l}|\mathbf{x})\) to decide whether to exit, we use the deterministic decision \(P(E^{l})>0.5\), which leads to more stable results.

**IM design \(f_{\theta}^{l}(\mathbf{z}_{i}^{\leq l})\):** The IMs \(f_{\theta}^{l}(\mathbf{z}_{i}^{\leq l})\) all have the same architecture. We reduce cost by limiting their input size and number of layers. The input to \(f_{\theta}^{l}\) is only \(\mathbf{z}_{i}^{l}\) and the IM is a single layer NN. The lightweight design allows us to insert exit branches at every layer. This gives our model greater exit flexibility, in contrast with many existing approaches (Chiang et al., 2021; Lahiany and Aperstein, 2022; Li et al., 2023; Ilhan et al., 2023), which typically use only a few exit branches.

**Training procedure:** Our training procedure consists of two phases: (i) warm-up; and (ii) bi-level optimization. In the warm-up phase the first \(L\)-\(1\) IMs are trained in parallel on all samples for a 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 optimizing the gate parameters \(\phi\) and the IM parameters \(\theta\). See Appendix 8.1.4 for more detail.
### Baselines

We compare our proposal with the following baselines:

* **BoostedNet**Yu et al. (2023) and **L2W-DEN**Han et al. (2022b) are state-of-the-art benchmarks with architecture-agnostic training procedures for general-purpose networks.
* **MSDNet**Huang et al. (2018) and **RANet**Yang et al. (2020) are representative EEDN-tailored architectures. We additionally compare with the improved training procedures in Yu et al. (2023) and Han et al. (2022b).

To obtain the uncertainty metrics, we calibrate the predicted \(\hat{\mathbf{p}}\) by first applying the post-calibration algorithm of temperature scaling from Guo et al. (2017) on the output logits using a validation set. Since all baselines also require a validation set to select gating thresholds, we split the validation set into two \(\mathcal{V}=\mathcal{V}^{1}\cup\mathcal{V}^{2}\); \(\mathcal{V}^{1}\) is used to compute the gating thresholds and \(\mathcal{V}^{2}\) is used for hyperparameter tuning, early stopping, setting the conformal threshold, and calibrating the output.

### Results

**Observation 1:** _Joint Training of the GMs and the IMs closes the train-test mismatch which leads to significant performance improvement._

Figure 1 depicts performance versus inference cost curves. Appendix 8.8 presents results on additional 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, it 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 CIFAR100-LT. The imbalance leads to noisier prediction, which makes it particularly challenging for the baselines. To provide further insight into why our model outperforms the baselines, we analyze the results in the subsequent sections. Appendix 8.6 provides an ablation study which demonstrates the value of learnable gates and joint training.

**Observation 2:** _Training the GMs leads to a better gate selection._

Because our GMs are jointly trained with the IMs and take into account the cost of inference, the GMs can avoid poorly performing gates or gates that provide only marginal performance improvement. Figure 2 shows which gates are 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 5 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 on the IMs that can perform satisfactorily for a given inference budget. In contrast, the threshold GMs employed by the baselines lead to 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-confident for some samples).

Figure 1: Accuracy vs Mul-Add of: **Left** CIFAR100 ((2t-14); **Middle** SVHN ((2t-7); and **Right** CIFAR100LT ((2t-14). The x-axes are scaled by the full model inference cost, Mul-Add (\(IC^{L}\)).
Moreover, the thresholding approach implicitly relies on the assumption that the IMs are well calibrated, which cannot be guaranteed (Guo et al., 2017; Minderer et al., 2021). Experimentally, we verify that it is not the case and that calibration is positively correlated with the depth of the IMs (see Figure 3). Consequently, calibration is also correlated with accuracy, which is in line with the findings of Minderer et al. (2021). A poor calibration leads to inaccurate predicted probabilities, which introduces noise in the inputs to the threshold GM. As a result, the exit selection is also more variable, leading to a less predictable accuracy-IC trade-off (as exhibited by the much wider CIs in Figure 1).

Observation 3:_Better gate selection concentrates training on better IMs._

Figure 2(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 at 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. 2(a)) than all samples (Fig. 2(b)). 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; for the baselines, they exit at gate 13. Despite having access to less informative features, JEI-DNN achieves a better average accuracy for these challenging samples at considerably lower inference cost.

Observation 4:_Learnable GMs leads to better uncertainty characterization capability_

Fig. 4**(a)** highlights that even with post-calibration, the baselines' predicted probabilities remain extremely poorly calibrated. JEI-DNN's calibration is much better and the calibration error correlates negatively with accuracy. 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-calibrated gates (Fig. 3); (ii) the baselines must increase the GM thresholds to achieve higher accuracy, but this leads to a bias -- samples only exit if the predictions are very confident.

Figure 3: ECE of the IMs averaged over all baselines on all datasets with t2t-7.

Figure 2: Decomposition of the contributions of the IMs to the final accuracies (depicted by dotted lines) for CIFAR100. The operational point is marked by a star in the left panel of Figure 1. **a)** Accuracy of each IM, \(f_{g}^{l}\), evaluated only on their exited samples (\(\mathcal{D}^{\text{I}}\)). **b)** Accuracy of each IM, \(f_{g}^{l}\), on the full dataset \(\mathcal{D}\). The size of a circle is proportional to the number of samples exited for a trial.
We now consider conformal intervals. These should circumvent the threshold GM bias as they are derived via a validation set and have probabilistic guarantees. Fig. 4**(b)** shows that JEI-DNN offers significantly tighter conformal intervals \(|\tilde{\mathcal{C}}|\) for a given constraint on empirical coverage \(\hat{\alpha}\). The empirical coverage values for JEI-DNN are also closer to the requested coverages \(\alpha\) (see Fig. 4**(c)**).

Observation 5:_JEI-DNN using a fixed backbone outperforms dedicated EEDN architectures_

We now compare with state-of-the-art dedicated EEDN approaches that employ a trainable backbone. We select MSDNet (Huang et al., 2018), RANet (Yang et al., 2020), BoostedNet (Yu et al., 2023), and L2W-DEN (Han et al., 2022b) as representative models. Figure 5 shows that JEI-DNN with a T2T-ViT-14 backbone achieves much better accuracy while keeping the inference cost low. We achieve an accuracy of 82% for 22% of the inference cost of the full trainable backbone architectures. By contrast, the trainable architectures can only achieve an accuracy of 73% at this cost. However, in the very low-cost regime, dedicated EE architectures like MSDNet and RANet are able to maintain higher accuracy. By this point, however, the reduction in accuracy is severe.

## 6 Limitations

Although our proposal outperforms the baselines, both with respect to the accuracy-IC trade-off and uncertainty characterization, it has some drawbacks. First, since we must specify the inference cost parameter \(\lambda\), we only obtain one accuracy/IC operational point per training. To obtain different values, we need to retrain the model with a changed \(\lambda\) value. By contrast, the baseline methods that employ threshold GMs can modify the inference time after deployment by changing the threshold. Second, we cannot directly target a specific accuracy or inference cost. In practice, if there is a constraint on one of the metrics, we either must train multiple models, adjusting \(\lambda\) until the constraint is met, or use an ad-hoc approach by applying different thresholds to the \(g_{\phi}^{\prime}\) values. In summary, if the desired trade-off between 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 architecture, JEI-DNN, that is compatible with off-the-shelf backbone architectures. By augmenting a 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. We show that this approach also outperforms state-of-the-art dedicated EEDN architectures. Furthermore, our approach provides significantly better uncertainty characterization. We show that the inference modules are better calibrated in the proposed JEI-DNN and tighter conformal intervals are obtained for a desired marginal coverage.