💬 NLP

Pass the Baton: Trajectory-Relayed On-Policy Distillation

The paper introduces Relay-OPD, a novel on-policy distillation method that mitigates prefix failure by dynamically triggering teacher interventions to correct student trajectory deviations, thereby achieving state-of-the-art performance on mathematical reasoning benchmarks while significantly reducing training trajectory length.

Original authors: Haolei Xu, Xiaowen Xu, Haiwen Hong, Zixuan Ni, Hongxing Li, Yiwen Qiu, Weiming Lu, Yongliang Shen

Published 2026-07-29

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Haolei Xu, Xiaowen Xu, Haiwen Hong, Zixuan Ni, Hongxing Li, Yiwen Qiu, Weiming Lu, Yongliang Shen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a world where giant, super-smart computers (called Large Language Models) are learning to solve complex puzzles, like math problems or logic riddles. To teach them, we often use a "teacher-student" setup. The teacher is a massive, powerful model that knows the answers, and the student is a smaller, faster model trying to learn from it. Usually, the teacher just writes out the perfect solution, and the student copies it. But there's a smarter way called On-Policy Distillation. Instead of just copying a finished essay, the student tries to write the solution itself. If it gets stuck or goes down the wrong path, the teacher steps in to guide the next word, right there in the moment. This is great because it teaches the student how to think, not just what to memorize.

However, this method has a tricky flaw. If the student makes a mistake early on—like taking a wrong turn in a maze—the teacher tries to guide the rest of the journey. But because the student is already lost, the teacher's guidance often gets confused or unreliable, and the student keeps wandering further off course. It's like trying to give someone directions while they are already driving in the wrong direction; the more they drive, the more confused they get. This paper tackles that specific problem: how do we stop the student from getting hopelessly lost before the teacher's help becomes useless?

The Relay Race of Reasoning

The authors of this paper, working with models from the Qwen family, realized that when a student model starts to go off the rails, the teacher and the student actually start to disagree on what to do next. The student, stubbornly committed to its wrong idea, wants to keep going straight. The teacher, seeing the error, wants to stop, rethink, and change direction.

They call this moment of disagreement a "handoff trigger." Think of it like a relay race. In a standard race, if the runner (the student) trips and starts running the wrong way, the coach (the teacher) might just shout instructions from the sidelines, but the runner keeps stumbling. In this new method, called Relay-OPD, the moment the coach sees the runner heading for a cliff, the coach jumps in, grabs the baton, and runs a few steps in the correct direction to reset the runner's path. Then, the coach hands the baton back, and the student continues the race, now on the right track.

How It Works: The "But" and "Wait" Tokens

The system is clever because it doesn't need a human to tell it when a mistake happens. It watches the probabilities. If the teacher is highly likely to say a "reflection" word like "But," "Wait," or "However" (which signal a change in thought), but the student is likely to say "So" or "Now" (which signal continuing the current thought), the system knows: We are off track!

At that exact moment, the teacher takes over for a brief "leg" of the journey. The teacher writes a short paragraph to correct the reasoning, then hands control back to the student. The student then finishes the rest of the problem, learning from the corrected path.

The researchers found that this intervention doesn't need to be long. In fact, the teacher only needs to write about 0.35% of the total words to make a huge difference. It's like a coach giving a single, sharp nudge to fix a runner's form, rather than running the whole race for them.

The Results: Smarter and Faster

The team tested this on eight different math benchmarks, using student models of different sizes (0.6 billion and 1.7 billion parameters) and a 4 billion parameter teacher. The results were impressive:

Better Scores: The Relay-OPD students scored significantly higher than those using standard methods. For the 1.7 billion parameter model, it improved the average score by 5.73% over the standard method and beat the next best competitor by 1.49%.
Less Wasted Time: Because the system stops the student from wandering down dead ends, the training process became much more efficient. The average length of the training paths was cut by more than 50%. The students learned the same amount in half the time.
The "Leg" Matters: The researchers proved that the teacher's short correction (the "teacher leg") was crucial. Simply stopping the student when they made a mistake (without the teacher fixing it first) didn't work as well. The teacher had to show the student how to get back on track.

What It's Not

The paper is careful to show what this method is not. It isn't just about cutting off long answers early (which other methods tried), nor is it about rewriting the student's whole essay after it's finished. It's a real-time, "in-the-moment" fix. It also isn't about the teacher taking over the whole race; the student still does most of the running. The teacher only steps in for those critical, short bursts when the student is about to make a permanent mistake.

In short, Relay-OPD is like a safety net that catches a falling student just before they hit the ground, gently redirects them, and lets them keep running. It turns a potential disaster into a learning moment, making AI models smarter, faster, and less likely to get lost in their own thoughts.

Technical Summary: Relay-OPD (Pass the Baton)

1. Problem Statement: Prefix Failure in On-Policy Distillation

On-policy distillation (OPD) improves knowledge transfer from strong teacher models to weaker student models by grounding token-level supervision in the student's own generated trajectories. However, this approach suffers from a critical flaw known as prefix failure. In long-chain reasoning tasks, if a student model commits to an incorrect reasoning direction early in generation, all subsequent tokens build upon this deviation. This results in long, misdirected continuations that:

Elicit unreliable and potentially harmful supervision from the teacher.
Waste substantial training compute on trajectories that cannot be corrected.

Existing remedies have structural limitations:

Fixed-length truncation (e.g., FastOPD, ESR): Cuts rollouts at rigid positions regardless of where the actual reasoning failure occurs.
Offline rewriting (e.g., TRD): Repairs trajectories only after completion, often leaving visible artifacts.
Token-level mixing (e.g., SKD): Switches based on generic distributional disagreement rather than explicit signals of reasoning failure.

