PaperlayerView exact revision on arXiv

Source-linked AI summary

Trajectory-Refined Distillation

Li Jiang, Haoran Xu, Yichuan Ding, Amy Zhang

arXiv:2606.08432v1cs.AI

TL;DR

OPD can fail when dense token-level supervision follows erroneous student prefixes, and token-loss fixes leave this structural problem unresolved. TRD instead refines student trajectories under teacher guidance before distillation, consistently improving accuracy and reasoning coverage across benchmarks and model scales.

  • Problem

    On-policy distillation can provide unreliable, fragmented supervision after erroneous student prefixes, while token-level interventions leave the sampled trajectory unchanged.

  • Method

    TRD samples an on-policy rollout, has the teacher refine it under guidance while preserving on-policy support, and distills from the refined trajectory.

  • Results

    Across five competition-math benchmarks and multiple Qwen3 scales, TRD consistently outperforms prior baselines, improving single-attempt accuracy and broadening reasoning coverage.

  • Takeaways & Limitations

    TRD provides a trajectory-level alternative to token-loss corrections for mitigating prefix failure and expanding valid reasoning paths under teacher guidance.

  • Takeaways & Limitations

    TRD requires an extra sampling budget and depends on the teacher’s ability to refine trajectories while keeping them close to the student’s on-policy distribution.

Abstract

from arXiv · show

On-policy distillation (OPD) has become a central post-training tool for large language models (LLMs), providing dense per-token teacher supervision along the student's own rollouts. In this work, we identify a common structural cause underlying OPD, which we call prefix failure. Under prefix failure, dense per-token supervision induces a bimodal teacher mixture and fragmented gradients that token-level loss truncation or reweighting fail to address. This observation motivates us to move beyond token-level loss interventions toward trajectory-level output corrections. We thus propose Trajectory-Refined Distillation (TRD), a trajectory-level correction method that revises the student's rollout under the teacher guidance while within on-policy support. By correcting problematic prefixes before distillation, TRD mitigates prefix failure at its source. Moreover, TRD improves the exploration by exposing the student to alternative valid derivations under teacher guidance, even when the original rolls are already correct. TRD can also be applied to on-policy self-distillation (OPSD), a parameter-sharing variant that uses the student model conditioned on privileged informations as the teacher. Across a wide range of benchmarks and base models at multiple scales, TRD consistently outperforms prior baselines, improving single-attempt accuracy and broadening reasoning coverage. Code is available at https://github.com/louieworth/trd

1 Introduction

OPD has become a prominent LLM post-training method, while OPSD provides a parameter-sharing alternative. However, vanilla OPD/OPSD can produce noisy or uninformative supervision, motivating TRD’s trajectory-level refinement of on-policy rollouts under teacher guidance.

  • Motivation: OPD computes per-token teacher supervision along student rollouts and is increasingly used in modern LLM post-training.Recent industry releases incorporate OPD alongside supervised fine-tuning and reinforcement learning with verifiable rewards.
  • Motivation: OPSD is a parameter-sharing variant of OPD in which the same model serves as teacher and student under different contexts.The student is conditioned only on the problem statement.
  • Problem: Vanilla OPD/OPSD recipes can yield noisy or uninformative supervision, while existing remedies reweight or clip token losses without changing sampled trajectories.These limitations motivate addressing failures at the trajectory level rather than only through token-loss interventions.
  • Trajectory-Refined Distillation: TRD refines raw on-policy rollouts under teacher guidance while retaining on-policy support, then uses the refined trajectories for subsequent distillation.For each problem-solution pair, the method samples a raw rollout and prompts the teacher to produce a refined version guided by reference solutions.

2 Related Work

