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

B站影视 日本电影 2025-06-09 19:30 2

摘要:知识蒸馏(KD)是一种将大模型(教师)的知识迁移到小模型(学生)的技术,学生通过模仿教师预测分布,充分利用软标签信息,通常优于传统监督微调,已在图像分类、文本生成等任务及最新工作(如 DeepSeek-R1、Qwen-3)中得到验证。其核心在于分布匹配方式的选

知识蒸馏(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)≈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 的有效性。

来源:机器之心Pro一点号

相关推荐