Title: Fast yet Safe: Early-Exiting with Risk Control
Venue: NeurIPS 2024
Accepted: yes
Keywords: Early-exit neural networks

risk control

efficiency

robustness
Abstract: Scaling machine learning models significantly improves their performance.  However, such gains come at the cost of inference being slow and resource-intensive.  Early-exit neural networks (EENNs) offer a promising solution: they accelerate inference by allowing intermediate layers to exit and produce a prediction early.  Yet a fundamental issue with EENNs is how to determine when to exit without severely degrading performance.  In other words, when is it 'safe' for an EENN to go 'fast'?  To address this issue, we investigate how to adapt frameworks of risk control to EENNs. Risk control offers a distribution-free, post-hoc solution that tunes the EENN's exiting mechanism so that exits only occur when the output is of sufficient quality.  We empirically validate our insights on a range of vision and language tasks, demonstrating that risk control can produce substantial computational savings, all the while preserving user-specified performance goals.

Fast yet Safe: Early-Exiting with Risk Control

Abstract

Scaling machine learning models significantly improves their performance. However, such gains come at the cost of inference being slow and resource-intensive. Early-exit neural networks (EENNs) offer a promising solution: they accelerate inference by allowing intermediate layers to 'exit' and produce a prediction early. Yet a fundamental issue with EENNs is how to determine when to exit without severely degrading performance. In other words, when is it 'safe' for an EENN to go 'fast'? To address this issue, we investigate how to adapt frameworks of risk control to EENNs. Risk control offers a distribution-free, post-hoc solution that tunes the EENN's exiting mechanism so that exits only occur when the output is of sufficient quality. We empirically validate our insights on a range of vision and language tasks, demonstrating that risk control can produce substantial computational savings, all the while preserving user-specified performance goals.

1 Introduction

As predictive models continue to grow in size, so do the costs of running them at inference time [12, 80]. This presents a challenge to domains ranging from mobile computing to smart appliances to autonomous vehicles - all of which require models that operate on resource-constrained hardware [62, 50, 47]. Even if computation is not limited by hardware, concerns over the energy usage and carbon footprint of large models motivates their efficient implementation [55, 44]. Additionally, since computational constraints can be dynamic, e.g ., due to variable loads in web traffic or energy demand, it is desirable for models to be able to adjust their computational needs to changing conditions [63].

Early-exit neural networks (EENNs) present a simple yet effective approach to such dynamic computation [70, 36]. Leveraging the neural network's compositional nature, EENNs can generate predictions at intermediate layers, thereby 'exiting' the computation 'early' when a stop condition is met. This early-exit ability has proven useful in settings ranging from vision and language to recommendations ([29], see § 4). Yet the flexibility of EENNs does not come for free: predictions generated at early exits are usually inferior to those produced by the full model. In turn, a dilemma arises in which the exit condition must balance computational savings with predictive performance.

In this work, we address the EENN's efficiency vs. performance trade-off via statistical frameworks of risk control (RC) [4, 8]. By tuning the EENN's exiting mechanism based on a user-specified notion of risk, RC aims to enhance the safety of early-exit outputs. We consider several risks that quantify the difference between the early-exit and full model's outputs, both in terms of prediction quality and uncertainty estimation. Moreover, we study RC frameworks that control the risk with varying degrees of stringency ( i.e ., in expectation vs. with high probability). We demonstrate the effectiveness of this light-weight, post-hoc solution across a range of tasks, including image classification, semantic segmentation, language modeling, image generation with diffusion, and speculative decoding in large language models. In particular, we make the following contributions:

We formalize EENNs as risk-controlling predictors, ensuring risk control is amenable to the early-exit setting by explicitly linking risk control and early-exit requirements (Prop. 1 & Prop. 2).

We propose risk functions to control early-exit performance both in terms of model predictions and their underlying predictive distributions (Eq. 6 & Eq. 8). Previous work has considered only prediction quality [65], not uncertainty quality (as we do).

We improve upon prior work for language modeling [65], demonstrating that our adaptions of risk control allow for less conservative early-exiting and result in larger efficiency gains (§ 3.3, § 5.3).

We apply, for the first time, risk control to early-exiting in image classification, semantic segmentation, and image generation, as well as speculative decoding in large language models (§ 5).

2 Background

Data. Let X × Y denote the sample space and assume a data-generating distribution P over it. We consider Y := { 1 , . . . , K } for classification and Y ⊆ R d for regression. Observed samples from P are split into disjoint train, calibration and test sets, denoted D train , D cal , and D test . We assume samples ( x , y ) in D cal and D test to be drawn i.i.d. from P , whereas D train is permitted to be drawn randomly from a different distribution (of same support).

Early-Exit Neural Networks. EENNs extend traditional static network models by dynamically adjusting computations during the model's forward pass ( e.g ., the number of evaluated network layers) on the basis of an input sample's complexity or 'difficulty'. More formally, we define an EENN as a sequence of probabilistic classifiers ˆ p ( y | x = x ; ϕ l , θ l ) , where l = 1 , . . . , L enumerates the model's exit layers, and ϕ l and θ l define the model's classification head and backbone parameters at the l -th exit, respectively. The final index L denotes the full model, i.e ., all layers are evaluated for a given input sample x . The obtained predictive distribution ˆ p ( y | x = x ; ϕ l , θ l ) at the l -th exit layer, denoted in short as ˆ p l ( y | x ) , permits to retrieve both a predicted class label ˆ y l = arg max y ∈Y ˆ p l ( y | x ) and an associated confidence score ˆ c l ∈ [0 , 1] , which aims to capture model's certainty about the exit's current prediction. One common choice for classification tasks is the maximum class probability ˆ c l = max y ∈Y ˆ p l ( y | x ) . However, different notions of confidence are possible, as we explore in § 5.2. At test-time, these confidence scores can be leveraged to determine the early-exit model's required computations for new samples via thresholding. For a given test input x , the EENN exits 2 at the first layer for which its confidence exceeds a pre-specified threshold value λ l ∈ [0 , 1] . For simplicity, a single threshold value λ l = λ, ∀ l ∈ { 1 , . . . , L -1 } is often fixed across exit layers, a setup we also consider here. The predictive distribution obtained from the model's early-exit mechanism is then given by

̸

The threshold parameter λ regulates the trade-off between the EENN's accuracy and efficiency gains. Lower values equate larger speed-ups by increasing the likelihood of an early exit (and vice versa ), at the cost of generally inferior predictions. Such marginally monotone behavior, where model performance improves on average across exits, is a core assumption for the practical use of EENNs (see also Fig. 1). We formalize it as

for some arbitrary loss function ℓ , and elaborate on its connection to risk control in § 3.3. It is common to determine early-exiting criteria by investigating these trade-offs between performance and efficiency on hold-out data, selecting thresholds that ensure the EENN meets a user's computational budget [36, 70] or performance goals [16, 21, 48]. A standard practice is to treat the EENN's predictive confidence ( e.g ., its softmax scores) as a heuristic for prediction quality. However, this is fallible, as EENNs can exhibit fluctuating or poorly-calibrated confidences [40, 37, 51], motivating more principled threshold selection.

Risk Control. Statistical frameworks for risk control (RC) [1, 4, 8] aim to improve prediction reliability by equipping threshold-based models with safety assurances. Specifically, consider a pre-trained prediction model ˆ fλ whose outputs depend on a threshold λ . For example, given a classification task, the set predictor ˆ fλ : X → 2 Y described by Bates et al. [8] includes a class label in the set if its probability exceeds the threshold, i.e ., ˆ f λ ( x ) := { y ∈ Y : ˆ p ( y | x ) ≥ λ } . Next, a notion of error for ˆ f λ is captured by defining a problem-specific loss function ℓ : Y × Y → R . For instance, a meaningful choice for the set predictor could be the miscoverage loss ℓ ( ˆ f λ ( x ) , y ) = ✶ [ y / ∈ ˆ f λ ( x )] , where ✶ [ · ] is the indicator function. The risk associated with a particular threshold λ ∈ Λ is then defined as the expected loss R ( λ ) := E ( x , y ) ∼P [ ℓ ( ˆ f λ ( x ) , y ) ] , (3)

Given such a risk, desired safety assurances may vary in strength. For a tolerated risk level ϵ ∈ (0 , 1) , risk control in expectation seeks to guarantee that with Λ the set of potential threshold candidates. RC frameworks leverage different probabilistic tools - which we detail further in § 3.3 - to determine a subset ˆ Λ ⊆ Λ for which the risk in Eq. 3 is guaranteed to be small. Note that ˆ Λ is retrieved in a post-hoc manner by leveraging the calibration set D cal sampled i.i.d. from P . Thus, R ( ˆ λ ) is a random quantity dependent on D cal for any ˆ λ ∈ ˆ Λ .

where the outer expectation is taken over randomly drawn calibration data of finite size |D cal | = n . A stronger statement on risk control with high probability requires additionally specifying a probability level δ ∈ (0 , 1) , and aims to ensure that

That is, rather than the average control over calibration data in Eq. 4, risk control according to Eq. 5 holds with high probability for any particular sampled set D cal . In both cases, we may refer to ˆ f ˆ λ for any ˆ λ ∈ ˆ Λ as a risk-controlling predictor . The risk level ϵ and probability level δ are user-specified parameters dictating how tightly the risk is controlled, and a particular choice has to consider the problem-specific setting and loss ℓ . For example, a reasonable choice for the stated miscoverage loss may be ( ϵ, δ ) = (0 . 05 , 0 . 1) . Observing ˆ Λ = ∅ implies that there is no risk-controlling predictor for the selected ( ϵ, δ ) , indicating overly stringent risk control conditions which ˆ f λ cannot satisfy.

The prediction guarantees obtained via RC are highly practical, since they are (i) distribution-free, i.e ., they do not impose any particular assumptions on the generating distribution P , (ii) are post-hoc applicable to any arbitrary choice of underlying predictor ˆ f λ , and (iii) hold in finite samples, thus not relying on asymptotic limit statements. Indeed, we experimentally find that the provided assurances hold even for remarkably small calibration sets ( n ≈ 100 , see § 5).

3 Safe Early-Exiting via Risk Control

We now detail our approach for early-exiting with safety guarantees based on risk control. We begin by outlining EENNs as risk-controlling predictors (§ 3.1). Next, we describe two general types of risk to measure performance drops resulting from early-exiting. Importantly, these risks can be employed to assess the quality of both predictions and predictive distributions (§ 3.2). Finally, we motivate and formalise how different risk control frameworks can be adapted to the early-exit setting (§ 3.3).

3.1 EENNs as Risk-Controlling Predictors

