Pass the Baton: Trajectory-Relayed On-Policy Distillation
2026-07-28 • Computation and Language
Computation and LanguageArtificial Intelligence
AI summaryⓘ
The authors identify a problem with on-policy distillation, where a student model following its own reasoning path can get stuck on mistakes, making further learning inefficient. To fix this, they create Relay-OPD, which lets a teacher model briefly take over when the student starts going off track, then returns control to the student. This targeted teacher intervention helps the student learn better reasoning early on, improving performance on math tasks while using less training time. Their method outperforms previous techniques across multiple benchmarks with smaller and larger student models.
On-policy distillationPrefix failureTeacher-student asymmetryRelay On-Policy DistillationTrajectoryMathematical reasoningStudent modelTeacher modelModel trainingIntervention
Authors
Haolei Xu, Xiaowen Xu, Haiwen Hong, Zixuan Ni, Hongxing Li, Yiwen Qiu, Weiming Lu, Yongliang Shen
Abstract
On-policy distillation (OPD) grounds token-level supervision in the student's own trajectory, yet suffers from prefix failure: once the student commits to a wrong reasoning direction, all subsequent generation builds on this deviation, producing misdirected continuations that elicit unreliable supervision and waste compute. We identify a teacher-student continuation asymmetry on failed prefixes, where the teacher tends to redirect while the student continues along the original direction, and convert it into a label-free handoff trigger in Relay On-Policy Distillation (Relay-OPD). During training, Relay-OPD constructs relay trajectories by letting the teacher briefly take over at detected trigger points to produce a teacher leg, after which the student resumes and is optimized on the resulting trajectory. A limited relay budget concentrates intervention on critical early positions while limiting departure from the student policy. With a Qwen3-4B-Instruct-2507 teacher and Qwen3-0.6B/1.7B-Non-Thinking students on eight mathematical reasoning benchmarks, Relay-OPD achieves the best or second-best results on every benchmark, outperforming standard OPD by +5.73% and the strongest baseline FastOPD by +1.49% on average for 1.7B, with consistent gains at 0.6B. Training trajectory length is reduced by over 50%.