无偏 Top-$k$ 估计用于 On-Policy Distillation

arXiv:2609.34447v2 公告类型：替换交叉版本 摘要：On-policy distillation（OPD）正逐渐成为大型语言模型（LLM）后训练的重要组成部分，用于将强教师 LLM 的推理能力迁移至较弱的学生 LLM。OPD 通过学生策略生成的 rollout，最小化教师与学生之间的反向 KL 散度来训练学生。然而，在 OPD 中估计反向 KL 散度的梯度仍然是一项挑战。仅使用学生生成 rollout 中采样到的 token 虽然计算成本低，但提供的分布监督有限，因而会降低准确率。此外，使用完整词表能够提供完整的分布监督，但计算成本高昂。因此，近期工作提出了 Top-$k$ OPD（TK-OPD），使用选定的 top-$k$ 个 token；与基于采样 token 的估计相比，它能够提供更丰富的分布监督，同时计算成本又显著低于基于完整词表的估计。然而，仅使用选定的 top-$k$ 个 token 会引入偏差，导致准确率下降，因为选定的 top-$k$ 个 token 之外的概率质量被丢弃。为解决 TK-OPD 的偏差问题，我们提出了 Tail-Corrected Top-$k$ On-Policy Distillation（TT-OPD）。它保留了 TK-OPD 的优势，包括丰富的分布监督和较低的计算成本，同时能够对反向 KL 散度的梯度提供无偏估计。TT-OPD 的关键思想是不仅使用选定的 top-$k$ 个 token，还使用学生生成 rollout 中采样到的 token，从而在期望意义下恢复被丢弃的概率质量，避免产生偏差。实验结果表明，TT-OPD 的表现显著优于其他经过测试的 OPD 变体。