As mentioned in § 2, risk control requires a predictor ˆ f λ whose outputs depend on a threshold λ ∈ Λ . The EENN's confidence-based thresholding behaviour (following Eq. 1) lends itself naturally to such a formulation. For a particular exit threshold λ ∈ [0 , 1] , the EENN ˆ p λ ( y | x ) will act as such a threshold predictor. To ensure that the EENN satisfies the user's control requirements, the risk-controlling threshold set ˆ Λ needs to be identified. Importantly, this can be done post-hoc using a pre-trained EENN with fixed weights, since only the exit threshold λ is modified. In order to maximize computational savings while ensuring that the user-defined risk is managed, we select ˆ λ := min ˆ Λ , since a low threshold encourages earlier exiting. If ˆ Λ = ∅ is empty, we default to ˆ λ = 1 , the equivalent of relying strictly on the full model output ˆ p L ( y | x ) .

3.2 Early-Exiting Risks

We next detail two types of risk which can be employed to guard against performance drops due to early-exiting. Similarly to Schuster et al. [65], our risks are defined in terms of relative exit performance, permitting their calculation for both labelled and unlabelled calibration data. Moving beyond their setting, we suggest these risks for controlling the quality of both model predictions and predictive distributions , from which confidence scores can be derived.

Performance Gap Risk. When calibration labels are present, these can be used to measure the early-exit performance through supervised losses. Let ˆ o l ( x ) denote a general EENN output for some input x . It takes the form ˆ y l for predictions and ˆ p l ( y | x ) for the underlying predictive distribution. The supervised performance gap risk is then defined as

where ˆ o λ ( x ) and ˆ o L ( x ) refer to early-exit and full model outputs, respectively. The choice of loss function ℓ is task-specific, and we outline relevant choices in § 5, such as the 0 -1 loss for image classification. For predictive distribution control, we suggest leveraging a squared distributional loss which, when averaged across samples, recovers the Brier score [13]. Specifically, we define such a 'Brier loss' for classification tasks as

where ˆ p l ( k | x ) denotes the predicted probability of a particular class k , and ✶ [ y = k ] its one-hot encoded label. The Brier score is a strictly proper scoring rule [26, 27], ensuring its suitability to assess probabilistic forecasts. Moreover, its mathematical formulation lends itself favorably to risk control when compared to other widely used probabilistic metrics. We defer further details to § A.3. Addressing risk control of the underlying predictive distribution ˆ p l ( y | x ) is a compelling extension, as confidence or uncertainty estimates are typically derived from it. Particularly in safety-critical scenarios where reliable uncertainties are essential [32, 9], such control can thus prove very useful.

Consistency Risk. In the case of unlabelled calibration data, an unsupervised version of Eq. 6 can be obtained by replacing the ground truth labels y with labels ˆ y L obtained from the full model. We define the unsupervised consistency risk as

where only input samples x ∼ P are required for its evaluation.

For regular prediction losses, Eq. 8 collapses to evaluating the per-sample loss ℓ (ˆ y λ , ˆ y L ) , since ℓ (ˆ y L , ˆ y L ) = 0 . For predictive distribution control, the loss difference remains, and we substitute y in Eq. 7 by sampling a label ˆ y L ∼ ˆ p L ( y | x ) from the EENN's final layer. Our reliance on the last layer's output is motivated by the EENN's marginal monotonicity (Eq. 2 and Fig. 1). Finally, we note that both the performance gap risk R G ( λ ) and consistency risk R C ( λ ) are quite agnostic to the EENN's actual predictive performance. In both cases, the risk formulation aims to ensure prediction consistency via the relative performance gap between exits, as opposed to absolute performance with respect to observed ground truth labels.

3.3 Risk Control Frameworks

0.60

0.70

0.80

↑

)

ImageNet

MSDNet

DViT

L2W-DEN

Dyn-Perc

1

2

3

4

5

Early Exit

0.30

0.40

0.50

Brier (

↓

)

After having defined our early-exit risks, we next outline how a desired risk-controlling exit threshold ˆ λ can be computed based on calibration data. We begin by considering a 'naive' empirical approach, followed by risk control in expectation (Prop. 1) and with high probability (Prop. 2).

Empirical Approach. For a tolerated risk level ϵ ∈ (0 , 1) , a 'naive' empirical threshold can be selected by picking the smallest threshold λ ∈ [0 , 1] from the candidate set Λ for which the risk on the calibration set D cal is controlled, i.e .,

Note that ˆ R ( λ ; D cal ) = 1 n ∑ n i =1 ℓ (ˆ o λ ( x i ) , y i ) is the empirical calibration risk, an approximation of the true risk in Eq. 3 computed on D cal (and likewise ˆ R ( λ ; D test ) denotes the empirical test risk). For the risks introduced in § 3.2 the threshold ˆ λ emp is always well-defined, since ˆ R ( λ = 1) is zero.

Risk Control in Expectation. The threshold ˆ λ emp is a straight-forward choice, but can fail to control the risk on test data if the approximation quality of R ( λ ) by ˆ R ( λ ; D cal ) is poor, e.g ., due to badly drawn calibration data. Perhaps surprisingly, only a slight modification of Eq. 9 is required to ensure risk control in expectation . Specifically, for a bounded loss function ℓ ≤ B where B > 0 , and assuming a monotone risk R ( λ ) , the threshold 3

guarantees Eq. 4, thus shielding against bad sample draws on average. It is easy to show that ˆ λ emp = lim n →∞ ˆ λ CRC, and since our losses are designed to be upper-bounded by B ∈ { 1 , 2 } , the two thresholds already coincide for small calibration sets ( n ≈ 100 ). We formalize risk control in expectation in the following proposition for our early-exit setting:

Proposition 1. Let ℓ : Λ → ( -∞ , B ] be a right-continuous bounded loss, and assume a marginally monotone EENN (Eq. 2). Then the exit threshold ˆ λ CRC ensures risk control in expectation, i.e ., it holds that E D cal ∼P n [ R ( ˆ λ CRC ) ] ≤ ϵ for any ϵ ∈ (0 , 1) .

Our proposition is an extension of Conformal Risk Control [4] (CRC) to the early-exit setting, and a proof can be found in § A.1. Our main technical insight is that risk control can be relaxed to assume monotone risks , rather than monotone losses as in the original formulation [4]. This relaxation is crucial for the early-exit setting, since we can relate monotone risks to assumptions of marginal monotonicity on the EENN (see Lemma 1). In contrast, monotone losses translate to assuming conditional monotonicity, a much stronger requirement suggesting the EENN's performance improves across exits per sample , and which has been shown to be violated in practice [40, 76, 37].

Risk control with High Probability. A stronger guarantee can be obtained by ensuring risk control with high probability for any drawn calibration set. We employ the Upper Confidence Bound (UCB) from Bates et al. [8] for this purpose. First, an empirical upper bound ˆ R + ( λ ; D cal ) is derived to bound the risk R ( λ ) with high probability. That is, for a probability level δ ∈ (0 , 1) it holds that

An exit threshold ensuring risk control according to Eq. 5 is then selected as

Similarly to Prop. 1, we can now formalize risk control with high probability for the early-exit setting: Proposition 2. Let ℓ : Λ → [ -B,B ] be a bounded loss, and assume a marginally monotone EENN (Eq. 2). Then the exit threshold ˆ λ UCB ensures risk control with high probability, i.e ., it holds that P D cal ∼P n ( R ( ˆ λ UCB ) ≤ ϵ ) ≥ 1 -δ for any ( ϵ, δ ) ∈ (0 , 1) 2 .

This reformulation of the main theorem from Bates et al. [8] (Thm. A.1) is proven in § A.1, and an algorithmic description is given in Appendix B. We employ their suggested Waudby-Smith-Ramdas bound [75] (WSR) to compute ˆ R + ( λ ; D cal ) , but relax the bounding requirements on the loss from ℓ ∈ [0 , 1] to ℓ ∈ [ -B,B ] for B > 0 . This change has important implications for the early-exit setting, since it admits 'rewarding' the EENN when an earlier exit performs better than the final exit for some samples, a phenomenon known as overthinking [40, 37]. In practice, this results in a better risk estimate and less conservative early-exiting. See § A.2 for more details on the effect of loss bounds.

Learn-then-Test and CALM [65]. Learnthen-Test [1] (LTT) is another framework for high-probability risk control, where threshold selection is framed as a multiple hypothesis testing problem. In contrast to UCB (Prop. 2), LTT does not require risk monotonicity, and can thus also be employed when the EENN is suspected to violate marginally monotone behaviour. LTT in the early-exit setting has been employed by Schuster et al. [65] (CALM), presumably motivated by the avoidance of this assumption. However, expecting an EENN to marginally improve across exits is a core requirement which implicitly underlies any practical implementation. Since this assumption is usually also empirically satisfied (Fig. 1), there is no obvious reason to explicitly avoid it. Furthermore, correcting for multiple testing in LTT via fixed sequence testing - as is done for CALM - will only yield practical savings if monotonicity is satisfied. We stress these observations since we find that UCB provides greater computational savings than LTT under the same guarantees (Fig. 2 and § 5.3), including

0.00

0.05

0.10

0.15

Test Risk

ˆ

R

CALM(T5-large), CNN/DM,

ncal

=

100

0.02

0.04

0.06

0.08

0.10

0.12

0.14

Risk Level

/epsilon1

5

10

15

↓

)

LTT

UCB

CRC

for small-sample regimes ( n ≈ 100 ), which are of high practical interest and were not explored by Schuster et al. [65]. Moreover, due to LTT's reliance on the Hoeffding-Bentkus bound [11], it cannot account for instances of model overthinking (see § A.2). Thus, unless Eq. 2 is known to be violated, we recommend UCB over LTT in the early-exit setting.

4 Related Work

Early-Exiting [70, 29] as a dynamic approach to accelerate model inference is both orthogonal and complementary to static model compression techniques such as pruning, quantization, and knowledge-distillation [6, 79, 68, 84, 25]. Its wide-ranging applicability has been demonstrated across numerous vision [36, 48, 15, 69, 22] and language tasks [21, 83, 65, 78, 5, 53]. While most prior work has focused on the trade-off between performance quality and computational savings, the safety of early-exit models has received less attention to date [64, 65, 51, 38]. Risk Control has gained traction due to an interest in efficient, post-hoc approaches with safety assurances for large models. Most related, conformal prediction [66] has been popularized as an effective method for uncertainty quantification with guarantees on the miscoverage risk [24, 3]. Recently, multiple proposals address the control of more general risk notions [4, 8, 1, 67, 54, 45], with explored applications ranging from imaging [71, 2, 81, 23, 43, 61, 10, 72] to language [86, 20, 45, 56] and beyond [39, 46]. Most closely related to our work is research by Schuster et al. [65], who first employ risk control for safe early-exiting in language modeling. We move beyond their setting by (i) controlling the quality of both model predictions and uncertainty estimates (§ 3.2), (ii) obtaining better efficiency gains through careful selection of our risk control framework (§ 3.3), and (iii) extending early-exit risk control to novel tasks (§ 5). Ringel et al. [60]'s work is also related: they apply risk control to exit early for a time series prediction task. Yet their emphasis is on exiting from a stream of input features, whereas we exit from the model itself ( i.e ., a stream of model layers).

