Fast yet Safe: Early-Exiting with Risk Control

Dear reviewer LhAs, we thank you for your time and comments, and address your raised point in the following.

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

We would politely point out here that our risk control approach is entirely post-hoc. As such, we do not change the performance of any underlying model but work with existing pretrained EENNs, and thus our model performances are equivalent to the numbers provided by the original authors. While we provide performance curves in Fig. 1 for some models, this was to stress the marginal monotonicity condition rather than testing model’s ability. Thus, we consider reporting the commonly seen performance vs. FLOP curves found in early-exit papers as redundant for our work, since we do not modify the underlying models.

Perhaps it's also helpful to add here that we may interpret our risk control threshold selection as a framework to navigate the performance vs. FLOP curve based on the user’s performance requirements, i.e. we find the point on the curve such that the EENN’s test performance is guaranteed to be at most $\epsilon$ (e.g., at most 5%) worse than the full model (i.e., last-layer exit). Hence, we believe that evaluating against different risk levels $\epsilon$ is the appropriate way to benchmark our procedure. Also note that provided “test risk” figures such as Fig. 2 and 3 in the paper are standard in the risk control domain (see, e.g., Fig. B.1 in [2]). Moreover, the fact that the test risk curves are close to the diagonal (e.g., see Figure 3, upper row) is encouraging, as it suggests that with our framework, there is very little efficiency cost for the added safety of having risk-control guarantees.

We thank you for the suggestion, highlighting the need to better explain provided test risk figures. We attempted to do so in §5 (L272-278), but shall expand upon it for the camera-ready version based on your feedback.

[2] Schuster, T., et al. Confident adaptive language modeling. NeurIPS 2022

Thanks for replying. I shall maintain my rating.

Dear reviewer q2CX, we thank you for your time and helpful comments, and address your two key concerns in the following.

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.

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.

Please consider our response under point (1) in the general rebuttal to this question, and kindly provide additional questions in case you feel it has not been sufficiently addressed.

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.

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.

However, the real performance of the evaluated EENNs should also be attached in the manuscript.

Thanks for raising these points as they made us realise there is some misunderstanding about the interpretation of our proposed risks, and their relationship to standard performance metrics. As we explicitly state in L164-168 and seen from the equations in §3.2, our risks are formulated as the mean relative performance difference between full (i.e., last-layer exit) and early-exit model, and provide a generic template in which we may plug in different labels, model outputs, confidence measures, and loss functions.

We explore a wide range of loss functions across tasks for prediction control, and in particular use standard task-specific performance losses (and precisely the ones you mention): 0-1 loss for classification, mIOU loss and miscoverage loss for segmentation, ROUGE-L loss and Token-F1 loss for language modelling, and LPIPS loss for image generation. Thus, the empirical test risk equates the difference in test accuracy between the full and early-exit model for our classification experiments in §5.1, the difference in mIOU between the full and early-exit model for some of our segmentation experiments in §5.2, etc., providing interpretable assessments of the EENN’s performance based on risk control.

We would also emphasis here that our risk control approach is entirely post-hoc. As such, we do not change the performance of any underlying model but work with existing pretrained EENNs, and thus our model performances are equivalent to the numbers provided by the original authors. Thus, we consider reporting the commonly seen performance vs. FLOP curves found in early-exit papers as redundant for our work, since we do not modify the underlying models (except for switching to a single exit threshold for some models like MSDNet, see our point (1) in meta rebuttal for a longer discussion on this).

Perhaps it's also helpful to add here that we may interpret our risk control threshold selection as a framework to navigate the performance vs. FLOP curve based on the user’s performance requirements, i.e. we find the point on the curve such that the EENN’s test performance is guaranteed to be at most $\epsilon$ (e.g., at most 5%) worse than the full model (i.e., last-layer exit). As such, evaluating against different risk levels $\epsilon$ is the appropriate way to benchmark our procedure, and provided “test risk” figures such as Fig. 2 and 3 in the paper are common in the risk control domain (see, e.g., Fig. B.1 in [2]). Moreover, the fact that the test risk curves are close to the diagonal (e.g., see Figure 3, upper row) is encouraging, as it suggests that with our framework, there is very little efficiency cost for the added safety of having risk-control guarantees.

