Source-linked AI summary
A Token-Level Analysis of Sampled-Token Reverse-KL On-Policy Distillation
Bing Shao, Jiazheng Zhang, Long Ma, Yujiong Shen, Senjie Jin, Xin Guo, Yuming Yang, Mingxu Chai, Zhiheng Xi, Tao Gui, Qi Zhang, Xuanjing Huang
TL;DR
Sampled-token reverse-KL OPD lacks a clear account of how its token-level loss allocates updates. The paper derives the student-logit gradient of the per-token K2 estimator and evaluates its allocation in Qwen3 math distillation, then tests SuRe as a bounded amplification rule. Low-student-probability tokens carry disproportionate gradient mass and larger teacher–student gaps, while SuRe improves several math metrics without clear degradation on selected out-of-domain benchmarks.
Problem
The paper asks which tokens receive large gradients in sampled-token reverse-KL OPD and why.
Method
The paper derives the K2 estimator’s gradient with respect to student logits and tests detached, bounded Surprise-aware Reweighting on Qwen3 math distillation.
Results
SuRe improves many, but not all, reported metrics over vanilla OPD, with the clearest gains on AMC23 and no clear degradation on selected out-of-domain benchmarks.
Takeaways & Limitations
Token-level gradient allocation is a useful lens for understanding sampled-token reverse-KL OPD, with SuRe serving as one lightweight empirical test.
Takeaways & Limitations
The study focuses on sampled token-level reverse-KL OPD, mainly mathematical reasoning datasets with limited-length traces, and a limited set of model scales and configurations.
Abstract
from arXiv · showhide
On-policy distillation (OPD) supervises a student on its own trajectories with token-level signals from a frozen teacher, yet how a sampled loss allocates updates across tokens remains poorly understood. We analyze the gradient of the per-token K2 estimator of reverse KL with respect to the student logits. The $\ell_1$ norm of this gradient factorizes into the absolute teacher--student log-probability gap and a student-side softmax factor that grows as the sampled token becomes less likely under the student. In our math-distillation runs, these per-token norms are highly non-uniform: low-student-probability tokens account for a disproportionate share of their sum and are also enriched in large teacher--student gaps. As a lightweight intervention suggested by this analysis, we study Surprise-aware Reweighting (SuRe), a detached, bounded weighting rule that further amplifies this existing allocation. Across two Qwen3 student scales, SuRe improves several math metrics over vanilla OPD and shows no clear degradation on the selected out-of-domain benchmarks. Our primary contribution is therefore a gradient-level characterization of reverse-KL OPD trained with the K2 estimator, with SuRe as one empirical instantiation.
1 Introduction
The paper characterizes how sampled-token reverse-KL OPD allocates gradient updates and tests SuRe, a bounded reweighting rule inspired by that analysis. In Qwen3 math distillation, low-student-probability tokens receive disproportionate gradient mass and overlap with large teacher–student gaps.
- OPD provides dense token-level feedback on student-generated trajectories while reducing distribution shift through on-policy training.
- Low-student-probability tokens account for a disproportionate share of summed gradient norms and are enriched in large teacher–student gaps.
- SuRe uses a detached, bounded per-token weight that mildly up-weights surprise tokens, approaches one for confident tokens, and recovers vanilla OPD at α = 0.
- The K2 gradient’s ℓ1 norm factorizes into the absolute teacher–student log-probability gap and the student-side factor 1 − πS(yt | ct).
- SuRe improves several metrics while controls do not fully separate exact surprise assignment from broader benefits of non-uniform weighting.
2 Preliminaries
The paper studies loss-based on-policy distillation with a frozen teacher, student-generated samples, and the per-token K2 estimator of reverse KL. The estimator is biased for the reverse-KL value but has an unbiased realized-loss logit gradient in expectation under current-student sampling at fixed context.
- OPD directly backpropagates the per-token K2 estimator as a loss, treating student rollouts as fixed samples within each update.
- The student generates responses token by token, with context ct = (x, y<t) and next-token distribution πS represented by student logits z.
- The distillation objective is reverse KL, DKL(πS(· | ct) ∥ πT(· | ct)), between student and frozen-teacher next-token distributions.
- K2 is a biased estimator of the reverse-KL value, but its realized-loss logit gradient is unbiased in expectation under current-student sampling at fixed context ct.
- The sampled token is held fixed during backpropagation, with no differentiation through the sampled trajectory or prefix distribution.
- Entropy and Jensen–Shannon divergence are used as diagnostics on next-token distributions, with entropy quantifying distributional uncertainty.
3 Token-Level Gradient Analysis
The analysis separates endpoint checkpoint shifts from the teacher–student residual used by the K2 loss, then derives and tests the K2 gradient-norm allocation. Gradient mass concentrates in low-student-probability tokens that also show larger absolute teacher–student gaps.
- 3.1 Post-training checkpoint-shift diagnostic: The checkpoint-shift diagnostic compares Base and final vanilla-OPD checkpoints on the same sampled tokens and prefixes, rather than measuring teacher endorsement or gradient concentration.
- 3.1 Post-training checkpoint-shift diagnostic: Only 8.5% of OPD-rollout tokens and 7.1% of Base-rollout tokens satisfy |∆OPD−Base| > 1.
- 3.2 An exact identity for the gradient norm of the K2 estimator: For the K2 estimator, the gradient norm identity separates the absolute teacher–student residual from the student-side softmax geometry term.
- 3.2 An exact identity for the gradient norm of the K2 estimator: Holding |∆log pt| fixed, tokens with smaller πS(yt | ct) receive larger gradient coefficients through the factor 1 − πS(yt | ct).
- 3.2 An exact identity for the gradient norm of the K2 estimator: Ranking by |∆log pt| captures 54.1%/74.6% of summed gradient norms in the top 5%/10% tokens, versus 28.4%/47.3% for JSD and 23.7%/42.4% for entropy.
- 3.2 An exact identity for the gradient norm of the K2 estimator: Lower-πS bins contain a larger share of gradient norms and are enriched in large absolute gaps, an empirical association not implied by 1 − πS alone.
4 An Analysis-Inspired Reweighting
The paper introduces SuRe, a simple student-side reweighting intervention motivated by the observed allocation of reverse-KL OPD gradients. It uses a detached, bounded weight that increases emphasis on sampled tokens with lower student probability.
- SuRe probes whether the observed concentration of gradient norms among student-surprised tokens is useful for optimization.
- SuRe uses a smooth, monotone, bounded weighting rule with wt ∈[1, 1 + α], recovering vanilla OPD through one dial.
- The weight uses detached sampled-token student probability, so it rescales the baseline per-token gradient without changing its gap factor.
- Because the denominator is not renormalized and wt ≥1, SuRe can increase the overall loss scale as well as relative emphasis on low-probability samples.
- The method applies a multiplicative weight to each per-token reverse-KL loss while retaining the unweighted token-mean denominator.
- The implementation change is isolated to the detached scalar wt.
5 Experiments
Experiments evaluate SuRe against vanilla OPD across two Qwen3 student scales, math benchmarks, selected OOD tasks, training dynamics, and controls. SuRe improves many but not all math metrics, while OOD results remain broadly comparable and ablations leave the exact source of gains unresolved.
- Setup: Evaluation covers AIME2024, AIME2025, AMC23, MATH-500, CRUX, IFEval, and MMLU-Pro with task-specific avg@k, pass@k, and pass@1 metrics.Table 1 reports in-domain performance, while selected OOD benchmarks measure code generation, instruction following, and general ability.
- Main Results: Vanilla OPD achieves higher avg@k than Base, KD, and SeqKD on every evaluated math benchmark across both student scales.KD and SeqKD sometimes fall below Base, including AMC23 at the 4B scale.
- Main Results: SuRe improves many, but not all, reported metrics over vanilla OPD, with clearest gains on AMC23: +3.7pp avg@8 and +7.5pp pass@8 at 1.7B, and +5.0pp pass@8 at 4B.MATH-500 changes are small, and both AIME25 metrics decrease at 4B.
- Main Results: On selected OOD tasks, OPD and SuRe are broadly comparable, and neither uniformly improves over Base; both are slightly below Base on 1.7B MMLU-Pro.The experiments therefore show no clear broad transfer effect on these selected tasks.
- Main Results: SuRe increases the early gradient norm while tracking OPD’s actor entropy, and its mean rollout score improves in the 1.7B training dynamics.Because unnormalized weights have mean above one, the larger norm alone cannot distinguish surprise alignment from increased loss scale.
- Ablations and Controls: On AMC23, α=1.0 is uniformly best for k∈{1, 2, 4}; α=2.0 weakens the small-k gain, so larger amplification is not uniformly better.The sweep supports a smooth but empirically non-monotonic dependence on amplification strength.
- Ablations and Controls: Orientation controls favor aligned surprise weighting over opposite weighting, but exact-shuffled comparisons remain statistically unresolved and do not isolate exact surprise assignment from generic non-uniform weighting.Mean-normalized SuRe preserves the MATH-500 improvement, while exact-shuffled and rank-reversed variants are numerically lower than the aligned normalized variant.
6 Related Work
Prior work studies distillation’s on-policy training rationale and token-level uncertainty or gradient concentration. This paper complements those perspectives by directly characterizing sampled OPD gradient allocation and using it to motivate token reweighting.
- Distillation: Offline distillation can suffer train-inference mismatch because students learn from teacher sequences but decode from their own distributions.OPD addresses this by training on student rollouts with dense teacher supervision.
- Reasoning Analysis and Optimization: Recent reasoning analyses report that RLVR gradients concentrate on high-entropy minority tokens and depend strongly on update direction.These findings motivate selecting or reweighting informative token positions.
7 Conclusion
The paper characterizes how sampled-token reverse-KL OPD allocates gradient norms and finds highly non-uniform concentration in low-student-probability tokens with large teacher–student gaps. SuRe tests amplifying this allocation and improves several metrics in the evaluated setting.
- The analysis measures how per-token reverse-KL gradients distribute across sampled positions in on-policy distillation.
- Largest gradient norms concentrate among low-student-probability samples that are also enriched in large teacher–student gaps.
- SuRe provides a lightweight test of amplifying the observed token-level gradient allocation.
- SuRe improves several metrics in the evaluated setting.
8 Limitations
The study’s scope is limited to sampled token-level reverse-KL OPD, primarily mathematical reasoning datasets with limited-length traces, and a restricted set of model and OPD configurations.
- Alternative objectives and designs, including full-vocabulary distillation and Jensen–Shannon divergence, are not systematically investigated.
- Experiments primarily use mathematical reasoning datasets with generated traces of limited length, which may restrict generality.
- Computational constraints limit the explored model and OPD configurations, excluding larger parameter scales.
A Detailed Gradient Derivation
The derivation analyzes the K2 reverse-KL loss for a fixed sampled token and expresses its student-logit gradient through the teacher–student gap and student probabilities. The resulting ℓ1 norm has symmetric sampled-token and competitor contributions, with distinct limits for confident and surprised tokens.
- The derivation fixes an on-policy sampled token and differentiates its scalar K2 loss with respect to student logits.
- The teacher log-probability is treated as a stopgradient constant, and the sampled trajectory is held fixed during backpropagation.
- The gradient-descent direction raises the sampled-token logit when the teacher favors that token more than the student, with signs reversing for the opposite gap.
- The gradient ℓ1 norm factorizes into 2|∆log p_t|(1 − p^S_t), combining the teacher–student gap with the student surprise factor.
- The sampled-token and competitor-coordinate channels contribute equally to the ℓ1 norm while pointing in opposite directions in logit space.
- As p^S_t approaches one the gradient norm vanishes, whereas as p^S_t approaches zero it approaches 2|∆log p_t|.
B Experimental Details
The experiments compare vanilla OPD and SuRe across Qwen3 student scales using shared DeepMath-hard training settings and aligned math prompts. Additional evaluations report best-of-k behavior, OOD metrics, larger-scale training dynamics, and a qualitative checkpoint comparison.
- Vanilla OPD and SuRe use Qwen3-8B as teacher, DeepMath’s 57K hard split, and Qwen3-1.7B-Base or Qwen3-4B-Base students.
- Both teacher and student receive the same training-time prompt, aligning contexts for the teacher–student gap.
- Evaluation uses the matching appended math instruction, with answers parsed from the final boxed expression and checked against benchmark ground truth.
- OOD appendix results report pass@1, pass@5, and pass@10, using strict accuracy for IFEval and exact match for MMLU-Pro.
- The only method difference in the corresponding scale runs is the per-token weight, while other listed hyperparameters are shared apart from memory-related settings.
- At the larger student scale, training-dynamics curves show qualitatively similar optimization behavior to the 1.7B analysis.
- The endpoint comparison finds most response tokens close between final OPD and Base, with sharp differences at a small number of local decision points.
C.4 Evaluation
The evaluation covers math and out-of-domain benchmarks with specified sampling protocols, while comparing SuRe against baselines and matched weighting controls. Results indicate gains on average metrics relative to GRPO, but the comparison remains an additional reference because supervision differs and no significance test is reported.
- Comparison with On-Policy RL (GRPO): SuRe and GRPO use the same DeepMath split, prompt template, hardware setup, decoding settings, and evaluation protocol for direct numerical comparison.The methods nevertheless use different supervision signals: reward-based answer matching versus teacher token-level distributions.
- Comparison with On-Policy RL (GRPO): SuRe matches or exceeds GRPO on every benchmark in avg@k, with roughly 1–3pp improvements across AIME24, AIME25, AMC23, and MATH-500.The pass@k comparison is mixed: SuRe ties or wins on three of four benchmarks but trails GRPO on AIME25.
- Comparison with On-Policy RL (GRPO): The GRPO comparison is treated as an additional reference rather than evidence that the objectives are equivalent.The authors cite different supervision and the absence of a significance test as reasons for this qualification.
- Matched Weighting Controls: Mean-normalized SuRe preserves the MATH-500 improvement, while shuffled and rank-reversed assignments are numerically lower than the aligned normalized variant.Because exact-shuffled still exceeds vanilla OPD, the gain is not attributed entirely to exact surprise assignment.