ABKD: Pursuing a Proper Allocation of the Probability Mass in Knowledge Distillation via $α$-$β$-Divergence

📄 arXiv: 2505.04560v3 📥 PDF

作者: Guanghui Wang, Zhiyong Yang, Zitai Wang, Shi Wang, Qianqian Xu, Qingming Huang

分类: cs.LG

发布日期: 2025-05-07 (更新: 2025-06-03)

备注: ICML 2025 Spotlight

🔗 代码/项目: GITHUB

💡 一句话要点

ABKD：通过α-β散度实现知识蒸馏中概率质量的合理分配

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

关键词: 知识蒸馏 模型压缩 α-β散度 Hardness-Concentration Confidence-Concentration 深度学习 模型优化

📋 核心要点

  1. 现有知识蒸馏方法在平衡Hardness-Concentration和Confidence-Concentration两种模式集中效应方面存在不足。
  2. ABKD框架通过引入α-β散度，实现了FKLD和RKLD之间的平滑插值，从而有效权衡了上述两种效应。
  3. 在多个语言和视觉数据集上的实验表明，ABKD框架在多种教师-学生模型设置下均表现出优越的性能。

📝 摘要（中文）

知识蒸馏(KD)通过最小化大型教师模型和小型学生模型输出分布之间的差异，将知识从教师模型传递到学生模型，通常使用前向Kullback-Leibler散度(FKLD)或反向KLD(RKLD)。由于教师分布提供了比one-hot标签更广泛的监督信息，KD已成为一种有效的训练范式。我们发现KD的核心挑战在于平衡两种模式集中效应： extbf{ extit{Hardness-Concentration}}效应，指的是关注具有较大误差的模式； extbf{ extit{Confidence-Concentration}}效应，指的是关注具有高学生置信度的模式。通过分析梯度更新期间概率的重新分配方式，我们观察到这两种效应在FKLD和RKLD中是纠缠在一起的，但以极端的形式存在。具体来说，FKLD中的两种效应都太弱，导致学生无法集中在目标类别上。相反，RKLD中的两种效应都太强，导致学生过度强调目标类别，而忽略了来自教师的更广泛的分布信息。为了解决这种不平衡，我们提出了ABKD，一个具有α-β散度的通用框架。我们的理论结果表明，ABKD提供了FKLD和RKLD之间的平滑插值，实现了这些效应之间的有效权衡。在17个语言/视觉数据集上进行的12个教师-学生设置的大量实验证实了其有效性。代码可在https://github.com/ghwang-s/abkd 获得。

🔬 方法详解

问题定义：知识蒸馏旨在将大型教师模型的知识迁移到小型学生模型，但现有方法如FKLD和RKLD在平衡Hardness-Concentration和Confidence-Concentration两种效应上存在缺陷。FKLD倾向于忽略目标类别，而RKLD则过度强调目标类别，导致学生模型无法充分学习教师模型的知识。

核心思路：论文的核心思路是通过引入α-β散度，构建一个更通用的知识蒸馏框架ABKD，该框架能够平滑地插值FKLD和RKLD，从而实现Hardness-Concentration和Confidence-Concentration两种效应之间的有效权衡。通过调整α和β参数，可以控制两种效应的强度，使学生模型既能关注困难样本，又能学习教师模型的整体分布。

技术框架：ABKD框架的核心是使用α-β散度作为学生模型学习的损失函数。整体流程与标准的知识蒸馏流程类似：首先，教师模型对输入数据进行预测，生成软标签；然后，学生模型也对输入数据进行预测，生成自己的预测分布；最后，使用α-β散度计算教师模型和学生模型预测分布之间的差异，并将其作为损失函数来训练学生模型。

关键创新：最重要的技术创新点在于使用α-β散度替代传统的FKLD或RKLD。α-β散度是一个更通用的散度度量，可以通过调整α和β参数来控制其行为。与FKLD和RKLD相比，α-β散度能够更灵活地平衡Hardness-Concentration和Confidence-Concentration两种效应，从而提高知识蒸馏的性能。

关键设计：ABKD框架的关键设计在于α和β参数的选择。论文中可能提供了关于如何选择这些参数的指导，例如通过实验或理论分析来确定最佳值。此外，损失函数的具体形式也需要仔细设计，以确保其能够有效地衡量教师模型和学生模型预测分布之间的差异。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文在17个语言/视觉数据集上进行了大量实验，并与12种教师-学生模型设置进行了对比。实验结果表明，ABKD框架在各种设置下均优于现有的知识蒸馏方法，证明了其有效性。具体的性能提升幅度可能在论文中给出，例如在特定数据集上，ABKD框架的准确率比基线方法提高了X%。

🎯 应用场景

ABKD框架可广泛应用于各种需要模型压缩和加速的场景，例如移动设备上的图像识别、自然语言处理等。通过知识蒸馏，可以将大型、复杂的模型压缩成小型、高效的模型，从而在资源受限的环境中实现高性能的推理。该方法在自动驾驶、智能安防等领域具有潜在的应用价值。

📄 摘要（原文）

Knowledge Distillation (KD) transfers knowledge from a large teacher model to a smaller student model by minimizing the divergence between their output distributions, typically using forward Kullback-Leibler divergence (FKLD) or reverse KLD (RKLD). It has become an effective training paradigm due to the broader supervision information provided by the teacher distribution compared to one-hot labels. We identify that the core challenge in KD lies in balancing two mode-concentration effects: the \textbf{\textit{Hardness-Concentration}} effect, which refers to focusing on modes with large errors, and the \textbf{\textit{Confidence-Concentration}} effect, which refers to focusing on modes with high student confidence. Through an analysis of how probabilities are reassigned during gradient updates, we observe that these two effects are entangled in FKLD and RKLD, but in extreme forms. Specifically, both are too weak in FKLD, causing the student to fail to concentrate on the target class. In contrast, both are too strong in RKLD, causing the student to overly emphasize the target class while ignoring the broader distributional information from the teacher. To address this imbalance, we propose ABKD, a generic framework with $α$-$β$-divergence. Our theoretical results show that ABKD offers a smooth interpolation between FKLD and RKLD, achieving an effective trade-off between these effects. Extensive experiments on 17 language/vision datasets with 12 teacher-student settings confirm its efficacy. The code is available at https://github.com/ghwang-s/abkd.