Does On-Policy Distillation Need to Complete Every Rollout?

📄 arXiv: 2610.21337v1 📥 PDF

作者: Zhengxin Zhang, Maosen Tang, Ning Wang, Changhao Li, Jingren Zhou, Bo Dai, Sainyam Galhotra, Alex Townsend, Claire Cardie

分类: cs.LG, cs.AI

发布日期: 2026-10-15

💡 一句话要点

提出随机截断的在线蒸馏方法以降低计算成本

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

关键词: 在线蒸馏 随机截断 块状逆概率加权 计算效率 模型训练加速

📋 核心要点

  1. 现有的在线蒸馏方法在生成完整回合和计算反馈时计算成本高，影响效率。
  2. 本文提出随机截断在线蒸馏（STOPD），通过随机截断回合并使用块状逆概率加权来修正偏差。
  3. STOPD在多个实验中显示出生成标记减少25-48%，同时保持与原始OPD相当的准确性，训练速度提升1.3-2.2倍。

📝 摘要（中文）

在线蒸馏（OPD）通过对学生生成的回合进行反馈，有效地将教师的能力转移给学生。然而，生成完整的学生回合并为每个标记计算教师反馈使得OPD的计算成本较高。本文提出随机截断在线蒸馏（STOPD），通过随机截断学生回合并使用块状逆概率加权（IPW）来修正采样偏差，从而加速OPD。在每个训练步骤中，我们在生成之前为每个回合采样一个停止位置，并让学生生成直到该位置或自然终止。我们为每个生成的标记获取教师反馈，并使用块状IPW来补偿保留后续标记的较低概率，同时控制权重。我们证明了这种加权在期望上保留了原始OPD的梯度和，并推导了SGD、Adam和Muon的一步更新误差界限。通过在多个教师-学生模型对、任务领域和优化器上进行评估，STOPD在保持与原始OPD相当的准确性的同时，减少生成标记25-48%，实现了1.3-2.2倍的训练加速。

🔬 方法详解

问题定义：本文旨在解决在线蒸馏（OPD）在生成完整学生回合和计算教师反馈时的高计算成本问题。现有方法需要为每个标记计算反馈，导致效率低下。

核心思路：论文提出随机截断在线蒸馏（STOPD），通过在生成之前随机选择回合的停止位置，允许学生在该位置或自然终止时生成，从而减少计算量。

技术框架：STOPD的整体架构包括三个主要阶段：首先，随机选择停止位置；其次，学生生成标记并获取教师反馈；最后，使用块状逆概率加权（IPW）来修正采样偏差。

关键创新：STOPD的主要创新在于引入了随机截断和块状IPW，显著降低了计算成本，同时保持了与原始OPD相当的性能。这一方法与传统的逐标记反馈机制本质上不同。

关键设计：在设计中，STOPD通过控制权重来补偿后续标记的低保留概率，确保了梯度和的期望保持不变。此外，推导了不同优化器的更新误差界限，为实际应用提供了理论支持。

🖼️ 关键图片

img_0
img_1
img_2

📊 实验亮点

STOPD在多个实验中表现出色，生成标记减少了25-48%，同时保持了与原始OPD相当的准确性，训练速度提升了1.3-2.2倍。这些结果表明STOPD在效率和性能之间取得了良好的平衡。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、强化学习和其他需要教师-学生模型的任务。通过降低计算成本，STOPD可以使得大规模模型的训练更加高效，推动相关领域的研究与应用发展，尤其是在资源受限的环境中。

📄 摘要（原文）

On-policy distillation (OPD) effectively transfers capabilities from a teacher to a student through feedback on student-generated rollouts. However, generating complete student rollouts and computing teacher feedback for every token make OPD computationally expensive. We introduce Stochastic Truncated On-Policy Distillation (STOPD), which accelerates OPD by randomly truncating student rollouts and correcting the resulting sampling bias with blockwise inverse-probability weighting (IPW). At each training step, we sample a stopping position for each rollout before generation and let the student generate until that position or natural termination. We obtain teacher feedback on every generated token and use blockwise IPW to compensate for the lower probability of retaining later tokens while controlling the weights. We prove that this weighting preserves Vanilla OPD's gradient sum in expectation and derive one-step update-error bounds for SGD, Adam, and Muon. We evaluate STOPD across multiple teacher-student model pairs, task domains, and optimizers. STOPD reduces generated tokens by 25-48% while maintaining accuracy comparable to Vanilla OPD, achieving 1.3-2.2x training speedups.