arXiv cs.AI· Zhengxin Zhang, Maosen Tang, Ning Wang, Changhao Li, Jingren Zhou, Bo Dai, Sainyam Galhotra, Alex Townsend, Claire Cardie·· 2 天前AI 评分24
STOPD:随机截断 rollout 的 on-policy 蒸馏方法
Does On-Policy Distillation Need to Complete Every Rollout?
AI 导读
研究者提出 Stochastic Truncated On-Policy Distillation(STOPD),通过随机截断学生模型的 rollout 并用分块逆概率加权(IPW)修正采样偏差,来加速 on-policy 蒸馏(OPD)。
正文 · 原文
Abstract: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.
| Subjects: | Machine Learning (cs.LG); Artificial Intelligence (cs.AI) |
| Cite as: | arXiv:2610.21337 [cs.LG] |
| (or arXiv:2610.21337v1 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2610.21337 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Sainyam Galhotra [view email]
[v1]
Thu, 15 Oct 2026 06:22:08 UTC (943 KB)
来源:arXiv cs.AI · arxiv.org