Pass the Baton: Trajectory-Relayed On-Policy Distillation

AI summary

Overview

Research area: Large language model post-training, specifically on-policy distillation (OPD) for mathematical reasoning, with a secondary contribution in efficient inference via speculative decoding.

Technical level: Advanced. The paper combines a reinforcement-learning distillation objective (reverse-KL single-sample advantage with PPO-style clipping) with a state-switched speculative decoding implementation. The prose summary here is written for a general technical reader.

One-sentence scope: Relay-OPD detects, without any label or verifier, the moment a student model's reasoning goes off track and briefly hands generation to the teacher before the student resumes, improving accuracy on eight math benchmarks while cutting training trajectory length by over 50%.

What This Paper Is About

On-policy distillation trains a small student model on trajectories it generates itself, with a larger teacher scoring each token. The problem is prefix failure: once the student commits to a wrong reasoning direction early in a long chain, every subsequent token builds on that mistake, so the teacher's token-level supervision becomes unreliable and a large amount of training compute is spent on misdirected continuations. Relay-OPD's goal is to intervene at exactly the point where the reasoning first goes wrong, correct it locally, then let the student continue on its own.

Key Contributions

1. Identification of a teacher–student continuation asymmetry. On failed reasoning prefixes, the teacher tends to redirect (for example by emitting a reflection token such as "But" or "Wait") while the student tends to continue in the same direction. The authors show through trajectory intervention experiments that correcting prefix failure requires only early and local teacher intervention, with teacher tokens as low as 0.35% of all generated tokens.

2. Relay-OPD, a training method built on that asymmetry. The divergence is converted into a label-free handoff trigger, and a budgeted relay rollout interleaves student legs with short teacher legs. The method uses no external verifier, no process supervision labels, and no answer correctness labels.

3. A single speculative-decoding engine that unifies student and teacher generation, using the student as draft model and the teacher as target model, so no switching between two generation pipelines is needed. The authors state the implementation reproduces the two-model relay process exactly.

4. Empirical validation across eight math benchmarks and two student scales, reporting gains over standard OPD and the strongest trajectory-intervention baseline while substantially shortening both training and inference trajectories.

Main Findings

• Best or second-best on every benchmark. With Qwen3-4B-Instruct-2507 as teacher and Qwen3-1.7B-Non-Thinking as student, Relay-OPD reaches an average accuracy of 46.96 across AIME 2024, AIME 2025, AIME 2026, MATH500, AMC 2023, OlympiadBench, HMMT February 2026 and HMMT November 2025. That is +5.73% over standard OPD and +1.49% over the strongest baseline, FastOPD. The teacher itself averages 61.30.

• Consistent gains at the smaller scale. For Qwen3-0.6B-Non-Thinking, Relay-OPD reaches 31.04 average accuracy, +3.01% over OPD and +0.62% over FastOPD.

• Largest per-benchmark jumps on the hardest sets. For the 1.7B student, improvements over OPD are +7.29% on AIME 2025 and +7.19% on AIME 2026; per-benchmark deltas range from +3.52 (OlympiadBench) to +7.29.

• Local correction alone moves the needle. Replacing only the single reflection token at each trigger, so that teacher tokens are merely 0.35% of all generated tokens, lifts accuracy from 27.73 to 34.96 (+7.23%) in the preliminary experiments on 128 DAPO-Math-17K English samples.

• Intervention value is front-loaded. Holding intervention length fixed but shifting it to later triggers instead of the earliest ones drops accuracy from 41.99 to 33.98 and then 29.49. The teacher–student gap narrows as generation proceeds, so a late takeover cannot redirect the trajectory.

• Longer teacher legs bring diminishing returns. Going from L=3 to L=6, accuracy plateaus around 41–44 while the teacher token ratio rises from 17.52% to 28.52%. Even at L=6 with a teacher token ratio of 28.52%, accuracy reaches only 43.55 versus 60.55 for the teacher generating independently.