5 Experiments

We empirically validate early-exiting via risk control on a suite of different tasks: image classification (§ 5.1), semantic segmentation (§ 5.2), language modelling (§ 5.3), image generation with diffusion (§ 5.4), and speculative decoding (§ 5.5). Our code is publicly available at https://github.com/ metodj/RC-EENN . We begin by outlining our general risk control design and evaluation metrics.

Risk control design. We target control of the performance gap and consistency risks defined in § 3.2. For predictions ˆ y l we employ target-specific losses, and, when applicable, for predictive

0.0

0.1

0.2

0.3

0.4

Test Risk

ˆ

R

R

G

(

ˆ

y

)(

0

-

1

)

R

G

(

ˆ

p

)

(Brier)

R

C

(

ˆ

y

)(

0

-

1

)

R

C

(

ˆ

p

)

(Brier)

0.0

0.1

0.2

0.3

0.4

1

2

3

4

5

↓

)

MSDNet

DViT

L2W-DEN

Dyn-Perc

0.0

0.1

0.2

0.3

0.4

CRC

UCB

0.0

0.1

0.2

0.3

0.4

0.0

0.1

0.2

0.3

0.4

Risk Level

/epsilon1

distributions ˆ p l ( y | x ) our Brier score formulation. We denote these four risks in short as R G (ˆ y ) , R G (ˆ p ) , R C (ˆ y ) and R C (ˆ p ) . Risk control requirements of different strength are assessed by varying the risk level ϵ . Note that our approach is entirely post-hoc , and our experiments leverage existing pretrained EENNs when possible. Thus the underlying models are typically not modified, and we omit the commonly seen performance vs. FLOP curves to instead appropriately benchmark against different risk levels. For risk control with high probability, we set δ = 0 . 1 ( i.e ., 90 % probability). Reported numbers are averaged across multiple trials of calibration and test splitting ( S = 100 ) to account for sampling effects. Additional results across experiments can be found in Appendix D.

Evaluation metrics. We evaluate our results based on obtained test risks and efficiency gains. We assess whether the guarantees stated in § 3.3 are satisfied by checking if the empirical test risk for a given risk-controlling threshold is controlled, i.e ., ˆ R ( ˆ λ ; D test ) ≤ ϵ holds. Ideally, the measured test risk should also approach ϵ from below, as to prevent overly conservative early-exiting. We measure efficiency gains by reporting the average exit layer across test samples, or its relative improvement over last-layer exiting (in %). Similar gains in terms of arithmetic operations (FLOPS) are reported in Appendix D. We favour approaches which, while controlling the test risk, exit as early as possible.

5.1 Image Classification

For image classification, we focus on the ImageNet dataset [19]. We employ four state-of-theart EENNs to demonstrate that our findings generalize across different models and architectures: MSDNet [36], DViT [74], L2W-DEN [30], and Dyn-Perc [31] (see § C.1 for details). We employ the standard 0 -1 loss for predictions, and the Brier loss formulation from Eq. 7 for predictive distributions. Fig. 3 displays results for the small-sample calibration regime ( n = 100 ). In line with our theoretical guarantees, the test risk remains controlled across all models, risk types, and risk levels ϵ ( top row). The steeply decreasing efficiency curves affirm that even under strict control requirements, substantial efficiency gains can be obtained ( bottom row). For example, controlling the prediction gap risk at a strict 5% ( R G (ˆ y ) for ϵ = 0 . 05 ) results in a model average of ∼ 61% less layers evaluated for control in expectation (CRC, Prop. 1), and ∼ 46% for control with high probability (UCB, Prop. 2), see Table 2. Naturally, UCB produces more cautious early-exiting due to its stronger safety assurance, but these differences decrease for larger calibration sets (see § D.1 for n = 1000 ). This highlights the practical benefit of allocating more calibration samples: a larger sample size can aid to mitigate the price paid by a high-probability guarantee in terms of obtained inference speed-ups.

5.2 Semantic Segmentation

For this task, we explore the effect of different confidence measures used in Eq. 1 on realizable speed-ups. We use the EENN with four exits proposed by Liu et al. [48] (ADP-C, see § C.2). ADP-C permits pixel-level early-exiting with per-pixel confidence scores. Since we desire to early-exit the

Ground Truth

First layer (l = 1)

Last layer (l = 4)

0.20

1.0

B

Early exit

0.15

0.8

Brier loss diff.

Confidence

0.10

0.6

0.05

0.4

Late exit

C

0.00

0.2

1

2

3

4

ExitLayer

0.0

entire image instead, we explore image-level aggregations alongside different confidence scores, which are briefly outlined below. As task-specific prediction losses, we consider the commonly used mean intersection-over-union (mIoU) and miscoverage for the labelled and unlabelled cases, respectively. For predictive distribution control, we employ pixel-averaged versions of the Brier loss in Eq. 7 (see § A.3). We evaluate our approaches on Cityscapes validation data (80% D cal , 20% D test ); in addition, we finetune and evaluate ADP-C on a subset of the GTA5 dataset [59] in § D.2.

We consider three pixel-level confidence scores: the top class softmax probability ( Top-1 ), the difference between top two class probabilities ( Top-Diff ), and the normalized entropy over a pixel's predictive distribution ( Entropy ). In addition, we consider three image-level aggregation strategies: the image's average pixel confidence ( Mean ), its 0 . 25 -th quantile ( Quantile ), and a patch-based approach ( Patch ), wherein a sliding window of fixed size ( e.g ., 50 × 50 pixels) computes the mean confidence over pixels per patch, and the min over such patch scores is retrieved. These aggregations consider both varying levels of prudence (Mean vs. Quantile) and granularity (Mean vs. Patch).

Table 1 displays obtained efficiency gains for risk control via UCB (Prop. 2) across different risk levels ϵ ∈ { 0 . 01 , 0 . 05 , 0 . 1 } . In line with Fig. 3, increased speed-ups are observed as the risk requirements are relaxed ( i.e ., ϵ increases). Notably, for a given ϵ the gains for Brier-based risks tend to be smaller than for prediction risks, affirming more challenging risk control. The differences between combinations of per-pixel confidence and image-level aggregation are most pronounced for small ϵ , where Patch records highest gains while Quantile is more conservative (see § D.2 for full results). In Fig. 4, we display a qualitative example of the model's exiting behaviour. For a sample which exits at the first layer ( top row), the EENN's confidence map remains fairly stable across subsequent layers, suggesting an accurate model assessment has been reached early on. In contrast, a sample which exits at the final layer ( bottom row) will see a substantial improvement in model certainty, justifying additional computations. Such behaviour is also visible when stratifying all samples across their respective model exits (Fig. 4, left ). For samples which exit later, the difference between distributional losses at the first and final layer increases, affirming that compute is spent meaningfully.

0.0

0.1

Test Risk

ˆ

R

'Full model'

/epsilon1

=

0.01

/epsilon1

=

0.05

/epsilon1

=

0.10

/epsilon1

=

0.15

0.00

0.10

Risk Level

/epsilon1

1

5

10

13

Exit (

↓

)

CRC

UCB

®

5.3 Language Modeling

For this task, we replicate the main experiments from Schuster et al. [65] (CALM), using their earlyexit version of the T5 model [57] for text summarization on CNN/DM [33] and question answering on SQuAD [58]. Recall that CALM makes use of the Learn-then-Test [1] (LTT) framework for early-exit prediction control, whereas we suggest the Upper Confidence Bound [8] (UCB). In contrast to their experiments which involve excessively large calibration sets ( n ≈ 8000 ), we explore more practical settings of low calibration sample counts with n ∈ { 100 , 1000 } . Our results for the performance gap risk (Eq. 6) based on task-specific losses ( ROUGE-L for CNN/DM and Token-F1 for SQuAD) are displayed in Fig. 2 and § D.3. In all cases, UCB exit thresholds provide larger computational savings over LTT, while ensuring the same risk control with high probability (Eq. 5). Since particularly pronounced for n = 100 , these results highlight the need for careful framework selection in order to minimize the cost of providing guarantees. Once more, risk control in expectation (CRC, Prop. 1) permits faster exiting due to its weaker safety assurance. Encouragingly, even with as few as n = 100 calibration samples, CRC exit thresholds reach near-optimal exiting, as indicated by their proximity to the ideal (diagonal) risk line. This suggests that even for modern language tasks, equipping an EENN with notions of safety does not necessitate a strong compromise on inference efficiency.

5.4 Image Generation with Early-Exit Diffusion

To demonstrate the wide-ranging applicability of our proposal, we lastly consider early-exiting for image generation with diffusion. We employ the DeeDiff model [69], which performs early-exiting on the denoising network at each sampling step during the reverse diffusion process 4 . We target control of the perceptual difference between images generated by the accelerated and full diffusion processes, which we measure with the LPIPS score [82], and where lower values indicate perceptually closer images. Our results on the CelebA dataset [49] are shown in Fig. 5 for both risk control via CRC (Prop. 1) and UCB (Prop. 2), asserting that the risk is controlled at all levels ϵ . The impact of the risk level on image generation is additionally visualized for two examples. For strict control requirements the early-exit generations perceptually resemble the full model, whereas generated image quality visibly deteriorates for larger ϵ (but remains controlled). For smaller ϵ , the speed-ups within each sampling step are relatively modest ( e.g ., ∼ 15% for ϵ = 0 . 05 ). However, such gains accrue over the large number of sampling steps in the image generation process ( ∼ 500), resulting in overall meaningful savings. Similar observations for CIFAR [42] are reported in § D.4.

5.5 Speculative Decoding for Large Language Models

While we primarily focus on early-exiting, our final experiment highlights how risk control can also be applied to other techniques for efficient inference. Here we consider accelerating the inference of large language models (LLMs) using the (soft) speculative decoding approach BiLD [41]. BiLD uses a small draft model to generate multiple tokens autoregressively while the original LLM is only employed for verification.

This step can be performed in a single forward pass for all proposed tokens, necessitating less computations from the larger, more expensive model. During verification the difference in token distributions between the models is computed, and the tokens generated by the draft model are rejected if the difference exceeds a tolerated threshold, triggering a 'rollback' (see Eq. 3 in [41]). We apply our risk control frameworks to this rollback threshold, which dictates (i) the similarity of generated text to the output produced solely by the original LLM, and (ii) the associated inference speed-ups in terms of sentences (or samples) per second. Our results in Fig. 6 for the performance gap risk (Eq. 6), as defined via the difference in sentence-level BLEU scores, corroborate our previous findings for language modeling (Fig. 2 and § 5.3). That is, our approaches via CRC (Prop. 1) and UCB (Prop. 2) provide meaningful speed-ups and improve upon the LTT method [1] used by CALM [65], while maintaining the desired risk control across test samples.

