重新思考逆向KL作为自适应熵蒸馏

# 将逆KL散度重新思考为自适应熵蒸馏
来源：https://arxiv.org/html/2608.14685
Shizhen Li1, Zhiyu Shen1, Yuyin Lu1, Yunhe Pang1, Jielin Song1, Yanghui Rao1, Fu Lee Wang2  
1中山大学计算机科学与工程学院，中国广州  
2香港都会大学科技学院，中国香港特别行政区  
\{lishzh57, shenzhy23, luyy37, pangyh8, songjlin6\}@mail2.sysu.edu.cn  
[email protected]  
[email protected]

###### 摘要
知识蒸馏（KD）被广泛用于将大型语言模型（LLM）的能力迁移给更小的学生模型，但现有目标函数往往难以在忠实模仿与鲁棒生成之间取得平衡。特别地，现有方法主要结合前向KL散度（FKL）和逆向KL散度（RKL），忽略了RKL本身已提供调整学生模仿强度的机制。受此启发，我们重新审视策略内的逆KL（RKL）蒸馏，并将其目标分解为教师拟合项和学生熵项，无需引入显式的FKL分支。我们从理论上证明，词元级别的最优学生分布对应于教师分布的一种温度化变体，其中自适应权重控制着模式寻求与不确定性保持之间的权衡。基于这一洞见，我们提出**自适应熵蒸馏（AED）**，利用教师的熵动态校准词元级别的模仿强度。在指令遵循和数学推理基准上的实验表明，AED取得了优越的整体性能，并普遍提升了教师-学生分布和熵的对齐程度。111代码：https://github.com/ShizhenL1/AED

## 1引言