Related work frames on-policy distillation as dense teacher supervision along student rollouts, with on-policy self-distillation removing the need for a separate teacher. Prior methods address reported failure modes mainly through token-level interventions, while this paper links empirical OPD failures partly to prefix failure.

  • On-Policy Distillation: On-policy distillation replaces fixed-corpus targets with per-token teacher supervision computed along the student’s own rollouts.OPD combines on-policy sampling with dense token-level learning through token-level KL loss.
  • On-policy Self-distillation: On-policy self-distillation instantiates teacher and student from the same model under different privileged contexts, enabling self-improvement without external supervision.This formulation removes the need for a separate teacher.
  • Common Failure Mode and Fix: Vanilla distillation methods often underperform in practice and exhibit failure modes including mode collapse, trajectory inflation, and misleading or vanishing supervision signals.Recent studies report these and other failure modes across on-policy distillation methods.
  • Common Failure Mode and Fix: Most reported failure modes are addressed through token-level interventions that select informative learning signals to stabilize training, with choices dictated by the divergence.The paper notes that empirical OPD failures may partly be attributable to prefix failure.

3 Preliminaries

The preliminaries define on-policy distillation as matching teacher outputs on trajectories sampled from the current student, and on-policy self-distillation as using shared parameters under different contexts. They also motivate reverse KL as a mode-seeking divergence for generative language-model distillation.

  • On-policy Distillation: OPD trains a student to match a teacher’s output distribution using trajectories sampled from the current student rather than fixed expert prefixes.The student samples an autoregressive rollout for a prompt from the training dataset.
  • On-policy Self-Distillation: OPSD uses teacher and student policies from the same model under different contexts, with the teacher receiving privileged information such as a reference answer or reasoning trace.The teacher and student share parameters θ.
  • Divergence Choice: Reverse KL is mode-seeking and generally preferred for generative language-model distillation, whereas forward KL is mode-covering.Reverse KL discourages the student from assigning probability to low-probability regions of the teacher.

4 Prefix Failure in Token-Level On-Policy Distillation

Prefix failure makes dense token-level on-policy distillation unreliable: the teacher can become bimodal or collapse onto the student’s wrong continuation, while per-token updates fragment the intended correction path. Empirical measurements show vanishing supervision on failed rollouts, motivating trajectory-level correction.

  • Failure mechanism: Prefix failure occurs when a student prefix contains reasoning errors that cannot reach the correct solution without retraction, making teacher supervision unreliable.On such prefixes, the teacher mixes a sequence-consistent wrong continuation with a pivot toward a correct continuation; degenerate prefixes can instead make guidance vanish.
  • Failure mechanism: Under prefix failure, forward KL emphasizes the correction-onset region and can force the student toward an OOD mode, whereas reverse KL emphasizes the wrong continuation and weakens recovery updates.The asymmetry follows from teacher-weighting in forward KL versus student-weighting in reverse KL, with the student assigning low probability to recovery tokens.
  • Gradient fragmentation: Dense per-token KL fragments correction supervision because later teacher contexts remain anchored to the failed trajectory instead of unfolding the correction path.The intended and delivered context-token pair sets diverge after one step and never re-intersect, leaving the correct information trapped in per-position marginals.
  • Gradient fragmentation: Loss-level interventions only reweight visited failed-prefix pairs rather than moving gradients onto correction-path pairs, motivating TRD’s refined-trajectory supervision.TRD generates a teacher-guided refined trajectory so supervision contexts grow along the correction path itself.
  • Empirical evidence: Both Dcorrect and Dincorrect stay near zero on student rollouts, indicating that teacher and student remain aligned and the teacher usually collapses onto failures.The observation holds under both forward and reverse KL in OPSD training.
  • Empirical evidence: 6–8 ‰ of teacher mass concentrates on 16 epistemic onset tokens, while the same metric on refined rollouts falls below 2 ‰.Student and teacher top-16 tokens already absorb 97–99% of total probability mass, making this allocation a dominant share of the residual budget.

5 Trajectory-Refined Distillation

