Unbiased Top-$k$ Estimation for On-Policy Distillation
作者: Linjian Meng, Siyuan Gan, YuHan Li, Xiran Wang, Ziyang Ding, Ditang Gou, Yiming Wu, Zhen Zhao
分类: cs.CL, cs.LG, stat.ML
发布日期: 2026-09-28
💡 一句话要点
提出TT-OPD以解决TK-OPD中的偏差问题
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 在线蒸馏 反向KL散度 Top-$k$方法 无偏估计 大语言模型 知识转移 模型训练
📋 核心要点
- 现有的Top-$k$在线蒸馏方法在计算效率上有所提升，但由于只使用Top-$k$令牌，导致了估计偏差，影响了模型的准确性。
- 本文提出的尾部校正Top-$k$在线蒸馏（TT-OPD）方法，通过结合Top-$k$令牌和样本令牌，恢复了被丢弃的概率质量，从而消除了偏差。
- 实验结果显示，TT-OPD在多个基准测试中表现优异，相比于其他OPD变体，准确性有显著提升。
📝 摘要（中文）
在大语言模型的后训练中，在线蒸馏（OPD）通过最小化教师与学生之间的反向KL散度来训练学生模型。然而，估计反向KL散度的梯度仍然是一个挑战。现有的Top-$k$ OPD方法虽然降低了计算成本，但会引入偏差，导致准确性下降。为了解决这一问题，本文提出了尾部校正的Top-$k$在线蒸馏（TT-OPD），它结合了选定的Top-$k$令牌和学生生成的样本令牌，从而提供无偏的梯度估计。实验结果表明，TT-OPD在性能上显著优于其他OPD变体。
🔬 方法详解
问题定义：本文旨在解决在线蒸馏（OPD）中反向KL散度梯度估计的偏差问题。现有的Top-$k$ OPD方法虽然降低了计算成本，但因只考虑Top-$k$令牌而丢失了部分概率信息，导致准确性下降。
核心思路：TT-OPD的核心思路是结合Top-$k$令牌和学生生成的样本令牌，以恢复丢失的概率质量，从而提供无偏的梯度估计。这种设计旨在兼顾计算效率和模型准确性。
技术框架：TT-OPD的整体架构包括两个主要模块：一是从学生生成的样本中选择Top-$k$令牌，二是结合这些令牌与样本令牌进行反向KL散度的梯度估计。
关键创新：TT-OPD的关键创新在于通过结合Top-$k$令牌和样本令牌，避免了仅使用Top-$k$令牌所带来的偏差。这一方法在保留计算效率的同时，显著提高了模型的准确性。
关键设计：在TT-OPD中，损失函数设计为结合了反向KL散度的无偏估计，参数设置上则考虑了Top-$k$的选择策略和样本令牌的采样方式，以确保模型训练的有效性和稳定性。
🖼️ 关键图片
📊 实验亮点
实验结果表明，TT-OPD在多个基准测试中显著优于其他OPD变体，准确性提升幅度达到X%（具体数据需根据实验结果填写），验证了其在无偏估计方面的有效性。
🎯 应用场景
该研究在大语言模型的后训练和知识蒸馏领域具有广泛的应用潜力。TT-OPD方法可以有效提升模型的推理能力，适用于对话系统、文本生成和其他自然语言处理任务。未来，该方法可能推动更高效的模型训练技术的发展。
📄 摘要（原文）
On-policy distillation (OPD) is becoming an important component of large language model (LLM) post-training for transferring the reasoning capability of a strong teacher LLM to a weaker student LLM. OPD trains the student by minimizing the reverse KL divergence between the teacher and the student via rollouts generated by the student's policy. However, estimating the gradient of the reverse KL divergence in OPD remains a challenge. Using only the sampled token from the student-generated rollout is computationally cheap but provides limited distributional supervision, which will degrade accuracy. In addition, using the full vocabulary provides complete distributional supervision but is computationally expensive. Therefore, recent works propose Top-$k$ OPD (TK-OPD) that use selected top-$k$ tokens, which provides richer distributional supervision than sampled-token estimation at substantially lower computational cost than full-vocabulary estimation. Unfortunately, using only the selected top-$k$ tokens induces bias, leading to accuracy degradation, as the probability mass outside the selected top-$k$ tokens is discarded. To address the bias of TK-OPD, we propose Tail-Corrected Top-$k$ On-Policy Distillation (TT-OPD). It preserves the advantages of TK-OPD, including rich distributional supervision and low computational cost, while providing an unbiased estimator of the gradient of the reverse KL divergence. The key insight of TT-OPD is to use not only the selected top-$k$ tokens, but also the sampled token from the student-generated rollout, thereby recovering the discarded probability mass in expectation, avoiding the bias. Experimental results demonstrate that TT-OPD significantly outperforms other tested OPD variants.