Does On-Policy Distillation Need to Complete Every Rollout?
作者: 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 (更新: 2026-10-16)
💡 一句话要点
提出STOPD以加速在政策蒸馏中的回合生成
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 在政策蒸馏 随机截断 块状逆概率加权 模型训练加速 教师-学生模型
📋 核心要点
- 现有的在政策蒸馏方法需要为每个令牌生成完整的学生回合，导致计算成本高昂。
- 本文提出随机截断的在政策蒸馏（STOPD），通过随机截断回合并使用块状逆概率加权来修正偏差。
- STOPD在多个模型和任务上表现出色，减少生成令牌数量25-48%，并实现1.3-2.2倍的训练速度提升。
📝 摘要（中文）
在政策蒸馏（OPD）中，通过对学生生成的回合进行反馈，有效地将教师的能力转移给学生。然而，为每个令牌生成完整的学生回合并计算教师反馈使得OPD计算成本高昂。本文提出随机截断的在政策蒸馏（STOPD），通过随机截断学生回合并使用块状逆概率加权（IPW）来修正采样偏差，从而加速OPD。每个训练步骤中，我们在生成之前为每个回合采样一个停止位置，并让学生生成直到该位置或自然终止。我们获得每个生成令牌的教师反馈，并使用块状IPW来补偿保留后续令牌的较低概率，同时控制权重。我们证明了这种加权在期望上保留了原始OPD的梯度和，并推导了SGD、Adam和Muon的一步更新误差界限。我们在多个教师-学生模型对、任务领域和优化器上评估了STOPD，结果表明STOPD减少生成令牌数量25-48%，同时保持与原始OPD相当的准确性，实现了1.3-2.2倍的训练加速。
🔬 方法详解
问题定义：本文解决的是在政策蒸馏（OPD）中，生成完整学生回合并计算教师反馈的高计算成本问题。现有方法在每个令牌上都需要进行完整的回合生成，导致效率低下。
核心思路：论文提出随机截断的在政策蒸馏（STOPD），通过在生成之前随机选择停止位置，允许学生在达到该位置或自然终止时生成回合，从而减少计算负担。
技术框架：STOPD的整体架构包括三个主要模块：随机停止位置采样、学生回合生成和块状逆概率加权。首先，在每个训练步骤中为每个回合采样停止位置，然后生成回合并获取教师反馈，最后通过块状IPW修正采样偏差。
关键创新：STOPD的主要创新在于引入了随机截断机制和块状逆概率加权，这与传统的OPD方法在生成完整回合的要求上有本质区别。通过这种设计，STOPD能够在保持准确性的同时显著提高训练速度。
关键设计：在STOPD中，关键参数包括停止位置的采样策略和块状IPW的权重控制。此外，损失函数的设计也考虑了教师反馈的有效利用，以确保生成令牌的质量。
🖼️ 关键图片
📊 实验亮点
STOPD在多个教师-学生模型对和任务领域的实验中表现优异，生成令牌数量减少了25-48%，同时保持了与传统OPD相当的准确性，训练速度提升达1.3-2.2倍，展示了其显著的性能优势。
🎯 应用场景
该研究的潜在应用领域包括自然语言处理、强化学习和其他需要教师-学生模型架构的任务。通过提高在政策蒸馏的效率，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.