6 Discussion

Our work addresses how to select a 'safe' exiting mechanism for early-exit neural networks (EENNs). We propose balancing the EENN's efficiency vs. performance trade-off via risk control, ensuring that accelerated inference does not compromise the quality of the early-exit outputs. We validate our light-weight, post-hoc solution on a variety of tasks and improve upon prior work [65] (see § 4).

Limitations and Future Work. A key limitation of our work is the reliance on a single shared exit threshold among layers (Eq. 1). While using a shared exit threshold is common [69, 48, 40, 76, 83, 77], relaxing this condition could lead to further efficiency gains. This, however, introduces new challenges both in terms of theory ( e.g ., defining monotonicity requirements) and practice ( e.g ., substantially larger search spaces). Overcoming them by adopting recently proposed risk control techniques for high-dimensional thresholds [60, 71] could prove promising for future extensions. Another (simple) workaround is to reduce the multi-dimensional problem back to a single threshold by use of a threshold function , as partially explored in [65]. Instead of working directly with multiple thresholds, our risk control framework can then be applied to this scalar parameter. Additionally, multiple risk control extensions provide natural avenues for future work. Firstly, risk control as used in our work is achieved only marginally across observations (Eq. 3), and one could aspire for more granular exit-conditional control [60]. Secondly, the employed risk control frameworks define risk in terms of the expected loss. One could instead aim to control the tails of the loss, e.g ., via specific quantiles of interest [67]. Lastly, relaxing the i.i.d assumption on calibration and test data could help extend risk-controlling EENNs to scenarios with test-time distribution shifts [86] or to online updating strategies [23].

Broader Impacts. EENNs provide a simple and effective approach to dynamic computation [29]. Their computational savings can reduce energy costs and the carbon footprint, as well as allow the model to be deployed on resource-constrained hardware. By incorporating a 'safe' exit mechanism into these models, we improve their trustworthiness and strengthen the reliability of their intermediate outputs, along with any decisions based on them. This facilitates safer model deployment in realworld applications and contributes to more responsible decision-making. While we do not foresee any direct negative consequences from our work, improper use of our risk control framework can lead to violations or misinterpretations of its provided guarantees. This, in turn, can risk instilling a false sense of security. Overall, we believe that our work outlines an effective approach to improve the reliability of EENNs and to safely balance their inherent efficiency vs. performance trade-off. In doing so, it contributes to the goal of developing models that are ultimately fast yet safe .

BiLD (Soft Speculative Decoding)

0.010

R

Risk

0.005

Test

0.000

E

3.6

CRC

UCB

LTT

3.4

0.000

0.002

0.004

0.006

0.008

0.010

Risk Level e

Acknowledgments and Disclosure of Funding

We thank Mona Schirmer, Rajeev Verma, Christoph-Nikolas Straehle, Patrick Forré and Stephen Bates for helpful discussions and clarifications. We are also grateful to the anonymous reviewers who helped improve the work with their constructive feedback. This project was generously supported by the Bosch Center for Artificial Intelligence. Eric Nalisnick did not utilize resources from Johns Hopkins University for this project.

References

Anastasios N. Angelopoulos, Stephen Bates, Emmanuel J. Candès, Michael I. Jordan, and Lihua Lei. Learn then Test: Calibrating Predictive Algorithms to Achieve Risk Control. arXiv Preprint (arXiv:2110.01052) , 2021. 3, 6, 9, 10, 11

Anastasios N. Angelopoulos, Amit Pal Kohli, Stephen Bates, Michael Jordan, Jitendra Malik, Thayer Alshaabi, Srigokul Upadhyayula, and Yaniv Romano. Image-to-Image Regression with Distribution-Free Uncertainty Quantification and Applications in Imaging. International Conference on Machine Learning , 2022. 6

Anastasios N Angelopoulos, Stephen Bates, et al. Conformal prediction: A gentle introduction. Foundations and Trends in Machine Learning , 2023. 6

Anastasios N Angelopoulos, Stephen Bates, Adam Fisch, Lihua Lei, and Tal Schuster. Conformal risk control. International Conference on Learning Representations , 2024. 1, 3, 5, 6, 2

Sangmin Bae, Jongwoo Ko, Hwanjun Song, and Se-Young Yun. Fast and robust early-exiting framework for autoregressive language models with synchronized parallel decoding. Conference on Empirical Methods in Natural Language Processing , 2023. 6, 7

Guangji Bai, Zheng Chai, Chen Ling, Shiyu Wang, Jiaying Lu, Nan Zhang, Tingwei Shi, Ziyang Yu, Mengdan Zhu, Yifei Zhang, et al. Beyond efficiency: A systematic survey of resource-efficient large language models. arXiv Preprint (arXiv:2401.00625) , 2024. 6

Fan Bao, Shen Nie, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: A vit backbone for diffusion models. Conference on Computer Vision and Pattern Recognition , 2023. 7

Stephen Bates, Anastasios Angelopoulos, Lihua Lei, Jitendra Malik, and Michael Jordan. Distribution-free, risk-controlling prediction sets. Journal of the ACM , 2021. 1, 3, 5, 6, 9, 2, 11

Edmon Begoli, Tanmoy Bhattacharya, and Dimitri Kusnezov. The need for uncertainty quantification in machine-assisted medical decision making. Nature Machine Intelligence , 2019. 4

Omer Belhasin, Yaniv Romano, Daniel Freedman, Ehud Rivlin, and Michael Elad. Principal uncertainty quantification with spatial correlation for image restoration problems. Transactions on Pattern Analysis and Machine Intelligence , 2023. 6

Vidmantas Bentkus. On hoeffding's inequalities. Annals of Probability , 2004. 6, 3

Tolga Bolukbasi, Joseph Wang, Ofer Dekel, and Venkatesh Saligrama. Adaptive neural networks for efficient inference. International Conference on Machine Learning , 2017. 1

Glenn W Brier. Verification of forecasts expressed in terms of probability. Monthly Weather Review , 1950. 4, 3

Mauro Cettolo, Marcello Federico, Luisa Bentivogli, Jan Niehues, Sebastian Stüker, Katsuitho Sudoh, Koichiro Yoshino, and Christian Federmann. Overview of the iwslt 2017 evaluation campaign. International Workshop on Spoken Language Translation , 2017. 10

Joud Chataoui, Mark Coates, et al. Jointly-learned exit and inference for a dynamic neural network. International Conference on Learning Representations , 2023. 6

Xinshi Chen, Hanjun Dai, Yu Li, Xin Gao, and Le Song. Learning to stop while learning to predict. International Conference on Machine Learning , 2020. 2

Marius Cordts, Mohamed Omran, Sebastian Ramos, Timo Rehfeld, Markus Enzweiler, Rodrigo Benenson, Uwe Franke, Stefan Roth, and Bernt Schiele. The cityscapes dataset for semantic urban scene understanding. Conference on Computer Vision and Pattern Recognition , 2016. 8, 6

A Philip Dawid. Probability forecasting. Encyclopedia of Statistical Sciences , 2004. 3

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. Conference on Computer Vision and Pattern Recognition , 2009. 7

Zhun Deng, Thomas Zollo, Jake Snell, Toniann Pitassi, and Richard Zemel. Distributionfree statistical dispersion control for societal applications. Advances in Neural Information Processing Systems , 2024. 6

Maha Elbayad, Jiatao Gu, Edouard Grave, and Michael Auli. Depth-adaptive transformer. International Conference on Learning Representations , 2020. 2, 6

Zhengcong Fei, Xu Yan, Shuhui Wang, and Qi Tian. Deecap: Dynamic early exiting for efficient image captioning. Conference on Computer Vision and Pattern Recognition , 2022. 6

Shai Feldman, Liran Ringel, Stephen Bates, and Yaniv Romano. Achieving risk control in online learning settings. Transactions on Machine Learning Research , 2023. 6, 10

Matteo Fontana, Gianluca Zeni, and Simone Vantini. Conformal prediction: a unified review of theory and new challenges. Bernoulli , 2023. 6

Amir Gholami, Sehoon Kim, Zhen Dong, Zhewei Yao, Michael W Mahoney, and Kurt Keutzer. A survey of quantization methods for efficient neural network inference. Low-Power Computer Vision , 2022. 6

Tilmann Gneiting and Adrian E Raftery. Strictly proper scoring rules, prediction, and estimation. Journal of the American Statistical Association , 2007. 4, 3

Tilmann Gneiting, Fadoua Balabdaoui, and Adrian E Raftery. Probabilistic forecasts, calibration and sharpness. Journal of the Royal Statistical Society Series B: Statistical Methodology , 2007. 4, 3

Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Weinberger. On calibration of modern neural networks. International Conference on Machine Learning , 2017. 3

Yizeng Han, Gao Huang, Shiji Song, Le Yang, Honghui Wang, and Yulin Wang. Dynamic neural networks: A survey. Transactions on Pattern Analysis and Machine Intelligence , 2021. 1, 6, 10

Yizeng Han, Yifan Pu, Zihang Lai, Chaofei Wang, Shiji Song, Junfeng Cao, Wenhui Huang, Chao Deng, and Gao Huang. Learning to weight samples for dynamic early-exiting networks. European Conference on Computer Vision , 2022. 7, 6

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. International Conference on Computer Vision , 2023. 7, 6

Geoffrey Heal and Antony Millner. Reflections: Uncertainty and decision making in climate change economics. Review of Environmental Economics and Policy , 2014. 4

Karl Moritz Hermann, Tomas Kocisky, Edward Grefenstette, Lasse Espeholt, Will Kay, Mustafa Suleyman, and Phil Blunsom. Teaching machines to read and comprehend. Advances in Neural Information Processing Systems , 2015. 9

Hans Hersbach. Decomposition of the continuous ranked probability score for ensemble prediction systems. Weather and Forecasting , 2000. 3

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 2020. 7

Gao Huang, Danlu Chen, Tianhong Li, Felix Wu, Laurens van der Maaten, and Kilian Weinberger. Multi-scale dense networks for resource efficient image classification. International Conference on Learning Representations , 2018. 1, 2, 6, 7

Metod Jazbec, James Allingham, Dan Zhang, and Eric Nalisnick. Towards anytime classification in early-exit architectures by enforcing conditional monotonicity. Advances in Neural Information Processing Systems , 2024. 2, 5, 3

Metod Jazbec, Patrick Forré, Stephan Mandt, Dan Zhang, and Eric Nalisnick. Early-exit neural networks with nested prediction sets. Conference on Uncertainty in Artificial Intelligence , 2024. 6

Ying Jin, Zhimei Ren, and Emmanuel J Candès. Sensitivity analysis of individual treatment effects: A robust conformal inference approach. Proceedings of the National Academy of Sciences , 2023. 6

Yigitcan Kaya, Sanghyun Hong, and Tudor Dumitras. Shallow-deep networks: Understanding and mitigating network overthinking. International Conference on Machine Learning , 2019. 2, 5, 10, 3

