Search papers, labs, and topics across Lattice.
This paper conducts a token-level analysis of on-policy distillation (OPD) using the K2 estimator of reverse KL divergence, revealing that the gradient updates are disproportionately allocated to low-probability tokens with significant teacher-student log-probability gaps. The authors introduce Surprise-aware Reweighting (SuRe), a novel weighting mechanism that enhances the allocation of updates, leading to improved performance metrics in math tasks without compromising out-of-domain benchmark results. This work not only clarifies the gradient dynamics of OPD but also provides a practical intervention that boosts the effectiveness of the distillation process.
Low-probability tokens disproportionately influence model updates, and a simple reweighting strategy can significantly enhance performance without sacrificing generalization.
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.
