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

On-policy distillation (OPD) is a standard technique for transferring knowledge from a large teacher model to a smaller student by training on the student's own trajectories. While effective, the underlying token-level update mechanism remains poorly understood. This paper investigates how the per-token K2 estimator of reverse KL distributes gradient norms across sampled tokens and whether this distribution can be optimized to improve student performance.

The authors derive a closed-form gradient identity for the K2 estimator of reverse KL with respect to student logits. They show that the L1 norm of this gradient factorizes into the absolute teacher-student log-probability gap and a student-side softmax factor that increases as the sampled token becomes less likely under the student. Using this insight, they propose Surprise-aware Reweighting (SuRe), a simple, detached, and bounded weighting rule that increases the relative importance of tokens where the student model assigns low probability.

Empirical analysis on Qwen3 models shows that gradient norms are highly non-uniform, with low-probability tokens accounting for a disproportionate share of the total update budget. SuRe, which amplifies these surprise-aware tokens, consistently outperforms vanilla OPD on several mathematical reasoning benchmarks, including AIME24 and AMC23. The authors also perform extensive ablations, showing that while the specific orientation of the weighting matters, the gains are robust to several implementation variations.

This work provides a formal gradient-level characterization of a widely used distillation technique. By identifying that current OPD methods implicitly focus on specific token types, the authors offer a principled, lightweight intervention that improves reasoning capabilities without requiring additional reference models or complex training architectures. This highlights the importance of token-level gradient dynamics in the post-training of large language models.