Skip to content

Beyond Greedy Exits: Improved Early Exit Decisions for Risk Control and Reliability

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

Early-Exit Deep Neural Networks (EEDNNs) suffer from overconfident wrong-class exits and static thresholds that fail under distribution shift. This paper proposes UAT, a UCB-based Multi-Armed Bandit algorithm that adapts exit thresholds online during inference using a reliability-calibrated confidence score. UAT achieves 1.70–2.10× speedup with less than 2% performance drop across classification, generation, and vision-language tasks.

Problem

Existing EEDNN strategies set a single confidence threshold on a held-out validation set and apply it greedily at inference time. This fails in two ways: (1) early layers can be overconfident on the wrong class (up to ~12% of samples on QNLI at layer 4 assign >0.7 probability to the wrong class), and (2) static thresholds do not adapt to deployment distribution shifts, undermining reliability and risk control.

Method

UAT operates in two phases. Offline: A lightweight reliability function g (single shared linear layer) is trained jointly with the EEDNN using a modified cross-entropy loss that couples the per-sample loss with g's output, forcing g to assign high scores to correctly classified, low-entropy samples and low scores to uncertain/misclassified ones. A coverage regularizer prevents the trivial solution g≡0.

Online inference: The set of candidate thresholds Ω (10 equally-spaced values in [0.5, 1.0]) are treated as arms in a MAB. At each step, UAT selects the threshold with the highest UCB index. For a sample at exit layer i, the effective signal is C_τ^i · (1 − C_g^i) — raw confidence multiplied by estimated correctness probability — minus a depth penalty ψ(i) = λ·i to discourage unnecessary deep processing. Rewards update running arm statistics. The global threshold λ is set to ϵ/L (where ϵ=0.01 is the risk tolerance), tying the efficiency–risk trade-off to a formal bound.

Theoretical guarantee: Under mild approximation assumptions, UAT achieves empirical risk ≤ ϵ + ϵ* (where ϵ* is the full-model risk floor) with high probability for sufficiently large T, via a sublinear regret bound.

Key Contributions

  • Reliability score g trained to approximate p(y=Å· | x, Å·) without ground-truth labels at inference time, decoupling overconfident wrong predictions from high-confidence correct ones.
  • Dynamic threshold selection via UCB-MAB; fully unsupervised at inference time, compatible with zero-shot settings.
  • Reward function that jointly maximizes reliability-weighted confidence and penalizes unnecessary depth, with a formal risk bound (Theorem 4.1).
  • Demonstrated system-level risk modeling where UAT can exceed full-model accuracy (via overthinking mitigation).

Results

  • BERT-large (GLUE/SQuAD): UAT reaches 95.8 SST-2 / 86.6 MNLI / 83.2 SQuAD F1 at 1.85× speedup, vs. full model 95.2 / 86.0 / 83.0 and best prior risk-based method EERC at 95.3 / 85.8 / 81.8 at 1.50×.
  • T5-large (summarization/QA/MT): UAT achieves 48.09 ROUGE-L on SamSum, 90.59 F1 on SQuAD, 38.62 BLEU on IWSLT at 2.10× speedup; full model is 48.82 / 90.63 / 39.19 at 1×.
  • BLIP-2-ViT-FlanT5-xl (VQA, captioning, VisDial): UAT scores 81.1 VQA accuracy, B4=43.3, CiDER=142.9 on COCO at 1.72× speedup; full model is 81.8 / 44.0 / 145.8; best prior risk method (EERS) is 80.0 / 42.5 / 140.3 at 1.38×.
  • NoCaps (out-of-domain captioning): UAT CiDER=119.5 / SPICE=15.6 full-dataset at 1.75×, vs. full model 121.6 / 15.8, showing strong robustness to domain shift.
  • Distribution robustness (Gaussian noise): At σ=1.0 noise, UAT B4=35.0, CiDER=125.5 at 1.81×; next best EERS achieves 32.7 / 116.9 at 1.68×.

Limitations

  • g is a global single linear layer shared across exits; per-exit local thresholds are not explored, leaving efficiency gains on the table.
  • Performance of UAT is coupled to the quality of the learned reliability function g; if g is poorly calibrated, exit decisions degrade.
  • Evaluation uses batch size 1 (per-instance inference), which is standard for EE baselines but limits throughput comparisons in batched settings.
  • Global threshold Ï„ is uniform across all exit layers; layer-adaptive thresholds are noted as future work.

Relevance to Vision-Language Models

UAT is directly evaluated on BLIP-2-based VLMs across image captioning (COCO, NoCaps), VQA (VQAv2), and visual dialogue (VisDial), establishing it as a practical inference-efficiency method for large multimodal models. For VLM researchers, the reliability-calibrated confidence is particularly relevant because VLM intermediate layers fuse heterogeneous vision and language representations that may be especially prone to early overconfidence. The framework's unsupervised online adaptation also makes it applicable to zero-shot or out-of-distribution VLM deployments where no labeled validation data is available — a common scenario in embodied or open-world settings.