Influence-Directed Distillation: Solving the Diversity Bottleneck in Sampled-Token On-Policy Distillation
Run Yang, Runpeng Dai, Jie Sun, Jielei Zhang, Fan Zhou, Hongtu Zhu, Peiyi Li, Longwen Gao
cs.CL, cs.LG
2026-08-30
采样token版on-policy蒸馏经常让pass@1上升、pass@k停滞,学生继承不到教师多样性。IDA-OPD按熵影响收缩降熵更新,四B模型AIME24的pass@16从78.7升到83.3。
On-policy蒸馏(OPD)已经进了GLM-5、Kimi、Qwen3这类后训练流水线。完整版要对齐教师在整个词表上的下一token分布,长轨迹上存logits太贵,于是工业上常用采样token版:学生自己滚,教师只回传被采到的那个token的对数概率。便宜的代价是多样性蒸馏失败:学生的pass@1上去了,pass@k却卡住,继承不到教师的多解能力。
已有补丁分两派。教师知情方法在高分歧位置再要top-K分布做正向KL,精度高、成本也高。学生侧加熵奖励或改优势形状更便宜,但干预太钝,降熵的更新和涨熵的更新一起被拧。
作者先把一次采样token更新对熵的一阶影响写成IH(y)=Ay·Dy。Ay是教师与学生在该token上的对数概率差,也就是这条更新的优势;Dy完全由学生当前的局部概率结构决定。同号的Ay可以涨熵也可以降熵:强化一个已经很高的token会把分布削尖,强化一个低概率候选则可能把分布摊平。真实训练位置上,Ay和一步熵变几乎不成单边关系,IH(y)的符号却能对上。
熵流失并不主要来自师生严重分歧的token。按归一化分歧δy分箱后,累积降熵峰值落在δy≈0:单次更新幅度很小,但这类token数量极大,堆起来才是主因。高分歧负尾也有一份,不是全部。
IDA-OPD因此只动IH(y)<0的位置。涨熵更新原样保留;降熵更新把优势乘上权重wy=|qy−py|/(qy+py)。师生已经接近时wy接近0,分歧大时wy接近1,方向不翻转。教师仍然只提供被采样token的一个对数概率,不需要全词表或top-K。
数学上两条线:Qwen3-8B/4B的RL-Math教师蒸到同尺寸Non-Thinking学生,数据是DeepMath103K里难度6。代码上用Qwen3-4B-RLCode教师。每个题采n=128条,用无偏估计报pass@1和pass@16。
| 设定 | 方法 | AIME24 pass@1/16 | AIME25 pass@1/16 |
| 8B | OPD | 61.7 / 79.1 | 47.9 / 70.7 |
| 8B | IDA-OPD | 63.3 / 83.3 | 50.0 / 76.7 |
| 8B | 教师 | 60.0 / 83.3 | 53.5 / 73.3 |
| 4B | OPD | 54.6 / 78.7 | 51.8 / 65.7 |
| 4B | IDA-OPD | 56.7 / 83.3 | 53.3 / 70.0 |
| 4B | 教师 | 53.3 / 80.0 | 53.5 / 73.3 |
4B上相对OPD,HMMT Feb的pass@16 +8.3,AIME24 +4.6,AIME25 +4.3。8B对应为+7.2、+4.2、+6.0。若干pass@16达到或超过教师,例如4B AIME24为83.3对教师80.0。需要top-K教师分布的AOPD/EOPD多样性也不错,IDA-OPD在表里的pass@16总体最高,成本仍是每位置一个标量。代码域MBPP+上pass@1/16从OPD的69.5/72.9到71.6/73.8,LiveCodeBench从26.8/54.2到28.1/55.2,方向相同,幅度更小。
消融里拿掉IH门控、把降熵更新直接置零、或改用sign(Ay)当门,4B AIME24的pass@1分别掉到53.8、52.1、54.2。线性wy比常数、开方、平方映射都更稳。
采样token OPD的效率是它能进大模型流水线的原因。这篇把多样性失败追到「低分歧、降熵、数量极大」的更新,并用学生侧已经算过的logits做门控,不必把全词表教师分布请回来。对已经在跑OPD、却发现pass@k不跟pass@1一起涨的团队,这是一个可插的改动。
pass@1多数时候只是略增或持平,主收益在多样性。若线上只看单次正确率,体感会偏小。
一阶分析假设小步logit梯度,作者用AdamW轨迹说明符号仍然可用,但IH(y)仍是代理,不是逐步无偏估计。主表几乎全是Qwen3非思考模型蒸Qwen3,跨家族、跨tokenizer没有主结果。代码域pass@16只涨大约1个点,故事主要建立在竞赛数学上。AIME每套30题,pass@k对个别题很敏感。没有报告对指令跟随或安全性的副作用。论文也没有单独的局限节,上述是读实验设置能直接看到的边界。