Learning beyond Teacher: Generalized On-Policy Distillation with Reward Extrapolation

在线策略蒸馏（OPD）通过在学生模型自主生成的轨迹上，使其 logits 分布与教师模型对齐，已在提升学生模型性能方面展现出显著的实证效果，且通常优于离线策略蒸馏及强化学习（RL）范式。本文首先从理论上证明：OPD 实质上是稠密 KL 约束强化学习（dense KL-constrained RL）的一种特例——其中奖励函数与 KL 正则项始终以相等权重进行加权，且参考模型可为任意模型。随后，我们提出广义在线策略蒸馏（G-OPD）框架：该框架在标准 OPD 目标函数基础上，引入一个灵活可选的参考模型，并增设一个奖励缩放因子（reward scaling factor），用于动态调节奖励项相对于 KL 正则项的相对权重。我们在数学推理与代码生成任务上开展了系统性实验，由此获得两项新发现：（1）将奖励缩放因子设为大于 1（即实施“奖励外推”，我们将其命名为 ExOPD），可在多种教师–学生模型尺寸组合下持续超越标准 OPD。尤其值得注意的是，在融合多个领域专家知识的设定中——这些专家由对同一学生模型分别施加领域专属强化学习所得到——ExOPD 能够将各领域专家的知识回传整合至原始学生模型，使学生性能不仅突破原有教师模型的性能上限，甚至反超各领域教师。（2）在“强→弱”蒸馏场景（即用较大教师模型蒸馏较小的学生模型）下，若进一步基于 ExOPD 框架，将参考模型设为教师模型在强化学习微调前的基座模型（base model），从而实施奖励校正（reward correction），则可获得更准确的奖励信号，并进一步提升蒸馏效果。然而，这一选择依赖于对教师模型强化学习前版本的访问权限，因而带来额外的计算开销。我们期望本工作能为未来在线策略蒸馏的相关研究提供新的理论视角与实践启示。

论文旨在解决监督式知识蒸馏在强化学习（RL）后大模型（如数学推理、代码生成教师模型）向学生模型迁移时性能受限的问题；特别关注标准On-Policy Distillation（OPD）为何有效、能否被理论统一，并进一步突破其性能边界——尤其是当学生弱于教师（strong-to-weak distillation）或需融合多领域专家知识时，如何超越教师性能上限。该问题虽植根于经典蒸馏与RL对齐，但首次系统性提出‘奖励外推’（reward extrapolation）和‘参考模型解耦’作为可证优的通用改进机制，具有理论新颖性。

核心创新在于：（1）首次从理论层面证明OPD等价于一种特殊形式的密集KL约束RL（即奖励与KL项恒等加权、参考模型任意），从而为蒸馏与RL建立严格统一框架；（2）据此提出广义OPD（G-OPD）——引入可学习的奖励缩放因子β和灵活参考模型，解耦奖励信号与正则化强度；（3）发现β>1（即ExOPD）能稳定提升性能，甚至实现学生反超教师；（4）在strong-to-weak场景中，将参考模型设为教师的预RL基模型（而非RL微调后版本）可校准奖励偏差，构成‘reward correction’新范式。

实验覆盖MATH、AMC、HumanEval、MBPP等权威数学与代码基准，在LLaMA-2/3、Qwen、Phi系列模型上验证；关键发现：① ExOPD在10+ teacher-student规模组合（如Qwen2-7B→Phi-3-mini）上一致优于标准OPD、PPO、Off-Policy Distillation；② 多专家融合（如数学+代码RL专家）经ExOPD蒸馏后，学生首次在MATH上达58.2%（超教师56.7%）；③ Reward correction需访问教师基模型（如Qwen2-7B-base），带来额外开销但显著提升稳定性；全文未开源代码，但实验细节完整、消融充分，未来工作可探索β自适应调度、免基模型的参考估计、及G-OPD与RAG/chain-of-thought的协同。

Recent related works include: 'Policy Distillation via On-Policy KL Regularization' (ICML 2023), 'Distilling RL Policies with Trajectory-Level Alignment' (NeurIPS 2022), 'Reward Modeling is Not All You Need: Aligning LMs via Direct Preference Optimization' (ICLR 2024), 'Self-Play Fine-Tuning for Math Reasoning' (ACL 2024), and 'Multi-Expert Mixture of Agents for Code Generation' (EMNLP 2023).