PaperlayerView exact revision on arXiv

Source-linked AI summary

Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes

Yuqian Fu, Haohuan Huang, Kaiwen Jiang, Jiacai Liu, Zhuo Jiang, Yuanheng Zhu, Dongbin Zhao

arXiv:2603.25562v2cs.LGcs.AIcs.CL

TL;DR

Sampled-token OPD offers favorable variance for long-horizon training but can provide brittle supervision when rollouts drift or tokenizations mismatch. The paper analyzes these failures and proposes teacher top-K local support matching with truncated reverse-KL and stabilization choices, reporting a +19.8% gain over standard sampled-token OPD in a multi-task setting.

  • Problem

    Sampled-token OPD can become brittle because supervision is imbalanced, teacher guidance is unreliable on student-generated prefixes, and tokenization or special-token mismatches distort comparisons.

  • Method

    The paper combines theoretical analysis with empirical failure-mode analysis and proposes teacher top-K local support matching using truncated reverse-KL, top-p rollouts, and special-token masking.

  • Results

    +19.8% average math score improvement is reported for the unmasked local-support variant over the sampled-token baseline in alternating multi-task training, while maintaining competitive ALFWorld performance.

  • Takeaways & Limitations

    Local support matching preserves local token-level updates while providing a less brittle signal and more stable optimization than sampled-token OPD.

  • Takeaways & Limitations

    The ablation is partial and preliminary, with possible effects from support-set construction and remaining off-policy effects.

Abstract

from arXiv · show

On-policy distillation (OPD) is increasingly used in LLM post-training because it can leverage a teacher model to provide dense supervision on student rollouts. The standard implementation, however, usually reduces distribution matching to a sampled-token log-ratio, which can make the learning signal fragile on long rollouts whose prefixes drift away from the teacher's typical support. We revisit this formulation from both theoretical and implementation perspectives. Theoretically, token-level OPD is biased relative to sequence-level reverse-KL minimization, but admits a substantially tighter worst-case variance bound; a controlled synthetic study further shows that stronger future-reward coupling increases gradient variance and destabilizes training. Empirically, we identify three failure modes of sampled-token OPD: imbalanced token-level supervision, unreliable teacher guidance on student-generated prefixes, and tokenizer or special-token mismatch. These findings motivate teacher top-K local support matching, a truncated reverse-KL objective that compares teacher and student distributions over a teacher-supported token set at each prefix, together with top-p rollout sampling and special-token masking. Across single-task reasoning and multi-task benchmarks spanning agentic and reasoning settings, this objective improves optimization stability and yields a +19.8% performance gain over standard sampled-token OPD baselines, providing a practical recipe for more stable on-policy distillation.

1 Introduction

The paper revisits OPD’s bias–variance trade-off and identifies practical brittleness in sampled-token supervision. It proposes local support matching and related stabilization choices to retain token-level efficiency while improving training stability.

  • Theoretical trade-off: Token-level OPD is biased relative to sequence-level OPD but has better worst-case variance scaling for long-horizon training.The paper frames this as a bias–variance trade-off motivating token-level supervision.
  • Theoretical trade-off: Stronger future-reward coupling increases gradient variance and reduces optimization stability in the toy experiment.
  • Failure modes: Sampled-token OPD can become brittle on long rollouts because its one-token signal is imbalanced and teacher guidance can fail on student-generated prefixes.
  • Failure modes: Tokenizer or special-token mismatch can further distort one-token comparisons.
  • Proposed revision: Teacher top-K local support matching compares teacher and student over a teacher-supported token subset, using truncated reverse-KL, top-p rollouts, and special-token masking.The proposed revision preserves local token-level updates while broadening supervision beyond the sampled token.
  • Empirical outcome: The proposed objective yields more stable optimization and stronger performance than sampled-token OPD in single-task reasoning and multi-task agentic-plus-reasoning training.

2 Understanding Sampled-Token OPD: Trade-offs and Failure Modes

The paper explains why token-level OPD is attractive for long-horizon training but why sampled-token implementation can fail in practice. It connects theoretical variance benefits to empirical failure modes and motivates broader local-support comparisons.

  • 2.1 From reverse-KL to token-level OPD: Token-level OPD removes future-reward coupling, producing bias relative to sequence-level reverse-KL but lower long-horizon variance.Under bounded rewards and gradients, its worst-case variance bound scales as O(T2), versus O(T4) for the sequence-level estimator.
  • 2.1 From reverse-KL to token-level OPD: Increasing future coupling raises gradient variance and destabilizes optimization in the toy experiment.γ = 0 recovers token-level OPD, while γ = 1 recovers the causal sequence-level estimator.
  • 2.2 Why sampled-token OPD is brittle in practice: Sampled-token OPD drives updates using the teacher–student log-ratio on one sampled token, making supervision highly imbalanced.Most sampled tokens receive negative reward, leaving optimization dominated by a small subset of locally positive tokens.
  • 2.2 Why sampled-token OPD is brittle in practice: Teacher guidance becomes unreliable when student-generated prefixes are uncommon under the teacher, allowing locally rewarded tokens during repetition or meaningless continuations.
  • 2.2 Why sampled-token OPD is brittle in practice: Tokenizer and special-token mismatch can make one-token comparisons confuse semantic disagreement with segmentation differences.The same raw text may receive different tokenizations, distorting the reward signal.
  • 2.2 Why sampled-token OPD is brittle in practice: Teacher-supported local support matching compares distributions over multiple next-token continuations while retaining token-level updates for stability.