TRD corrects problematic student rollouts at the trajectory level before standard on-policy distillation, targeting prefix failure while preserving on-policy support. It also expands exploration by exposing the student to teacher-guided alternative derivations and applies to on-policy self-distillation.

  • Method: TRD first samples an on-policy rollout and asks the teacher to construct a refined trajectory used for the subsequent OPD update.The procedure is yo ∼πθ(· | x), yr ∼πT(· | x, yo), followed by updating θ on yr.
  • Method: Conditioning the refined trajectory on the original rollout anchors it within demonstrated policy support, while the teacher rewrites erroneous portions to mitigate prefix failure.In OPSD, the backbone additionally conditions on the reference solution y∗ when constructing yr.
  • Exploration: TRD surfaces alternative valid derivations on correct rollouts, broadening exploration beyond the high-probability paths standard OPD tends to reinforce.These teacher-suggested paths are rarely sampled from πθ(· | x) alone.
  • Empirical behavior: 3x less: the teacher’s epistemic onset mass decreases 3x less on yr than on yo, while KL and the perplexity gap decrease as epistemic onset concentration fades.The refined trajectory also restores the incremental supervision associated with the teacher-student perplexity gap.

6 Experiments

TRD is evaluated against four dense-KL baselines across OPD and OPSD on math and code benchmarks, using Qwen3 teachers and students at multiple scales. It consistently preserves or improves performance, with trajectory refinement particularly benefiting harder math benchmarks and improving verifier-pass rates before distillation.

  • Experimental setup: TRD trains on refined trajectories yr, whereas Forward KL, clipped Forward KL, Reverse KL, and Top-K Reverse KL train on raw rollouts yo.The evaluation covers both OPD and OPSD settings across math and code benchmarks.
  • Experimental setup: The experiments use Qwen3-8B as the separate OPD teacher, while OPSD shares the Qwen3 backbone and derives teacher distributions through privileged conditioning.OPD students are Qwen3-1.7B and Qwen3-4B-Instruct-2507; OPSD uses Qwen3-4B-Instruct-2507 and Qwen3-8B.
  • OPD results: TRD is best or tied-best on seven of eight OPD Avg@16 benchmarks in each student block, with the smaller student gaining +4.6% on AIME24.Raw-rollout OPD variants generally fail to match the Qwen3-4B-Instruct-2507 base, whereas refined-trajectory training preserves its capabilities.
  • OPD results: TRD improves OPD Pass@16 on AMOBench by +5.1% for Qwen3-1.7B and +12.8% for Qwen3-4B-Instruct-2507.AIME24 and AIME25 are mostly saturated; on code, TRD matches the best HumanEval+ value and is best on MBPP+, while every method trails the base on LiveCodeBench.
  • OPSD results: In OPSD Avg@16, TRD is best on every benchmark at both scales and never falls below base, while three of four dense-KL baselines regress on at least one benchmark.For OPSD Pass@16, TRD separates most clearly from per-token interventions; on Qwen3-8B, it reports 50% relative gain on AMOBench and 15% on HMMT25.
  • Trajectory analysis: Refinement raises the OPSD verifier-pass rate from 65.8% on yo to 81.4% on yr and compresses median trajectory length roughly 9×, from 7.7K to 0.88K.The compressed trajectories move toward the reference median of ∼0.49K; under OPD AMOBench rollouts, TRD is slightly shorter on average, 18.9K →18.5K characters, with similar successful-rollout length distributions.

7 Conclusion · Appendix: Trajectory-Refined Distillation

