发表机构
East China Normal University; Hugging Face(华东师范大学; Hugging Face)
机构由 AI 辅助整理,请以论文原文为准。AI 中文总结
针对Direct-OPD中低散度状态监督无效的问题,提出选择性监督方法S$^2$D-OPD,仅保留前10%高散度状态进行蒸馏,在多个模型和基准上提升准确率且无需额外计算。
AI 中文摘要
直接在线策略蒸馏(Direct-OPD)通过使用强化学习前后检查点之间的token级对数比率作为学生自身轨迹上的密集监督,将强化学习诱导的策略改进从一个小模型迁移到一个更大的学生模型。这种迁移在每个状态上奖励策略偏移,然而对数比率仅衡量相对变化:即使两个检查点分配给学生候选token的概率质量消失,它也可能保持不变。通过精确构造,我们证明当检查点之间的Jensen-Shannon散度(JSD)和两个KL方向随该质量消失时,Direct-OPD的奖励及其更新可以保持不变,并指出小的JSD限制了教师行为变化的程度。受此分析启发,我们提出了面向Direct-OPD的选择性监督(S$^2$D-OPD),该方法根据教师参考JSD对学生采样状态进行排序,并在低散度状态下掩蔽Direct-OPD监督,仅保留每个响应中前10%的状态。在两组教师配对和四个从1.7B到8B参数的学生模型上,S$^2$D-OPD在AIME和HMMT基准上,在八个设置中的七个中提高了保留集上的准确率,并在第八个设置中与之持平,且无需额外的前向传播。我们的代码可在https URL获取。
英文摘要
Direct On-Policy Distillation (Direct-OPD) transfers reinforcement-learning-induced policy improvements from a small model to a larger student by using the token-level log-ratio between post-RL and pre-RL checkpoints as dense supervision on the student's own rollouts. This transfer rewards the policy shift at every state, yet the log-ratio measures only relative change: it can stay fixed even as the probability mass that both checkpoints assign to the student's candidate tokens vanishes. Through an exact construction, we show that the Direct-OPD reward and its update can remain unchanged while the Jensen-Shannon divergence (JSD) and both KL directions between the checkpoints vanish with this mass, and we note that a small JSD bounds how much the teacher's behavior changed. Motivated by this analysis, we propose Selective Supervision for Direct-OPD (S$^2$D-OPD), which ranks student-sampled states by their teacher-reference JSD and masks Direct-OPD supervision at low-divergence states, retaining only the top 10% of states per response. Across two teacher pairs and four student models ranging from 1.7B to 8B parameters, S$^2$D-OPD improves held-out accuracy over dense Direct-OPD on AIME and HMMT benchmarks in seven of eight settings and matches it in the eighth, without extra forward passes.
Comments19 pages. Yibo Zhao and Zixuan Yang are equal contributors and may list their names in either order on their CVs