• Shorter training trajectories and earlier optimum. For the 1.7B student, Relay-OPD reaches its best checkpoint at step 35 (OPD: step 55; FastOPD: step 45) with an average rollout response length of 2,296 tokens, a 50.7% reduction from OPD's 4,658 and shorter than FastOPD's 2,709. For the 0.6B student, average training trajectory length is 2,490 tokens, a 63.9% reduction from OPD.

• Shorter inference responses than FastOPD. Relay-OPD reduces mean response length by 17.9%, 14.2% and 28.3% on AIME 2025, AIME 2026 and HMMT February 2026 respectively versus FastOPD, while improving accuracy by +2.39%, +4.17% and +1.14% on those same sets.

• Higher pass@k than OPD. The paper reports that Relay-OPD achieves higher pass@k than standard OPD across sampling budgets on HMMT February 2026 and HMMT November 2025.

• Baselines fall short in distinct ways. TRD degrades against standard OPD on both students (30.69 vs. 41.23 at 1.7B; 18.07 vs. 28.03 at 0.6B), with the authors observing visible rewriting artifacts. SKD marginally beats OPD at 1.7B (42.35 vs. 41.23) and drops to 24.38 at 0.6B, reportedly struggling to break repetitive generation patterns. FastOPD concentrates signal at the sequence front via fixed truncation but cannot demonstrate how to recover from failed prefixes.

• The teacher leg matters beyond early stopping. With M=1 for both variants, Trigger-stop (terminate at the first trigger, no teacher tokens) scores 43.48 average, while Relay-OPD with L=3 scores 46.25, a +2.77% improvement.

• The relay-token objective beats alternatives. Using the actually generated relay tokens scores 46.96, versus 44.56 for the Student draft token variant and 44.08 for Teacher FKL (k=128). The authors attribute the Teacher FKL decline to its mode-covering nature forcing the student to match unreliable teacher guidance.

• Teacher intervention shrinks as the student improves. Over the first 60 training steps, the fraction of trajectories exhausting the relay budget falls from roughly 75–85% to roughly 50–60%, and the teacher token ratio starts near 13% then stabilizes at 2–3% after about 20 steps. Policy entropy stays consistently higher than OPD and FastOPD.

• Moderate intervention is best. Varying the relay budget, average accuracy rises from 44.31 to 47.10 as L goes from 0 to 4, then drops to 46.47 at L=5. For M, M=1 and M=2 give 46.25 and 46.96, while M=4 drops to 44.01. For the handoff top-K, K=1, K=5 and K=10 give 44.27, 46.96 and 43.14 respectively, all above the 41.23 of standard OPD (equivalent to K=|V|).

Methodology in Plain English

Standard on-policy distillation has the student generate a full answer, then asks the teacher to score each token the student produced. If the student went wrong early, all the tokens after that point get scored too, even though the reasoning has already derailed.

Relay-OPD changes the rollout itself. While the student generates, the method watches one specific signal: is the teacher's single most likely next token a "reflection" word from a fixed list (Wait, But, Hmm, Actually, Hold, However, Yet, Oh, Alternatively, No, Ah, Oops, Well, plus case and leading-space variants), while the student's top-K candidate list contains no reflection word at all? When that happens, the two models disagree about direction, not just about probabilities. That moment is called a handoff trigger.

When a trigger fires, the teacher briefly takes over, producing a short "teacher leg." The leg starts with the reflection token that triggered the handoff and continues for L paragraphs (paragraphs are delimited by "\n"; a paragraph averages 23.2 tokens in their setting). Control then returns to the student, which continues from the extended prefix. A relay budget (M, L) caps the number of takeovers (M) and leg length (L), concentrating correction at the early positions where failure originates. When the M-th teacher leg ends, the rollout terminates.

The student is then trained on this relay trajectory with a reverse-KL-style single-sample advantage, computed on the tokens that were actually generated, and clipped in the PPO style. In the default configuration K=5 and (M, L)=(2, 3).

