A Token-Level Analysis of Sampled-Token Reverse-KL On-Policy Distillation

• URL: http://arxiv.org/abs/2608.25643v2
• Date: Thu, 27 Aug 2026 06:20:00 GMT
• Title: A Token-Level Analysis of Sampled-Token Reverse-KL On-Policy Distillation
• Abstract summary: We analyze the gradient of the per-token K2 estimator of reverse KL with respect to the student logits.<n>Our primary contribution is therefore a gradient-level characterization of reverse-KL OPD trained with the K2 estimator, with SuRe as one empirical instantiation.
• Score: 59.43605629311855
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• Mismatch Matters: On-Policy Distillation Beyond Token Agreement [27.396104103072844]
  We show that mismatch tokens can be mainly categorized into two types: student-excess tokens and student-deficit tokens.<n>We propose TIDE, which applies bounded Hellinger shaping to suppress the most severe sampled excesses and an analytic teacher top-$K$ injection to restore deficient probability mass.
  arXiv  Detail & Related papers  (2026-08-10T16:53:14Z)
• Distill Where You Fail: Recovering Learning Signals of Negative RL-Groups from Adaptive Teacher Guidance [50.77757355236912]
  We propose RSTG (Recovering Learning Signals via Adaptive Teacher Guidance), which applies distillation selectively and precisely where it matters most.<n> Experiments demonstrate that RSTG substantially outperforms naive GRPO+OPD by +4.02% on math and+3.05% on code.
  arXiv  Detail & Related papers  (2026-08-01T17:29:14Z)
• Rethinking Reward Supervision: Rubric-Conditioned Self-Distillation [60.55792673956761]
  We propose textbfRubric-Conditioned Self-Distillation, a framework that incorporates rubrics as structured, fine-grained feedback for on-policy self-distillation.<n>We evaluate on a diverse suite of science reasoning benchmarks and results show that rubric-conditioned self-distillation effectively converts rubric-level criteria into token-level guidance.
  arXiv  Detail & Related papers  (2026-06-17T17:54:04Z)
• OmniOPD: Logit-Free On-Policy Distillation via Speculative Verification [21.095025092772257]
  On-Policy Distillation (OPD) trains a student model on its own generative trajectories under dense token-level feedback from a stronger teacher.<n>We introduce OmniOPD, a novel framework that addresses both limitations through a logit-free, chunk-level supervision signal.
  arXiv  Detail & Related papers  (2026-05-31T22:31:15Z)
• Your Teacher Can't Help You Here: Combating Supervision Fidelity Decay in On-Policy Distillation [81.10000755917712]
  On-policy distillation transfers reasoning capabilities by training a student model on its own generated trajectories using token-level feedback from a teacher.<n>As student-generated prefixes lengthen, the teacher's next-token distribution becomes less confident and less discriminative.<n>To mitigate SFD, we introduce textbfLookahead Group Reward (ours).
  arXiv  Detail & Related papers  (2026-05-29T04:39:20Z)
• Cornerstones or Stumbling Blocks? Deciphering the Rock Tokens in On-Policy Distillation [4.624042537090342]
  We investigate high-loss tokens, a token type that--as the most direct signal of student-teacher mismatch under On-Policy Distillation's per-token KL objective--should progressively diminish as training converges according to existing studies.<n>We find that these tokens provide negligible functional contribution to the model's actual reasoning performance.
  arXiv  Detail & Related papers  (2026-05-10T01:41:43Z)
• Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes [31.95045602299568]
  On-policy distillation (OPD) is appealing for large language model (LLM) post-training because it evaluates teacher feedback on student-generated rollouts rather than fixed teacher traces.<n>We revisit OPD from the estimator and implementation sides.<n>We identify three failure modes of sampled-token OPD: an imbalanced one-token signal, unreliable teacher guidance on student-generated prefixes, and distortions caused by tokenizer or special-token mismatch.
  arXiv  Detail & Related papers  (2026-03-26T15:35:59Z)
• ToDi: Token-wise Distillation via Fine-Grained Divergence Control [9.958797874295355]
  Token-wise Distillation (ToDi) is a novel method that adaptively combines Forward KL and Reverse KL per token using a sigmoid-based weighting function.<n>ToDi consistently outperforms recent distillation baselines using uniform or less granular strategies.
  arXiv  Detail & Related papers  (2025-05-22T06:51:16Z)
• Shrinking Class Space for Enhanced Certainty in Semi-Supervised Learning [59.44422468242455]
  We propose a novel method dubbed ShrinkMatch to learn uncertain samples. For each uncertain sample, it adaptively seeks a shrunk class space, which merely contains the original top-1 class. We then impose a consistency regularization between a pair of strongly and weakly augmented samples in the shrunk space to strive for discriminative representations.
  arXiv  Detail & Related papers  (2023-08-13T14:05:24Z)
• Unbiased Risk Estimators Can Mislead: A Case Study of Learning with Complementary Labels [92.98756432746482]
  We study a weakly supervised problem called learning with complementary labels. We show that the quality of gradient estimation matters more in risk minimization. We propose a novel surrogate complementary loss(SCL) framework that trades zero bias with reduced variance.
  arXiv  Detail & Related papers  (2020-07-05T04:19:37Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.