SAFE-KD: Risk-Controlled Early-Exit Distillation for Vision Backbones

2026-02-03

AI Summary

This work addresses the challenge of enabling early-exit networks to reduce computational cost while rigorously controlling misclassification risk. The authors propose SAFE-KD, a novel framework that integrates conformal risk control (CRC) with decoupled knowledge distillation (DKD) to introduce a multi-exit architecture with explicit risk constraints in vision backbones. By leveraging lightweight exit heads, hierarchical distillation, and exit consistency constraints, SAFE-KD guarantees—under finite-sample settings—that the selective misclassification risk remains below a user-specified threshold. Experiments across multiple datasets and architectures demonstrate that SAFE-KD achieves superior accuracy-computation trade-offs, enhanced calibration, and robustness to corrupted inputs, all while providing theoretically grounded finite-sample risk guarantees.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Computer Vision: Adversarial Attacks & Robustness
• Natural Language Processing: Safety and Robustness

Application Category

• Security and Privacy: Large-scale security measurements
• Responsible Web: Human-perceived consequences of algorithmic deployment on the web
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Early-exit networks reduce inference cost by allowing ``easy''inputs to stop early, but practical deployment hinges on knowing \emph{when} early exit is safe. We introduce SAFE-KD, a universal multi-exit wrapper for modern vision backbones that couples hierarchical distillation with \emph{conformal risk control}. SAFE-KD attaches lightweight exit heads at intermediate depths, distills a strong teacher into all exits via Decoupled Knowledge Distillation (DKD), and enforces deep-to-shallow consistency between exits. At inference, we calibrate per-exit stopping thresholds on a held-out set using conformal risk control (CRC) to guarantee a user-specified \emph{selective} misclassification risk (among the samples that exit early) under exchangeability. Across multiple datasets and architectures, SAFE-KD yields improved accuracy compute trade-offs, stronger calibration, and robust performance under corruption while providing finite-sample risk guarantees.

Problem

• early-exit
• risk control
• selective classification
• conformal prediction
• vision backbones

Innovation

• conformal risk control
• early-exit networks
• decoupled knowledge distillation
• risk-guaranteed inference
• vision backbones