The paper identifies prefix failure as a structural limitation of on-policy (self)-distillation and proposes Trajectory-Refined Distillation (TRD), which refines trajectories under privileged context before per-token KL supervision. TRD has additional sampling and depends on teacher-guided refinement that remains close to the student's on-policy distribution.

  • 7 Conclusion: Prefix failure makes frozen-rollout per-token KL induce a bimodal teacher mixture and fragmented gradients.The paper states that loss-level fixes leave this structural problem intact.
  • 7 Conclusion: TRD addresses prefix failure through trajectory-level refinement under privileged context.The refined trajectory is denoted yr.
  • 7 Conclusion: TRD supervises per-token KL along the refined trajectory yr rather than the student's original frozen rollout.This changes the supervision trajectory used for distillation.
  • 7 Conclusion: The proposed refinement recovers the ideal supervision-pair structure while remaining within the intended on-policy setting.The passage describes this as the mechanism for addressing the structural limitation.
  • 7 Conclusion: TRD requires one extra sampling budget to construct yr.This is identified as the first limitation.
  • 7 Conclusion: On Qwen3-8B, faster KL training on shorter refined trajectories partially offsets the extra sampling overhead.The passage says total wall-clock nearly matches dense-KL baselines, as reported in Appendix C.3.
  • 7 Conclusion: TRD depends on the teacher guiding refinement to mitigate prefix failure while keeping refined trajectories close to the student's on-policy distribution.The passage identifies this teacher-guidance requirement as a limitation.
  • 7 Conclusion: This teacher-guidance limitation is less severe when the teacher is stronger.The supplied passage states this condition without further quantification.

A Derivation of the OPD Policy Gradient · B Additional Experiments

The appendix derives the dense policy-gradient form of on-policy distillation, including its causal per-token surrogate and OPSD correspondence. It also outlines complementary experiments analyzing training-corpus properties, test-time difficulty buckets, and KL-objective ablations.

  • A Derivation of the OPD Policy Gradient: The derivation starts from sequence-level reverse KL and expresses OPD using the log-ratio δt between student and teacher token probabilities.Here, δt := log πθ(yt | x, y<t) − log πT (yt | x, y<t).
  • A Derivation of the OPD Policy Gradient: Differentiating the KL expectation requires accounting for θ-dependence in both the sampling distribution and the integrand.The derivation applies the product rule and score-function trick.
  • A Derivation of the OPD Policy Gradient: Causality makes cross terms with earlier tokens vanish after conditioning on the relevant prefix.For t′ < t, δt′ is measurable with respect to (x, y<t).
  • A Derivation of the OPD Policy Gradient: Applying discount factor 0 retains only the current-token term from the return-to-go, yielding the per-token surrogate.The resulting surrogate is identified as the gradient of the per-token KL loss.
  • A Derivation of the OPD Policy Gradient: With privileged-context substitution, the same derivation connects the OPD surrogate to the OPSD loss.The substitution is πT (· | x, y<t)=πθ(· | x, y⋆, y<t).
  • B Additional Experiments: Appendix B adds three analyses beyond dataset-averaged results in Tables 1–4.These analyses cover corpus characterization, AMOBench test-time decomposition, and KL-objective ablations.
  • B Additional Experiments: Appendix B.1 compares yo and yr on DeepScaleR across length, verifier accuracy, and joint outcome.The comparison uses Qwen3-4B and Qwen3-8B with-CoT and without-CoT subsets.
  • B Additional Experiments: Appendix B.2 decomposes AMOBench Avg@16 and Pass@16 by base-difficulty bucket, while B.3 ablates Forward-KL, Reverse-KL, and TRD.The stated purpose is to localize TRD’s test-time gains and compare training objectives.

B.1 Training-Trajectory Analysis: yo vs. yr