3 Method

The method replaces sampled-token supervision with a truncated reverse-KL comparison over teacher-supported candidate tokens at each rollout prefix. It combines local support matching with support renormalization, top-p rollout sampling, and special-token masking to make token-level updates less brittle while retaining local efficiency.

  • Teacher top-K local support matching: The objective compares teacher and student next-token distributions over a teacher-defined local support set instead of a single sampled token.The support set contains the K highest-probability teacher tokens at each prefix.
  • Teacher top-K local support matching: The method averages a truncated reverse-KL objective across rollout positions and renormalizes both distributions within the support set.Separate support-set softmax normalization makes probability masses directly comparable and prevents gradients from propagating directly to tokens outside the support.
  • Teacher top-K local support matching: Relative to sampled-token OPD, distribution-level matching avoids making the update depend entirely on one sampled-token log-ratio while remaining cheaper than full-vocabulary KL.The comparison is performed over teacher-supported candidates at the same prefix.
  • Teacher top-K local support matching: Top-p rollout sampling keeps trajectories closer to typical continuations, improving teacher-signal reliability on student-generated prefixes.Unconstrained sampling can produce very low-probability tokens and less informative prefixes.
  • Teacher top-K local support matching: Special-token masking reduces false negatives caused by incompatible tokenization conventions and is less necessary for the local support objective than for sampled-token OPD.The paper presents masking as an orthogonal practical fix that materially helps sampled-token OPD.

4 Experiments

The experiments evaluate local support matching across single-task math, alternating math-and-agentic training, and component and support-construction ablations. Results indicate that local support matching improves performance and stability when combined with appropriate rollout sampling and renormalization, while support construction matters especially in multi-task training.

  • Experimental setup: Experiments cover single-task math reasoning, alternating ALFWorld-and-math training, and ablations of rollout, support, and masking choices.The evaluation reports pass@1 on five math benchmarks and success rate on ALFWorld.
  • Single-task math reasoning: 40.7 average score follows special-token masking of sampled-token OPD, up from 36.4 without masking and 28.2 for the student.Local support matching further improves over sampled-token OPD in single-task math reasoning.
  • Multi-task training: +19.8% raises average math score from 34.8 to 41.7 in alternating multi-task training, while ALFWorld performance remains competitive.The masked variant achieves the best ALFWorld score at 97.7 but gives up part of the math improvement.
  • Component ablations: Under identical top-p rollout conditions, teacher top-K local support matching improves AIME24 avg@32 from 21.6 to 23.6.Renormalization is essential, while very small supports and fully unconstrained rollouts destabilize training; performance is less sensitive once K is sufficiently large.
  • Support variants: Support construction matters: teacher top-K, student top-K, and teacher top-K plus the sampled token are broadly comparable in single-task training but diverge substantially in multi-task results.The preliminary ablation does not justify over-interpreting the ranking among variants, especially because remaining off-policy effects may contribute.

5 Conclusion

The paper revisits OPD theoretically and empirically, identifying why sampled-token supervision becomes brittle and proposing teacher top-K local support matching as a more stable alternative. Across math and mixed agentic-reasoning settings, the modification improves optimization stability and downstream performance while leaving teacher matching an imperfect proxy for task success.

  • Conclusion: Token-level OPD offers better worst-case variance scaling than sequence-level reverse-KL, while sampled-token implementation suffers imbalanced, unreliable, and mismatch-sensitive supervision.The theoretical and empirical analyses motivate replacing one-token supervision with a local distribution-level comparison.
  • Conclusion: Teacher top-K local support matching preserves local token-level updates while replacing one-token supervision with truncated distribution-level comparison.The approach improves optimization stability and downstream performance over sampled-token OPD across single-task math and alternating agentic-plus-reasoning training.

A Discussion and Limitations

The discussion characterizes support matching as a truncated estimator with unresolved limitations involving sampling mismatch, teacher-task alignment, and broader distribution-shift settings.

  • Estimator limitations: Support-set KL restricts gradient contributions to the teacher top-K token subset, introducing bias relative to full-vocabulary reverse-KL.The authors describe this as an estimator property rather than a settled benefit or drawback.
  • Training limitations: Rollout prefixes are generated with top-p sampling while training updates omit correction for that sampling process, potentially creating training–inference mismatch.The authors leave this issue to ongoing work.
  • Task-alignment limitations: Teacher matching can diverge from task success because locally teacher-preferred continuations may remain rewardable on unhelpful or harmful trajectories.The authors note that stronger rollout control, distribution-shift handling, uncertainty modeling, or outcome-verifiable rewards may be needed.
  • Estimator limitations: Top-p truncation offers an alternative adaptive support by including tokens until a prescribed cumulative probability mass is reached.
  • Future directions: The paper identifies matched-task, matched-compute OPD-versus-RL comparisons and continual-learning evaluations as open directions for studying transfer and approximation error.

