Trust-Region Behavior Blending for On-Policy Distillation
作者: Daniil Plyusov, Alexey Gorbatovski, Alexey Malakhov, Nikita Balagansky, Boris Shaposhnikov, Daria Korotyshova, Daniil Gavrilov
分类: cs.LG, cs.AI
发布日期: 2026-05-29
💡 一句话要点
提出Trust-Region Behavior Blending，提升On-policy蒸馏的早期训练效果
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: On-policy蒸馏 知识蒸馏 行为策略 信任区域 数学推理
📋 核心要点
- On-policy蒸馏在训练初期，学生策略较弱，导致教师监督信号作用于低质量数据。
- Trust-Region Behavior Blending (TRB)在KL信任区域内，用更接近教师的行为策略替代学生策略，进行预热。
- 实验表明，TRB在数学推理蒸馏任务中，相较于其他方法，取得了更优的平均性能。
📝 摘要（中文）
On-policy蒸馏(OPD)通过匹配更强的教师模型，利用学生自身策略采样的前缀来训练学生模型。虽然这解决了离线蒸馏的前缀不匹配问题，但早期学生模型的rollout效果可能较差，导致教师模型的监督作用于质量较弱或低质量的前缀上。我们提出了Trust-Region Behavior Blending (TRB)，一种预热方法，在以学生为中心的KL散度信任区域内，用最接近教师模型的行为策略替换早期rollout策略，同时保持每个前缀的反向KL OPD损失不变。KL预算逐渐退火至零，因此训练在预热后恢复为纯粹的学生模型rollout。在两个数学推理蒸馏设置中，TRB在所比较的方法中获得了最强的平均性能。
🔬 方法详解
问题定义：On-policy distillation (OPD) 旨在利用学生自身的策略生成的数据来训练学生模型，从而避免离线蒸馏中由于学生和教师策略差异导致的前缀不匹配问题。然而，OPD 的一个关键痛点在于，在训练初期，学生模型的策略通常较弱，生成的 rollout 数据质量较低。这导致教师模型的监督信号作用于这些低质量的前缀上，从而影响学生模型的学习效率和最终性能。
核心思路：TRB 的核心思路是在训练初期，用一个更接近教师模型的行为策略来替代学生模型的 rollout 策略。具体来说，TRB 维护一个以学生模型为中心的 KL 散度信任区域，并在该区域内寻找最接近教师模型的行为策略。通过使用该行为策略进行 rollout，可以生成更高质量的训练数据，从而更好地利用教师模型的监督信号。随着训练的进行，KL 散度信任区域逐渐缩小，最终恢复为纯粹的学生模型 rollout。
技术框架：TRB 的整体框架可以分为以下几个阶段：1. 初始化： 初始化学生模型和教师模型。2. 行为策略选择： 在以学生模型为中心的 KL 散度信任区域内，选择最接近教师模型的行为策略。3. Rollout： 使用选定的行为策略生成 rollout 数据。4. 损失计算： 计算每个前缀的反向 KL OPD 损失。5. 模型更新： 使用计算得到的损失更新学生模型。6. KL预算退火： 逐渐减小 KL 散度信任区域的大小。7. 重复步骤2-6，直到 KL 预算为零。
关键创新：TRB 的关键创新在于引入了 Trust-Region 的概念，并将其应用于 On-policy 蒸馏的早期训练阶段。通过在 KL 散度信任区域内选择行为策略，TRB 能够有效地利用教师模型的监督信号，同时避免了由于学生模型策略过弱导致的问题。与传统的 OPD 方法相比，TRB 能够更有效地提升学生模型的学习效率和最终性能。
关键设计：TRB 的关键设计包括：1. KL 散度信任区域： 使用 KL 散度来衡量学生模型和行为策略之间的差异，并设置一个 KL 预算来限制行为策略的偏离程度。2. 行为策略选择： 可以使用不同的方法来选择行为策略，例如，可以使用一个额外的神经网络来学习行为策略，或者可以使用教师模型的策略作为行为策略。3. KL 预算退火： 可以使用不同的退火策略来逐渐减小 KL 散度信任区域的大小，例如，可以使用线性退火或指数退火。
🖼️ 关键图片
📊 实验亮点
实验结果表明，在两个数学推理蒸馏任务中，TRB 方法相较于其他基线方法，取得了显著的性能提升。具体来说，TRB 在平均性能上优于所有比较方法，证明了其在提升 On-policy 蒸馏早期训练效果方面的有效性。这些结果表明，TRB 是一种有前景的知识蒸馏方法，具有广泛的应用前景。
🎯 应用场景
该研究成果可应用于各种需要知识迁移和模型压缩的场景，尤其是在计算资源有限或对模型推理速度有较高要求的环境中。例如，可以将大型语言模型蒸馏到小型设备上，或者将复杂的强化学习策略迁移到机器人控制系统中，从而实现更高效、更灵活的智能应用。
📄 摘要（原文）
On-policy distillation (OPD) trains a student on prefixes sampled from its own policy while matching a stronger teacher. This addresses the prefix mismatch of offline distillation, but early student rollouts can still be poor, placing teacher supervision on weak or low-quality prefixes. We propose Trust-Region behavior Blending (TRB), a warmup method that replaces the early rollout policy with the closest-to-teacher behavior policy inside a student-centered KL trust region, while keeping the per-prefix reverse-KL OPD loss unchanged. The KL budget is annealed to zero, so training returns to pure student rollouts after warmup. Across two math-reasoning distillation settings, TRB attains the strongest average among the compared methods.