Early-Exit Neural Networks with Nested Prediction Sets

TL;DR

This paper proposes using anytime-valid confidence sequences for early-exit neural networks to generate nested, reliable uncertainty estimates, improving safety and efficiency in adaptive inference.

Contribution

It introduces AVCSs as a novel method for uncertainty quantification in EENNs, addressing the non-nested set issue of traditional techniques.

Findings

AVCSs produce nested prediction sets across exits

Traditional methods like conformal prediction are unsuitable for EENNs

AVCSs enhance safety and efficiency in early-exit models

Abstract

Early-exit neural networks (EENNs) enable adaptive and efficient inference by providing predictions at multiple stages during the forward pass. In safety-critical applications, these predictions are meaningful only when accompanied by reliable uncertainty estimates. A popular method for quantifying the uncertainty of predictive models is the use of prediction sets. However, we demonstrate that standard techniques such as conformal prediction and Bayesian credible sets are not suitable for EENNs. They tend to generate non-nested sets across exits, meaning that labels deemed improbable at one exit may reappear in the prediction set of a subsequent exit. To address this issue, we investigate anytime-valid confidence sequences (AVCSs), an extension of traditional confidence intervals tailored for data-streaming scenarios. These sequences are inherently nested and thus well-suited for an EENN's sequential predictions. We explore the theoretical and practical challenges of using AVCSs in EENNs and show that they indeed yield nested sets across exits. Thus our work presents a promising approach towards fast, yet still safe, predictive modeling