跳到主要内容
返回时间线
arXiv来源发表:

只用1%的token就能追平全量蒸馏:IER用梯度估计可靠性重排稀疏监督

核心概要

该工作把在线策略蒸馏(OPD)中的token选择问题重新表述为固定前缀下的梯度估计可靠性问题,在信息几何下把单样本反向KL梯度分解为信号与采样噪声,提出信息效率比(IER)作为最优标量基线下的信噪比,并用top-K候选集近似来排序token;在数学与医学推理任务上,IER单独作为选择器在0.1%–1%的极小token预算下接近或超过全量OPD,与已有有用性分数结合后(IER-OR、IER-AND)在多个设置下匹配甚至超过全量OPD。

AI-generated editorial illustration: 1% of Tokens Can Be Enough: On Gradient Estimation in On-Policy Distillation

深度剖析

论文指出稀疏OPD的一个被忽略维度:即使某个token的教师监督是有用的,用单个采样到的下一个token去估计反向KL梯度时仍可能产生高噪声更新,因此token选择应同时考虑“有用性”和“梯度估计可靠性”。 此前OPD的token选择方法(Prefix、Entropy、TIP、TA-OPD、CA-SoftOR等)主要围绕有用性、重要性、可教性或位置设计指标,作者称这些工作“mostly emphasize different aspects of token usefulness and ignore the estimation reliability”,本文把可靠性作为互补的第二个轴提出。 该判断由理论推导支撑:在固定前缀下引入与动作无关的标量基线,给出梯度信号与均方估计方差的分解,并求出最小化方差的最优基线;作者同时说明高IER并不代表监督高有用性。

作者提出信息效率比(IER):在最优标量基线下,期望梯度信号与最小化采样误差之比,其倒数对应单样本token梯度估计的相对均方误差。 与vOPD等针对单样本OPD梯度方差做控制变量/方差缩减的工作不同,本文从信息几何(Fisher度量下的自然梯度)角度分析单样本反向KL梯度,并给出信号-噪声分解与最优基线闭式解。 结论以定理形式给出(Theorem 1),并在附录中给出证明;推论1进一步把IER的倒数与独立重复采样平均估计的相对MSE联系起来,指出达到给定相对MSE所需的采样复杂度。

为降低计算成本,作者用学生与教师top-K logits并上采样token构成候选集来近似IER,得到排序代理,并可与已有有用性分数通过软OR(IER-OR)和软AND(IER-AND)组合,同时保留采样的反向KL训练目标。 这把一个全词表统计量转化为可在每个前缀上实际计算的排序分数,使可靠性信号能直接插入现有稀疏OPD流程,而不是替换训练目标。 实现细节在附录中给出:top-K取K=10,对候选集中缺失logit赋小值并重归一化,对对数似然比做裁剪以保证数值稳定;训练时每个rollout batch含4个prompt、每个prompt 16条学生回复,选择器分数在batch内归一化并共享一个全局预算,每条回复至少保留一个token。

在数学推理(JustRL-Nemotron-1.5B→OpenMath-Nemotron-1.5B的强到弱蒸馏、JustRL-Qwen3-4B→Qwen3-1.7B的大到小蒸馏)与医学推理(ClinAlign-4B→Qwen3-4B)上,IER及其组合在0.1%–1%的极小token预算下匹配或超过全量OPD,且更高预算并不总是带来更好表现。 作者报告IER单独使用时在0.1%预算下AIME26为58.9对全量OPD 59.9、HMMT26为34.1对34.7,并在另一模型对上于四个基准中的三个超过全量OPD;医学推理中0.1%预算下IER的HealthBench总分为45.25,接近全量OPD的45.77,而Prefix在该预算下“provides essentially no improvement”。 证据来自多任务、多模型对、多预算(0.1%、1%、10%,并扫描5%、20%、50%、80%)与thinking-on/thinking-off两种模式的实验;作者也明确说明增益依赖选择器与设置,基于Entropy的组合在数学推理上有时反而下降,且两种分数都是近似,不保证选择更优。

启示与展望

该结果面向使用教师模型对学生自生成轨迹逐token打分的在线策略蒸馏场景,适用于数学这类可验证推理任务与医学这类开放式推理任务,覆盖强到弱(同规模、教师经过后训练)与大到小(教师更大更强、共享分词器)两类蒸馏设置,以及thinking-on与thinking-off两种模式。它给出的是一套可插入现有选择器的可靠性排序信号:IER可单独作为选择器,也可通过IER-OR(有用或可靠)与IER-AND(既有用又可靠)与Prefix、Entropy、TIP、TA-OPD、CA-SoftOR等有用性分数组合,训练目标仍保持采样的反向KL。作者指出,基于IER的截断式设计(当已积累足够多可靠梯度token时提前剪枝轨迹)是提升训练效率的可行方向,留作未来工作。

作者自己指出,如何最好地选择与加权token仍是开放问题:增益随选择器与设置变化,基于Entropy的组合在数学推理上有时下降,且有用性分数与IER都是近似,不保证选出更优token。IER衡量的是相对估计精度而非梯度幅度,作者明确说明高IER不保证梯度幅度、任务对齐或参数更新的收益,因此不能把IER当作有用性的替代。此外,本文实现仍生成完整轨迹并额外计算token评分与选择,稀疏监督并不带来成比例的计算节省;训练效率表显示IER及其组合相对全量OPD增加约2.2%–2.5%的步时与最多2.17 GiB(0.92%)的峰值显存。医学推理的评分由gpt-oss-120B作为评判模型完成,其在HealthBench元评估上的宏F1为0.6614,低于GPT-4.1,作者说明这是出于成本考虑。最后,本次读取的markdown中数学推理主结果表(Table 2)的数值单元格为空,因此无法逐项核对各选择器在AIME25/26与HMMT25/26上的具体分数,只能依据正文叙述与医学推理表格;若需要精确对比,应回到原文表格。

来源