PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yuxuan Jiang, Francis Ferraro

arXiv:2606.00305v2cs.CLcs.AI

TL;DR

OPD’s token-level loss does not reliably identify or repair trajectory-level reasoning divergence. TOPD uses near-future continuations and short-window OT guidance to detect divergent states and guide multiple future tokens. The paper reports that TOPD consistently improves standard OPD on mathematical reasoning benchmarks.

  • Problem

    OPD’s token-level reverse-KL supervision does not reliably detect trajectory-level reasoning errors or repair the student’s future reasoning path.

    Source: §1 · pp. 1–2

  • Method

    TOPD compares teacher and student near-future continuations to identify divergent states and uses short-window OT alignment to distribute teacher trajectory guidance across multiple tokens.

    Source: §1 · p. 2

  • Results

    TOPD consistently improves standard OPD on competitive mathematics benchmarks, and mechanistic analysis finds it more effectively reduces local trajectory divergence.

    Source: §7 · p. 9

  • Takeaways & Limitations

    The findings suggest that OPD reasoning failures are trajectory-level phenomena rather than isolated token mismatches.

    Source: §7 · p. 10

    The experiments are mainly on mathematical reasoning benchmarks, so validation on domains such as code generation and open-ended reasoning remains needed.

    Source: §8 · p. 10

Abstract

from arXiv · show

On-Policy Distillation (OPD) improves large language model reasoning by training a student model on trajectories sampled from its own policy under teacher supervision. Although OPD operates on trajectories, its learning signal remains token-level: it identifies deviations through high-loss tokens and repairs them through local reverse-KL correction. We show that this "trajectory-sampled but token-learned" mechanism cannot reliably bridge student trajectories toward teacher trajectories. About 30% of high-loss tokens fall into the low-divergence regime, indicating that many are surface-form mismatches rather than real reasoning forks. Moreover, even truly divergent tokens are difficult to repair with isolated token-level supervision, since reasoning failures often unfold as short-horizon distributional drift. We propose Trajectory-aware OPD (TOPD), which uses near-future trajectory information to identify real divergent states and distribute guidance across multiple future tokens. Experiments show that suppressing non-divergent high-loss tokens improves standard OPD from 47.8% to 48.2% average accuracy, while TOPD further improves performance to 52.2%, with gains on AIME24 from 60.0% to 63.3% and AIME25 from 46.7% to 53.3%.

1 Introduction

Standard OPD uses token-level reverse-KL correction to train on student-generated trajectories, but token mismatches do not reliably identify or repair reasoning-path divergence. TOPD uses near-future trajectory information to detect divergent states and distribute teacher guidance across multiple future tokens.

1 Introduction

  • Single-token reverse-KL correction struggles to repair an entire reasoning trajectory: in the 5+2×3 example, local tokens change, yet the model still generates “is 7,”.Source: §1 · p. 2
  • TOPD identifies genuinely divergent states from teacher and student short-window continuations, then uses OT alignment to guide multiple future tokens toward the teacher path.This extends supervision from isolated token correction to short-window trajectory correction.Source: §1 · pp. 2–3
  • Token-level supervision does not reliably provide the trajectory-level error detection and path repair required for multi-step reasoning.The section frames this as a mismatch between OPD’s token-level reverse-KL objective and trajectory-level reasoning correction.Source: §1 · p. 2
  • About 30% of high-loss tokens are low-divergence false alarms, so local loss can reflect surface-form mismatches rather than reasoning forks.The passage reports that ignoring these false alarms improves OPD performance.Source: §1 · pp. 1–2

2 When Token-Level OPD Fails to Redirect Reasoning Trajectories

Standard OPD’s token-level supervision is a blunt signal for trajectory divergence and can fail to redirect reasoning even after correcting a divergent token. The section examines these two failure modes through short-window probing and a local correction case study.

2 When Token-Level OPD Fails to Redirect Reasoning Trajectories

  • The section identifies token-level supervision as structurally mismatched with the need to detect divergence and then correct the reasoning path.Its two failure modes are noisy high-loss targets and insufficient trajectory repair from isolated token correction.Source: §2 · p. 3, §2.1 · p. 3, §2.2 · p. 4

2.1 Informative but Blunt: High-Loss Tokens Are Not Always Trajectory-Critical

  • High-loss tokens can mark stylistic or equivalent-form mismatches without indicating genuine future reasoning divergence.The short-window analysis characterizes token loss as informative but noisy, with substantial overlap between high-loss and low-divergence regions.Source: §2.1 · p. 3
  • 29.23% of high-loss tokens fall within the typical low-divergence range, despite high-loss tokens having larger median OT distance than random or low-loss tokens.Median OT distances are 0.566 for high-loss, 0.361 for random, and 0.315 for low-loss tokens; loss-to-distance correlation is limited.Source: §2.1 · p. 3