Sehoon Kim, Karttikeya Mangalam, Suhong Moon, Jitendra Malik, Michael W Mahoney, Amir Gholami, and Kurt Keutzer. Speculative decoding with big little decoder. Advances in Neural Information Processing Systems , 2024. 9, 10

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Toronto, ON, Canada , 2009. 9, 12

Gilad Kutiel, Regev Cohen, Michael Elad, and Daniel Freedman. What's behind the mask: Estimating uncertainty in image-to-image problems. arXiv Preprint (arXiv:2211.15211) , 2022. 6

Loïc Lannelongue, Jason Grealey, and Michael Inouye. Green algorithms: quantifying the carbon footprint of computation. Advanced Science , 2021. 1

Bracha Laufer-Goldshtein, Adam Fisch, Regina Barzilay, and Tommi S Jaakkola. Efficiently controlling multiple risks with pareto testing. International Conference on Learning Representations , 2022. 6

Bracha Laufer-Goldshtein, Adam Fisch, Regina Barzilay, and Tommi Jaakkola. Risk-controlling model selection via guided bayesian optimization. arXiv Preprint (arXiv:2312.01692) , 2023. 6

Shiya Liu, Dong Sam Ha, Fangyang Shen, and Yang Yi. Efficient neural networks for edge devices. Computers & Electrical Engineering , 2021. 1

Zhuang Liu, Zhiqiu Xu, Hung-Ju Wang, Trevor Darrell, and Evan Shelhamer. Anytime dense prediction with confidence adaptivity. International Conference on Learning Representations , 2021. 2, 6, 7, 10

Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. International Conference on Computer Vision , 2015. 9

Arnab Neelim Mazumder, Jian Meng, Hasib-Al Rashid, Utteja Kallakuri, Xin Zhang, Jae-Sun Seo, and Tinoosh Mohsenin. A survey on the optimization of neural network accelerators for micro-ai on-device inference. IEEE Journal on Emerging and Selected Topics in Circuits and Systems , 2021. 1

Lassi Meronen, Martin Trapp, Andrea Pilzer, Le Yang, and Arno Solin. Fixing overconfidence in dynamic neural networks. Winter Conference on Applications of Computer Vision , 2024. 2, 6

Allan H Murphy. A new vector partition of the probability score. Journal of Applied Meteorology and Climatology , 1973. 3

Xuchen Pan, Yanxi Chen, Yaliang Li, Bolin Ding, and Jingren Zhou. Ee-tuning: An economical yet scalable solution for tuning early-exit large language models. arXiv Preprint (arXiv:2402.00518) , 2024. 6

Sangdon Park, Osbert Bastani, Nikolai Matni, and Insup Lee. Pac confidence sets for deep neural networks via calibrated prediction. International Conference on Learning Representations , 2020. 6

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

Victor Quach, Adam Fisch, Tal Schuster, Adam Yala, Jae Ho Sohn, Tommi S Jaakkola, and Regina Barzilay. Conformal language modeling. International Conference on Learning Representations , 2023. 6

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research , 2020. 9

Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text. Conference on Empirical Methods in Natural Language Processing , 2016. 9

Stephan R Richter, Vibhav Vineet, Stefan Roth, and Vladlen Koltun. Playing for data: Ground truth from computer games. European Conference on Computer Vision , 2016. 8, 6

Liran Ringel, Regev Cohen, Daniel Freedman, Michael Elad, and Yaniv Romano. Early time classification with accumulated accuracy gap control. International Conference on Machine Learning , 2024. 6, 10

Swami Sankaranarayanan, Anastasios Nikolas Angelopoulos, Stephen Bates, Yaniv Romano, and Phillip Isola. Semantic uncertainty intervals for disentangled latent spaces. Advances in Neural Information Processing Systems , 2022. 6

Kavya Saravanan and Abbas Z Kouzani. Advancements in on-device deep neural networks. Information , 2023. 1

Simone Scardapane, Michele Scarpiniti, Enzo Baccarelli, and Aurelio Uncini. Why should we add early exits to neural networks? Cognitive Computation , 2020. 1

Tal Schuster, Adam Fisch, Tommi Jaakkola, and Regina Barzilay. Consistent accelerated inference via confident adaptive transformers. Conference on Empirical Methods in Natural Language Processing , 2021. 6

Tal Schuster, Adam Fisch, Jai Gupta, Mostafa Dehghani, Dara Bahri, Vinh Tran, Yi Tay, and Donald Metzler. Confident adaptive language modeling. Advances in Neural Information Processing Systems , 2022. 2, 4, 6, 9, 10, 3, 7, 11

Glenn Shafer and Vladimir Vovk. A tutorial on conformal prediction. Journal of Machine Learning Research , 2008. 6

Jake Snell, Thomas P Zollo, Zhun Deng, Toniann Pitassi, and Richard Zemel. Quantile risk control: A flexible framework for bounding the probability of high-loss predictions. International Conference on Learning Representations , 2022. 6, 10

Max Sponner, Bernd Waschneck, and Akash Kumar. Adapting neural networks at runtime: Current trends in at-runtime optimizations for deep learning. ACM Computing Surveys , 2024. 6

Shengkun Tang, Yaqing Wang, Caiwen Ding, Yi Liang, Yao Li, and Dongkuan Xu. Deediff: Dynamic uncertainty-aware early exiting for accelerating diffusion model generation. arXiv Preprint (arXiv:2309.17074) , 2023. 6, 9, 10, 7, 12

Surat Teerapittayanon, Bradley McDanel, and Hsiang-Tsung Kung. Branchynet: Fast inference via early exiting from deep neural networks. International Conference on Pattern Recognition , 2016. 1, 2, 6

Jacopo Teneggi, Matthew Tivnan, Web Stayman, and Jeremias Sulam. How to trust your diffusion model: A convex optimization approach to conformal risk control. International Conference on Machine Learning , 2023. 6, 10

Alexander Timans, Christoph-Nikolas Straehle, Kaspar Sakmann, and Eric Nalisnick. Adaptive bounding box uncertainties via two-step conformal prediction. European Conference on Computer Vision , 2024. 6

Jingdong Wang, Ke Sun, Tianheng Cheng, Borui Jiang, Chaorui Deng, Yang Zhao, Dong Liu, Yadong Mu, Mingkui Tan, Xinggang Wang, et al. Deep high-resolution representation learning for visual recognition. Transactions on Pattern Analysis and Machine Intelligence , 2020. 6

Yulin Wang, Rui Huang, Shiji Song, Zeyi Huang, and Gao Huang. Not all images are worth 16x16 words: Dynamic transformers for efficient image recognition. Advances in Neural Information Processing Systems , 2021. 7, 6

Ian Waudby-Smith and Aaditya Ramdas. Estimating means of bounded random variables by betting. Journal of the Royal Statistical Society Series B: Statistical Methodology , 2024. 5, 1, 2

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

Ji Xin, Raphael Tang, Jaejun Lee, Yaoliang Yu, and Jimmy Lin. Deebert: Dynamic early exiting for accelerating bert inference. Proceedings of the Association for Computational Linguistics , 2020. 10

Canwen Xu and Julian McAuley. A survey on dynamic neural networks for natural language processing. Findings of the Association for Computational Linguistics , 2023. 6

Mengwei Xu, Wangsong Yin, Dongqi Cai, Rongjie Yi, Daliang Xu, Qipeng Wang, Bingyang Wu, Yihao Zhao, Chen Yang, Shihe Wang, et al. A survey of resource-efficient llm and multimodal foundation models. arXiv Preprint (arXiv:2401.08092) , 2024. 6

Xiaowei Xu, Yukun Ding, Sharon Xiaobo Hu, Michael Niemier, Jason Cong, Yu Hu, and Yiyu Shi. Scaling for edge inference of deep neural networks. Nature Electronics , 2018. 1

Yunpeng Xu, Wenge Guo, and Zhi Wei. Conformal risk control for ordinal classification. Conference on Uncertainty in Artificial Intelligence , 2023. 6

Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as a perceptual metric. Conference on Computer Vision and Pattern Recognition , 2018. 9, 7

Wangchunshu Zhou, Canwen Xu, Tao Ge, Julian McAuley, Ke Xu, and Furu Wei. Bert loses patience: Fast and robust inference with early exit. Advances in Neural Information Processing Systems , 2020. 6, 10

Zixuan Zhou, Xuefei Ning, Ke Hong, Tianyu Fu, Jiaming Xu, Shiyao Li, Yuming Lou, Luning Wang, Zhihang Yuan, Xiuhong Li, et al. A survey on efficient inference for large language models. arXiv Preprint (arXiv:2404.14294) , 2024. 6

Shlomo Zilberstein. Using anytime algorithms in intelligent systems. AI Magazine , 1996. 2

Thomas P Zollo, Todd Morrill, Zhun Deng, Jake Snell, Toniann Pitassi, and Richard Zemel. Prompt risk control: A rigorous framework for responsible deployment of large language models. International Conference on Learning Representations , 2023. 6, 10

APPENDIX

The appendix is organized as follows:

Appendix A contains mathematical details, namely proofs for our propositions (§ A.1), elaboration on bounding conditions of the loss function (§ A.2), and further details on our Brier score formulation (§ A.3).

Appendix B contains algorithmic descriptions of risk control via the Upper Confidence Bound [8] (UCB, Prop. 2) with the Waudby-Smith-Ramdas bound [75] (WSR, Prop. 3) in Algorithm 1 and Algorithm 2.

Appendix C contains various implementation details of our four experiments: image classification (§ C.1), semantic segmentation (§ C.2), language modeling (§ C.3), and image generation (§ C.4).

Appendix D contains additional results across our suite of experiments (in the same order as above), including risk control and efficiency curves for varying risk levels ϵ (§ D.1, § D.2, § D.3, § D.4), as well as additional efficiency gain tables (§ D.1, § D.2).

A Mathematical Details

A.1 Proofs for Risk Control

We begin by formalizing the connection between marginal monotonicity requirements on the EENN (Eq. 2) and the monotonicity of risks (Eq. 3) in Lemma 1 below.

Lemma 1. A marginally monotone EENN satisfying Eq. 2 for some arbitrary loss function ℓ implies monotone decreasing risks of the form in § 3.2, i.e ., we have that R ( λ 1 ) ≥ R ( λ 2 ) for λ 1 ≤ λ 2 .

Proof. For a given test sample x , denote the exit layers corresponding to exit thresholds λ 1 and λ 2 as l 1 and l 2 . From the EENN's confidence-based exiting mechanism in Eq. 1 it follows that l 1 ≤ l 2 , i.e ., a smaller exit threshold λ 1 will result in exits that are earlier or equal to the larger threshold λ 2 . From our risk formulation in terms of relative exit performances in § 3.2 and marginal monotonicity according to Eq. 2 it then follows that

