教師を超えた学習：報酬外挿による一般化された方策蒸留

要旨

オン方針蒸留（OPD）は、生徒モデルが生成した軌道上で教師モデルのロジット分布に合わせる手法であり、生徒の性能向上において強力な実証的効果を示し、オフ方針蒸留や強化学習（RL）のパラダイムをしばしば上回る。本研究ではまず、OPDが密なKL制約付きRLの特殊ケースであることを理論的に示す。具体的には、報酬関数とKL正則化が常に等しく重み付けられ、参照モデルは任意のモデルでよいという特性を持つ。次に、標準的なOPD目的関数を拡張した一般化オン方針蒸留（G-OPD）フレームワークを提案する。これは柔軟な参照モデルと、報酬項のKL正則化に対する相対的重みを制御する報酬スケーリング係数を導入するものである。数学推論とコード生成タスクにおける包括的実験を通じて、二つの新たな知見を得た：（1）報酬スケーリング係数を1より大きく設定する「報酬外挿」（ExOPDと命名）は、様々な教師-生徒サイズ組み合わせにおいて標準OPDを一貫して改善する。特に、同一の生徒モデルに分野特化RLを適用して得た複数専門家の知識を元の生徒に統合する設定では、ExOPDにより生徒は教師の性能限界を超え、分野専門教師を凌駕する性能を発揮した。（2）ExOPDを基盤とし、強→弱蒸留設定（大規模教師から小規模生徒への蒸留）では、参照モデルとしてRL適用前の教師のベースモデルを選択する「報酬補正」を行うことで、より正確な報酬信号が得られ蒸留性能がさらに向上する。ただしこの選択は教師の事前RL版へのアクセスを必要とし、計算コストも増大する。本研究がOPDに関する将来の研究に新たな示唆をもたらすことを期待する。

English

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.

教師を超えた学習：報酬外挿による一般化された方策蒸留

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

要旨

Support