ReTaCo：用于在线策略蒸馏的残差目标控制

中文摘要 AI 辅助

在线策略蒸馏（OPD）在由学生自身生成的前缀上训练学生，并利用词元级别的教师反馈，但传输或存储教师在每个词元上的全词汇分布代价高昂。熵感知的在线策略蒸馏（EOPD）在反向KL散度上增加前向监督，以帮助学生恢复其低估的合理词元，仅使用教师的top-$k$概率以限制成本。由于EOPD对这些概率进行重新归一化，其目标对省略的词汇不分配质量。我们证明，即使学生在top-$k$集合内匹配了教师的相对概率，所得损失仍会持续推动学生的top-$k$质量趋向于1，因此只要被省略的词元具有正的教师概率，教师本身就不是一个不动点。我们提出ReTaCo（残差目标控制），它单独保留top-$k$词元，并将剩余词元分组为一个残差符号，同时将此前向目标与一个单样本估计器配对，该估计器的期望等于全词汇反向KL散度。对于教师的top-$k$质量$m$，残差目标为$(1-\beta)(1-m)$，其中$\beta\in[0,1]$：$\beta=0$保留教师的质量，更大的$\beta$将更多质量转移到top-$k$词元上而不改变其相对概率。在固定前缀下，我们证明总体目标具有唯一最优解，其top-$k$质量位于$m$和$m+\beta(1-m)$之间，并随$\beta$单调增加；在$\beta=0$时，被低估的top-$k$词元仍获得不消失的恢复梯度。数值优化证实了这些预测，并且在三组教师-学生配对中，ReTaCo在大多数数学和代码基准上优于EOPD。

英文摘要

On-policy distillation (OPD) trains a student on its own generated prefixes with token-level teacher feedback, but transmitting or storing the teacher's full-vocabulary distribution at every token is costly. Entropy-aware OPD (EOPD) adds forward supervision to reverse KL to help the student recover plausible tokens it underestimates, using only the teacher's top-$k$ probabilities to limit cost. Because EOPD renormalizes these probabilities, its target assigns no mass to the omitted vocabulary. We prove that the resulting loss keeps pushing the student's top-$k$ mass toward one even after the student matches the teacher's relative probabilities within the top-$k$ set, so the teacher itself is not a stationary point whenever the omitted tokens have positive teacher probability. We propose ReTaCo (Residual-Target Control), which keeps the top-$k$ tokens individually and groups the remaining tokens into one residual symbol, and pairs this forward target with a single-sample estimator whose expectation equals the full-vocabulary reverse KL. With teacher top-$k$ mass $m$, the residual target is $(1-β)(1-m)$ for $β\in[0,1]$: $β=0$ preserves the teacher's mass, and larger $β$ moves more mass onto the top-$k$ tokens without changing their relative probabilities. At a fixed prefix, we prove that the population objective has a unique optimum whose top-$k$ mass lies between $m$ and $m+β(1-m)$ and increases monotonically with $β$; at $β=0$, underestimated top-$k$ tokens still receive non-vanishing recovery gradients. Numerical optimization confirms these predictions, and across three teacher-student pairs, ReTaCo outperforms EOPD on most mathematics and code benchmarks.

发表机构

• Zhejiang University（浙江大学）
• Georgia Institute of Technology（佐治亚理工学院）
• Boston College（波士顿学院）
• Yale University（耶鲁大学）
• University of California, San Diego（加利福尼亚大学圣迭戈分校）
• Tsinghua University（清华大学）
• Shanghai Innovation Institute（上海创新研究院）
• Massachusetts Institute of Technology（麻省理工学院）

机构由 AI 辅助整理，请以论文原文为准。