β-OPSD: 方策最適化による導出、自己蒸留による訓練

要旨

オン方策自己蒸留（OPSD）は、推論言語モデルを改善する有望な手法であるが、実際には依然として脆く、確実に動作させるにはしばしばかなりの工学的労力が必要である。我々は、この困難さの構造的な原因を特定する。バニラOPSDは、より広範な方策最適化ファミリーにおけるβ=1の構成要素にほかならない。ここでβは、生徒モデルを参照方策に固定するKLペナルティの重みである。この等価性により、βは暗黙的に1に固定された値から、制御可能な正則化パラメータへと転換され、参照方策への近さと特権的教師によるガイダンスとの間のトレードオフを実現する、より一般的な定式化が得られる。我々はβ-OPSDを導入し、その最適方策を参照方策と特権的教師との間の幾何学的補間として導出する。しかしながら、この目的を強化学習で直接最適化することは、コストが高く、高分散になるであろう。我々はRL目的を直接最適化する代わりに、その閉形式解を蒸留ターゲットへと変換する。各βの値は、参照から教師への経路に沿ったターゲットを選択し、これは両者のトークンレベルのロジットを混合することで効率的に実装される。このようにして、低コストな蒸留は、高コストな方策最適化の解を近似する。さらに、return-to-goクレジット割り当ては、OPSDの単純さを保ちながら、トークン更新を系列レベルの目的と整合させる。数学的推論ベンチマークにおける実験は、β-OPSDがバニラOPSDを一貫して上回り、最適化の安定性と下流の推論性能を改善することを示している。我々の結果は、OPSDを実用的にする効率性を犠牲にすることなく、自己蒸留から方策最適化へ、そしてその逆へと向かう原理に基づく経路を提供する。

English

On-policy self-distillation (OPSD) is a promising approach to improve reasoning language models, but it remains brittle in practice: making it work reliably often requires substantial engineering effort. We identify a structural source of this difficulty: vanilla OPSD is precisely the β=1 member of a broader policy-optimization family, where β weights the KL penalty anchoring the student to a reference policy. This equivalence turns β from an implicit value fixed at one into a controllable regularization parameter, yielding a more general formulation that trades off proximity to a reference policy against privileged teacher guidance. We introduce β-OPSD and derive its optimal policy as a geometric interpolation between the reference policy and the privileged teacher. Directly optimizing this objective with reinforcement learning, however, would be costly and high-variance. Rather than optimize the RL objective directly, we turn its closed-form solution into a distillation target. Each value of β selects a target along the reference-to-teacher path, which we implement efficiently by mixing their token-level logits. In this way, inexpensive distillation approximates the solution of expensive policy optimization. Return-to-go credit assignment further aligns token updates with the sequence-level objective while retaining the simplicity of OPSD. Experiments on mathematical reasoning benchmarks show that β-OPSD consistently outperforms vanilla OPSD, improving optimization stability and downstream reasoning performance. Our results provide a principled route from self-distillation to policy optimization and back without sacrificing the efficiency that makes OPSD practical.