Fast yet Safe: Early-Exiting with Risk Control

Authors: metod  jazbec
alexander  timans
tin hadži veljković
kaspar  sakmann
dan  zhang
christian a naesseth
eric  nalisnick
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 18:09:48

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.

Main paper: 
 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). 
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 pl (y|x), permits to retrieve both a predicted class label ŷl = arg max y∈Y pl (y | x) and an associated confidence score ĉ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 ĉl = max y∈Y pl (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 exits2 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 pλ (y|x) := pe (y|x), where e = minEQUATIONThe 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 asEQUATIONfor 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) = 1[y / ∈ fλ (x)],where 1[•] 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) 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 λ ∈ Λ.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 thatEQUATIONwhere 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 thatP D cal ∼P n (R( λ) ≤ ϵ) ≥ 1 -δ ∀ λ ∈ Λ.(5) 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). 
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).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 pL (y|x).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 ôl (x) denote a general EENN output for some input x. It takes the form ŷl for predictions and pl (y|x) for the underlying predictive distribution. The supervised performance gap risk is then defined asEQUATIONwhere ôλ (x) and ô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 asEQUATIONwhere pl (k|x) denotes the predicted probability of a particular class k, and 1[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 pl (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 ŷL obtained from the full model. We define the unsupervised consistency risk asEQUATIONwhere only input samples x ∼ P are required for its evaluation. Marginally monotone performance trends (Eq. 2) are generally observed across models, with last-layer exits performing best.For regular prediction losses, Eq. 8 collapses to evaluating the per-sample loss ℓ( ŷλ , ŷL ), since ℓ( ŷL , ŷL ) = 0. For predictive distribution control, the loss difference remains, and we substitute y in Eq. 7 by sampling a label ŷL ∼ pL (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.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., λemp := min{λ ∈ Λ : R(λ; D cal ) ≤ ϵ}.Note that R(λ; D cal ) = 1 n n i=1 ℓ( ôλ (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 threshold3 λCRC := min λ ∈ Λ :EQUATIONguarantees 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 thatEQUATIONAn exit threshold ensuring risk control according to Eq. 5 is then selected asEQUATIONSimilarly 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 thatP 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 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. 
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). 
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 ŷl we employ target-specific losses, and, when applicable, for predictive 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.0.0 0.1 0.2 0.3 0.4 Test Risk R R G ( ŷ)(0-1) R G ( p) (Brier) R C ( ŷ)(0-1) R C ( p) (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 ( ŷ) 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.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 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. 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.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.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. [41] for a machine translation task (De-En) on IWSLT [14] (with n = 500). We fix the fallback threshold to 0.5, and apply risk control to the rollback threshold. Our adaptation of UCB again outperforms LTT by yielding larger efficiency gains under the same guarantees. Shading denotes the standard deviation across S = 100 calibration/test splits. 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. 
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).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.procedures. Thus, the min is well-defined. Next, consider the empirical risk Rn+1 (λ; D) computed over D using the n + 1 available samples, and define the following threshold choice: λn+1 := min{λ ∈ Λ : Rn+1 (λ; D) ≤ ϵ}.(13) λn+1 is always well-defined, since Rn+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 haveRn+1 (λ; D) = n n + 1 R(λ; D cal ) + ℓ n+1 (λ) n + 1 ≤ n n + 1 R(λ; D cal ) + B n + 1 ,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 thatEQUATION) By Eq. 3 and our loss definition above, we can rewrite for a general λ the expressionE D cal ∼P n [R(λ)] = E D cal ∼P n [E (x,y)∼P [ℓ n+1 (λ)]] i.i.d = E D∼P n+1 [ℓ n+1 (λ)],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:E[ℓ n+1 ( λCRC )] (Eq. 14) ≤ E[ℓ n+1 ( λn+1 )] l.o.t.e. = E E[ℓ n+1 ( λn+1 ) | D ] Unif = E 1 n + 1 n+1 i=1 ℓ i ( λn+1 ) (Eq. 13) & r.c. ≤ E[ϵ] = ϵ.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 thatP 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. 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/2B, 2 log(1/δ)nσ 2 i-1}. Since the second term is always non-negative, it follows that ν i ∈ [0, 1/2B]. For the loss ℓ i ∈ [-B, B], we then have that(ℓ i - E[ℓ i ]) ∈ [-2B, 2B]. 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/2B 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) := ℓ ôλ (x), yℓ ô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)) ∈ [-2B, B].Hence, κ i will be a non-negative martingale even when the upper bound 1/B is used for ν i instead of 1/2B.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 formEQUATIONwith ℓ(λ, L) denoting the relative early-exit loss. Recall that ôλ and ôL are based on pλ and pL , 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 pL .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 ℓ(λ, L) := max{0, ℓ ôλ (x), yℓ ôL (x), y }.(16) 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.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. Exit Layer (↓)[0, 1] [-1, 1]Figure 7 : We display risk curves based on relative early-exit losses ℓ(λ, L) for the CALM model [65] on CNN/DM (left). The risk based on the zero-bounded losses ( ; Eq. 16) naturally upper-bounds the one without ( ; Eq. 15). Since LTT requires ℓ(λ, L) ∈ [0, 1], threshold tuning is performed based on the larger risk ( ), despite aiming to control the actual risk ( ). This results in overly conservative early-exiting, as indicated by the LTT test risk ( ) deviating furthest from the optimal (diagonal) risk line (middle), and its lowest empirical gains (right). Since CRC and UCB permit negative losses, we display their results for both zero-bounded ( ; Eq. 16) and unrestricted ( ; Eq. 15) settings. Also here, the unrestricted setting results in controlled risks with larger efficiency gains, highlighting its relevance for the early-exit setting with relative losses of the form ℓ(λ, L).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 haveBrier(p l ) = 1 N N n=1 ℓ B,n (p l ) = 1 N N n=1 ℓ B (p l (y|x n ), y n ) (Eq. 7) = 1 N N n=1 K k=1 pl (k|x n )-1[y n = k] 2 ,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 RG (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:RG (p) = 1 N N n=1 ℓ B (p λ (y|x n ), y n ) -ℓ B (p L (y|x n ), y n ) = 1 N N n=1 ℓ B,n (p λ ) - 1 N N n=1 ℓ B,n (p L ) = Brier(p λ ) -Brier(p L ).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 asBrier(p l ) = 1 N HW N n=1 (i,j)∈H×W ℓ B,n (p l,i,j ) = 1 N HW N n=1 (i,j)∈H×W ℓ B (p l (y|x n,i,j ), y n,i,j ),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 HW ), 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.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).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.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 ( ŷ)(0-1) R G ( p) (Brier) R C ( ŷ)(0-1) R C ( p) (EQUATIONEQUATIONEQUATIONWe 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 ( ŷ) (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 ϵ. Left: The quality of generated images is directly related to the targeted risk control level ϵ.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.Such model-driven exiting is distinct from anytime settings, where exits are environment-driven[85].In Eq. 10 and Eq. 12, we default to λ = 1 if Λ = ∅, implying that risk control does not permit early-exiting.Since the code for DeeDiff is not publicly available, we implement it ourselves (see § C.4).See https://github.com/liuzhuang13/anytimeSee https://github.com/raymin0223/fast_robust_early_exitSee https://github.com/google-research/t5x/tree/main/t5x/contrib/calm 
 
 Appendix:  


Decision: ### Summary
This paper tries to improve the efficiency of deep models by exploring a type of Early-Exit Neural Network (EENN). EENNs accelerate inference by allowing intermediate layers to exit and produce a prediction early. However, this paper improves them by using approaches inspired by leveraging a method called as Risk Control (RC.)  RC manually sets the confidence thresholds for early exiting. The paper presents promising experimental results on vision and language tasks.

### Decision

The paper is well-written and clear. The approach is interesting, and the problem is important and potentially impactful. There were several criticisms directed toward the paper during the rebuttal period. However, the authors did a good job replying to them during the rebuttal period. The authors provided clarifications for many of the questions raised by the reviewers, and the proposed approach has merit on its own, as it studies a very interesting and often overlooked aspect of EE models and offers a holistic approach that demonstrates its effectiveness across models and tasks. Thus, I would recommend this paper for acceptance. 

However, I recommend that the authors to read the rebuttals carefully and address the reviewers' concerns in the camera-ready of the paper.

Reviews: review: paper_summary: paper_summary: This 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).
strength_weakness: strength_weakness: - 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). 
- 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. 
- 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
questions: questions: - Please consider the comments raised on the Weaknesses section and share any insights you may have on them.
- 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?

