Pass the Baton: Trajectory-Relayed On-Policy Distillation (Relay-OPD)

Paper: Pass the Baton: Trajectory-Relayed On-Policy Distillation — arXiv: 2607.26057 Authors: Haolei Xu, Xiaowen Xu, Haiwen Hong, et al. (Zhejiang University / Alibaba Yuvion) Code: https://github.com/zju-real/Relay-OPD Project page: https://zju-real.github.io/Relay-OPD

The Relay Race Metaphor

Training a student model with standard on-policy distillation is like coaching a young relay runner by only reviewing the tape after they finish: once they run off course, the whole leg accumulates error, and late feedback can't fix it. The smarter approach — the one Relay-OPD takes — is to take the baton at the moment the runner strays, run a few corrective steps, and hand it back.

The Problem: Prefix Failure

On-Policy Distillation (OPD) lets the student generate reasoning trajectories while the teacher provides token-level probability supervision, keeping training and inference distributions aligned. But OPD suffers from prefix failure: once the student goes wrong early, everything afterward is a long misguided continuation. The teacher's token-level feedback on such trajectories is unreliable, and compute is wasted.

Prior fixes fall short:

• Fixed-length truncation (ESR, FastOPD): cuts at a fixed position regardless of where failure actually occurs.
• Offline rewriting (TRD): teacher rewrites finished trajectories, but the rewritten text carries visible "editing traces."
• Token-level mixing (SKD): switches between teacher and student tokens by distribution difference — too coarse a signal, since distribution difference ≠ direction error.

What's missing is an online mechanism for detecting "the student went off-track."

Key Insight: Teachers Turn, Students Persist

On failing prefixes, the teacher tends to generate reflection tokens ("But", "Wait", "However") while the student tends to continue ("So", "Therefore"). In one example (Figure 1a), the teacher gives "But" 74.4% probability at a handoff point while the student gives "So" 50.6% and "But" nearly zero.

This asymmetry is detectable without external labels or verifiers. The handoff criterion:

\[\phi(h) = \mathbb{1}[a^T(h) \in \mathcal{R}] \cdot \mathbb{1}[\mathcal{K}_S(h) \cap \mathcal{R} = \emptyset]\]

where \(a^T(h)\) is the teacher's top-1 token, \(\mathcal{K}_S(h)\) is the student's top-K support set, and \(\mathcal{R}\) is a predefined set of reflection words. Crucially, this detects teacher-student directional divergence, not student error per se — ensuring interventions are precise.

Relayed Trajectories

When a handoff triggers:

1. The student runs the first leg up to the handoff point. 2. The teacher takes over briefly (reflection token + corrective reasoning). 3. Control returns to the student until trajectory end. 4. Standard OPD training runs on the mixed trajectory.

A relay budget (2–3 teacher interventions per trajectory) concentrates interventions at critical early junctures and keeps trajectories close to the student's own distribution. The whole rollout runs inside a speculative decoding engine, so teacher handoffs only require parameter switches — cost is near single-model inference.

Results

Setup: Qwen3-4B-Instruct-2507 teacher; Qwen3-0.6B / 1.7B Non-Thinking students; eight math benchmarks (AIME24/25/26, MATH, AMC23, Olympiad Bench, HMMT Feb26, HMMT Nov25).

| Method | Avg Accuracy | Trajectory Length | |--------|-------------|-------------------| | Student baseline | 24.84 | — | | SFT | 33.20 | 4262 | | KD | 33.75 | 4262 | | GRPO | 34.42 | 2558 | | OPD | 41.23 | 4658 | | TRD | 30.69 | 2785 | | FastOPD | 45.47 | 2709 | | SKD | 42.35 | 4753 | | Relay-OPD | 46.96 | 2296 |

• +5.73% over OPD; +1.49% over FastOPD; +7.29% and +7.19% on AIME25/AIME26.
• Training trajectory length reduced 50.7% (4658 → 2296).
• vs FastOPD: responses 17.9%/14.2%/28.3% shorter on AIME25/AIME26/HMMT Feb26 while being more accurate.
• 0.6B student: +3.01% over OPD.

TRD underperforms OPD (rewriting artifacts), and SKD gains little or regresses — it cannot break established repetitive generation patterns. Relay-OPD intervenes at the moment of directional divergence rather than patching afterward.

Limitations

1. The reflection-word set ℛ is hand-defined — effective for math reasoning, questionable for code generation or open dialogue. 2. Only validated on math benchmarks. 3. Relay budget is a fixed hyperparameter (2–3), with no adaptive adjustment by task difficulty. 4. The asymmetry assumption may not generalize across model families — some teachers may also persist rather than reflect.

Takeaway

The deeper insight: the difference between strong and weak models isn't just what they know, but when they turn back. Weak models' core failure is not lacking knowledge — it's not knowing when to reverse direction. Supervision, accordingly, should concentrate at directional turning points, not be given only at the end or sprinkled uniformly across every token.

Links:

• Paper: https://arxiv.org/abs/2607.26057
• HTML: https://arxiv.org/html/2607.26057v1
• Code: https://github.com/zju-real/Relay-OPD
• Project page: https://zju-real.github.io/Relay-OPD

FAQ

Q: Who is this for? Researchers, engineers, and students interested in LLM distillation and reasoning training.

Q: What's the core contribution? An online, verifier-free handoff mechanism that detects teacher-student directional divergence and lets the teacher briefly correct the trajectory, fixing prefix failure in on-policy distillation.

Q: Is code available? Yes: https://github.com/zju-real/Relay-OPD