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

arXiv: https://arxiv.org/abs/2608.25643 · HTML: https://mlsanity.com/paper/2608.25643/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Placed in On-policy and self-distillation for LLMs, one of 12 directions in Knowledge distillation.
- Builds on 13 of its 26 references that are in the map.
- First posted on arXiv 26 Aug 2026.

## Placed in

- Knowledge distillation / [On-policy and self-distillation for LLMs](https://mlsanity.com/distill/on-policy-and-self-distillation-for-llms/)

## Builds on (most cited first, 13 of 13)

- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [MMLU-Pro: A More Robust and Challenging Multi-Task Language Understanding Benchmark](https://mlsanity.com/paper/2406.01574/) (2024)
- [On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes](https://mlsanity.com/paper/2306.13649/) (2023)
- [Sequence-Level Knowledge Distillation](https://mlsanity.com/paper/1606.07947/) (2016)
- [Beyond the 80/20 Rule: High-Entropy Minority Tokens Drive Effective Reinforcement Learning for LLM Reasoning](https://mlsanity.com/paper/2506.01939/) (2025)
- [CRUXEval: A Benchmark for Code Reasoning, Understanding and Execution](https://mlsanity.com/paper/2401.03065/) (2024)
- [Rethinking On-Policy Distillation of Large Language Models: Phenomenology, Mechanism, and Recipe](https://mlsanity.com/paper/2604.13016/) (2026)
- [DeepMath-103K: A Large-Scale, Challenging, Decontaminated, and Verifiable Mathematical Dataset for Advancing Reasoning](https://mlsanity.com/paper/2504.11456/) (2025)
- [Learning beyond Teacher: Generalized On-Policy Distillation with Reward Extrapolation](https://mlsanity.com/paper/2602.12125/) (2026)
- [Entropy-Aware On-Policy Distillation of Language Models](https://mlsanity.com/paper/2603.07079/) (2026)
- [Scaling Reasoning Efficiently via Relaxed On-Policy Distillation](https://mlsanity.com/paper/2603.11137/) (2026)
- [On the Direction of RLVR Updates for LLM Reasoning: Identification and Exploitation](https://mlsanity.com/paper/2603.22117/) (2026)

## Built on by (most cited first, 6 of 6)

- [What Shared Prefixes Hide: Trajectory Dropout for On-Policy Distillation](https://mlsanity.com/paper/2609.33455/) (2026)
- [Do We Really Need KL Divergence for On-Policy Distillation of Large Language Models?](https://mlsanity.com/paper/2609.33791/) (2026)
- [Teacher-Student Gaps Are Not Enough: Outcome-Guided On-Policy Distillation for Multi-Turn Autonomous Agents](https://mlsanity.com/paper/2609.35319/) (2026)
- [Dr. OPD: Learning What to Follow for Optimal On-Policy Distillation of Large Language Models](https://mlsanity.com/paper/2609.38025/) (2026)
- [DIAL-OPD: Learning More from Fewer Tokens in On-Policy Distillation](https://mlsanity.com/paper/2610.11659/) (2026)
- [Beyond Domain-Routed Supervision: Cross-Teacher Guarding for Multi-Teacher On-Policy Distillation](https://mlsanity.com/paper/2610.13521/) (2026)

## Abstract (arXiv)

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 ℓ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.