**Post-Rebuttal Edit:**
Provisionally increasing my score from WA to A, given the thorough clarifications provided by the authors that adequately addressed most of my concerns.
limitations: limitations: Limitations and future work are adequately discussed in the manuscript.
paper_summary: paper_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. ​
strength_weakness: strength_weakness: - 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: questions: Refer to Weaknesses.
limitations: limitations: N/A
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: Please see weaknesses.
limitations: limitations: yes
paper_summary: paper_summary: This 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.

-------------after rebuttal----------

The authors address most of my concerns. Although the single exiting threshold seems somehow a little flaw, the proposed new method is interesting and new.
strength_weakness: strength_weakness: The 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. 

The 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.

It is also suggested to review the related important works in early exiting neural network areas, such as:
[1] Improved techniques for training adaptive deep networks. in ICCV, 2019.
[2] Resolution Adaptive Networks for Efficient Inference, in CVPR, 2020.
[3] Learning to Weight Samples for Dynamic Early-exiting Networks, in ECCV, 2022.
questions: questions: See 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.
limitations: limitations: See 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.

Comments: comment: Thanks for your acknowledgment of our rebuttal. We would be interested to know which concern that you brought up made you decide to keep your current score? Based on your review and response, this is not clear to us at the moment, and getting a clarity on this would help us a lot in improving our work. We are particularly interested in which part of our experimental setting do you find unsatisfactory and would appreciate any concrete suggestions for improving it.

