Rethinking Calibration for Early-Exit Neural Networks
Abstract
Early-exit neural networks (EENNs) accelerate inference by allowing intermediate classifiers to stop computation once predictions are confident enough. Most methods rely on confidence thresholds for exiting, and consequently, improving classifier calibration is widely assumed to improve performance. In this work, we challenge this assumption and show that calibration alone is not sufficient for EENNs to exploit adaptive computation. To address this insufficiency, we introduce Early-Exit Failure Prediction (EEFP), which accounts for both prediction correctness and the cost of further computation. We also propose a lightweight, EEFP-motivated procedure to improve the intermediate classifiers, which can directly replace calibration in EENNs. Extensive experiments demonstrate that our approach achieves superior cost-accuracy trade-offs compared to calibration, and EEFP more reliably reflects overall EENN performance. Our code is available at https://github.com/gmum/rethinking-calibration-for-eenns.
Lay Summary
Deep learning models can be made faster and more energy-efficient by using "early-exit" checkpoints that stop processing once they are confident in an answer. Traditionally, these systems rely on confidence to decide when to stop. However, we show that this has a major flaw: if an input is highly confusing or impossible to identify, the model will have low confidence at every checkpoint. It then wastes valuable energy running all the way to the end of the network, only to make an incorrect prediction anyway. To solve this, we adapted the concept of failure prediction for early-exit models. We developed a lightweight confidence-correction procedure that evaluates whether spending more computing power is actually useful. If the model determines that additional effort is futile, it stops immediately to save resources. In our experiments, this approach achieved better trade-offs between computing cost and accuracy than traditional methods. Our code is openly available to help other researchers make their AI models more efficient.