Trajectory-Refined Distillation

TRD: Beyond Token-Level Tweaks—Fixing the "Prefix Failure" in LLM Distillation

Trajectory-Refined Distillation

2026-06-01

Li Jiang, Haoran Xu, Yichuan Ding, Amy Zhang

Summary

Problem

Method

Results

Takeaways

Abstract

The paper introduces Trajectory-Refined Distillation (TRD), a post-training method for Large Language Models that moves beyond token-level supervision to trajectory-level correction. TRD leverages teacher models (or self-distillation with privileged info) to rewrite student-generated reasoning paths, achieving significant SOTA improvements such as a 50% relative gain on the AMOBench competition-math benchmark for Qwen3-8B.

TL;DR

On-policy distillation (OPD) is the current gold standard for LLM post-training, but it has a structural flaw: Prefix Failure. When a model starts a reasoning path incorrectly, standard per-token supervision becomes chaotic. Trajectory-Refined Distillation (TRD) solves this by having the teacher "rewrite" the student's bad drafts before training. The results? A massive 50% relative performance jump on elite math benchmarks and 60% faster training on 8B models.

The Problem: The Curse of the Broken Prefix

In standard OPD, we let a student model (the one we are training) generate a response ( $y_{o}$ ). We then ask a smarter teacher model to look at every single token the student wrote and say "what would I have said here?".

The authors identify a fatal flaw they call Prefix Failure:

Bimodal Teachers: If the student starts a math problem with a wrong assumption, the teacher is torn between continuing the student's flawed logic (to be consistent) or jumping to the right answer (which looks like a nonsensical "teleportation" to the student).
Fragmented Gradients: Because the training happens on the student's frozen, broken path, the teacher's advice never "unfolds" into a coherent correction. It's like a coach giving advice at every step of a race, but the runner is stuck in a muddy swamp—no amount of "step better" advice helps if the runner is in the wrong place entirely.

Prefix Failure Visualization Figure: Under prefix failure, the teacher's distribution splits into two modes, making the learning signal noisy and ineffective.

The Solution: Trajectory-Refined Distillation (TRD)

Instead of trying to fix the training loss (as others have done with clipping or Top-K filtering), TRD fixes the data.

The Workflow:

Drafting ( $y_{o}$ ): The student generates a raw rollout.
Refining ( $y_{r}$ ): The teacher (or the student itself using privileged info like the ground-truth answer) rewrites the rollout. Crucially, the teacher is shown the student's draft so the new version stays "close" to the student's style but corrects the logic.
Distilling: The student is trained to match the teacher's distribution along the refined path ( $y_{r}$ ), not the broken draft ( $y_{o}$ ).

This recovers what the authors call the "Ideal Gradient". The context for the student now grows along a successful reasoning path, allowing the model to learn the multi-step "unfolding" of a correction.

Algorithm Overview Figure 1: TRD refines student trajectories $y_{o}$ into $y_{r}$ , leading to superior performance across all scales.

Experimental Results: High-Stakes Math & Code

TRD was tested on intense competition-level benchmarks like AIME and AMOBench.

AMOBench (The Ultimate Test): For the Qwen3-8B model, TRD improved the "Pass@16" score from a baseline of 41% to a staggering 61.5%.
Efficiency Gains: One of the most surprising findings is that $y_{r}$ (the refined paths) are significantly shorter than $y_{o}$ (the student's ramblings). On the 8B model, this compressed training time by 60%, effectively paying for the extra cost of the teacher's "refinement" step.

Training Trajectory Analysis Figure: Analysis showing TRD finds shorter, more efficient solution paths compared to baseline models.

Deep Insights: Exploitation vs. Exploration

TRD doesn't just fix errors; it improves "Exploration" even when the student is already correct. By asking the teacher to refine a correct student rollout, the teacher often provides an alternative, more elegant derivation. This expands the student's repertoire of correct reasoning patterns, making it more robust during test time.

Conclusion and Future Outlook

The success of TRD suggests that the LLM community has been focusing too much on the probability math of distillation and not enough on the structural integrity of the reasoning paths. By shifting the focus to trajectory-level correction, TRD provides a blueprint for making smaller models think like "reasoning" models (like o1 or DeepSeek-R1).

Limitations: The method does require a stronger teacher or access to reference solutions during training. However, given the massive gains in both accuracy and training efficiency, this "extra step" is likely to become a standard part of the LLM post-training recipe.

Find Similar Papers

Try Our Examples

Examine recent literature on "prefix failure" or "exposure bias" in on-policy distillation of Large Language Models to find alternative trajectory-level correction methods.
Trace the theoretical origins of "On-policy Self-distillation" (OPSD) and identify how TRD specifically modifies the sequence-level KL divergence optimization compared to the original SDR (Self-Distilled Reasoner) paper.
Explore the application of trajectory refinement techniques for multi-step reasoning in domains beyond Math and Code, such as complex planning in Robotics or strategic decision-making in Reinforcement Learning.

Contents

TRD: Beyond Token-Level Tweaks—Fixing the "Prefix Failure" in LLM Distillation

1. TL;DR

2. The Problem: The Curse of the Broken Prefix

3. The Solution: Trajectory-Refined Distillation (TRD)

3.1. The Workflow:

4. Experimental Results: High-Stakes Math & Code

5. Deep Insights: Exploitation vs. Exploration

6. Conclusion and Future Outlook