There is a lack of a mechanism that acts online, correcting failures as they emerge, and intervenes based specifically on the reasoning state itself.

2. Methodology: Relay On-Policy Distillation (Relay-OPD)

The authors propose Relay-OPD, a method that interleaves student generation with brief, targeted teacher takeovers at detected failure points.

2.1 Core Insight: Teacher-Student Continuation Asymmetry

The method relies on an observed asymmetry on failed prefixes:

The Student: Tends to continue along the current (wrong) direction.
The Teacher: Tends to stop, re-examine, and redirect the reasoning (often using "reflection" tokens like But, Wait, However).

This divergence serves as a label-free handoff trigger. When the teacher's most probable next token is a reflection token, but the student's top- $K$ support set contains no such tokens, a handoff is triggered.

2.2 Relay Trajectory Construction

When a trigger fires, the teacher briefly takes over to produce a teacher leg, after which the student resumes.

Handoff Trigger ( $\phi(h)$ ): Defined as $1[a_T(h) \in \mathcal{R}] \cdot 1[K_S(h) \cap \mathcal{R} = \emptyset]$ , where $\mathcal{R}$ is the set of reflection tokens, $a_T$ is the teacher's top token, and $K_S$ is the student's top- $K$ set.
Relay Budget ( $M, L$ ): To prevent the trajectory from drifting too far from the student's policy, a budget caps the intervention:
- $M$ : Maximum number of teacher takeovers.
- $L$ : Number of additional paragraphs generated by the teacher after the reflection token.
Process: The student generates autoregressively. Upon detecting a trigger (and if budget remains), the teacher generates $L$ paragraphs (starting with the reflection token). The student then resumes generation from the corrected prefix.

2.3 Optimization Objective

Relay-OPD optimizes the student policy on the resulting relay trajectory $z$ . Unlike standard OPD which uses the student's original tokens, Relay-OPD uses the actually generated tokens $z_t$ from the relay trajectory (including teacher-generated tokens).
The objective uses a reverse-KL-style single-sample advantage:
$A_t^{Relay} = \log \pi_T(z_t | h_t^z) - \log \bar{\pi}_\theta(z_t | h_t^z)$
This is optimized via a clipped PPO-style loss, allowing the student to selectively absorb the teacher's corrective signals rather than fully fitting the teacher's distribution.

2.4 Efficient Implementation

To avoid the overhead of switching between two generation pipelines, Relay-OPD unifies the process within a single speculative decoding engine:

Student acts as the draft model.
Teacher acts as the target model.
The engine dynamically switches the target policy ( $\pi_{tgt}$ ) between the student and teacher based on the current state (Student leg vs. Teacher leg).
Teacher legs perform standard speculative decoding against the teacher, while student legs accept all drafts. This ensures the distributional identity of the relay trajectory is preserved exactly.

3. Key Contributions

Identification of Asymmetry: The paper identifies that teacher-student divergence on failed prefixes is a reliable, label-free signal for intervention. Experiments show that correcting prefix failure requires only early, local intervention, with teacher tokens comprising as little as 0.35% of total generated tokens.
Relay-OPD Framework: A novel method that converts this asymmetry into a budgeted relay rollout, unifying student and teacher generation in a single speculative decoding engine.
Empirical Superiority: Demonstrated across eight mathematical reasoning benchmarks and two student scales (0.6B and 1.7B).

4. Experimental Results

The authors evaluated Relay-OPD using Qwen3-4B-Instruct-2507 as the teacher and Qwen3-0.6B/1.7B-Non-Thinking as students on eight benchmarks (AIME 2024/2025/2026, MATH500, AMC 2023, OlympiadBench, HMMT Feb/Nov).

Performance:
- 1.7B Student: Achieved the best or second-best results on all benchmarks.
  - Average accuracy: 46.96% (vs. 41.23% for standard OPD, +5.73% improvement).
  - Outperformed the strongest baseline, FastOPD, by +1.49%.
- 0.6B Student: Consistent gains, improving average accuracy by +3.01% over standard OPD.
Efficiency:
- Training Length: Reduced average training trajectory length by >50% (from 4,658 to 2,296 tokens for 1.7B).
- Convergence: Reached the best checkpoint in fewer steps (Step 35 vs. Step 55 for OPD).
- Inference: Reduced mean response length by 14–28% compared to FastOPD while improving accuracy.
Ablation Studies:
- Teacher Leg Necessity: Simply stopping at the trigger (Trigger-stop) yielded 43.48% accuracy; adding the teacher leg ( $L=3$ ) raised it to 46.25%, proving the value of corrected context.
- Objective: The reverse-KL style objective on generated tokens outperformed both "Student Draft Token" objectives and "Teacher FKL" objectives, confirming the benefit of selective absorption of corrective signals.

5. Significance and Claims

The paper claims that prefix failure is a fundamental bottleneck in on-policy distillation that existing truncation or rewriting methods fail to address dynamically. Relay-OPD addresses this by:

Acting Online: Correcting failures as they emerge rather than post-hoc.
Minimizing Intervention: Using a "relay" approach where the teacher only intervenes briefly to redirect, keeping the trajectory close to the student's policy distribution.
Efficiency: Significantly reducing the compute cost of training by shortening trajectories and avoiding wasted computation on unrecoverable errors.

The authors conclude that Relay-OPD offers a robust, label-free mechanism for strong-to-weak distillation in reasoning tasks, achieving state-of-the-art results while drastically reducing training costs.