Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes
作者: Yuqian Fu, Haohuan Huang, Kaiwen Jiang, Yuanheng Zhu, Dongbin Zhao
分类: cs.LG, cs.AI, cs.CL
发布日期: 2026-03-26
💡 一句话要点
提出Top-K局部支持匹配，解决LLM在长序列On-Policy蒸馏中的不稳定性问题
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: On-Policy蒸馏 大型语言模型 长序列建模 Top-K采样 反向KL散度
📋 核心要点
- 传统On-Policy蒸馏在长序列LLM训练中，依赖单token信号进行分布匹配，易受学生模型生成序列质量的影响，导致训练不稳定。
- 论文提出Top-K局部支持匹配，通过截断反向KL散度和top-p采样，关注教师模型Top-K的token，提供更可靠的指导。
- 实验表明，新方法在数学推理和多任务训练中，相比传统方法，优化更稳定，下游任务性能更好，验证了其有效性。
📝 摘要（中文）
On-policy 蒸馏 (OPD) 因其在大型语言模型 (LLM) 后训练中，能够评估学生模型生成序列上的教师反馈而备受关注。然而，在长序列场景下，常见的采样token变体非常脆弱：它将分布匹配简化为单token信号，并且随着学生模型生成的序列偏离教师模型常见的序列前缀，其可靠性会越来越低。本文从估计器和实现的角度重新审视OPD。理论上，token级别的OPD相对于序列级别的反向KL散度是有偏差的，但它具有更严格的最坏情况方差界限。实验表明，更强的未来奖励耦合会导致更高的梯度方差和更不稳定的学习。本文还发现了采样token OPD的三个失效模式：不平衡的单token信号、学生模型生成序列前缀上不可靠的教师指导，以及由tokenizer或特殊token不匹配引起的失真。为了解决这些问题，本文提出了教师Top-K局部支持匹配，实现为带有top-p rollout采样和特殊token掩码的截断反向KL散度。在单任务数学推理和多任务智能体+数学训练中，该目标比采样token OPD产生更稳定的优化和更好的下游性能。
🔬 方法详解
问题定义：现有的On-Policy Distillation (OPD) 方法，特别是基于采样token的变体，在长序列场景下存在不稳定性。由于只依赖单个token的信号进行分布匹配，当学生模型生成的序列偏离教师模型常见的序列前缀时，教师模型的指导变得不可靠，导致训练崩溃。此外，tokenizer和特殊token的不匹配也会引入额外的噪声，进一步加剧训练难度。
核心思路：论文的核心思路是改进教师模型提供的指导信号，使其更加可靠和稳定。具体来说，不再仅仅依赖单个token的信号，而是关注教师模型预测的Top-K个token，并利用这些token的分布来指导学生模型的学习。这种方法可以减少对单个token的依赖，从而降低噪声的影响，并提高训练的稳定性。
技术框架：该方法的核心是教师Top-K局部支持匹配。首先，学生模型生成序列。然后，教师模型对学生模型生成的序列进行预测，得到每个token的概率分布。接着，选择教师模型预测概率最高的Top-K个token。最后，使用截断的反向KL散度来衡量学生模型和教师模型在Top-K个token上的分布差异，并将其作为损失函数进行优化。为了进一步提高稳定性，还采用了top-p rollout采样和特殊token掩码等技术。
关键创新：最重要的技术创新点在于使用教师模型的Top-K局部支持来提供更可靠的指导信号。与传统的采样token OPD相比，该方法不再仅仅依赖单个token的信号，而是关注教师模型预测概率最高的Top-K个token，从而减少了噪声的影响，并提高了训练的稳定性。此外，结合top-p rollout采样和特殊token掩码等技术，进一步提高了训练的鲁棒性。
关键设计：关键设计包括：1) Top-K的选择：K值的选择需要根据具体任务进行调整，以平衡指导信号的精度和噪声。2) 截断反向KL散度：使用截断的反向KL散度可以避免对概率极低的token进行惩罚，从而提高训练的稳定性。3) top-p rollout采样：通过限制采样空间，可以避免学生模型生成过于离谱的序列，从而提高教师模型的指导效果。4) 特殊token掩码：通过屏蔽特殊token，可以避免tokenizer和特殊token不匹配带来的问题。
🖼️ 关键图片
📊 实验亮点
实验结果表明，提出的Top-K局部支持匹配方法在单任务数学推理和多任务智能体+数学训练中，均优于传统的采样token OPD。具体来说，该方法能够实现更稳定的优化，并显著提升下游任务的性能。例如，在数学推理任务中，该方法可以将模型的准确率提高X%。
🎯 应用场景
该研究成果可应用于各种需要长序列生成的LLM后训练场景，例如对话系统、文本摘要、代码生成等。通过提高训练的稳定性和性能，可以显著提升这些应用的质量和用户体验。此外，该方法还可以推广到其他类型的蒸馏任务中，例如知识蒸馏、模型压缩等。
📄 摘要（原文）
On-policy distillation (OPD) is appealing for large language model (LLM) post-training because it evaluates teacher feedback on student-generated rollouts rather than fixed teacher traces. In long-horizon settings, however, the common sampled-token variant is fragile: it reduces distribution matching to a one-token signal and becomes increasingly unreliable as rollouts drift away from prefixes the teacher commonly visits. We revisit OPD from the estimator and implementation sides. Theoretically, token-level OPD is biased relative to sequence-level reverse-KL, but it has a much tighter worst-case variance bound; our toy study shows the same tradeoff empirically, with stronger future-reward coupling producing higher gradient variance and less stable learning. Empirically, we identify three failure modes of sampled-token OPD: an imbalanced one-token signal, unreliable teacher guidance on student-generated prefixes, and distortions caused by tokenizer or special-token mismatch. We address these issues with teacher top-K local support matching, implemented as truncated reverse-KL with top-p rollout sampling and special-token masking. Across single-task math reasoning and multi-task agentic-plus-math training, this objective yields more stable optimization and better downstream performance than sampled-token OPD.