Problem
Research questions and friction points this paper is trying to address.
On-policy Distillation
Sampled Loss
Token-level Analysis
Reverse KL
A Token-Level Analysis of Sampled-Token Reverse-KL On-Policy Distillation
🤖 AI Summary
本文分析了基于采样令牌的反向KL在策略蒸馏中的梯度分配问题，并提出了一种名为SuRe的方法来改进这一过程，从而提高数学指标。
📝 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.
Innovation
Methods, ideas, or system contributions that make the work stand out.
On-policy Distillation
Reverse-KL
K2 Estimator
Surprise-aware Reweighting
Token-level Analysis
🔎 Similar Papers
No similar papers found.
💼 Related Jobs