D Bias and variance analysis of token-level versus sequence-level OPD

The analysis formalizes a bias–variance trade-off: token-level OPD drops future-reward coupling and is biased, while sequence-level OPD is closer to the trajectory objective but has worse long-horizon variance scaling.

  • Bias analysis: Token-level OPD keeps only the immediate-reward term, removing future-reward coupling and making it generally biased relative to the sequence-level objective.
  • Variance analysis: O(T2) is the worst-case variance scaling for the token-level estimator under bounded rewards and gradients.
  • Variance analysis: O(T4) is the corresponding worst-case variance scaling for the sequence-level estimator.
  • Bias–variance trade-off: Sequence-level estimation is closer to the exact trajectory objective but couples each score term with many future rewards, producing quartic rather than quadratic variance growth.
  • Experimental setup: The toy study uses a two-task one-dimensional continuous-control environment with a roughly 4K-parameter three-layer MLP student.
  • Experimental setup: Gradient variance is estimated from eight micro-batches formed from batches of 64 trajectories, using output-layer gradient vectors.

E.3 Additional Results of Toy Experiments

Across three random seeds, the toy experiments show that stronger future-reward coupling generally raises gradient variance and can move learned trajectories toward sub-optimal regions.

  • Experimental comparison: Figures A1–A3 compare gradient variance and state visitation for γ ∈ {0.0, 0.25, 0.5, 0.75, 1.0} across three random seeds.
  • Gradient variance: Larger γ typically remains at a higher variance level later in training, with γ = 0.75 or γ = 1.0 sometimes one to several orders of magnitude above smaller γ values.
  • State visitation: Token-level OPD at γ = 0 consistently learns trajectories moving toward target states for both tasks across runs.
  • State visitation: When γ approaches the sequence-level case at γ = 1.0, learned trajectories often deviate from the desired direction and stabilize around sub-optimal regions.

F Experiment Setups

The experiments use established verl-based infrastructure and evaluate local support matching beyond the main math and ALFWorld settings with a single-task WebShop test.

  • Infrastructure: Experiments run on verl and verl-agent using a node with 8 NVIDIA H100 GPUs.
  • Training settings: Table A1 records the training and evaluation settings, including that rollout top-p is used for local support matching but not for the sampled-token baseline.
  • Additional evaluation: 57.8 is the WebShop success rate for local support matching, compared with 50.0 for sampled-token OPD.
  • Additional evaluation: The WebShop evaluation uses Qwen2.5-1.5B-Instruct as the student and omits a masking variant because tokenizer-mismatch effects are less central in this setting.

G.2 Training Dynamics and Alignment

The appendix examines how local support matching affects optimization, alignment, and variant performance across single-task and multi-task training. It finds more stable optimization and strong performance for teacher top-K support, while treating variant rankings as diagnostic rather than definitive.

  • Learning curves: Training reward and evaluation performance improve across most of math-reasoning training in both single-task and alternating multi-task settings.The improvement is not limited to the final checkpoint.
  • Optimization stability: Local support matching produces smaller gradient norms, fewer clipping-boundary hits, and sufficient policy entropy, indicating more stable optimization.Special-token masking substantially reduces clipping-boundary fractions for sampled-token OPD during early and middle training but has minor effects on the proposed method.
  • Alignment: The teacher–student sampled-token log-probability gap moves closer to zero under truncated local support matching.This suggests improved alignment even under the sampled-token diagnostic used for the baseline.
  • Variant design: The five variants differ in local support construction and sampled-token handling, using either renormalized local reverse-KL or an EMA-PG head-plus-tail correction.Variants 1–3 use renormalized reverse-KL, whereas Variants 4–5 use a truncated head term with sampled-token correction.
  • Variant comparison: Teacher top-K performs reasonably well in single-task training and remains strong in multi-task training, whereas EMA-PG variants do not improve empirical performance here.The comparison is presented as a diagnostic result rather than a definitive ranking of support-set objectives.

H Qualitative OPD reward-hacking case study

The case study traces sampled-token OPD from valid answers into over-continuation, hesitation loops, and malformed off-distribution text. Throughout, the local teacher signal can remain positive on low-information or degenerate continuations.

  • Over-continuation: After a valid answer is available, sampled-token OPD can continue rewarding generic fillers and connective tokens instead of clean termination.The same pattern appears on prefixes such as “confirm,” where additional verification remains favored.
  • Hesitation loops: Repeated “wait” tokens, punctuation-heavy continuations, and other weak fillers can remain locally rewardable after the trajectory becomes unproductive.These hesitation loops are consistent with the paper’s repetition-loop discussion.
  • Off-distribution drift: After the student drifts off-distribution, malformed non-English text and degenerate continuations can still receive high teacher probability.The local signal therefore does not necessarily self-correct the deteriorating trajectory.
Loading 2603.25562v2…