Not Every Token Is Worth Distilling: Selective Supervision for Direct-OPD

📄 arXiv: 2609.29142v1 📥 PDF

作者: Yibo Zhao, Zixuan Yang, Yunshi Lan, Xiang Li

分类: cs.LG, cs.AI

发布日期: 2026-09-24

备注: 19 pages. Yibo Zhao and Zixuan Yang are equal contributors and may list their names in either order on their CVs

💡 一句话要点

提出选择性监督方法S$^2$D-OPD以提升直接策略蒸馏效果

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 直接策略蒸馏 选择性监督 强化学习 Jensen-Shannon散度 策略转移

📋 核心要点

  1. 现有的Direct-OPD方法在处理低散度状态时，可能导致无效的监督信号，从而影响策略学习效果。
  2. 本文提出的S$^2$D-OPD方法通过选择性监督，优先关注高散度状态，以提高策略蒸馏的有效性。
  3. 实验结果显示，S$^2$D-OPD在多个模型和基准测试中，准确率在七个设置中超过了密集Direct-OPD，且在第八个设置中持平。

📝 摘要（中文）

直接策略蒸馏（Direct-OPD）通过使用策略改进的token级对数比作为密集监督，将小模型的强化学习策略转移到大模型。然而，该方法的对数比只反映相对变化，可能在概率质量消失时保持不变。本文提出选择性监督方法S$^2$D-OPD，根据教师参考的Jensen-Shannon散度（JSD）对学生采样状态进行排名，并在低散度状态下屏蔽Direct-OPD监督，仅保留前10%的状态。实验表明，S$^2$D-OPD在多个基准测试中相较于密集Direct-OPD显著提高了准确率。

🔬 方法详解

问题定义：本文旨在解决Direct-OPD方法在低散度状态下的监督信号无效问题，导致策略学习效果不佳。现有方法未能有效利用教师模型的行为变化。

核心思路：提出选择性监督方法S$^2$D-OPD，通过根据教师参考的JSD对学生状态进行排名，屏蔽低散度状态的监督，保留高散度状态的有效信息，从而提升策略蒸馏效果。

技术框架：S$^2$D-OPD的整体架构包括状态采样、JSD计算和监督筛选三个主要模块。首先从学生模型中采样状态，然后计算与教师模型的JSD，最后根据JSD值筛选出前10%的状态进行监督。

关键创新：S$^2$D-OPD的核心创新在于引入了选择性监督机制，通过JSD评估状态的重要性，从而优化了策略蒸馏过程。这一方法与传统的密集监督方法相比，能够更有效地利用教师模型的信息。

关键设计：在参数设置上，S$^2$D-OPD保留了前10%的高散度状态，损失函数设计上与Direct-OPD相似，但通过选择性监督提高了训练效率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在实验中，S$^2$D-OPD在AIME和HMMT基准测试中，在七个设置中超越了密集Direct-OPD，且在第八个设置中表现持平，显示出显著的性能提升，验证了选择性监督的有效性。

🎯 应用场景

该研究的潜在应用领域包括强化学习模型的训练，尤其是在需要高效策略转移的场景，如机器人控制、游戏智能体和自动驾驶等。通过优化策略蒸馏过程，S$^2$D-OPD能够在资源有限的情况下提升模型性能，具有重要的实际价值和未来影响。

📄 摘要（原文）

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. Our code is available at https://anonymous.4open.science/r/S2D-OPD-8868.