Pass the Baton: Trajectory-Relayed On-Policy Distillation

2026-07-28

AI Summary

This work addresses the challenge in policy distillation where early reasoning errors by the student model lead to persistent trajectory deviations, unreliable supervision signals, and inefficient computation. To mitigate this, the authors propose Relay-OPD, a novel method that employs an unsupervised trigger mechanism to detect student deviation and selectively delegates critical segments to the teacher model, thereby constructing relay trajectories to refine the student policy. Relay-OPD uniquely leverages behavioral asymmetry between teacher and student on erroneous prefixes to enable precise, early intervention within a constrained computational budget. Experimental results demonstrate that Relay-OPD achieves state-of-the-art or runner-up performance across eight mathematical reasoning benchmarks, with a 1.7B-parameter student model outperforming standard OPD by 5.73% and FastOPD by 1.49% on average, while reducing training trajectory length by over 50%.

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%.

Problem

• on-policy distillation
• prefix failure
• trajectory
• teacher-student asymmetry
• mathematical reasoning

Innovation

• On-policy distillation
• trajectory relay
• prefix failure
• teacher-student asymmetry
• label-free handoff