Next, we prove Prop. 1, our adaptation of the Conformal Risk Control [4] (CRC) guarantee on risk control in expectation for the early-exit setting. Our proof closely follows the proof for Thm. 1 in Angelopoulos et al. [4], but we relax the condition on monotone losses to that on monotone risks , which implies assuming marginal monotonicity on the EENN according to Lemma 1. We restate our proposition from the main paper for completeness first.

Proposition 1. Let ℓ : Λ → ( -∞ , B ] be a right-continuous bounded loss, and assume a marginally monotone EENN (Eq. 2). Then the exit threshold ˆ λ CRC ensures risk control in expectation, i.e ., it holds that E D cal ∼P n [ R ( ˆ λ CRC ) ] ≤ ϵ for any ϵ ∈ (0 , 1) .

Proof. Consider the calibration set D cal = { ( x i , y i ) } n i =1 ∼ P n and some test sample ( x , y ) ∼ P drawn i.i.d from P , and denote their union set as ˜ D := D cal ∪ ( x , y ) . Additionally, define ℓ i ( λ ) := ℓ (ˆ o λ ( x i ) , y i ) as the loss of the EENN's early-exit output for the i -th sample. In particular, ℓ n +1 ( λ ) now denotes the test sample's loss. Let us first recall the definition of ˆ λ CRC (Eq. 10):

Note that Λ is a discrete grid of values over [0 , 1] , e.g ., equidistant values { 0 . 01 , 0 . 02 , . . . , 1 } , and ff ˆ Λ = ∅ and the risk control condition is never met, we default to ˆ λ = 1 for all threshold selection

procedures. Thus, the min is well-defined. Next, consider the empirical risk ˆ R n +1 ( λ ; ˜ D ) computed over ˜ D using the n +1 available samples, and define the following threshold choice:

ˆ λ n +1 is always well-defined, since ˆ R n +1 ( λ = 1) is zero for the risks introduced in § 3.2. As we assume a bounded loss function ℓ ≤ B , we observe that for any λ ∈ Λ we have

which implies that ˆ λ n +1 ≤ ˆ λ CRC. Since we assume a marginally monotone EENN, by Lemma 1 it follows that the risk is monotone decreasing and by R ( ˆ λ n +1 ) ≥ R ( ˆ λ CRC ) we also have that

By Eq. 3 and our loss definition above, we can rewrite for a general λ the expression

in short E [ ℓ n +1 ( λ )] . The remainder of the proof follows Thm. 1 in Angelopoulos et al. [4]. Assume a particular set ˜ D is given. Then the threshold ˆ λ n +1 is void of randomness and a constant, and by the i.i.d condition we also have that ℓ n +1 ( λ ) | ˜ D ∼ Unif ( { ℓ 1 , . . . , ℓ n +1 } ) is uniformly distributed. Combining these observations with the law of total expectation (l.o.t.e.) and right-continuity (r.c.) of the loss, the final result follows:

Next, we sketch a proof for Prop. 2, our adaptation of the Upper Confidence Bound [8] (UCB) guarantee on risk control with high probability for the early-exit setting. Our main change includes modifying the Waudby-Smith-Ramdas bound [75] (WSR) to relax the bounding condition on the loss from ℓ ∈ [0 , 1] to ℓ ∈ [ -B,B ] (Prop. 3). We first restate our proposition from the main paper.

Proposition 2. Let ℓ : Λ → [ -B,B ] be a bounded loss, and assume a marginally monotone EENN (Eq. 2). Then the exit threshold ˆ λ UCB ensures risk control with high probability, i.e ., it holds that P D cal ∼P n ( R ( ˆ λ UCB ) ≤ ϵ ) ≥ 1 -δ for any ( ϵ, δ ) ∈ (0 , 1) 2 .

Proof. Our result follows almost directly from the proofs in Bates et al. [8] (for their Thm. A.1), where we leverage the required risk monotonicity by Lemma 1. We omit the technical requirement on risk continuity from the original proof, since a relaxation to non-continuous risks is permitted ([8], Remark 3). A proof that the WSR bound can be used to construct a valid upper confidence bound can be found in Bates et al. [8], Sec. 3.1.3. However, an assumption on losses bounded to ℓ ∈ [0 , 1] is made, which is overly restrictive for the early-exit setting. We relax this assumption to ℓ ∈ [ -B,B ] in Prop. 3 below, concluding the result.

Since our risk definitions in § 3.2 can naturally assume negative values, and we thus want to account for the occurrence of earlier exits performing better, we relax the bounding condition on the loss function for the Waudby-Smith-Ramdas bound [75] (WSR) in the following proposition.

Proposition 3. A valid upper confidence bound (Eq. 11) based on the Waudby-Smith-Ramdas bound can be constructed for losses ℓ ∈ [ -B,B ] with B > 0 .

Proof. Observe the definitions of individual components ( µ i , σ 2 i , ν i , κ i ) for the WSR bound in Algorithm 2. In particular, define ν i = min { 1 / 2 B, √ 2 log(1 /δ ) nσ 2 i -1 } . Since the second term is always non-negative, it follows that ν i ∈ [0 , 1 / 2 B ] . For the loss ℓ i ∈ [ -B,B ] , we then have that ( ℓ i -E [ ℓ i ]) ∈ [ -2 B, 2 B ] . Hence, it follows that 1 -ν i ( ℓ i -E [ ℓ i ]) ≥ 0 , which implies that κ i = ∏ i j =1 { 1 -ν j ( ℓ j -E [ ℓ j ]) } is a non-negative martingale. The rest of the proof follows Prop. 5 from Bates et al. [8], concluding the result.

Observe that the proof for Prop. 3 utilizes the fraction 1 / 2 B in its definition of ν i which can make the WSR bound less tight. However, in the case of a marginally monotone EENN (Lemma 1) we can relax this fraction to 1 /B instead, a comment we formalize in the following remark.

Remark 1. Note that in the case of a marginally monotone EENN, the bound from Prop. 3 can be further optimized. For a bounded loss ℓ ∈ [0 , B ] , the relative early-exit loss ℓ ( λ, L ) := ℓ ( ˆ o λ ( x ) , y ) -ℓ ( ˆ o L ( x ) , y ) will fall in the [ -B,B ] range (see also § 3.2 and § A.2). Additionally, due to the marginal monotonicity assumption (Eq. 2) the risk based on the relative loss will be non-negative, i.e ., R ( λ, L ) = E [ ℓ ( λ, L )] ∈ [0 , B ] . This implies that ( ℓ ( λ, L ) -R ( λ, L )) ∈ [ -2 B,B ] . Hence, κ i will be a non-negative martingale even when the upper bound 1 /B is used for ν i instead of 1 / 2 B .

A.2 On Bounding of the Loss Function

The risks outlined in § 3.2 rely on an early-exit loss definition in terms of relative exit performance. That is, our risks take the general form

with ℓ ( λ, L ) denoting the relative early-exit loss. Recall that ˆ o λ and ˆ o L are based on ˆ p λ and ˆ p L , respectively. For a bounded loss ℓ ∈ [0 , B ] , we then have that ℓ ( λ, L ) ∈ [ -B,B ] . In our early-exit setting, negative losses have an intuitive interpretation. The associated test samples are cases where the EENN overthinks [40, 37], i.e ., the early-exit ˆ p λ performs better than the final exit ˆ p L .

Risk control via CRC (Prop. 1) or UCB (Prop. 2) with the relaxed WSR bound (Prop. 3) conveniently account for such occurrences. In contrast, this presents a challenge for the Learn-then-Test [1] (LTT) framework employed by Schuster et al. [65], since the underlying Hoeffding-Bentkus [11] (HB) bound requires ℓ ( λ, L ) ∈ [0 , 1] . As a workaround, Schuster et al. [65] instead impose a lower loss limit of zero, i.e ., they use

While solving their technical requirement, it introduces a key drawback in that the risk control procedure cannot account for samples where the risk requirement is satisfied 'for free'. This introduces substantially more conservative early-exiting (see Fig. 7), since an upper bound on the empirical calibration risk is used for threshold tuning.

In addition, observe that the Brier score is naturally bounded by [0 , 2] for the multi-class setting [13]. Thus, our relative early-exit Brier losses assume values in the range of [ -2 , 2] . While acceptable for risk control via CRC and UCB (by setting B = 2 ), it once again does not align with the LTT requirement on ℓ ( λ, L ) ∈ [0 , 1] . Thus, applying LTT requires additional restrictions such as scaling ( e.g ., with a 1 / 2 term) and Eq. 16 to satisfy the bounds. This highlights another drawback where the intuitive risk interpretation as a Brier score difference is partially lost (see § A.3 below).

Note that while CRC and UCB are amenable to ℓ ( λ, L ) ∈ [ -B,B ] , it can still be beneficial to ensure non-negative losses ( e.g ., by Eq. 16) in order to improve the marginal monotonicity of the EENN (Eq. 2). Namely, it can happen that an EENN is not marginally monotone for an unrestricted loss (Eq. 15), but is for its zero-bounded counterpart (Eq. 16). Hence, such approaches might be useful when there is reason to believe that EENN violates its marginal monotonicity assumption, though in such cases a practitioner might better opt for the pruning of unnecessary model layers instead.

A.3 Brier score formulation

Brier score motivation. The Brier score is a strictly proper scoring rule [26, 27], ensuring its suitability to assess probabilistic forecasts. This can be demonstrated by its decomposition structure, which highlights that both calibration and sharpness properties of the forecaster are taken into account [52, 18]. Moreover, its mathematical formulation lends itself favorably to risk control when compared to other widely used probabilistic metrics. These include the expected calibration error (ECE [28]), which requires binning and thus cannot be defined at a per-sample level; the negative log-likelihood, which is less interpretable and unbounded; the ranked probability score (RPS), which does not treat class distances equally; and the continuous ranked probability score (CRPS [34]) or f -divergences like the Hellinger distance, which can be overly conservative by aiming to control the (potentially long) tail of the distribution, and require access to the full (ground truth) distribution. While amenable to our overall risk control framework, such probabilistic metrics, or any derived top-k uncertainty measures, seem either less practical or less principled.

0.00

0.25

0.50

0.75

1.00

λ

0.00

0.05

0.10

Risk

ˆ

R

/lscript

(

λ

)

-

/lscript

(

1

)

{

/lscript

(

λ

)

-

/lscript

(

1

)

, 0

}

0.05

0.10

0.15

Risk Level

/epsilon1

0.00

0.05

0.10

0.15

Test Risk

ˆ

R

CRC

UCB

LTT

0.05

0.10

0.15

Risk Level

/epsilon1

5

10

15

↓

)

[

0, 1

]

[

-

1, 1

]