We thank you for raising these points, highlighting the need to better explain the provided test risk figures. We attempted to do so in §5 (L272-278), but shall expand upon it for the camera-ready version. Likewise, we will aim to clarify that model performance is not affected by our post-hoc risk control procedure.

[2] Schuster, T., et al. Confident adaptive language modeling. NeurIPS 2022

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.

We appreciate the additional references. We would politely point out that reference [3] is already contained in our image classification experiments, where we refer to it as L2W-DEN (see Figure 3 in the paper). Following your advice, we have also added the suggested reference [1] to our experiments. We refer to the model as IMTA and observe similar results to the other models, see Figure A of the attached rebuttal PDF. We will incorporate this model alongside your other suggested reference [2] (RANet) for the camera-ready version (the ImageNet training of the RANet unfortunately didn't finish before the rebuttal deadline), either into the main figure or into an appendix figure to preserve readability.

We thank you again for your comments and would appreciate you considering raising your score if you believe they have been addressed sufficiently.

Dear reviewer c7rs, we thank you for your time and comments, and address your two raised concerns in the following.

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.

We politely disagree with this statement. Our experiments incorporate a wide variety of models, confidence measures and risk types across multiple machine learning tasks, datasets and data modalities (image classification on ImageNet with MSDNet, DViT, L2W-DEN and Dyn-Perc for 0—1 and Brier losses; semantic segmentation on Cityscapes & GTA5 with ADP-C for mIoU, miscoverage and Brier losses and several confidence measures; text summarization on CNN/DM and question answering on SQuAD with T5 for ROUGE-L and Token-F1 losses; image generation with early-exit diffusion on CelebA and CIFAR for LPIPS loss). Moreover, our extensive experiments have been highlighted as a key strength by all other reviewers (qAvq: "experimentally demonstrated on a wide range of predictive ... and generative tasks"; LhAs: "extensive experimental evidence across multiple domains"; q2CX: "The experiments on different vision tasks can be strong evidence to show the effectiveness…”).

Lastly, we are surprised since your review also lists our experimental evaluation as a strength (”The paper provides empirical validation of their approach on a range of vision and language tasks, demonstrating substantial computational savings while preserving performance goals”); perhaps one of the mentions of our experiments was in error?

If you have any concrete additional experiments in mind, please let us know and we would be happy to consider them.

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.

We politely point to §6, which is titled “Conclusion, Limitations and Future Work”. There we state our main limitation (a single exit threshold) and point to several limitations and opportunities for future work that are rooted in risk control itself, rather than our adaptation to the early-exit setting. Furthermore, we are explicit about our marginal monotonicity requirement for the EENN (e.g., in our propositions) and discuss it throughout the paper. The presence of Limitations has also been acknowledged by the other reviewers (qAvq: "Limitations and future work are adequately discussed in the manuscript."). Lastly, we will use the additional page of the camera-ready version to further expand on the limitations.

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Based on your current review, it is not entirely clear to us why you are recommending rejection. We would very much appreciate further elaboration of your concerns, or a consideration to raise your score if you believe they have been addressed sufficiently.

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 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!

Dear reviewer qAvq, we thank you for your time and engaged questions.

The reliance of the proposed approach to marginal monotonicity between performance and depth can be a limiting [...] (see Elhoushi et al, LayerSkip, 2024 - Fig.2).

This is a great point and we fully agree that for a particular sample the predictive accuracy can fluctuate across exits. We refer to this effect as overthinking (following [5, 12]) and mention it several times throughout the paper (L77-80, L196-198, L211-214). However, note that our approaches do not require such strong assumptions on conditional monotonicity, i.e., monotonic behaviour per sample. Rather, we only require the milder assumption of marginal monotonicity, i.e., that predictive accuracy is monotone across exits on average over samples (Eq. 2 in the paper). Such a marginal condition is often implicitly assumed in the early-exit literature, and validated by the commonly shown performance vs. FLOP curves which monotonically grow as the budget increases. Indeed, Fig. 2 in Elhoushi et al. [13] is showcasing a lack of conditional monotonicity, since the model outputs are shown for a particular input prompt. Looking at Fig. 6 and 8 from the same paper, we observe that performance is generally monotonic w.r.t. exits on average (marginally). In fact, the difference between conditional and marginal monotonicity serves as a motivation for one of our main technical contributions - in the original formulation, CRC requires conditional monotone risks, while we extend this to marginal monotone ones (see Prop. 1). As we state in the paper (L192-198), such an extension is crucial to make CRC amenable to EENNs because of the overthinking issue.

As stated in the limitations, the adoption of a single and static threshold across exits also significantly reduces the search space [...] inefficiency imposed by this design choice has not been quantified... .

Please consider our response under point (1) in the general rebuttal, and kindly provide additional questions in case you feel it has not been sufficiently addressed.

how the proposed approach can be applied to the emerging line of works with learnable exit policies

We assume you refer to papers such as JEI-DNN [14], where exits are determined via direct modeling of exit probabilities. If so, it is correct that our procedure does not directly transfer since we require a threshold-based exiting mechanism. We will ensure to better define the scope of our contribution, and we thank you for pointing this out. Note, however, that models with learnable exit policies are less adaptive compared to their threshold counterparts, as they require retraining the model from scratch for every new computational budget. In contrast, threshold-based models can be adapted in a post-hoc fashion by simply tuning the threshold.

Let us also add here that thresholding is not used only in early-exiting, but also in some other emerging lines of work like (soft) speculative decoding [15], where the rollback policy is governed by thresholding. Our risk control framework could be applied to such settings as well, and we will try to add an experiment to the camera-ready version to better exemplify the scope of our contribution.

whether it will maintain its performance on different confidence metrics, such as state-saturation, used in LLMs

We note that our approach is agnostic to the choice of confidence measure used when early-exiting. We aim to demonstrate this in the paper with our segmentation experiment in §5.2 (see Table 1; and Table 3 in §D.2), where we consider different confidence measures and aggregators. Based on your comment, we have extended our language modeling experiment to include additional confidence measures: (i) hidden-state saturation and (ii) meta-classifiers, as proposed by [2]. We observe in the Figure B of the attached rebuttal PDF that our employed risk control frameworks based on CRC and UCB continue to outperform LTT across all confidence measures, and will include these results in the camera-ready version.

Can the proposed Consistency Risk formulation be used for online adaptation [...], but input distribution shifts may occur?

We are not entirely sure what form of online adaptation is meant here, so kindly inform us if the question was not properly addressed.

For distribution shifts between training and calibration data (i.e., $P_{train} \neq P_{cal}$) risk control continues to be applicable as long as the model’s marginal monotonicity is not violated. Our post-hoc approach treats the EENN as a black-box and makes no assumptions on the training distribution (as stated in §2, L50-51). For shifts between calibration and test data (i.e., $P_{cal} \neq P_{test}$) the provided guarantees cease to hold, since an i.i.d. assumption on those samples is made. However, under benign shifts the empirical test risks may continue to be controlled.

Since we guarantee risk control even for small calibration set sizes ($n \approx 100$) and are computationally light-weight, a naive online update could see a periodical re-calibration of the exit threshold based on collected test samples, which, in the case of our consistency risk, would indeed not even require any test labels. We fully agree that it is an interesting direction for future work to explore proper online updating.

We thank you for your comments and would appreciate you considering raising your score if you believe they have been addressed sufficiently.

References

[12] Jazbec, M., et al. Towards anytime classification in early-exit architectures by enforcing conditional monotonicity. NeurIPS 2023

[13] Elhoushi, Mostafa, et al. Layer skip: Enabling early exit inference and self-speculative decoding. arXiv preprint (2024)

[14] Regol, F., et al. Jointly-learned exit and inference for a dynamic neural network: Jei-dnn. ICLR 2024

[15] Kim, S., et al. Speculative decoding with big little decoder. NeurIPS 2023