在策略蒸馏中是否需要完成每次 rollout？

arXiv:2610.21337v1 发布类型：new 摘要：On-policy distillation（OPD）通过对学生模型生成的 rollout 提供反馈，有效地将能力从教师模型迁移到学生模型。然而，生成完整的学生模型 rollout，并为每个 token 计算教师反馈，使 OPD 的计算成本高昂。我们提出 Stochastic Truncated On-Policy Distillation（STOPD），通过随机截断学生模型 rollout，并使用分块逆概率加权（IPW）校正由此产生的采样偏差，从而加速 OPD。在每个训练步骤中，我们在生成前为每个 rollout 采样一个停止位置，并让学生模型生成至该位置或自然终止。我们为每个生成的 token 获取教师反馈，并使用分块 IPW 补偿后续 token 被保留概率较低的影响，同时控制权重大小。我们证明，这种加权在期望意义下保留了 Vanilla OPD 的梯度总和，并推导出适用于 SGD、Adam 和 Muon 的单步更新误差界。我们在多个教师模型—学生模型组合、任务领域和优化器上评估了 STOPD。STOPD 在保持与 Vanilla OPD 相当准确率的同时，将生成的 token 数量减少了 25%–48%，实现了 1.3–2.2 倍的训练加速。