On-policy distillation speeds training by cutting rollout lengths randomly
Does On-Policy Distillation Need to Complete Every Rollout?
Machine LearningArtificial Intelligence
Summary
Training a student AI using examples generated from the student’s own attempts can be slow because the process requires completing long sequences each time. The authors propose a method called STOPD that randomly stops these sequences early and adjusts the learning so that it remains fair and accurate despite seeing fewer steps. They prove mathematically that their adjustment keeps learning unbiased and show experimentally that their method reduces the work needed by up to nearly half without hurting performance. This means training can be done significantly faster while keeping the results similar.
What this means in practice
- •For machine learning engineers: Speed up training of student models that learn from teacher feedback by shortening rollout sequences without accuracy loss.
- •For ai infrastructure teams: Reduce computational resource use during on-policy distillation by applying stochastic rollout truncation and bias correction techniques.
Authors
Zhengxin Zhang, Maosen Tang, Ning Wang, Changhao Li, Jingren Zhou, Bo Dai, Sainyam Galhotra, Alex Townsend, Claire Cardie
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.