2.2 Local Token Correction Does Not Guarantee Trajectory Correction

  • Correcting the divergent token “(” still leaves the student on an erroneous continuation instead of the teacher-consistent trajectory.In the 5+2×3 case, the student generates “)” and returns to an erroneous path producing 2 = 7.Source: §2.2 · p. 4, Figure 3 · p. 4
  • The token-by-token learning trap arises because reverse-KL supervision edits one position at a time without sufficiently reshaping near-future transition dynamics.The local correction setting fixes a prefix and compares regenerated student continuations against the teacher continuation.Source: §2.2 · p. 4

3 Methodology

TOPD identifies student states associated with near-future reasoning drift, then adds trajectory-level guidance to the training objective. It uses optimal-transport alignment of teacher and student continuations to guide multiple future-token predictions while retaining token-level correction.

Divergence Detection

  • TOPD treats a position as genuinely divergent only when it has both high token-level loss and high short-window OT divergence.The teacher and student continuations are generated from the same student prefix; TOPD uses their trajectory discrepancy to identify real divergent points.Source: §3 · pp. 4–5

Trajectory-level Teacher Signal Injection

  • At each detected point, TOPD uses the OT transport plan to form a trajectory-aware soft target for each student future position.The target is based on soft alignment between student and teacher future positions, and the student prediction distribution is trained to match it.Source: §3 · p. 5

Training Objective

  • TOPD adds trajectory-level supervision to the original OPD loss, with λ controlling its strength.This preserves token-level correction while encouraging the student to follow a teacher-consistent short-window reasoning path.Source: §3 · p. 5

4 Empirical Experiments

The experiments train Qwen3-4B through off-policy then on-policy distillation and evaluate it on competitive mathematics benchmarks. TOPD outperforms standard OPD across benchmarks, reaching 52.2 average accuracy.

4.1 Training Setup

  • The student is initialized with forward-KL and cross-entropy distillation on 20k teacher solutions, then trained with reverse-KL on four student rollouts for each of 50k prompts.The two stages use kd_ratio values of 0.5 and 1.0, respectively.Source: §4.1 · p. 5
  • TOPD adds modest computation overhead: training takes 1.41× as long as standard OPD, with OT computed only on selected high-loss short windows.Source: §4.1 · p. 5

4.2 Evaluation

  • Evaluation uses zero-shot Pass@1 averaged over five runs on AIME 24, AIME 25, and HMMT 25-Feb, with 90 problems aggregated as the primary score.The evaluation uses the LM-EVALUATION-HARNESS framework.Source: §4.2 · pp. 5–6

4.3 Main Results

  • 52.2 average accuracy: TOPD consistently outperforms standard OPD across three competitive mathematics benchmarks.Standard OPD reaches 47.8, compared with 38.9 for the offline-distilled model.Source: §4.3 · p. 6
  • TOPD improves AIME24 from 60.0 to 63.3 and AIME25 from 46.7 to 53.3.Source: §4.3 · p. 6

5 Ablations and Analysis

Downweighting low-OT high-loss tokens slightly improves OPD, whereas suppressing matched high-OT tokens harms it; TOPD achieves the strongest reported accuracy. A local correction analysis also finds that TOPD lowers short-window trajectory divergence more than standard OPD.

5.1 Trajectory-Aware Divergence Detection

  • High-loss tokens span a broad range of OT distances, motivating selective downweighting of low-OT tokens rather than treating every high-loss position as a real divergence.Selected tokens have their loss contribution capped at the sample-level median; the comparison downweights the bottom 30% of high-loss tokens or the same number sampled from the high-OT region.Source: Figure 4 · p. 6, §5.1 · p. 6
  • 52.2% accuracy: TOPD performs best, versus 47.8% for standard OPD, 48.2% after low-OT downweighting, and 42.1% after matched high-OT downweighting.The contrast indicates that low-OT high-loss tokens can be noisy signals while high-OT tokens carry useful correction signals.Source: §5.1 · p. 6, Table 2 · p. 7

5.2 Does TOPD Improve Local Trajectory Correction?

  • The local correction test fixes the prefix at a real divergence point, targets the next 50 tokens, and measures OT distance between regenerated student and teacher continuations.Source: §5.2 · p. 7
  • TOPD lowers mean short-window OT distance from 0.538 before updates to 0.204, compared with 0.350 after standard OPD updates.The experiment selects 1,000 real divergence points and applies five consecutive local updates on the same prefix-window pair.Source: §5.2 · pp. 7–8, Figure 5 · p. 7

6 Related Works

Related work positions OPD as a post-training approach for reasoning models, while noting limitations of token-level objectives for cross-step trajectory alignment. The paper distinguishes TOPD by integrating trajectory-aware transport signals into reverse-KL optimization.

6.1 On-Policy Distillation for LLM Post-Training

  • Prior work presents OPD as a post-training paradigm that combines trajectory exploration with token-level guidance, and studies factors such as teacher complementarity, reasoning-style overlap, and adaptive teacher selection.Other cited directions examine reasoning efficiency, compression, and performance-efficiency trade-offs.Source: §6.1 · p. 8

