Bridging Reasoning Trajectories in On-Policy Distillation via Near-Future Guidance

Technical Summary: Bridging Reasoning Trajectories in On-Policy Distillation via Near-Future Guidance

Problem Statement
On-Policy Distillation (OPD) is a cornerstone of modern post-training for Large Language Models (LLMs), designed to improve reasoning by training a student model on trajectories sampled from its own policy under teacher supervision. Standard OPD operates on a "trajectory-sampled but token-learned" mechanism: it identifies deviations via high-loss tokens and repairs them using local reverse-KL divergence correction.

The authors identify a fundamental mismatch in this approach, arguing that token-level supervision is structurally misaligned with the requirements of trajectory-level reasoning correction. This misalignment manifests in two critical failure modes:

1. Noisy False Alarms: High-loss tokens do not always correspond to genuine reasoning divergence. Empirical analysis reveals that approximately 30% of high-loss tokens fall into a low-divergence regime, representing surface-form mismatches (e.g., stylistic preferences like "and" vs. "then") rather than logical forks. Treating these as critical errors provides noisy supervision that can actively interfere with reasoning correction.
2. The Token-by-Token Learning Trap: Even when genuine divergence points are identified, single-token reverse-KL correction is often insufficient to redirect the entire reasoning path. Reasoning failures typically evolve as short-horizon distributional drift. Correcting one token locally does not guarantee the model will align with the teacher's future trajectory; the model often drifts back into erroneous branches immediately after the local correction.

Methodology: Trajectory-aware OPD (TOPD)
To address these limitations, the authors propose Trajectory-aware OPD (TOPD), which bridges the gap between local token correction and global trajectory evolution by incorporating near-future trajectory information. The method consists of two primary stages:

1. Divergence Detection:

   • For high-loss candidate positions, the system steps back to the preceding token to establish a shared prefix.
   • Both the teacher and student generate short-window continuations (length $K=50$) from this prefix.
   • The system measures trajectory-level discrepancy using Optimal Transport (OT) distance between the teacher and student continuations.
   • Real divergent points are identified as positions exhibiting both high token-level loss and high short-window OT divergence. This filters out benign surface-level mismatches.

2. Trajectory-level Teacher Signal Injection:

   • Instead of optimizing only the current token, TOPD extends supervision to the entire short-window trajectory.
   • Using the OT transport plan ($\gamma$), the method constructs a "trajectory-aware soft target" for each student future position. This target aligns the student's future tokens with the teacher's trajectory based on the optimal transport mapping.
   • The student is trained to match these OT-aligned soft targets across the window, effectively learning how to steer its near-future reasoning path toward the teacher's path, not just how to correct the current token.
   • The final objective combines the standard OPD loss with a trajectory loss term ($L_{TOPD} = L_{OPD} + \lambda L_{traj}$).

Key Contributions

• Identification of Mismatch: The work establishes that token-level supervision is a "blunt proxy" for trajectory-level reasoning correction, failing to distinguish between genuine logical forks and benign surface variations.
• Empirical Characterization of Failure Modes: The authors provide evidence that ~30% of high-loss tokens are low-divergence false alarms and demonstrate the "token-by-token learning trap" where local corrections fail to restore the overall reasoning path.
• Trajectory-Level Guidance Principle: The paper introduces a principle where near-future trajectory information is distributed across multiple tokens. By using short-window OT alignment, TOPD injects teacher-student path discrepancies directly into the loss function, enabling the student to learn trajectory shifts rather than isolated token edits.

Experimental Results
Experiments were conducted using Qwen3 models (distilling a 30B teacher into a 4B student) on competitive mathematics benchmarks (AIME24, AIME25, HMMT25-Feb).

• Performance Gains: Standard OPD improved the warm-start model's average accuracy from 38.9% to 47.8%. TOPD further increased this to 52.2%.
• Specific Benchmarks: TOPD improved AIME24 accuracy from 60.0% to 63.3% and AIME25 from 46.7% to 53.3%.
• Ablation Studies:
  • Suppressing low-OT high-loss tokens (false alarms) in standard OPD yielded a modest gain (+0.4%), confirming these tokens are noisy.
  • Suppressing high-OT high-loss tokens (real divergences) caused a significant performance drop, confirming their necessity.
  • Local trajectory correction analysis showed that while standard OPD reduced short-window OT distance partially (from 0.538 to 0.350), TOPD achieved a substantial reduction (to 0.204), demonstrating superior ability to redirect near-future trajectories.
• Efficiency: TOPD introduces a modest computational overhead, with training time being 1.41$\times$ that of standard OPD, as OT is computed only on selected high-loss windows.

Significance and Claims
The paper claims that reasoning failures in OPD are fundamentally trajectory-level phenomena rather than isolated token mismatches. By shifting from point-wise token editing to global, transition-aware alignment, TOPD enables more effective reasoning alignment. The authors position this work as a step toward more trajectory-aware distillation objectives, arguing that successful reasoning correction requires the student to learn not just which token to correct, but how to navigate the subsequent reasoning path to remain consistent with the teacher.

Limitations
The authors acknowledge that TOPD incurs additional computational costs due to OT calculations, focuses primarily on short-horizon divergence (potentially missing long-context errors), and is currently validated mainly on mathematical reasoning benchmarks rather than code or open-ended tasks. Additionally, OT distance measures trajectory proximity but may not capture semantic equivalence where different paths lead to the same correct answer.