Brier loss and Brier score. For convenience (but easily transferable), consider the supervised setting where we target risk control of the performance gap risk (Eq. 6) for predictive distributions, which we denote in short as R G (ˆ p ) . This requires computing the Brier loss ℓ B (ˆ p l ( y | x ) , y ) , which we define for the multi-class setting in Eq. 7. Now consider a dataset D ∼ P of size N ( e.g ., the calibration set D cal ). The associated Brier score [13] denoted Brier (ˆ p l ) for the l -th exit layer is then defined as the mean Brier loss across samples, i.e ., we have

where ℓ B,n (ˆ p l ) is an abbreviation for the n -th sample's Brier loss. The risk R G (ˆ p ) that we aim to control is approximated by its empirical equivalent ˆ R G (ˆ p ) on D , and can then be interpreted as the difference in Brier scores between our EENN with threshold λ and the full (last-layer exit) model:

Mean-pixel Brier score. Since our semantic segmentation experiment (§ 5.2) is a pixel-level classification task, we obtain pixel-level predictive distributions, and thus compute per-pixel Brier losses. For an image of height H and width W , we can denote by ℓ B,n (ˆ p l,i,j ) the n -th sample's Brier loss at the pixel location ( i, j ) ∈ H × W . Since we target image-level early-exiting, these per-pixel Brier losses are averaged across pixels to compute a per-image Brier loss, which in turn is averaged across samples to obtain the l -th exit layer's Brier score. That is, the layer's associated Brier score Brier (ˆ p l ) is defined as

and can be interpreted as the mean-pixel Brier score of the particular exit layer. Note that if we interpret every image pixel as an individual sample ( i.e ., define ˜ N := NHW ), the Brier score formulation as a sample-averaged Brier loss continues to hold. Similar to above, the targeted risk is then interpreted as the mean-pixel Brier score difference.

B Algorithmic Details

Here, we sketch the algorithm for computing the risk-controlling threshold ˆ λ UCB (Eq. 12) via the Upper Confidence Bound [8] (UCB, Prop. 2) with the Waudby-Smith-Ramdas bound [75] (WSR, Prop. 3). The algorithm is an adaptation of the approach presented in Bates et al. [8] to the early-exit setting. Note that our practical code implementation differs slightly from the pseudo-code presented here, as we omit here some code optimization steps such as vectorization to improve readability.

Algorithm 1: Risk control for EENNs via UCB (Prop. 2)

C Implementation Details

Since our approach is post-hoc , the required compute resources needed are minimal. The primary requirement is sufficient memory to process the required calibration and test data, which can be quite large ( e.g ., images of size 2048 × 1024 pixels for Cityscapes). All our experiments can be performed and replicated on a single A100 GPU with experiment runtimes of < 1 day. We detail further case-specific implementations below for each experiment.

C.1 Image Classification

Model details. The models we consider are: the multi-scale dense network ( MSDNet ; [36]), an adaptation of traditional convolutional NNs for the early-exit setting; the dynamic vision transformer ( DViT ; [74]), which comprises multiple transformers with an increasing number of input patches; an enhanced MSDNet model that weights easy and hard examples differently during training ( L2WDEN ; [30]); and a recently proposed dynamic perceiver ( Dyn-Perc ; [31]), which decouples feature extraction and early prediction tasks via a novel dual-branch architecture. For all models, we either work with the publicly available pretrained checkpoints or train the models ourselves, closely following the original implementation details.

C.2 Semantic Segmentation

Model details. We consider the EENN proposed by Liu et al. [48] (ADP-C), which adds three intermediate exit heads to the HRNet segmentation model [73] (for a total of four exits) and is trained end-to-end on Cityscapes [17]. The model comes in two sizes, small (W18) and large (W48). We focus on the larger model (ADP-C-W48), but find results hold equivalently for the smaller one (in fact, larger gains can be obtained). Across experiments we employ the publicly available model checkpoints from the original implementation 5 .

Model finetuning. Since the model is trained on Cityscapes, we consider additional finetuning to evaluate ADP-C on GTA5 [59]. We take the available pre-trained model checkpoint (APD-CW48) and finetune the model for 50 epochs on the GTA5 training set ( ∼ 12000 images). For this purpose, we employ the original training script and training parameters ( e.g ., learning rate, batch size, etc .). However, we find that our finetuned model does not perform on par with the original, i.e ., performance on GTA5 is substantially inferior to that on Cityscapes. In particular, the performance improvement across subsequent exits on GTA5 is marginal, resulting in an EENN that is less suitable for early-exiting (see also § D.2). Yet, we find that risk control frameworks still apply, highlighting their robust model-agnostic properties even in light of an inferior underlying predictor.

Image-level aggregation. ADP-C provides an exiting mechanism following Eq. 1 on pixel-level , which is less useful for down-stream applications and decision making. For details on the exact mechanism, we refer to Liu et al. [48]. Rather than exiting only for selected image pixels, we instead want to early-exit the entire image whilst ensuring risk control. Thus, alongside different per-pixel confidence scores ˆ c l,i,j , ( i, j ) ∈ H × W , we also consider confidence aggregations ϕ ( · ) which produce a single image-level confidence measure ˆ c l to perform image-level risk control. Note that our prediction losses mean intersection-over-union and miscoverage already aggregate from pixel- to image-level, whereas our distributional loss (Eq. 7) is adapted to additionally average over pixels, resulting in a mean-pixel Brier score interpretation (see § A.3).

Risk control evaluation. We evaluate the original ADP-C-W48 on Cityscapes validation data, with a split of 80% D cal ( i.e ., 400 images) and 20% D test ( i.e ., 100 images). Similarly, we randomly select a subset of 1000 images from the GTA5 validation set and evaluate our finetuned model using 80% D cal ( i.e ., 800 images) and 20% D test ( i.e ., 200 images). In both cases, we average risk control results over 100 trials of random calibration and test splits.

C.3 Language Modeling

Model details. For our language modeling experiments, we employ the early-exit pretrained model based on T5-large (770M parameters) from Bae et al. [5] 6 . While this model closely follows the implementations in Schuster et al. [65], we found it easier to work with than the original code 7 . Note that Schuster et al. [65] report results for T5-small (60M parameters) and T5-base (220M parameters), whereas we use the larger T5-large . For risk control evaluation, we follow their exact exiting mechanism. Specifically, we compute softmax-based confidences at every exit and deploy their threshold decay mechanism, where early-exiting is more conservative for initial tokens and becomes progressively more permissive ([65], Eq. 5).

C.4 Image Generation with Early-Exit Diffusion

Model details. For our image generation experiment, we re-implement the early-exit diffusion model proposed by Tang et al. [69] (DeeDiff), since the original code is not publicly available. We model our training procedure as closely as possible to the original. As suggested in the paper, we use the U-ViT transformer [7] as a backbone denoising network. Early-exiting in DeeDiff is performed on the denoising network at each sampling step. Specifically, for every sampling step t and exit layer e = 1 , . . . , L , a per-pixel confidence map c e,t is obtained. Then, c e,t is used to compute the global (scalar) confidence score c e,t by averaging the confidence across all pixels. If the scalar confidence score satisfies the exit condition c e,t ≥ λ , we proceed to the next denoising step t +1 , employing the output ( i.e ., the predicted noise) of the e -th exit layer at the t -th sampling step. The model is trained using a standard diffusion denoising loss [35] and two uncertainty-aware losses, closely following the approach described in Tang et al. [69].

LPIPS metric. Our task-specific prediction loss measures the perceptual difference between earlyexit and full model image generations. For this, we employ the LPIPS metric [82], which computes the similarity between activations of image patches for a selected pre-trained network. LPIPS values are in the range of [0 , 1] , with smaller values indicating perceptually more similar images.

D Further Experimental Results

D.1 Image Classification

Additional test risk and efficiency curves for calibration set size n = 1000 on ImageNet are displayed in Fig. 8, while tables with efficiency gain values on ImageNet are given in Table 2.

0.0

0.1

0.2

0.3

0.4

Test Risk

ˆ

R

R

G

(

ˆ

y

)(

0

-

1

)

R

G

(

ˆ

p

)

(Brier)

R

C

(

ˆ

y

)(

0

-

1

)

R

C

(

ˆ

p

)

(Brier)

0.0

0.1

0.2

0.3

0.4

1

2

3

4

5

↓

)

MSDNet

DViT

L2W-DEN

Dyn-Perc

0.0

0.1

0.2

0.3

0.4

CRC

UCB

0.0

0.1

0.2

0.3

0.4

0.0

0.1

0.2

0.3

0.4

Risk Level

/epsilon1

(a) UCB and

n

= 100

D.2 Semantic Segmentation

We report full tables with efficiency gains across risks and confidence measure combinations for Cityscapes (Table 3) and GTA5 (Table 4) in terms of mean exit layer and GFLOPS. In addition, we display test risk and efficiency curves on both Cityscapes (Fig. 9) and GTA5 (Fig. 10) for risk control via CRC (Prop. 1) and UCB (Prop. 2). The figures are for the simplest combination of Top-1 pixel-level confidence and mean image-level aggregation. Note that the figures for other confidence combinations are similar and thus omitted, with the test risk being controlled in all cases.

GTA5 results. We observe that for GTA5 the performance gap risk for prediction control R G (ˆ y ) (mIoU) seems particularly easy to control, with high gains reached even for very strict ϵ = 0 . 01 . This relates to the underlying predictor's generally inferior performance due to finetuning (see § C.2). The obtained model records lower performance and marginal improvements across exit layers, resulting in a small risk that is easy to control. Intuitively, the price paid by exiting early is marginal, since the early-exit layer performs almost on par with the final layer. Thus, the scale of the risk deviates from that of other risks, and more meaningful risk control should consider both improving the underlying predictor, and selecting a different scale of risk levels ϵ .

RG(9)(mIoU)

RG(p)(Brier)

RC(g)(Miscov.)

RC(p)(Brier)

0.25

0.25

Risk

0.20

Risk

0.20

0.15

0.15

0.10

Test

0.10

0.05

0.05

0.00

0.00

0.00

0.05

0.10

0.15

0.20

0.25

0.00

0.05

0.10

0.15

0.20

0.25

0.00

0.05

0.10

0.15

0.20

0.25

0.00

0.05

0.10

0.15

0.20

0.25

4

UCB

4

UCB

3

CRC

3

CRC

2

2

0.00

0.05

0.10

0.15

0.20

0.25

0.00

0.05

0.10

0.15

0.20

0.25

0.00

0.05

0.10

0.15

0.20

0.25

0.00

0.05

0.10

0.15

0.20

0.25

Risk Level e

Risk Level e

RG(9)(mIoU)

RG(p)(Brier)

RC(9)(Miscov.)

RC(p)(Brier)

0.030

0.25

0.25

Risk

0.025

0.20

Risk

0.20

0.020

0.15

0.15

Test

0.015

0.010

0.10

Test

0.10

0.005

0.05

一

0.05

0.000

0.00

0.00

0.000 0.005 0.010 0.015 0.020 0.025 0.030