On DeepScaleR, TRD refines raw student rollouts into substantially shorter trajectories with higher verifier accuracy across both reference-quality regimes. The correction is predominantly beneficial: it recovers many failed prefixes while rarely disrupting already-correct rollouts.

  • Training-trajectory analysis: TRD analyzes raw rollouts yo and refined trajectories yr across a with-CoT subset (n=4,419) and an answer-only subset (n=35,826).The analysis reports Qwen3-4B-Instruct-2507 and Qwen3-8B results across both regimes.
  • Training-trajectory analysis: Refinement compresses trajectories in both regimes, with the 8B compression factor remaining ∼9× even for answer-only references.With CoT, yr is 0.85K / 0.88K from yo’s 2.2K / 7.7K; without CoT, yr is 0.93K / 0.83K from yo’s 2.1K / 7.5K.
  • Training-trajectory analysis: Verifier accuracy rises after refinement: with CoT, yr improves over yo by +8.9% / +15.6%, while without CoT the gains are +10.8% / +10.0%.With CoT, yo passes on 66.8% / 65.8% and yr on 75.7% / 81.4%; without CoT, yo passes on 69.5% / 69.8% and yr on 80.3% / 79.8%.
  • Training-trajectory analysis: Fail→pass outcomes exceed pass→fail outcomes by ∼80× on 4B-with-CoT and ∼44× on 8B-with-CoT, indicating predominantly corrective refinement.Comparable ratios hold on the without-CoT subset, consistent with refinement correcting dead-end prefixes more often than disturbing correct ones.

B.2 Test Rollout Analysis

Test-time rollouts on AMOBench show that TRD improves medium-difficulty performance, expands reachable support on questions the base model never solves, and finds shorter correct solution paths without longer incorrect reasoning.

  • Setup: The analysis evaluates 39 AMOBench questions using K=16 independent completions from the same Qwen3-8B checkpoints used for Tabs. 3 and 4.Generation uses temperature 0.6, top-p = 0.95, and a 38,912-token response budget.
  • Difficulty buckets: Questions are grouped by the base model’s 16-sample pass count into B0 (n = 23), B1–8 (n = 12), and B9–16 (n = 4).The buckets represent all-failure, medium-difficulty, and easy questions, respectively, and remain fixed when scoring Base, +Forward KL, and TRD.
  • Rollout lengths: TRD’s correct-rollout lengths are bimodal, including a pronounced low-length mode around 104 tokens absent from Base and +Forward KL.Incorrect-rollout distributions remain similar and near the generation cap, so TRD’s accuracy gains do not require longer reasoning.
  • Pass@16 by bucket: On B0, TRD achieves Pass@16= 0.39 versus +Forward KL at 0.22 across 23 questions the base model fails on all 16 attempts.Positive Pass@16 in B0 reflects expansion beyond the base policy’s reachable support at K=16.

B.3 Ablation: Trajectory-Subset for OPSD on Qwen3-8B … C.2 Evaluation Protocol

The ablation shows that full-corpus training is the appropriate default: filtering removes complementary trajectory signals, while TRD particularly depends on corpus breadth. The experiments use DeepScaleR/TACO training data, specified OPD/OPSD teacher setups, and K=16-completion evaluation across math and code benchmarks.

  • B.3 Ablation: Trajectory-Subset for OPSD on Qwen3-8B: Both fail and succ halves are necessary for coverage, making the full corpus the right default because filtered wins incur regressions elsewhere.The filters partition 40,245 tuples: fail has n=12,318, succ has n=27,927, and fail→succ has n=4,372.
  • B.3 Ablation: Trajectory-Subset for OPSD on Qwen3-8B: Filtering changes which trajectories enter D rather than optimizing Eq. (6), and each single-half filter drops a complementary teacher signal.succ-only loses hard failed problems; fail-only loses alternative-path signal from easy problems.
  • B.3 Ablation: Trajectory-Subset for OPSD on Qwen3-8B: Forward KL is the most data-sensitive, with vanilla +Forward KL regressing Base on four of five Avg@16 benchmarks while filtering recovers selected AIME25 and BeyondAIME results.Filtering still leaves AIME24 and HMMT25 below Base under every filter and mostly trades down Pass@16.
  • B.3 Ablation: Trajectory-Subset for OPSD on Qwen3-8B: Reverse KL is almost flat under filtering: both subset rows remain within ±1.0 of no-filter Avg@16, with mild Pass@16 tradeoffs rather than noticeable quality improvement.This stability is attributed to mode-seeking reverse KL discounting low-probability regions.
  • B.3 Ablation: Trajectory-Subset for OPSD on Qwen3-8B: TRD is hurt most by filtering: every subset row loses Pass@16 on every benchmark, with AMOBench losses of −10 to −15.The fail→succ subset also fails to outperform the full corpus, indicating that TRD benefits from the breadth of the refined corpus.
  • C Experiment Details: The appendix covers training data, trajectory construction, evaluation protocols, hardware budgets, refinement prompts, training metrics, and OPD/OPSD consistency checks.These materials are organized across Appendices C.1–C.6.
  • C.1 Training Data and Trajectory Construction: Training uses DeepScaleR for math and TACO for code; OPD uses a separate Qwen3-8B teacher, whereas OPSD shares the backbone and gives the teacher privileged access to y⋆.Stage 1 constructs yo for all methods, while Stage 2 constructs yr only for TRD.
  • C.2 Evaluation Protocol: Evaluation samples K=16 independent completions per question, reports Avg@16 and Pass@16, and uses boxed-answer verification for math and unit tests for code.Math uses final boxed-answer extraction; HumanEval+ and MBPP+ use EvalPlus, and LiveCodeBench uses lcb_runner release version 6.

