揭示了采样反KL训练中令牌更新的不均衡机制，提出新重加权方法提升效果。
A Token-Level Analysis of Sampled-Token Reverse-KL On-Policy Distillation
- 从梯度层面分析反KL损失对每个令牌的更新分配机制。
- 低概率令牌贡献了大部分梯度，且与教师-学生差异大。
- 提出SuRe重加权法，轻量高效，提升数学任务表现。
在线策略蒸馏（OPD）使用冻结教师模型的令牌级信号监督学生模型在自身轨迹上的学习，但采样损失如何在令牌间分配更新仍不清楚。本文分析了反KL损失的每令牌K2估计器梯度，其ℓ₁范数可分解为教师-学生对数概率差的绝对值和一个随学生对采样令牌概率降低而增大的软最大因子。在数学蒸馏实验中，这些每令牌范数极不均匀：低学生概率令牌占总和比例过高，且教师-学生差距更大。基于此分析，我们提出一种轻量干预方案——惊喜感知重加权（SuRe），一种独立、有界的权重规则，进一步放大该分配。在两个Qwen3学生规模上，SuRe在多个数学指标上优于基线OPD，且在选定域外基准上无明显退化。主要贡献是首次以梯度层级刻画采用K2估计器的反KL OPD，SuRe为其一个具体实现。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。