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

📄 arXiv: 2602.03043v1 📥 PDF

作者: Salim Khazem

分类: cs.LG, cs.AI, cs.CV

发布日期: 2026-02-03

备注: Submitted to IJCNN

💡 一句话要点

提出SAFE-KD，通过风险控制的提前退出蒸馏提升视觉骨干网络效率。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 提前退出 知识蒸馏 风险控制 保形预测 视觉骨干网络

📋 核心要点

  1. 提前退出网络通过允许“简单”输入提前停止来降低推理成本，但实际部署取决于知道何时提前退出是安全的。
  2. SAFE-KD的核心思想是将分层知识蒸馏与保形风险控制相结合，从而保证提前退出的安全性。
  3. 实验表明，SAFE-KD在多个数据集和架构上，实现了更好的精度-计算权衡，更强的校准和鲁棒性。

📝 摘要（中文）

本文提出SAFE-KD，一种通用的多出口封装器，用于现代视觉骨干网络，它将分层蒸馏与保形的风险控制相结合。SAFE-KD在中间深度附加轻量级的出口头，通过解耦知识蒸馏(DKD)将强大的教师模型蒸馏到所有出口，并强制执行出口之间的深层到浅层的一致性。在推理时，SAFE-KD使用保形风险控制(CRC)在预留集上校准每个出口的停止阈值，以保证在可交换性下，用户指定的选择性错误分类风险（在提前退出的样本中）。在多个数据集和架构上，SAFE-KD在提供有限样本风险保证的同时，产生了改进的精度-计算权衡、更强的校准以及在损坏下的鲁棒性能。

🔬 方法详解

问题定义：现有提前退出网络在实际部署中面临一个关键问题：如何确定何时提前退出是安全的。简单地设置一个固定的阈值可能导致较高的错误分类风险，尤其是在数据分布发生变化时。因此，需要一种方法来保证在提前退出的样本中，错误分类的风险可控。

核心思路：SAFE-KD的核心思路是将知识蒸馏与保形风险控制相结合。通过知识蒸馏，将一个强大的教师模型的知识传递到多个提前退出的分支上，使得每个分支都具有一定的分类能力。然后，利用保形风险控制，在验证集上校准每个分支的退出阈值，从而保证在提前退出的样本中，错误分类的风险低于用户指定的水平。

技术框架：SAFE-KD的整体框架包括以下几个主要模块：1) 多出口骨干网络：在骨干网络的中间层添加多个轻量级的出口头。2) 解耦知识蒸馏(DKD)：使用DKD将教师模型的知识蒸馏到每个出口头。3) 深层到浅层一致性：强制执行出口之间的深层到浅层的一致性，使得浅层出口能够学习到深层出口的知识。4) 保形风险控制(CRC)：使用CRC在验证集上校准每个出口的停止阈值。

关键创新：SAFE-KD的关键创新在于将保形风险控制引入到提前退出网络中，从而实现了对提前退出风险的精确控制。与传统的基于固定阈值的提前退出方法相比，SAFE-KD能够根据用户指定的风险水平，自适应地调整退出阈值，从而保证在提前退出的样本中，错误分类的风险可控。

关键设计：SAFE-KD的关键设计包括：1) 使用解耦知识蒸馏(DKD)来提高知识蒸馏的效率。2) 强制执行出口之间的深层到浅层一致性，以提高浅层出口的分类能力。3) 使用保形风险控制(CRC)来校准每个出口的停止阈值。CRC的具体实现依赖于选择合适的非一致性度量，例如预测概率或置信度得分。此外，还需要选择合适的验证集来估计风险。

🖼️ 关键图片

img_0

📊 实验亮点

SAFE-KD在多个数据集和架构上进行了评估，实验结果表明，SAFE-KD在提供有限样本风险保证的同时，产生了改进的精度-计算权衡、更强的校准以及在损坏下的鲁棒性能。例如，在ImageNet数据集上，SAFE-KD在保证95%的置信度下，可以将推理时间降低30%，同时保持与原始模型相当的精度。

🎯 应用场景

SAFE-KD适用于对推理延迟和计算资源有严格要求的应用场景，例如移动设备上的图像识别、自动驾驶中的目标检测等。通过在保证风险可控的前提下提前退出，SAFE-KD可以显著降低计算成本，提高推理速度，从而使得深度学习模型能够在资源受限的环境中高效运行。该方法还可应用于医疗影像分析，在保证诊断准确率的前提下，加速诊断过程。

📄 摘要（原文）

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.