0.00

0.05

0.10

0.15

0.20

0.25

0.00

0.05

0.10

0.15

0.20

0.25

0.00

0.05

0.10

0.15

0.20

0.25

4

UCB

4

UCB

3

CRC

3

CRC

2

2

2

0.000 0.005 0.010 0.015 0.020 0.025 0.030

0.00

0.05

0.10

0.15

0.20

0.25

0.00

0.05

0.10

0.15

0.20

0.25

0.00

0.05

0.10

0.15

0.20

0.25

Risk Level e

Risk Level e

(b) Efficiency gains in terms of GFLOPS improvement.

(a) Efficiency gains in terms of mean exit layer improvement.

(b) Efficiency gains in terms of GFLOPS improvement.

D.3 Language Modeling

0.00

0.05

0.10

0.15

0.20

Test Risk

ˆ

R

CNN/DM(

ncal

=

100)

0.00

0.05

0.10

0.15

0.20

CNN/DM(

ncal

=

1000)

0.00

0.05

0.10

0.15

0.20

SQuAD (

ncal

=

100)

0.00

0.05

0.10

0.15

0.20

SQuAD (

ncal

=

1000)

0.10

0.20

5.0

7.5

10.0

12.5

15.0

↓

)

0.10

0.20

5.0

7.5

10.0

12.5

15.0

0.10

0.20

1.5

2.0

2.5

3.0

0.10

0.20

1.5

2.0

2.5

3.0

CRC

UCB

LTT

State-Saturation

Meta-Classifier

Softmax

0.15

0.15

0.15

0.10

0.10

0.10

Test

0.05

0.05

0.05

0.00

0.00

0.00

Layer

20

20

20

CRC

UCB

LTT

Exit

10

10

10

0.05

0.10

0.15

0.05

0.10

0.15

0.05

0.10

0.15

Risk Level e

D.4 Image Generation with Early-Exit Diffusion

0.00

0.05

0.10

Test Risk

ˆ

R

'Full model'

/epsilon1

=

0.01

/epsilon1

=

0.05

/epsilon1

=

0.10

/epsilon1

=

0.15

0.00

0.05

0.10

Risk Level

/epsilon1

1

5

10

13

Exit (

↓

)

CRC

UCB

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We provide all theoretical (see § 3) and experimental (see § 5) results backing up our claims.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We clearly list out our key limitations in § 6.

3. Theory Assumptions and Proofs

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

Answer: [Yes]

Justification: We state our theoretical results in § 3.3 and provide the proofs in § A.1.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We describe our experiments in § 5 and provide further implementation details in Appendix C.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: Our code is publicly available at https://github.com/metodj/RC-EENN . In § 5 and Appendix C, we describe our experimental setup in detail. All employed datasets are publicly available.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We describe our experiments in § 5 and provide further implementation details in Appendix C.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We display 1-sigma standard deviations across multiple trials of random calibration and test data splitting in all figures with empirical test risk and efficiency curves. The exceptions are Fig. 3 and Fig. 8 where we omit them to improve readability, without affecting any interpretation of results.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Since our approach is post-hoc , the required compute resources needed are minimal. The primary requirement is sufficient memory to jointly load the required calibration and test data, which can be quite large ( e.g ., for 2048 × 1024 images on Cityscapes). All our experiments can be performed/replicated on a single A100 GPU with runtimes of < 1d. We mention this in Appendix C.

9. Code Of Ethics

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

Answer: [Yes]

Justification: We checked the Code of Ethics carefully and our work complies with all the points listed there.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: We clearly mention what we believe to be the impacts of our work in § 6.

11. Safeguards

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

Answer: [NA]

Justification: Our work doesn't introduce any new foundation or generation models, hence additional safeguards are not needed.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We made sure to properly cite all the related work and datasets used.

13. New Assets

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

Answer: [Yes]

Justification: We describe our theory in § 3.3. We also additionally provide pseudo-code for our main algorithm in Appendix B, and reference all underlying papers. Our code is publicly available at https://github.com/metodj/RC-EENN .

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: No experiments involving crowd-sourcing or human experiments were conducted in our research.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

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

Answer: [NA]

Justification: No experiments involving crowd-sourcing or human experiments were conducted in our research.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigate how to adapt frameworks of risk control to early-exit neural networks (EENN). The new framework is designed to exit without severely degrading performance as well as to guarantee the reliability of the prediction during early exiting.  The applied risk control technique offers a distribution-free, post-hoc solution that tunes the EENN’s exiting mechanism so that exits only occur when the output is of sufficient quality. Experiments on a series of tasks, including image classification, semantic segmentation, language modelling, and image generation show the effectiveness of the proposed method.\n\n-------------after rebuttal----------\n\nThe authors address most of my concerns. Although the single exiting threshold seems somehow a little flaw, the proposed new method is interesting and new.\n\nSTRENGTHS:\n1. The experiments on different vision tasks can be strong evidence to show the effectiveness of the proposed method. \n2. The theoretical analysis seems reasonable to be.\n\nWEAKNESSES:\nThe main weakness of this paper is under the strong assumption that all exiting thresholds are set to the same values. I admit that relaxing this assumption by adopting RC techniques for a high-dimensional threshold can lead to some difficulties in theoretical analysis, while using the same threshold for all exits may harm the overall performance of EENN. Only empirical test risk is used in all experiments. The general performance of EENN models in different tasks is also an important evaluation metric for EENN. \n\nThe applicability of the proposed RC method can be impeded if the prediction performance is unsatisfactory when using the same exiting threshold for every exit classifiers, even though it has a lower risk in the prediction.\n\nIt is also suggested to review the related important works in early exiting neural network areas, such as:\n[1] Improved techniques for training adaptive deep networks. in ICCV, 2019.\n[2] Resolution Adaptive Networks for Efficient Inference, in CVPR, 2020.\n[3] Learning to Weight Samples for Dynamic Early-exiting Networks, in ECCV, 2022.\n\nQUESTIONS:\nSee weakness. I think more evaluation metrics, such as test accuracy for classification tasks, or the IOU for segmentation tasks, should be included to show the effectiveness of the proposed method.\n\nLIMITATIONS:\nSee weakness. I think it is acceptable to use the same threshold for theoretical analysis. However, the real performance of the evaluated EENNs should also be attached in the manuscript.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-26"}

comments: SUMMARY:
This paper introduces a risk control framework for early-exit neural networks to balance the trade-off between inference efficiency and model performance. Compared to the conventional methods that manually set the confidence thresholds for early exiting, this work proposes a method to determine exit points from the perspective of risk control. Experiments on vision and language tasks demonstrate the effectiveness.

STRENGTHS:
* Novelty: The integration of risk control into EENNs is novel, and the studied problem is an interesting and important challenge in deploying early-exiting models;
* Empirical Validation: The paper provides extensive experimental evidence across multiple domains, demonstrating the practical applicability of the proposed methods;
* Soundness: The risk control framework is well-developed, offering both theoretical foundations and practical algorithms for determining exit thresholds.

WEAKNESSES:
Overall I like the paper. A small suggestion is to demonstrate the accuracy-FLOPs curves as in the compared methods (MSDNet, L2W-DEN, and Dyn-Perceiver), because the current x-axis, "risk level", might be insufficiently straightforward.

QUESTIONS:
Please see weaknesses.

LIMITATIONS:
yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-06-26

comments: SUMMARY:
This paper proposes a method for improving the efficiency of early-exit neural networks (EENNs) while maintaining performance. ​ EENNs allow for predictions to be made at intermediate layers, resulting in faster inference. ​ However, the challenge is determining when it is safe to exit without sacrificing performance. ​ The authors address this issue by adapting risk control frameworks to EENNs. ​ They propose risk functions to measure the performance drop resulting from early-exiting and demonstrate the effectiveness of their approach on various vision and language tasks. ​

STRENGTHS:
- The paper addresses an important problem in machine learning - improving the efficiency of inference without sacrificing performance. ​
- The authors propose a novel approach of adapting risk control frameworks to EENNs, which provides a post-hoc solution for determining when it is safe to exit. ​
- The paper provides empirical validation of their approach on a range of vision and language tasks, demonstrating substantial computational savings while preserving performance goals. ​
- The authors consider both prediction quality and uncertainty estimation in their risk control framework, which is important for safety-critical scenarios. ​

WEAKNESSES:
- The empirical validation of the approach is limited to a range of vision and language tasks. It would be beneficial to see the application of the method to other domains as well.
- The paper does not discuss the limitations or potential drawbacks of the proposed approach. It would be helpful to have a discussion on the potential trade-offs or challenges in implementing the risk control framework in practice.

QUESTIONS:
Refer to Weaknesses.

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-11

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis manuscript revisits the confidence-based thershold tuning for early-exit networks following a risk-control formulation. The proposed approach aims to provide an enhanced mechanism to identify confident predictions early-on during the computation, and exit early to improve inference efficiency with minimal effect on accuracy. The underlying trade-off between accuarcy and latency is balanced by a post-hoc solution, that determines the thershold value for the exits, through a risk control formulation considering a labelled calibration set (Performance Gap Risk), or the predictions of the final exit of the original model (Consistency Risk).\n\nSTRENGTHS:\n- The manuscript studies a very interesting and timely problem, often overlook by the booming filed of early-exit models. \n- The proposed approach is post-hoc and can be applied to existing approaches without the need of re-training/finetunning the model, and does not add any computational overhead at inference time. \n- The effectiveness of the proposed approach is experimentally demonstrated on a wide range of predictive (image classification and segmentation) and generative tasks (language modelling and image generation), proving the generality of the proposed approach. \n- The paper is well-written and offers adequate background information and a nice formulation.\n\nWEAKNESSES:\n- The reliance of the proposed approach to marginal monotonicity between performance and depth can be a limiting factor in some applications. E.g. recent works on EE-based LLMs indicate the predictive accuracy fluctuates notably across exits (see Elhoushi et al, LayerSkip, 2024 - Fig.2). \n- As stated in the limitations, the adoption of a single and static threshold across exits also significantly reduces the search space of the proposed approach. The inefficiency imposed by this design choice has not been quantified, but is likely to be significant without any calibratuon technics on exit confidence or special training of the model. \n- It is unclear how the proposed approach can be applied to the emerging line of works with learnable exit policies, as well as whether it will maintain its performance on different confidence metrics, such as state-saturation, used in LLMs etc\n\nQUESTIONS:\n- Please consider the comments raised on the Weaknesses section and share any insights you may have on them.\n- Can the proposed Consistency Risk formulation be used for online adaptation of the thresholds, as during runtime calibration labels are not available, but input distribution shifts may occur?\n\n**Post-Rebuttal Edit:**\nProvisionally increasing my score from WA to A, given the thorough clarifications provided by the authors that adequately addressed most of my concerns.\n\nLIMITATIONS:\nLimitations and future work are adequately discussed in the manuscript.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}