Implementation-wise, the whole process runs in one speculative decoding engine: the student is the draft model and the teacher is the target model. In student legs, every draft is accepted (the acceptance ratio equals 1), so generation is ordinary student decoding. In teacher legs, standard speculative rejection sampling runs against the teacher. The only exception is the first position of each teacher leg, where the engine emits the trigger token directly without verification. Because of the correctness property of speculative sampling, the authors state this single engine reproduces the two-model relay process exactly, without switching pipelines, and the teacher logits computed during verification supply the trigger criterion at no extra cost. After each parameter update, student weights are synchronized to the draft model; teacher parameters stay frozen.

Experiments use Qwen3-4B-Instruct-2507 as teacher and Qwen3-0.6B-Non-Thinking and Qwen3-1.7B-Non-Thinking as students, trained on the English subset of DAPO-Math-17K for one epoch with a maximum response length of 16,384 tokens, on 8 H100 GPUs, implemented on verl and vLLM 0.21.0.

Why This Matters

Impact on research. The paper reframes prefix failure not as something to truncate or repair after the fact, but as a state that can be detected online from the reasoning content itself. It also provides evidence that efficient correction is extremely sparse (0.35% teacher tokens in the minimal case) and front-loaded, which is a directly usable design principle for other distillation and imitation pipelines. The label-free trigger sidesteps the cost of verifiers, process reward models, or correctness labels.

Real-world applications (as implied by the work):

• Training small, cheaply servable reasoning models for mathematics from much larger teachers, where the smaller model is the one that actually gets deployed.
• Cutting inference cost in production reasoning systems, since Relay-OPD reduced mean response length by 17.9%, 14.2% and 28.3% versus FastOPD on AIME 2025, AIME 2026 and HMMT February 2026 respectively while improving accuracy.
• Reducing training compute budgets, given a 50.7% reduction in average rollout length for the 1.7B student and 63.9% for the 0.6B student relative to OPD.
• Serving as a general recipe for long-chain reasoning domains where early directional errors compound, beyond the mathematical benchmarks evaluated here.

Industry relevance. The method is built on standard open infrastructure (verl, vLLM) and uses off-the-shelf Qwen3 checkpoints, so it fits existing post-training pipelines. The combination of shorter training trajectories, earlier best checkpoints (step 35 versus 55 for OPD at 1.7B), and shorter inference responses speaks directly to cost.

Future Directions

• Extending beyond mathematical reasoning. All eight evaluation benchmarks are math competitions (AIME 2024/2025/2026, MATH500, AMC 2023, OlympiadBench, HMMT February 2026, HMMT November 2025); the paper does not report results on coding, general instruction following, or other long-chain reasoning domains.

• Replacing the handcrafted reflection-token list. The trigger depends on a fixed set of 13 base words plus case and leading-space variants. Whether the criterion generalizes to domains where redirection is expressed differently, or whether the token set could be learned, is not addressed.

• Scaling the student and teacher sizes. Results are reported only for Qwen3-0.6B-Non-Thinking and Qwen3-1.7B-Non-Thinking students with a Qwen3-4B-Instruct-2507 teacher. Behavior at larger student scales, or with a weaker teacher, is not reported.

• Improving the teacher-leg objective. The Teacher FKL variant scored 44.08 versus 46.96 for the relay-token objective, indicating the teacher leg's supervision remains imperfect. The authors note that even at L=6 with a 28.52% teacher token ratio, accuracy reaches only 43.55 versus 60.55 for the teacher alone, leaving a gap between relay intervention and teacher-level performance.

Target Audience

This paper is most useful to machine learning engineers and researchers working on model distillation, reinforcement-learning-style post-training, and small-model deployment for reasoning tasks. It will also interest practitioners building efficient inference pipelines, since the single speculative-decoding engine is central to the method, and readers tracking how to reduce training and inference cost for long-chain reasoning models. A working familiarity with policy-gradient objectives and speculative decoding helps, but the core idea, letting the teacher step in briefly when the student's reasoning first goes wrong, is accessible without it.

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