ICML Spotlight 2025丨追求概率质量的帕累托最优：基于广义α-β散度引导的知识蒸馏框架ABKD
仅用于站内搜索，没有排版格式，具体信息请跳转上方微信公众号内链接
知识蒸馏（KD）是一种将大模型（教师）的知识迁移到小模型（学生）的技术，学生通过模仿教师预测分布，充分利用软标签信息，通常优于传统监督微调，已在图像分类、文本生成等任务及最新工作（如DeepSeek-R1、Qwen-3）中得到验证。其核心在于分布匹配方式的选择，主流方法多用前向KL散度（FKLD）或反向KL散度（RKLD），但FKLD易导致输出过度平滑，难以聚焦关键类别，RKLD则使学生过度自信、降低多样性。实验证明，两者在多任务中表现有限，且目前缺乏系统方法揭示其深层问题，阻碍了更通用高效KD框架的发展。因此，一个自然的问题产生了：
究竟是什么潜在因素导致了FKLD和RKLD的次优表现？
为了解答这个问题，我们通过追踪对数质量比（LogR），分析不同散度在训练中如何影响学生分布的概率分配。进一步分析表明在温和假设下，LogR与损失函数对logits的梯度成正比，这使我们将问题转化为分析：不同散度算法如何影响∣LogR∣下降。
标题：ABKD:PursuingaProperAllocationoftheProbabilityMassinKnowledgeDistillationviaα-β-Divergence
论文：https ://arxiv. org/abs/2505. 04560
代码：https ://github. com/ghwang-s/abkd
在此框架下，我们识别出两种关键的「模式集中效应」：难度集中与置信集中。
「难度集中」指的是更关注于在教师分布p与学生分布q之间误差较大的模式。
「置信集中」指的是更关注于学生分布q本身高度自信的模式。
进一步研究发现，FKLD和RKLD的局限性正源于对这两种集中效应的极端利用。
FKLD集中效应较弱，对所有类别误差一视同仁，导致学生难以聚焦目标类别，易出现错误预测（见图1d）。
RKLD集中效应过强，兼顾高误差和高置信度类别，易陷入「平凡解」，即学生只关注目标类别，忽略教师分布的其他知识（见图1e）。
揭示这些局限性后，我们提出一个关键问题：我们能否找到一种通用且理论上有依据的方法，来平衡「难度集中」与「置信集中」效应？
为此，我们引入了α-β散度。
其广义统一了FKLD、RKLD及Hellinger距离等多种散度。
正如下一节理论表明，α-β散度可通过调节α和β在FKLD与RKLD间平滑插值，实现对难度集中和置信集中效应的灵活权衡。这一机制带来更合理的概率分配，据此我们提出了通用蒸馏框架ABKD，形式为最小化：
ABKD提供了一个统一空间权衡难度集中与置信集中。为解释这一点，回到对数质量比（LogR）。下列命题解释了超参数α和β如何影响∣LogR∣的减小。
在(a)和(a1)中，α-β散度通过幂次形式体现置信集中效应：
当β→1时，退化为RKLD的效应。
当β→0时，退化为FKLD的效应。
较大的β值会增强置信集中效应，使匹配性能更加专注于模型最有信心的类别（见图1c）。
在(b)和(b1)中，使用以下方式衡量难度集中效应：
当α→1时，退化为FKLD的效应。
当α→0时，退化为RKLD的效应。
较小的α会放大难度集中效应，在困难类别上进行更强的误差惩罚，从而实现更激进的匹配（见图1b）。
通过调节α和β，ABKD实现对两种效应的精细平衡，避免了传统方法（FKLD/RKLD）的极端情况（见图1g）。
α-散度是ABKD框架的一个特例（当β=1−α时成立）。
虽然它也能实现FKLD（α→1）和RKLD（α→0）之间的转换，但根据命题4. 2，为了减小α，必须相应地增大β，以满足α+β=1的限制条件。这种不必要的限制阻碍了其性能的进一步提升，具体如图1(a)和图1(f)所示。
另一种方案是对FKLD和RKLD取加权和（WSD）：
不幸的是，这种组合度量会过度强调p和q中概率较小的模式。当q(k)≈0,p(k)>0或p(k)≈0,q(k)>0时，有：
因此，算法必须关注极端情况以最小化目标函数，导致概率分配不合理。此外这种情况下梯度范数也会过度增长，导致参数更新中可能含有噪声，破坏优化稳定性。
最后一种方法是使用Jensen-Shannon散度：
其中m=(p+q)/2。它的缺点是当分布p和q相距较远时（这是训练初期的常见情况），会出现梯度消失，阻碍模型收敛。
我们在五个指令跟随基准上评估方法，使用ROUGE-L指标，对比了SFT、FKLD、GKD和DISTILLM等方法。
实验结果表明仅修改蒸馏目标，ABKD在不同数据集上均优于或匹配FKLD、SFT。对比基于SGO的方法（如GKD、DISTILLM）效果提升明显，ABKD在不公平对比下依然表现优越。
损失函数影响
α-β散度优于α或β散度及WSD。在五个数据集上相较基线有稳定提升。训练过程中的优势也体现在图2。
在12个常用图像识别数据集上进行实验，对比方法比较了多种先进蒸馏方法：KD、DKD、LSD和TTM等。
结果表明ABKD能根据不同教师-学生组合自动选择合适的蒸馏目标。此外ABKD可作为简单的插件工具，优化现有方法的损失函数，带来额外性能提升。
α对难度集中的影响：图6(a)(b)展示了不同α下的训练表现。在CIFAR-100中，较小α提升有限；在Dolly等大输出空间任务中，适当减小α有助于避免陷入局部最优。
β对置信集中的影响：图6(c)(d)显示了β对输出分布的Shannon熵和Self-BLEU的影响。较小β使模型更关注低置信度类别，提升输出分布的平滑性和多样性（熵更高，Self-BLEU更低）。
本文指出，知识蒸馏的核心在于平衡「难度集中」和「置信集中」两种效应，而传统FKLD和RKLD只覆盖两个极端，导致概率分配不合理。为此，我们提出基于α-β散度的ABKD框架，统一并推广了现有方法，实现两种效应的灵活权衡。理论与大量实验均验证了ABKD的有效性。
©THEEND
转载请联系本公众号获得授权
投稿或寻求报道：liyazhou@jiqizhixin. com
