🤖 AI

Trajectory-Refined Distillation

This paper introduces Trajectory-Refined Distillation (TRD), a novel on-policy distillation method that mitigates the structural "prefix failure" problem by applying trajectory-level corrections to student rollouts under teacher guidance, thereby improving reasoning accuracy and coverage across various benchmarks and model scales.

Original authors: Li Jiang, Haoran Xu, Yichuan Ding, Amy Zhang

Published 2026-06-09

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Li Jiang, Haoran Xu, Yichuan Ding, Amy Zhang

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 you are teaching a student (the AI) how to solve a complex math problem. You have a brilliant teacher (a more advanced AI) who knows the right answer.

In the standard method described in the paper, called On-Policy Distillation (OPD), the process works like this:

The student tries to solve the problem on their own.
As the student writes down their steps, the teacher watches and says, "For this specific word you just wrote, here is the probability of what comes next."
The student tries to adjust their brain to match the teacher's probabilities for every single word.

The Problem: "Prefix Failure" (The Wrong Turn)

The paper identifies a major flaw in this standard method called Prefix Failure.

Imagine the student is walking down a path to solve the problem. At step 3, they take a wrong turn. They are now on a dead-end street.

The Teacher's Dilemma: The teacher sees the student is lost. The teacher knows the correct path starts with "Wait, actually..." and then goes a completely different direction.
The Confusion: But the teacher is also forced to look at the student's current location (the dead end). So, the teacher's advice becomes a confusing mix: "Stay on this dead end because you're already here" AND "Turn around immediately."
The Result: The student gets a jumbled signal. They try to learn from the teacher, but because the student is stuck on the wrong path, the teacher's advice gets fragmented. The student learns to be "good at being lost" rather than learning how to get back on track. The paper calls this a "bimodal mixture"—the teacher is trying to teach two opposite things at once, and the student gets confused.

Existing solutions tried to fix this by just tweaking the "grade" (the loss function) for specific words, but they didn't fix the fact that the student was still stuck on the wrong path.

The Solution: Trajectory-Refined Distillation (TRD)

The authors propose a new method called Trajectory-Refined Distillation (TRD). Instead of just grading the student's messy draft word-by-word, they change the workflow:

The Draft: The student still takes a first shot at the problem (the "raw rollout").
The Rewrite: Before grading, the teacher takes that messy draft and rewrites it. The teacher fixes the wrong turns, corrects the logic, and produces a clean, perfect version of the solution based on the student's starting point.
The Lesson: The student then learns from this clean, corrected version, not the messy original.

The Analogy:

Old Way: You write a story with a plot hole in the middle. Your editor reads it and says, "For the word 'cat' in paragraph 3, you should have written 'dog'." But you are still stuck in paragraph 3 trying to figure out why you wrote 'cat' in the first place. It's confusing.
TRD Way: You write the story with the plot hole. Your editor takes your draft, fixes the plot hole, and rewrites the whole middle section so the story flows perfectly. Then, you study the editor's rewritten version to learn how to write better stories next time.

Why This Works Better

Fixes the Root Cause: By correcting the "wrong turn" before the student tries to learn from it, the teacher's advice is no longer confused. It's a clear, single path.
Explores New Paths: Even if the student gets the answer right the first time, the teacher might show them a different, shorter, or smarter way to solve it. This teaches the student to be more flexible.
Efficiency: The paper found that the "rewritten" paths were often much shorter and clearer than the student's original rambling attempts, which actually made training faster.

The Results

The researchers tested this on difficult math competitions (like AIME and HMMT) and coding challenges. They found that:

The new method (TRD) consistently beat the old methods.
It was especially good at helping the AI solve the hardest problems where it usually got stuck.
It worked whether the teacher was a separate, bigger AI or the same AI acting as its own "smart" teacher (using a technique called "privileged information").

In short, the paper argues that to teach an AI effectively, you shouldn't just grade its mistakes word-by-word while it's still making them. Instead, you should let the teacher fix the whole path first, and then let the student learn from the corrected journey.

Technical Summary: Trajectory-Refined Distillation (TRD)

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

The paper identifies a fundamental structural limitation in On-Policy Distillation (OPD) and On-Policy Self-Distillation (OPSD) for Large Language Models (LLMs), termed prefix failure.

In standard OPD/OPSD, a student model generates a rollout (trajectory) $y_o$ , and a teacher model provides dense, per-token supervision based on that specific rollout. The authors argue that when the student's rollout takes an incorrect reasoning path (a "failed prefix"), the teacher's supervision signal becomes structurally compromised:

Bimodal Teacher Mixture: When a prefix is incorrect, the teacher distribution becomes a mixture of two modes: one that continues the failed prefix (for sequence consistency) and another that pivots toward the correct solution.
Fragmented Gradients: Standard per-token loss functions (KL divergence) evaluate supervision along the frozen failed trajectory. Consequently, the gradient updates are fragmented. The teacher repeatedly recommends a correction token (e.g., "Wait" or "Actually") conditioned on the wrong context (the failed prefix), rather than conditioning on the unfolding correction path.
Ineffectiveness of Token-Level Fixes: Existing remedies (e.g., loss clipping, top-K truncation, reweighting) operate at the token-loss level. They attempt to select "informative" tokens or suppress "noisy" ones but fail to address the root cause: the failed prefix itself remains structurally intact. The supervision pairs delivered to the student diverge from the ideal correction path, preventing the model from learning to recover from errors effectively.

