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

TL;DR

SAFE-KD tackles safe, latency-aware vision inference by introducing a universal multi-exit wrapper for backbones, trained with hierarchical distillation via Decoupled Knowledge Distillation (DKD) and cross-exit consistency. It calibrates per-exit thresholds through Conformal Risk Control (CRC) to bound the selective misclassification risk $R_j(τ)$ at each exit with a slack of $O(1/|D_{cal}|)$ under exchangeability, enabling early exits without sacrificing safety. Empirically, SAFE-KD delivers favorable accuracy-vs-compute trade-offs, enhanced calibration, and robustness under corruption across diverse backbones and datasets, while providing finite-sample risk guarantees. The approach offers a practical, risk-bounded framework for deploying early-exit vision models in latency- or energy-constrained systems, with principled guarantees that distinguish it from heuristic gating.

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.