C.3 Hardware and Compute … C.8 Common Optimization Hyperparameters

The implementation uses a single 8× H100 node and task-specific two-stage prompting, with diagnostics, model configurations, and optimization settings specified for OPD and OPSD. TRD adds a refinement sampling pass but can nearly match Vanilla OPSD wall-clock because refined trajectories are shorter.

  • C.3 Hardware and Compute: All runs use one node of 8× H100 80GB GPUs with FSDP2 sharding via verl, and each pipeline run includes generation, training, LoRA merging, and evaluation when enabled.Each row of Tables 1–4 corresponds to one offline pipeline run.
  • C.3 Hardware and Compute: 9:20 vs. 9:40: on Qwen3-8B, TRD and Vanilla OPSD have nearly matched total wall-clock despite TRD’s extra sampling pass.The overhead is partly offset by faster KL training because refined trajectories are much shorter than yo.
  • C.4 Initial and Refinement Prompts: Stage-1 prompts generate raw rollout yo, while task-specific refinement prompts differ between OPD and OPSD in whether the reference solution y⋆ is shown.OPD hides y⋆ and uses a separate teacher; OPSD includes y⋆ and uses the shared model under privileged conditioning.
  • C.5 Training Metrics for Fig. 3: OPSD diagnostics measure per-token KL by rollout outcome, epistemic-token mass, and the teacher–student perplexity gap on student rollouts.The curves use OPSD to control for teacher–student model mismatch, with token-weighted masking and aggregation described for each metric.
  • C.6 Models and Distillation Setup: OPD uses a frozen separate Qwen3-8B teacher with Qwen3-1.7B or Qwen3-4B-Instruct-2507 students, whereas OPSD shares Qwen3-4B-Instruct-2507 and Qwen3-8B backbones under privileged conditioning.In OPSD, LoRA updates only the student branch while the same base model supplies privileged teacher logits.
  • C.7 Method-Specific Hyperparameters: Direct OPD/OPSD rows use full-vocabulary KL over |V| ≈152K, T=1.0, AdamW, bfloat16, gradient checkpointing, one trainer epoch, and LoRA rank 64 / alpha 128.The default per-GPU batch is 1 with gradient accumulation 16, yielding effective batch 128 unless overridden.
  • C.8 Common Optimization Hyperparameters: Clipping caps each per-token KL at c, with c=0.1 for OPD direct clip01 rows and c=0.06 for canonical OPSD clipped-forward rows; Top-K instead uses the teacher’s top-32 support.Top-K renormalizes teacher and student distributions on that support before evaluating KL.
Loading 2606.08432v1…