Escaping the KL Agreement Trap in On-Policy Distillation

📄 arXiv: 2606.09471v1 📥 PDF

作者: Haoran Xin, Anhao Zhao, Ying Sun, Jin Li, Xiaoyu Shen, Hui Xiong

分类: cs.LG, cs.CL

发布日期: 2026-06-08

备注: 13 pages, 8 figures


💡 一句话要点

提出KAT以解决在线策略蒸馏中的低KL一致性陷阱问题

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

关键词: 在线策略蒸馏 KL一致性 动态阈值 强化学习 监督信号

📋 核心要点

  1. 现有的在线策略蒸馏方法在学生模型进入低KL一致性陷阱时,无法提供有效的纠正信号,导致训练效果不佳。
  2. 论文提出KAT,通过动态阈值检测低KL一致性,及时终止无效的训练过程,从而提升监督信号的质量。
  3. 实验结果显示,KAT在多个数学基准上显著提高了模型的准确率和通过率,同时减少了回合长度,提升了训练效率。

📝 摘要(中文)

在线策略蒸馏(OPD)通过教师对学生生成的回合进行评分,提供了密集的标记级监督。然而,当学生进入不可恢复的前缀时,教师可能会对退化状态产生局部一致性,导致低反向KL但缺乏有效的纠正训练信号。我们将这种持续状态称为低KL一致性陷阱。进一步分析表明,在此类陷阱期间及之后的标记产生的监督信号效用较低。为此,我们提出了KAT(KL一致性陷阱终止),这是一种在线OPD终止规则,通过动态训练自适应阈值检测持续的低KL一致性。通过过滤来自退化一致性的弱监督,KAT在四个数学基准上提高了avg@k准确率2.66%和pass@k 3.43%,同时将平均回合长度减少了59.73%。

🔬 方法详解

问题定义:论文要解决的问题是在线策略蒸馏中出现的低KL一致性陷阱,这种情况导致教师与学生模型在低效状态下达成一致,缺乏有效的监督信号。现有方法在此情况下无法提供有效的纠正训练信号,影响模型性能。

核心思路:论文的核心解决思路是引入KAT机制,通过动态调整的阈值来检测和终止低KL一致性状态,从而避免学生模型在无效状态下继续训练,提升监督信号的有效性。

技术框架:整体架构包括教师模型和学生模型的交互过程,KAT机制作为监测模块,实时评估KL一致性并决定是否终止训练。主要阶段包括模型训练、低KL检测和监督信号过滤。

关键创新:最重要的技术创新点在于KAT的动态阈值设计,使得模型能够根据训练过程中的反馈自适应调整,区别于传统的静态阈值方法,从而更有效地识别和终止低效训练。

关键设计:在KAT中,设置了动态阈值以适应不同训练阶段,损失函数设计上强调对有效监督信号的强化,同时优化了教师与学生模型的交互机制,以提高整体训练效率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,KAT在四个数学基准上显著提高了avg@k准确率2.66%和pass@k 3.43%。同时,KAT还将平均回合长度减少了59.73%,展示了其在提升训练效率方面的显著效果。

🎯 应用场景

该研究的潜在应用领域包括强化学习、机器人控制和自动驾驶等需要高效训练的场景。通过提高训练过程中的监督信号质量,KAT能够显著提升模型的学习效率和性能,具有广泛的实际价值和未来影响。

📄 摘要(原文)

On-policy distillation (OPD) provides dense token-level supervision by asking a teacher to score student-generated rollouts. However, when the student drifts into an unrecoverable prefix, the teacher may locally agree with the degraded state, producing low reverse KL but little corrective training signal. We identify this persistent regime as a low-KL agreement trap. Further analyses show that tokens during and after such traps produce less useful supervision signals. We propose KAT (KL Agreement Trap Termination), an online OPD termination rule that detects persistent low-KL agreement with a dynamic training-adaptive threshold. By filtering weak supervision from degenerate agreement, KAT improves avg@k accuracy by 2.66% and pass@k by 3.43% across four mathematical benchmarks, while reducing average rollout length by 59.73%.