追求概率质量的帕累托最优：广义α-β散度引导知识蒸馏框架ABKD教师

文章图片

文章图片

文章图片

文章图片

文章图片

文章图片

文章图片

文章图片

文章图片

知识蒸馏（KD）是一种将大模型（教师）的知识迁移到小模型（学生）的技术，学生通过模仿教师预测分布，充分利用软标签信息，通常优于传统监督微调，已在图像分类、文本生成等任务及最新工作（如 DeepSeek-R1、Qwen-3）中得到验证。其核心在于分布匹配方式的选择，主流方法多用前向 KL 散度（FKLD）或反向 KL 散度（RKLD），但 FKLD 易导致输出过度平滑，难以聚焦关键类别， RKLD 则使学生过度自信、降低多样性。实验证明，两者在多任务中表现有限，且目前缺乏系统方法揭示其深层问题，阻碍了更通用高效 KD 框架的发展。因此，一个自然的问题产生了：

究竟是什么潜在因素导致了 FKLD 和 RKLD 的次优表现？

为了解答这个问题，我们通过追踪对数质量比（LogR），分析不同散度在训练中如何影响学生分布的概率分配。进一步分析表明在温和假设下， LogR 与损失函数对 logits 的梯度成正比，这使我们将问题转化为分析：不同散度算法如何影响 ∣LogR∣ 下降。

标题：ABKD: Pursuing a Proper Allocation of the Probability Mass in Knowledge Distillation via α-β-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∣ 的减小。

当 β→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) 所示。

与WSD的比较

另一种方案是对 FKLD 和 RKLD 取加权和（WSD）：

不幸的是，这种组合度量会过度强调 p 和 q 中概率较小的模式。当 q(k)≈0p(k)>0 或 p(k)≈0q(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 更低）。

结语

【追求概率质量的帕累托最优：广义α-β散度引导知识蒸馏框架ABKD】本文指出，知识蒸馏的核心在于平衡「难度集中」和「置信集中」两种效应，而传统 FKLD 和 RKLD 只覆盖两个极端，导致概率分配不合理。为此，我们提出基于 α-β 散度的 ABKD 框架，统一并推广了现有方法，实现两种效应的灵活权衡。理论与大量实验均验证了 ABKD 的有效性。