Are Full Rollouts Necessary for On-Policy Distillation?

📄 arXiv: 2605.31490v1 📥 PDF

作者: Yaocheng Zhang, Jiajun Chai, Songjun Tu, Yuqian Fu, Xiaohan Wang, Wei Lin, Guojun Yin, Qichao Zhang, Yuanheng Zhu, Dongbin Zhao

分类: cs.CL

发布日期: 2026-05-29

备注: 14 pages, 16 figures


💡 一句话要点

提出渐进式和截断式策略,提升On-policy蒸馏在长序列推理中的训练效率。

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)

关键词: On-policy蒸馏 长序列推理 训练效率 rollout长度控制 知识蒸馏

📋 核心要点

  1. 传统On-policy蒸馏(OPD)在长序列推理中计算成本高昂,且学生模型在训练初期易受后期不可靠教师反馈的影响。
  2. 论文提出渐进式OPD(POPD)和截断式OPD(TOPD),通过控制rollout长度,提升训练效率并减少资源消耗。
  3. 实验表明,POPD可将OPD训练效率提升3倍,TOPD仅用10%的rollout长度即可达到原有OPD的性能。

📝 摘要(中文)

On-policy 蒸馏 (OPD) 提供由学生模型生成的 rollout 过程中的密集教师反馈,已成为长序列推理中一种有前景的后训练范式。然而,标准的 OPD 通常在训练期间生成完整的 rollout,这在计算上是昂贵的,并且可能使学生在 rollout 后期位置暴露于不可靠的教师反馈,尤其是在早期训练期间。我们确定 rollout 长度是 OPD 的一个关键瓶颈,它会显著影响训练效率。与具有可验证奖励的强化学习 (RLVR) 不同,OPD 不需要完整的轨迹或最终答案奖励来提供学习信号。这一观察结果表明,完整的 rollout 可能并不总是对有效的 OPD 都是必要的。受此启发,我们提出了两种简单的 rollout 长度控制策略:渐进式 OPD (POPD),它在训练期间逐渐扩展 rollout 长度;截断式 OPD (TOPD),它永久地对可靠的截断 rollout 执行蒸馏。在数学推理上的实验表明,POPD 将 OPD 的训练效率提高了高达 3 倍,而 TOPD 仅使用 10% 的 rollout 长度即可匹配 OPD 的性能,从而显著减少了时间和内存。这些结果表明,控制 rollout 长度为更高效的 OPD 提供了一条简单而实用的途径。

🔬 方法详解

问题定义:论文旨在解决On-policy蒸馏(OPD)在长序列推理任务中训练效率低下的问题。现有OPD方法通常需要生成完整的rollout,这导致计算资源消耗巨大,并且在训练初期,学生模型容易受到rollout后期不可靠的教师反馈的干扰,影响学习效果。

核心思路:论文的核心思路是控制rollout的长度,认为不需要完整的rollout也能实现有效的蒸馏。通过限制或逐步增加rollout长度,可以减少计算负担,并使学生模型专注于更可靠的早期反馈。

技术框架:论文提出了两种策略: 1. 渐进式OPD (POPD):在训练初期使用较短的rollout长度,随着训练的进行,逐步增加rollout长度。这使得学生模型可以先学习更容易的部分,然后再逐步学习更复杂的部分。 2. 截断式OPD (TOPD):始终使用截断的rollout进行蒸馏。选择一个合适的截断长度,保证教师反馈的可靠性,并减少计算量。

关键创新:论文的关键创新在于认识到在On-policy蒸馏中,完整的rollout并非总是必要的。通过控制rollout长度,可以在保证性能的同时,显著提高训练效率。与传统的OPD方法相比,POPD和TOPD更加灵活,可以根据任务的特点和计算资源进行调整。

关键设计: 1. Rollout长度的控制策略:POPD采用线性或指数方式逐步增加rollout长度。TOPD则需要选择一个合适的截断长度,可以通过实验或先验知识确定。 2. 损失函数:POPD和TOPD仍然使用标准的蒸馏损失函数,例如KL散度或均方误差,用于衡量学生模型和教师模型输出之间的差异。 3. 训练过程:POPD和TOPD的训练过程与标准的OPD类似,都是通过最小化蒸馏损失来更新学生模型的参数。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,POPD可以将OPD的训练效率提高高达3倍,这意味着在相同的计算资源下,可以训练出性能更好的模型。TOPD仅使用10%的rollout长度即可达到与标准OPD相当的性能,从而显著减少了时间和内存消耗。这些结果验证了控制rollout长度对于提高OPD训练效率的有效性。

🎯 应用场景

该研究成果可应用于各种需要长序列推理的任务,例如数学问题求解、代码生成、自然语言理解等。通过提高On-policy蒸馏的训练效率,可以降低模型训练的成本,并加速模型部署。此外,该方法还可以扩展到其他类型的蒸馏训练中,例如离线蒸馏和跨模态蒸馏。

📄 摘要(原文)

On-policy distillation (OPD) provides dense teacher feedback along rollouts generated by the student and has emerged as a promising post-training paradigm for long-horizon reasoning. However, standard OPD typically generates full rollouts during training, which is computationally expensive and may expose the student to unreliable teacher feedback at late rollout positions, especially during early training. We identify the rollout horizon as a key bottleneck in OPD that substantially impacts training efficiency. Unlike Reinforcement Learning with Verifiable Rewards (RLVR), OPD does not require a complete trajectory or a final answer reward to provide learning signals. This observation suggests that full rollouts may not always be necessary for effective OPD. Motivated by this insight, we propose two simple horizon-control strategies: Progressive OPD (POPD), which gradually expands the rollout horizon during training, and Truncated OPD (TOPD), which permanently performs distillation on reliable truncated rollouts. Experiments on mathematical reasoning show that POPD improves the training efficiency of OPD by up to 3$\times$, while TOPD matches OPD performance using only 10\% of the rollout horizon, leading to substantial wall-clock and memory reductions. These results demonstrate that controlling the rollout horizon offers a simple and practical path to more efficient OPD.