Learning beyond Teacher: Generalized On-Policy Distillation with Reward Extrapolation
作者: Wenkai Yang, Weijie Liu, Ruobing Xie, Kai Yang, Saiyong Yang, Yankai Lin
分类: cs.LG, cs.AI, cs.CL
发布日期: 2026-02-12
备注: Work in progress. Github repo: https://github.com/RUCBM/G-OPD
💡 一句话要点
提出广义On-Policy蒸馏框架G-OPD,通过奖励外推提升学生模型性能,甚至超越教师模型。
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)
关键词: On-Policy蒸馏 奖励外推 知识融合 模型压缩 强化学习 KL散度 代码生成 数学推理
📋 核心要点
- 现有On-Policy蒸馏方法中,奖励函数和KL正则化权重固定,限制了学生模型的学习能力。
- 提出广义On-Policy蒸馏框架G-OPD,通过引入奖励缩放因子和灵活参考模型,解耦奖励和正则化。
- 实验表明,奖励外推(ExOPD)能有效提升学生模型性能,甚至超越教师模型,尤其在知识融合场景下。
📝 摘要(中文)
本文提出广义On-Policy蒸馏(G-OPD)框架,该框架扩展了标准OPD目标,引入了灵活的参考模型和奖励缩放因子,用于控制奖励项与KL正则化之间的相对权重。理论上,本文首先证明OPD是密集KL约束强化学习的一个特例,其中奖励函数和KL正则化始终等权重,且参考模型可以是任何模型。通过在数学推理和代码生成任务上的综合实验,本文得出两个新颖的见解:(1)将奖励缩放因子设置为大于1(即奖励外推),称为ExOPD,在各种教师-学生模型大小配对中,始终优于标准OPD。特别是,在将来自不同领域专家的知识合并回原始学生模型时,ExOPD使学生甚至能够超越教师的性能边界,并优于领域教师。(2)在强到弱的蒸馏设置中,通过选择教师的RL前模型作为参考模型进行奖励校正,可以产生更准确的奖励信号,并进一步提高蒸馏性能。然而,这种选择需要访问教师的预RL变体,并产生更多的计算开销。希望本文的工作能为OPD的未来研究提供新的见解。
🔬 方法详解
问题定义:现有On-Policy蒸馏(OPD)方法将学生模型与教师模型的logit分布对齐,在提升学生模型性能方面表现出色。然而,标准OPD方法中,奖励函数和KL正则化项的权重被固定为相等,限制了学生模型探索更优策略的能力,尤其是在教师模型并非最优的情况下。此外,参考模型的选择也较为单一,缺乏灵活性。
核心思路:本文的核心思路是通过解耦奖励函数和KL正则化项的权重,并引入更灵活的参考模型,来扩展标准OPD。具体而言,通过引入奖励缩放因子,允许对奖励项进行放大或缩小,从而使学生模型能够更积极地探索潜在的更优策略(奖励外推)。同时,允许使用不同的参考模型,例如教师模型的预训练版本,以提供更准确的奖励信号。
技术框架:G-OPD框架的核心在于修改了标准的OPD目标函数。原始OPD目标函数可以看作是KL约束的强化学习,其中奖励和KL正则化权重相等。G-OPD通过引入奖励缩放因子λ来调整奖励项的权重,目标函数变为优化学生模型的策略,使其在考虑与参考模型的KL散度的同时,最大化缩放后的奖励。此外,G-OPD允许使用任意模型作为参考模型,而不仅仅是教师模型。
关键创新:本文最重要的技术创新点在于提出了奖励外推(ExOPD)的概念,即通过将奖励缩放因子设置为大于1的值,来放大奖励信号,鼓励学生模型探索超出教师模型策略范围的更优策略。这种方法在知识融合场景下尤为有效,可以使学生模型超越多个教师模型的性能。
关键设计:G-OPD的关键设计包括:1) 奖励缩放因子λ的选择:实验表明,λ>1的奖励外推通常能带来更好的性能。2) 参考模型的选择:在强到弱的蒸馏设置中,使用教师模型的预训练版本作为参考模型可以提供更准确的奖励信号。3) 损失函数:G-OPD使用KL散度作为正则化项,以约束学生模型的策略与参考模型的策略之间的差异。
🖼️ 关键图片
📊 实验亮点
实验结果表明,ExOPD在数学推理和代码生成任务上均优于标准OPD。特别是在知识融合场景下,ExOPD使学生模型能够超越教师模型的性能边界。例如,在特定任务上,使用ExOPD的学生模型性能提升超过5%,甚至超越了多个领域专家教师的平均水平。
🎯 应用场景
该研究成果可应用于各种需要知识迁移和模型压缩的场景,例如:将大型语言模型的知识迁移到小型模型,提升移动端设备的AI能力;融合多个领域专家的知识,训练出更强大的通用模型;通过奖励外推,使学生模型超越教师模型,实现更高效的知识学习。
📄 摘要(原文)
On-policy distillation (OPD), which aligns the student with the teacher's logit distribution on student-generated trajectories, has demonstrated strong empirical gains in improving student performance and often outperforms off-policy distillation and reinforcement learning (RL) paradigms. In this work, we first theoretically show that OPD is a special case of dense KL-constrained RL where the reward function and the KL regularization are always weighted equally and the reference model can by any model. Then, we propose the Generalized On-Policy Distillation (G-OPD) framework, which extends the standard OPD objective by introducing a flexible reference model and a reward scaling factor that controls the relative weight of the reward term against the KL regularization. Through comprehensive experiments on math reasoning and code generation tasks, we derive two novel insights: (1) Setting the reward scaling factor to be greater than 1 (i.e., reward extrapolation), which we term ExOPD, consistently improves over standard OPD across a range of teacher-student size pairings. In particular, in the setting where we merge the knowledge from different domain experts, obtained by applying domain-specific RL to the same student model, back into the original student, ExOPD enables the student to even surpass the teacher's performance boundary and outperform the domain teachers. (2) Building on ExOPD, we further find that in the strong-to-weak distillation setting (i.e., distilling a smaller student from a larger teacher), performing reward correction by choosing the reference model as the teacher's base model before RL yields a more accurate reward signal and further improves distillation performance. However, this choice assumes access to the teacher's pre-RL variant and incurs more computational overhead. We hope our work offers new insights for future research on OPD.