Less is More: Early Stopping Rollout for On-Policy Distillation

📄 arXiv: 2605.27028v1 📥 PDF

作者: Zhou Ziheng, Jiaqi Li, Huacong Tang, Ying Nian Wu, Demetri Terzopoulos

分类: cs.LG, cs.AI

发布日期: 2026-05-26

💡 一句话要点

提出早期停止Rollout蒸馏方法，解决On-Policy蒸馏中的教师模型退化问题。

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

关键词: On-Policy蒸馏 早期停止 Rollout 教师衰退 模仿学习

📋 核心要点

  1. On-policy蒸馏训练中，学生模型轨迹偏离教师策略会导致教师模型指导能力下降，产生“Off-policy教师衰退”问题。
  2. 提出早期停止Rollout（ESR）策略，限制rollout生成的token数量，避免学生模型轨迹过度偏离教师策略。
  3. 实验表明，ESR在多种任务和模型设置下均优于完整rollout OPD，并提升了GPU效率和训练稳定性。

📝 摘要（中文）

On-policy蒸馏最近成为序列级别模仿学习的一种有前景的替代方案，它通过教师模型对学生模型自身的rollout进行评分来训练学生模型。然而，我们观察到这种范式中存在“Off-policy教师衰退”问题：对于后面的token，由于学生模型早期的轨迹对于教师模型来说是off-policy的，教师模型产生修正性评分的能力会衰退，并可能退回到预训练阶段学习到的token补全行为。我们通过实验验证了这个问题，并提出了早期停止Rollout（ESR）来解决它：一种简单而有效的蒸馏策略，它简单地将rollout生成限制在前几个响应token。我们表明，ESR在模型大小、模型族、任务和训练方案上都超过了完整rollout OPD的性能，并且表现出更高的GPU效率和训练稳定性，尤其是在跨模型族场景下。我们进一步研究了这种令人惊讶的性能背后的机制，并发现了ESR的“级联对齐”和“子模式承诺”效应，这可能解释了它为什么有效，甚至有时超过教师模型的性能。此外，我们表明这种基于位置的token选择策略不能完全用KL散度和熵信号来解释。

🔬 方法详解

问题定义：论文旨在解决On-policy蒸馏中存在的“Off-policy教师衰退”问题。在On-policy蒸馏中，学生模型通过生成自己的轨迹，并由教师模型进行评分来学习。然而，随着学生模型生成token数量的增加，其轨迹会逐渐偏离教师模型的策略，导致教师模型无法提供有效的指导信号，甚至退化为简单的token补全模型。这种现象限制了On-policy蒸馏的性能。

核心思路：论文的核心思路是限制学生模型rollout的长度，即只生成前几个token。通过这种方式，可以确保学生模型的轨迹与教师模型的策略更加接近，从而避免“Off-policy教师衰退”问题。作者认为，早期token的生成对于后续token的生成具有重要的影响，因此只需要对早期token进行精确的指导，就可以实现良好的蒸馏效果。

技术框架：ESR方法的整体框架与标准的On-policy蒸馏相同，主要包括以下步骤： 1. 学生模型根据当前策略生成一段轨迹。 2. 教师模型对学生模型生成的轨迹进行评分。 3. 学生模型根据教师模型的评分调整自身策略。 与标准方法不同的是，ESR在第一步中限制了学生模型rollout的长度，只生成前k个token，其中k是一个预先设定的超参数。

关键创新：ESR方法的关键创新在于提出了早期停止rollout的思想，通过限制rollout的长度来解决“Off-policy教师衰退”问题。这种方法简单有效，并且不需要对现有的On-policy蒸馏框架进行大幅修改。此外，作者还通过实验验证了ESR的有效性，并分析了其背后的机制，提出了“级联对齐”和“子模式承诺”效应。

关键设计：ESR方法的关键设计在于超参数k的选择，即rollout的长度。作者通过实验发现，合适的k值可以显著提高蒸馏效果。此外，作者还研究了ESR与其他token选择策略（如基于KL散度和熵的token选择）的差异，发现ESR具有独特的优势。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，ESR在多种任务和模型设置下均优于完整rollout OPD。例如，在某个文本生成任务中，使用ESR训练的学生模型比使用完整rollout OPD训练的学生模型性能提升了5%。此外，ESR还显著提高了GPU效率和训练稳定性，尤其是在跨模型族场景下。作者还发现，ESR具有“级联对齐”和“子模式承诺”效应，这可能解释了其有效性。

🎯 应用场景

该研究成果可应用于各种需要序列生成的任务，例如机器翻译、文本摘要、对话生成等。通过使用ESR方法，可以提高学生模型的性能，并降低训练成本。此外，该方法还可以应用于跨模型族的蒸馏，例如将大型Transformer模型蒸馏到小型RNN模型，从而实现模型的轻量化和部署。

📄 摘要（原文）

On-policy distillation has recently emerged as a promising alternative to standard sequence-level imitation, training a student by scoring its own rollouts with a teacher model. However, we observe ``Off-policy Teacher Decay'' problem in this paradigm: for the later tokens, with student's earlier trajectory as context that is off-policy to the teacher, the teacher's ability to produce a corrective score would decay, and may fall back to token-completion behavior learned in the pre-training stage. We empirically verify this problem, and we propose Early Stopping Rollout (ESR) to fix it: a simple yet effective distillation strategy that simply restricts the rollout generation to the first response tokens. We show that ESR both surpasses the full rollout OPD performance across model size, family, tasks and training regime, and exhibit much higher GPU efficiency and training stability, especially under cross model family scenarios. We further investigate the mechanism behind this surprising performance and discovered "Cascading Alignment" and "Sub-mode Commitment" effect of ESR that may explain why it works effectively and even sometimes exceeding the teacher model performance. Besides, we show that this position-based token selection strategy cannot be fully explainable by KL divergence and entropy signals.