Thanks in advance!
Thanks for the response. I will keep my rating.
Dear reviewer c7rs, given that your recommendation is in conflict with the other reviewers' and we believe to have addressed the weaknesses raised in your original review, an acknowledgment or response to our rebuttal would be much appreciated. We are very much open to address any concerns in the remaining discussion period.
Thanks for updating your review, we are glad to hear we were able to address most of your concerns. We also appreciate you increasing your score and will make sure to include a longer discussion on the single-threshold assumption in the camera-ready version. In case any additional questions arise during the discussion period, we would be happy to address them.
Dear reviewer LhAs, we appreciate your acknowledgement and consideration. In case any additional questions arise during the discussion period or can be clarified in order to improve your rating, we would be happy to address those.
Dear reviewer qAvq, we are glad that we were able to resolve your key concerns and appreciate you raising your score. We will ensure that a more thorough discussion on the single vs. multi-threshold discussion along the lines of the rebuttal text is included in the camera-ready version, and thank you for pointing out this improvement. In case any additional questions arise during the discussion period, we would be happy to address them.
Thank you very much for the thorough and insightful reply. Most of my raised concerns have been adequately addressed, and as such, I am inclined to provisionally increase my score to Accept, pending the upcoming reviewer discussion. 

I would suggest that the discussion about the applicability of the proposed approach to multi-exit models with different thresholds in each exit (where the search space is notably harder to explore) is extended in the manuscript, along the lines of the relevant discussion in the rebuttal.
Thanks for replying. I shall maintain my rating.