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

arXiv: https://arxiv.org/abs/2602.03043 · HTML: https://mlsanity.com/paper/2602.03043/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Early exit and edge inference, one of 12 directions in Efficient LLM inference and serving.
- Builds on 18 of its 27 references that are in the map.
- First posted on arXiv 3 Feb 2026.

## Placed in

- Efficient LLM inference and serving / [Early exit and edge inference](https://mlsanity.com/inference/early-exit-and-edge-inference/)

## Builds on (most cited first, 18 of 18)

- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Benchmarking Neural Network Robustness to Common Corruptions and Perturbations](https://mlsanity.com/paper/1903.12261/) (2019)
- [FitNets: Hints for Thin Deep Nets](https://mlsanity.com/paper/1412.6550/) (2014)
- [Paying More Attention to Attention: Improving the Performance of Convolutional Neural Networks via Attention Transfer](https://mlsanity.com/paper/1612.03928/) (2016)
- [Relational Knowledge Distillation](https://mlsanity.com/paper/1904.05068/) (2019)
- [DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification](https://mlsanity.com/paper/2106.02034/) (2021)
- [Selective Classification for Deep Neural Networks](https://mlsanity.com/paper/1705.08500/) (2017)
- [A Gentle Introduction to Conformal Prediction and Distribution-Free Uncertainty Quantification](https://mlsanity.com/paper/2107.07511/) (2021)
- [BranchyNet: Fast Inference via Early Exiting from Deep Neural Networks](https://mlsanity.com/paper/1709.01686/) (2017)
- [Weight Uncertainty in Neural Networks](https://mlsanity.com/paper/1505.05424/) (2015)
- [Dynamic Neural Networks: A Survey](https://mlsanity.com/paper/2102.04906/) (2021)
- [DeeBERT: Dynamic Early Exiting for Accelerating BERT Inference](https://mlsanity.com/paper/2004.12993/) (2020)
- [Conformal Risk Control](https://mlsanity.com/paper/2208.02814/) (2022)
- [BERT Loses Patience: Fast and Robust Inference with Early Exit](https://mlsanity.com/paper/2006.04152/) (2020)
- [PolygoNet: Leveraging Simplified Polygonal Representation for Effective Image Classification](https://mlsanity.com/paper/2504.01214/) (2025)
- [TopoLoRA-SAM: Topology-Aware Parameter-Efficient Adaptation of Foundation Segmenters for Thin-Structure and Cross-Domain Binary Semantic Segmentation](https://mlsanity.com/paper/2601.02273/) (2026)

## Abstract (arXiv)

Early-exit networks reduce inference cost by allowing ``easy'' inputs to stop early, but practical deployment hinges on knowing when early exit is safe. We introduce SAFE-KD, a universal multi-exit wrapper for modern vision backbones that couples hierarchical distillation with 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 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.