2. Methodology: Trajectory-Refined Distillation (TRD)

To mitigate prefix failure, the authors propose Trajectory-Refined Distillation (TRD), a method that shifts the intervention from the token level to the trajectory level.

Core Mechanism

Instead of using the raw student rollout $y_o$ directly for distillation, TRD introduces a refinement step before the loss calculation:

Sampling: The student model $\pi_\theta$ generates a raw on-policy rollout $y_o \sim \pi_\theta(\cdot | x)$ .
Refinement: The teacher model $\pi_T$ $π_{T}$ (or the student itself in OPSD, conditioned on privileged information like the ground truth $y^*$ $y^{*}$ ) generates a refined trajectory $y_r$ $y_{r}$ .
- In OPD: $y_r \sim \pi_T(\cdot | x, y_o)$ .
- In OPSD: $y_r \sim \text{sg}[\pi_\theta(\cdot | x, y_o, y^*)]$ .
- Crucially, the teacher is prompted to revise the student's rollout guided by the reference solution, effectively "rewriting" the erroneous portions of the prefix while maintaining the student's reasoning patterns where they are valid.
Distillation: The refined trajectory $y_r$ is used as the supervision target for the standard distribution-matching loss (Forward KL with full vocabulary).

Theoretical Justification

TRD addresses the "fragmented gradient" problem by ensuring the supervision context grows along the correction path rather than the failed prefix.

Ideal Gradient Recovery: By supervising along $y_r$ , the context for each token includes the corrected prefix, allowing the teacher to provide a coherent, multi-step correction signal ( $g_{ideal}$ ) rather than a repetitive, context-mismatched signal ( $g_{frag}$ ).
Exploration: Even when $y_o$ is correct, TRD exposes the student to alternative valid derivations suggested by the teacher, broadening the distribution of correct reasoning paths the student learns.

Algorithm

The process is summarized in Algorithm 1:

Sample $y_o \sim \pi_\theta(\cdot | x)$ .
Generate $y_r \sim \pi_T(\cdot | x, y_o)$ .
Update $\theta$ using the KL divergence between $\pi_\theta$ and the teacher's distribution over $y_r$ .

3. Key Contributions

Identification of Prefix Failure: The paper formalizes "prefix failure" as a structural cause of degradation in on-policy distillation, explaining why token-level interventions fail to recover from reasoning errors.
Trajectory-Level Correction: TRD is proposed as the first trajectory-level optimization design that explicitly targets the recovery of failed prefixes while respecting the on-policy support constraint, without prohibitive computational overhead.
Dual-Regime Improvement: TRD improves performance in two regimes:
- Error Correction: It mitigates prefix failure by correcting problematic prefixes before distillation.
- Exploration: It broadens reasoning coverage by exposing students to alternative valid derivations even when the original rollout was correct.
Applicability to Self-Distillation: The method is naturally extended to OPSD, where the student acts as its own teacher conditioned on privileged information, removing the need for a separate teacher model.

4. Experimental Results

The authors evaluated TRD across five competition-math benchmarks (AIME24/25, HMMT25, BeyondAIME, AMOBench) and code generation benchmarks (HumanEval+, MBPP+, LiveCodeBench) using Qwen3 models (1.7B, 4B, 8B).

Key Findings

Superior Performance: TRD consistently outperforms prior baselines (Forward/Reverse KL, with/without clipping/truncation) across both OPD and OPSD settings.
Significant Gains on Hard Tasks: The improvements are most pronounced on the hardest benchmarks.
- AMOBench: TRD achieved a ~50% relative improvement in Pass@16 for Qwen3-8B under the OPSD setup compared to the base model.
- AIME24/25: TRD showed consistent gains, particularly for smaller models (e.g., +4.6% on AIME24 for Qwen3-1.7B).
Stability: Unlike some baselines that degrade on certain benchmarks (e.g., Reverse KL dropping performance on AIME24), TRD never dropped below the base model performance.
Efficiency: In the OPSD setting, the refined trajectories $y_r$ were significantly shorter (median length compressed by ~9x) than raw rollouts $y_o$ , reducing training wall-clock time by ~60% despite the extra sampling step.
Code Generation: TRD matched or exceeded baselines on HumanEval+ and MBPP+, though gains on LiveCodeBench were limited, suggesting the teacher's refinement capability may be bounded by task difficulty.

5. Significance and Claims

The paper claims that TRD represents a necessary shift in post-training strategy for LLMs. The authors argue that:

Token-level interventions are insufficient: The structural nature of prefix failure requires trajectory-level corrections.
On-policy support is preserved: TRD achieves these corrections without violating the on-policy constraint (unlike using a fixed expert policy), making it compatible with standard compute-efficient learning paradigms.
Self-Improvement is viable: The success of TRD in the OPSD setting demonstrates that models can effectively refine their own reasoning paths using privileged information (ground truth) without requiring a separate, larger teacher model, offering a lightweight path to self-improvement.

The work suggests that future distillation methods must account for the structural integrity of the reasoning trajectory, not just the statistical alignment of individual tokens.