Tail-Aware Top-$k$ On-Policy Distillation

技术摘要：尾部感知 Top-k 在线蒸馏 (TA-OPD)

问题陈述
在线蒸馏 (On-policy distillation, OPD) 通过在学生模型自身的轨迹上，训练其下一个 token 的分布与教师模型的分布对齐，从而解决离线蒸馏中固有的训练-推理差异问题。一种普遍的方法是 归一化 Top-k OPD (Normalized Top-k OPD)，它最小化教师分布中前 $k$ 个最高概率 token 范围内的反向 Kullback-Leibler (KL) 散度。

本文指出这种归一化目标函数存在一个关键缺陷：它丢弃了关于尾部概率（即 top-$k$ token 之外的总概率质量）的信息。理论分析和实证证据表明，最小化该归一化目标会导致一种系统性偏差，使学生的尾部概率和熵稳步增加。随着学生尾部概率的增长，它会越来越多地采样到教师 top-$k$ 集合之外的 token，从而漂移到教师监督不可靠的前缀路径上。这种现象会导致下游准确率的下降，尤其是在学生与教师之间存在较大能力差距时。

方法论：尾部感知 Top-k OPD (TA-OPD)
为了解决尾部概率信号丢失的问题，作者提出了 TA-OPD。该方法通过引入一个显式的尾部 token ($v_{tail}$) 来恢复这些缺失的信息。

1. 增强 Token 集： TA-OPD 不再仅在 top-$k$ 个 token 上进行概率归一化，而是定义了一个增强集 $S^+_t = S^k_t \cup \{v_{tail}\}$，其中 $S^k_t$ 代表教师的 top-$k$ 个 token。
2. 概率分配：
   • 属于 $S^k_t$ 的 token 保留其原始概率 $p_t(v)$。
   • 尾部 token $v_{tail}$ 被分配总尾部概率：$p_t(v_{tail}) = p^{tail}_t = \sum_{v \notin S^k_t} p_t(v)$。
   • 同理，教师的分布也通过 $q^{tail}_t$ 进行增强。
3. 损失函数： 该方法最小化在该增强集上的反向 KL 散度：
   $$\ell^{TA}_t = \sum_{v \in S^k_t} p_t(v) \log \frac{p_t(v)}{q_t(v)} + p^{tail}_t \log \frac{p^{tail}_t}{q^{tail}_t}$$
   这一公式显式地比较了学生与教师的尾部概率，驱动学生的尾部概率向教师趋近。

理论特性

• 梯度动力学： 与提供零梯度的归一化目标（对于 top-$k$ 集合之外的 logits）不同，TA-OPD 通过尾部对数比来更新尾部 logits。此外，top-$k$ logits 上的梯度不再求和为零，这使得该方法能够通过显式提升 top-$k$ logits 相对于尾部 logits 的比例，来应对学生过度加权尾部的情况。
• 下界： 证明了 TA-OPD 损失是全词表反向 KL 散度的紧下界。两者的差值是一个与尾部概率成正比的非负残差项。在实践中，随着 TA-OPD 将尾部概率推向零，该残差也会随之消失，使目标函数成为一个紧密的近似。
• 无偏变体 (SC-TA-OPD)： 当采样 token 的概率可用时，作者推导出了 样本修正的 TA-OPD (Sample-Corrected TA-OPD)。该变体增加了一个基于采样 token 的偏差修正项，从而在数值和梯度上都能实现全词表反向 KL 散度的无偏估计。

核心贡献

1. 识别尾部概率增加问题： 本文从理论和实证上证明了归一化 top-k OPD 会导致学生尾部概率和熵的持续增加，进而导致性能下降。
2. 提出 TA-OPD： 一种新型蒸馏方法，通过引入尾部 token 恢复尾部概率信号，防止了与 top-k 归一化相关的熵增现象。
3. 理论保证： 证明了 TA-OPD 是全词表目标的紧下界，并推导出了无偏的样本修正变体。
4. 实验优越性： 大量实验表明 TA-OPD 优于基准方法，特别是在学生与教师能力差距较大的场景下。

实验结果
作者在数学推理基准测试（MATH500, Minerva, OlympiadBench, AMC, AIME24/25）和分布外 (OOD) 任务（ARC-c, MMLU-Pro）上评估了 TA-OPD，使用了三组学生-教师对（例如，由 OpenThinker3-7B 蒸馏的 Qwen2.5-7B-Instruct）。

• 性能提升： TA-OPD 一致取得了最佳的平均准确率。显著的改进包括：
  • Llama-3.1-8B 在数学基准测试平均分上提升了 8.05 个百分点（相比于归一化 top-k OPD）。
  • Llama-3.1-8B 在 MATH500 上提升了 19.53 个百分点（从 31.30% 提高到 50.83%）。
  • 在 OOD 任务中也有显著收益（例如，Qwen2.5-7B-Instruct 的 ARC-c 准确率从 23.07% 提升至 76.30%）。
• 训练动力学： TA-OPD 使训练更加稳定，使学生的熵保持在 1.5 以下，尾部概率接近 0；而归一化 top-k OPD 则导致熵上升至 ~6，且尾部概率上升至 ~0.6。
• 能力差距敏感性： 当学生与教师之间的能力差距较大时，该方法最为有效。当差距较小时，性能提升会减弱。
• 参数敏感性： TA-OPD 在 $k$ 值较小（如 $k=2$）时依然有效，且一旦 $k \ge 2$ 后对 $k$ 的取值表现出不敏感性。

意义与主张
本文声称 TA-OPD 解决了现有 top-k 蒸馏方法中的一个根本缺陷，即通过保留尾部概率信号来维持学生与教师的对齐。该方法被认为是一种实用的改进，除了 top-$k$ 概率外不需要额外的教师查询，且易于实现。作者指出，虽然 TA-OPB 在大差距场景下能带来实质性收益，但随着 $k$ 的增加或能力差距的缩小，其优势会减弱。由于计算限制，该研究目前仅限于 8B 参数以下的模型。