6.2 Optimal Transport and Structured Alignment

  • Token-level forward- and reverse-KL objectives can treat prediction steps independently, limiting their modeling of cross-step dependencies in reasoning trajectories.TOPD differs by integrating trajectory-aware transport signals into reverse-KL optimization to represent short-horizon shifts and path dependencies.Source: §6.2 · pp. 8–9

6.3 Efficient Reasoning and Agentic Systems

  • The related-work section also surveys efficiency, retrieval, evaluation, agentic, and graph-based reasoning studies across domain-specific settings.Source: §6.3 · p. 9

7 Conclusion

The conclusion identifies noisy high-loss false alarms and insufficient isolated token-level correction as limitations of standard OPD. TOPD combines divergence detection with short-window OT supervision, and the analysis reports reduced local trajectory divergence.

7 Conclusion

  • Standard OPD has two stated limitations: high-loss tokens may be false alarms, and isolated token corrections may not redirect the student’s future reasoning trajectory.Source: §7 · p. 9
  • TOPD combines near-future divergence detection with short-window OT guidance, and mechanistic analysis finds reduced local divergence as the student absorbs teacher trajectory information more effectively.Source: §7 · pp. 9–10
  • The conclusion frames reasoning failures in OPD as trajectory-level phenomena rather than isolated token mismatches.It presents trajectory-aware distillation objectives as a direction for reasoning language models.Source: §7 · p. 10

8 Limitations

The paper identifies computational cost, short-horizon analysis, mathematical-benchmark coverage, and limits of OT distance as limitations of TOPD.

8 Limitations

  • TOPD adds computation, focuses on short-horizon divergence and mathematical benchmarks, and OT distance may not capture semantic equivalence between reasoning paths.The authors call for validation on code generation and open-ended reasoning; different reasoning paths can reach the same correct answer.Source: §8 · p. 10

9 Ethics

The work uses public datasets and benchmarks under their released terms, and its experiments focus on English mathematical reasoning. The authors note that they cannot guarantee the data are free of undesirable content.

9 Ethics

  • Experiments use publicly available datasets and open-access benchmarks, without accessing proprietary training data or processing personal or sensitive user data.Experiments use standard inference and optimization procedures.Source: §9 · p. 10
  • Datasets and benchmarks are used under their released terms and intended research purposes; raw datasets and proprietary model outputs are not redistributed.Derived artifacts are intended for research and evaluation.Source: §9 · p. 10
  • The artifacts focus on English mathematical reasoning benchmarks, and the authors cannot guarantee datasets are free from biased, toxic, or otherwise undesirable content.The benchmarks primarily cover symbolic and multi-step reasoning rather than demographic or user-centered data; ChatGPT is used only for grammar correction and language polishing.Source: §9 · p. 10

A Additional correlation analysis

The additional correlation analysis finds that teacher token loss is informative but insufficiently selective for identifying trajectory-critical divergence points.

A Additional correlation analysis

  • Teacher token loss has a weak, dispersed relationship with near-future OT distance, so high loss does not reliably distinguish genuine trajectory divergence.Many high-loss tokens have small OT distances, some moderate-loss tokens have strong divergence, and the binned mean rises gradually rather than separating sharply.Source: A Additional correlation analysis · p. 14

B.1 off policy and on policy training

The appendix describes the two-stage KDFlow training setup and TOPD’s added trajectory-processing cost. TOPD’s overhead remains manageable in the reported setup, at 1.41× standard OPD training time.

B.1 off policy and on policy training

  • Training distills from Qwen3-30B-A3B-Instruct-2507 into Qwen3-4B-Instruct-2507 with thinking mode disabled, using KDFlow and two training stages.The setup uses a single node of 4× H100 GPUs with FSDP2, bf16, and gradient checkpointing.Source: §B.1 · pp. 14–15
  • Teacher token loss has a weak, highly dispersed correlation with near-future OT distance at high-loss positions.Figure 6 indicates token-level loss alone is an imprecise predictor of trajectory-level divergence.Source: Figure 6 · p. 15
  • Stage 1 trains the student on 20k teacher responses sampled from OpenThoughts3, using forward-KL distillation combined with cross-entropy at kd_ratio = 0.5.Training lasts one epoch with learning rate 2 × 10−5 and global batch size 128.Source: §B.1 · p. 15
  • Stage 2 initializes from the Stage-1 checkpoint and distills teacher token distributions on four student rollouts per prompt using pure reverse-KL distillation.The rollouts use a separate 10k-prompt slice; training lasts one epoch with learning rate 2 × 10−6 and global batch size 4.Source: §B.1 · p. 15

B.2 Computational Overhead.

  • TOPD applies trajectory-level correction only at selected high-loss positions, with probing overhead approximately m/N rather than processing every token.In the reported setting, m = 218 and N = 9688; each position uses a length-50 continuation and a 50 × 50 OT cost matrix.Source: §B.2 · p. 15
  • TOPD increases total training time to 1.41× standard OPD under the same training setup.The authors characterize this as moderate additional cost.Source: §B.2 · p. 15

Related papers

Loading 2606.00305v2…