On-policy distillation sped up by stopping early and restarting smartly
Know When to Stop, Where to Restart: Accelerating Multi-Turn Agentic On-Policy Distillation
Machine LearningComputation and Language
Summary
Training smaller, faster AI agents using examples from bigger, smarter ones takes a lot of time because each step has to be carefully replayed. The authors studied how the helpful teaching signals vary over time and found that important guidance mostly appears at the start of each step, and mistakes are clear from the very first action gone wrong. Using these insights, they created a method called STRIDE that stops teaching when the guidance is unreliable and restarts from good points, making training faster without losing quality. This method speeds up training significantly across different tests while keeping performance high.
What this means in practice
- •For ai model trainers: Accelerate training of compact AI agents by stopping rollouts early when teacher guidance weakens and restarting from reliable points.
- •For game development teams: Reduce training time for game AI characters that rely on multi-turn decision making using the proposed adaptive stopping and prefix buffering method.
Authors
Zhiyu Gui, Kexin Huang, Jia Guo, Junkang Wu, Zihao Wang, Zhiqiang Zhang, Jun Zhou, Jiancan Wu, Xiang Wang
Abstract
On-policy distillation (OPD) has become a standard approach for transferring capabilities from large teachers to compact students. Its cost, however, is dominated by autoregressive student rollouts and scales poorly in multi-turn agentic settings. Existing acceleration methods truncate or relocate the supervision signal according to fixed, offline budgets, despite substantial variation in teacher-signal reliability both within and across trajectories. Our empirical analysis on $τ^2$-bench reveals a clear structure in this variation: informative supervision is concentrated in the prefix of each turn, and, most importantly for multi-turn agentic training, the cross-turn loss of teacher endorsement is temporally locked to the student's first erroneous action rather than accumulating gradually over turns. Building on these findings, we propose STRIDE (Stop-and-Restart on-policy Distillation acceleration), which combines two complementary techniques: adaptive early stopping, which terminates a rollout once the cumulative teacher log-probability falls below an out-of-distribution threshold, and a prefix buffer, which caches high-quality prefixes and restarts generation at the weakest correct turn. Together, these mechanisms induce a data-driven curriculum that progressively extends coverage to later turns. On $τ^2$-bench retail, our method matches full-trajectory OPD and exceeds the 30B teacher at a $3.73\times$ speedup, surpasses the baseline itself at $2.34\times$, and retains a $4.51\times$ speedup under cross-domain multi-teacher training. As a supplementary generalization test beyond the agentic setting, STRIDE outperforms full OPD on AIME 2025 at a $5.10\times$ speedup and on AIME 2024 at a $3.08\times$ speedup; averaged across the two evaluations, both fixed-budget truncation baselines remain below full OPD.