> 参见标题图1：固定强度蒸馏与AED的比较：固定强度采用均匀模仿，而AED适应教师的不确定性。
大语言模型提供了强大能力，但伴随着高昂的计算和部署成本。知识蒸馏（KD）（Hinton et al., 2015 (https://arxiv.org/html/2608.14685#bib.bib1)）通过压缩大型教师模型为轻量级学生模型来解决这个问题。传统的序列级KD（Kim and Rush, 2016 (https://arxiv.org/html/2608.14685#bib.bib2)）优化教师生成轨迹上的前向KL散度，但这种离策略训练存在暴露偏差（Bengio et al., 2015 (https://arxiv.org/html/2608.14685#bib.bib3); Ranzato et al., 2016 (https://arxiv.org/html/2608.14685#bib.bib4)）。近期工作转向使用学生采样轨迹的策略内RKL蒸馏（Gu et al., 2024 (https://arxiv.org/html/2608.14685#bib.bib5)），以减少训练-推理的不匹配。

尽管如此，充分利用策略内蒸馏的益处仍然具有挑战性，因为FKL的概率覆盖与RKL的概率集中之间存在固有的权衡（Gu et al., 2024 (https://arxiv.org/html/2608.14685#bib.bib5)）。现有框架通常通过静态插值（Agarwal et al., 2024 (https://arxiv.org/html/2608.14685#bib.bib6)）或类似Skewed RKL的变体（Ko et al., 2024 (https://arxiv.org/html/2608.14685#bib.bib7)）来解决。然而，这些全局的、固定权重策略在整条序列上均匀缩放目标函数，未能考虑词元级别的教师置信度差异。为了克服这一点，近期方法如AKL（Wu et al., 2025 (https://arxiv.org/html/2608.14685#bib.bib16)）和ToDi（Jung et al., 2025 (https://arxiv.org/html/2608.14685#bib.bib17)）引入了动态平衡策略，通过词元级机制避免均匀缩放。

虽然现有框架主要关注自适应地结合FKL和RKL，但RKL本身的内部结构已经提供了一种控制学生模仿教师强度的机制。

因此，我们从分布几何的角度重新审视策略内蒸馏。我们表明，RKL目标本身包含了集中与覆盖的互补力，通过重新加权其两个内部成分，可以实现仅使用RKL的公式化。在这种重新加权的公式下，分析最优解对应于一个温度化的教师分布，为如何控制学生的模仿强度提供了理论表征。

基于这些理论发现，我们开发了AED，一个理论驱动、不确定性感知的策略内蒸馏框架，如图1 (https://arxiv.org/html/2608.14685#S1.F1) 所示。AED的动机源于此理论最优解，并在词元级别自适应地调整模仿强度。为了平衡RKL内的两项，AED引入了自适应加权机制。我们通过将教师的熵映射到权重来实例化该机制，鼓励学生内化教师的分布。我们在指令蒸馏和数学推理基准上评估了AED，并进一步分析了其词元级别的分布行为。

基于对RKL的几何分析，AED重新加权其内在成分以产生仅含RKL的目标，并使用教师熵来校准词元级别的模仿强度。实验证明，这种设计提升了任务性能和教师-学生分布对齐，突显了该理论表征的实用价值。

## 2相关工作

### 2.1大语言模型知识蒸馏的演进
传统KD是一种主要的模型压缩技术，通常遵循离策略范式，学生模型从静态的教师生成目标（如软逻辑值、中间表示或推理链）中学习（Hinton et al. (2015 (https://arxiv.org/html/2608.14685#bib.bib1)); Sanh et al. (2019 (https://arxiv.org/html/2608.14685#bib.bib8)); Jiao et al. (2020 (https://arxiv.org/html/2608.14685#bib.bib9))）。对于大语言模型，它已被用于通过思维链推理蒸馏（Hsieh et al. (2023 (https://arxiv.org/html/2608.14685#bib.bib10)); Fu et al. (2023 (https://arxiv.org/html/2608.14685#bib.bib11))）和从合成数据集中蒸馏指令遵循行为（Jiang et al. (2023 (https://arxiv.org/html/2608.14685#bib.bib12))）。然而，离策略KD存在暴露偏差和累积分布偏移（Bengio et al. (2015 (https://arxiv.org/html/2608.14685#bib.bib3)); Ranzato et al. (2016 (https://arxiv.org/html/2608.14685#bib.bib4))），因为学生模型并未在其自身自回归错误诱导的状态上进行训练。

为了解决这个问题，近期的策略内蒸馏方法让学生模型从其自身分布中采样，同时接受教师的监督（Gu et al. (2024 (https://arxiv.org/html/2608.14685#bib.bib5)); Agarwal et al. (2024 (https://arxiv.org/html/2608.14685#bib.bib6))）。MiniLLM（Gu et al. (2024 (https://arxiv.org/html/2608.14685#bib.bib5))）使用RKL散度提供词元级别的纠正信号，而GKD（Agarwal et al. (2024 (https://arxiv.org/html/2608.14685#bib.bib6))）通过替代的散度目标泛化了策略内蒸馏。然而，这些方法通常将RKL视为固定目标；我们进一步分解和分析RKL，以更好地理解其蒸馏行为。

### 2.2生成建模中的FKL-RKL权衡
散度度量的选择强烈地塑造了学生的行为。先前的工作探索了策略内蒸馏的替代目标，包括DistilLLM中的Skewed RKL（Ko et al. (2024 (https://arxiv.org/html/2608.14685#bib.bib7))）、GKD中的JS散度（Agarwal et al. (2024 (https://arxiv.org/html/2608.14685#bib.bib6))）以及对称KL散度和广义f散度（Wen et al. (2023 (https://arxiv.org/html/2608.14685#bib.bib13))）。这些目标旨在平衡RKL的模式寻求行为与更广泛的分布覆盖。

除了静态的散度选择，近期工作进一步探索了自适应平衡策略。相关的自适应加权方法已在计算机视觉领域得到研究（Zheng and Yang, 2024 (https://arxiv.org/html/2608.14685#bib.bib14); Amara et al., 2022 (https://arxiv.org/html/2608.14685#bib.bib15)），而自回归序列生成则强调词元级别的动态性来平衡不同的散度行为。AKL（Wu et al. (2025 (https://arxiv.org/html/2608.14685#bib.bib16))）基于教师分布头部和尾部区域的不同拟合行为自适应地调整FKL和RKL，而ToDi（Jung et al. (2025 (https://arxiv.org/html/2608.14685#bib.bib17))）则根据它们在策略内采样期间增加或减少词元概率的梯度来协调它们。最近，EOPD（Jin et al. (2026 (https://arxiv.org/html/2608.14685#bib.bib18))）建立在RKL基础上，为高熵教师词元引入了一个额外的FKL项。相比之下，AED不进行FKL和RKL目标的插值，也不激活额外的FKL分支。相反，它利用RKL两个内部项的不同拟合作用，产生一个自适应的目标分布和一个简单的仅含RKL的策略内蒸馏目标。

## 3预备知识与分析
本节首先介绍问题设置、符号和基于KL的蒸馏目标，然后检查这些目标诱导的训练行为，这反过来激发了后续开发的自适应方法。

### 3.1问题设置与符号
令 \(P\) 表示教师模型，\(Q_{\theta}\) 表示参数化为 \(\theta\) 的学生模型。给定一个输入提示 \(x\) 和一个输出序列 \(y = (y_1, \dots, y_T)\)，教师和学生定义了自回归序列分布：
\[
p(y \mid x) = \prod_{t=1}^{T} p(y_t \mid x, y_{<t}), \quad q_{\theta}(y \mid x) = \prod_{t=1}^{T} q_{\theta}(y_t \mid x, y_{<t}),
\]
其中 \(y_{<t} = (y_1, \dots, y_{t-1})\) 表示步骤 \(t\) 之前的前缀。

在策略内蒸馏中，响应从学生模型采样：\(y \sim q_{\theta}(\cdot \mid x)\)。对于给定的前缀 \(y_{<t}\)，我们分别用 \(p_t\) 和 \(q_t\) 表示教师和学生在步骤 \(t\) 整个词表 \(\mathcal{V}\) 上的条件分布。具体地，对于任何词元 \(v \in \mathcal{V}\)：
\[
p_t(v) = p(v \mid x, y_{<t}), \quad q_t(v) = q_{\theta}(v \mid x, y_{<t}).
\]
在此符号下，序列级目标可以沿着采样轨迹分解为词元级项。

### 3.2基于KL蒸馏的几何分析
前向和逆向KL的不同行为可以通过它们诱导的训练目标来理解。为了分析策略内设置中的优化动态，我们将序列级目标重新表述为学生生成轨迹上的词元级期望。对于RKL，最小化 \(D_{\mathrm{KL}}(q_{\theta} \parallel p)\) 等价于最小化以下策略内目标，该目标在从学生模型采样的响应上评估：
\[
\mathcal{L}_{\mathrm{RKL}} = \mathbb{E}_{y \sim q_{\theta}(\cdot \mid x)} \left[ \sum_{t=1}^{T} \sum_{v \in \mathcal{V}} q_t(v) \log \frac{q_t(v)}{p_t(v)} \right].
\]
这个目标揭示了训练中的两种竞争力量。对数比自然可以分解为 \(-\sum_{v \in \mathcal{V}} q_t(v) \log p_t(v)\) 和 \(\sum_{v \in \mathcal{V}} q_t(v) \log q_t(v)\)。第一项 \(-\sum_{v \in \mathcal{V}} q_t(v) \log p_t(v)\) 起着源自教师的交叉熵损失的作用。当其最小时，它严厉惩罚学生将非可忽略的概率质量分配给教师认为不太可能的词元，从而驱动学生分布向教师的高概率模式靠拢。第二项 \(\sum_{v \in \mathcal{V}} q_t(v) \log q_t(v)\) 对应于学生分布的负熵，它内在地惩罚过度集中。

因此，基于RKL的训练可以解释为在学生采样轨迹上平衡教师驱动的模式寻求与熵诱导的概率扩散。

类似地，虽然标准的FKL是离策略的（即在教师生成的序列上评估），但近期的策略内蒸馏方法通常采用词元级FKL表述来缓解暴露偏差（Agarwal et al., 2024 (https://arxiv.org/html/2608.14685#bib.bib6)）。在此设置中，前缀从学生模型采样，并且学生被强制在每个步骤通过词元级FKL匹配教师的词表分布。相应的策略内FKL目标可以写为：
\[
\mathcal{L}_{\text{On-FKL}} = \mathbb{E}_{y \sim q_{\theta}(\cdot \mid x)} \left[ \sum_{t=1}^{T} \sum_{v \in \mathcal{V}} p_t(v) \log \frac{p_t(v)}{q_t(v)} \right].
\]
通过分离对数比，词元级目标可以分解为两项：\(-\sum_{v \in \mathcal{V}} p_t(v) \log q_t(v)\) 和 \(\sum_{v \in \mathcal{V}} p_t(v) \log p_t(v)\)。在这种分解中，第一项 \(-\sum_{v \in \mathcal{V}} p_t(v) \log q_t(v)\) 是交叉熵项。在教师加权的目标下，每当学生对教师支持的词元分配不足的概率质量时，该项会诱导大的修正梯度，因此迫使学生尽可能广泛地覆盖教师分布的支撑。

相比之下，第二项 \(\sum_{v} p_t(v) \log p_t(v)\) 等于 \(-H(p)\)，仅依赖于教师分布，因此相对于学生参数是常数。它不对学生优化贡献任何梯度。因此，基于FKL的训练完全由交叉熵项的支撑覆盖压力驱动，当学生族的表达能力不如教师时，这可能导致均值寻求行为。

然而，RKL具有不同的内部结构。上述分解表明它已经包含了两个互补项：一个教师拟合项和一个学生熵项。在标准的RKL中，这两项具有相等的相对权重。这自然引出了以下问题：如果我们明确地重新加权它们会发生什么？增加教师拟合项的权重会使学生分布向教师的高概率区域集中，而增加学生熵项的权重则会产生更平滑的分布，具有更广泛的概率覆盖。这表明，可以直接通过RKL的内在成分来控制概率集中与覆盖之间的权衡，这启发了下一节引入的自适应重新加权公式。

> 参见标题图2：固定权重蒸馏与AED的概念比较。

## 4方法
在本节中，我们介绍AED，一个理论驱动、不确定性感知的策略内蒸馏框架...（后续内容翻译）