Complementing reinforcement learning with SFT through logit averaging in the post training of LLMs

技術サマリー：ログイット平均化による強化学習の教師あり微調整（SFT）による補完

1. 問題定義

大規模言語モデル（LLM）のポストトレーニングは、通常、教師あり微調整（SFT）と検証可能な報酬を伴う強化学習（RLVR）という 2 つのパラダイムに依存しています。SFT はフォーマットやスタイルを学習するためにトークンごとの密な監督を提供しますが、トレーニングデモンストレーションの分布によって制限されます。一方、近接方策最適化（PPO）やグループ相対方策最適化（GRPO）などの RL 手法は、これらの分布の限界を超えられますが、スパースでノイズの多い報酬信号に悩まされます。

標準的なアプローチは、これらを逐次的に組み合わせるものです：まず SFT を適用し、その後、学習可能な方策 $\pi_\theta$ を SFT の参照 $\pi_{\text{sft}}$ に固定するための Kullback-Leibler（KL）正則化項 $\beta \text{KL}(\pi_\theta \| \pi_{\text{sft}})$ を伴う RL を行います。著者らは、この「KL 固定」レシピに根本的な緊張関係があることを特定しました：KL ペナルティは $\pi_\theta$ を $\pi_{\text{sft}}$ へと引き寄せ、方策を SFT チェックポイントの限界（例えば、代数的誤りや反復ループ）に制約し、RL 中に発展した推論能力を十分に活用するのを妨げる可能性があります。実務家はしばしばこれを緩和するために小さな $\beta$ を選択しますが、著者らは正則化に対する構造的な代替案を提案します。

2. 手法

本論文は、KL 正則化をログイット平均化で置き換える手法を導入します。乖離にペナルティを課す代わりに、学習可能な方策 $\pi_\theta$ と凍結された参照方策 $\pi_{\text{ref}}$（例えば SFT）を、そのログイットを平均化して混合方策 $\pi^{\text{mix}}_{\theta, \alpha}$ を形成することで結合します。

中核メカニズム:
重み $\alpha \in [0, 1]$ と状態 $s$ が与えられたとき、混合ログイット $z^{\text{mix}}_{\theta, \alpha}(s)$ は以下のように定義されます：
$$z^{\text{mix}}_{\theta, \alpha}(s) = (1 - \alpha) z_\theta(s) + \alpha z_{\text{ref}}(s)$$
混合方策は $\pi^{\text{mix}}_{\theta, \alpha}(a | s) = \text{softmax}(z^{\text{mix}}_{\theta, \alpha}(s))$ となります。

トレーニング目的:
著者らは、この混合方策に対して GRPO 目的関数を適応させます：

• ロールアウト: 軌跡は、純粋な学習可能方策 $\pi_\theta$ ではなく、混合方策 $\pi^{\text{mix}}_{\theta, \alpha}$ からサンプリングされます。
• アドバンテージ推定: グループ相対アドバンテージは、これらの混合ロールアウトからの報酬に基づいて計算されます。
• 最適化: 学習可能なパラメータ $\theta$ は、混合方策下での報酬を最大化するように更新されます。参照ログイット $z_{\text{ref}}$ は全体を通じて凍結されたままです。
• 重み付けスキーム:
  1. 固定重み: $\alpha$ はトレーニング全体を通じて一定に保たれます（例：$\alpha=0.5$）。
  2. 適応的重み: $\alpha$ はホールドアウト検証セットに基づいて動的に更新されます。更新則は、純粋方策と混合方策の正解性を比較する「ネットゲイン」指標を使用し、混合がネットゲインをもたらす場合は $\alpha$ を 1 へ、そうでない場合は 0 へと押しやります。

RLVR との比較:
クリティックと KL 正則化を使用する標準的な RLVR 手法（PPO、GRPO、DAPO）とは異なり、このアプローチには KL 項もクリティックも存在しません。参照方策はログイット平均化を通じて探索に参加し、損失関数内の重要度比は、現在の混合方策と過去の混合方策との間で計算されます。

3. 主要な貢献と知見

相補的スキル（パターン 1 & 2）:
MATH、cn-k12、MMLU データセット全体の実証分析により、SFT と RL 学習済み方策が相補的スキルを有することが明らかになりました：

• SFT はフォーマット（出力テンプレートへの準拠）において優れていますが、複雑な推論では失敗する可能性があります。
• RL 方策（$\pi_\theta$） は推論（正しい数値答えを導出）において優れていますが、しばしばフォーマットの規律を失います。
• 相補性: 問題の相当な割合（22%〜38%）は、2 つの方策のちょうど片方によって正しく解決されます。ログイット平均化は「専門家積（Product of Experts）」として機能し、SFT のフォーマット能力と RL 方策の推論能力を組み合わせます。

性能結果:
Qwen2.5-Instruct モデル（1.5B、3B、7B）を 3 つのデータセットで評価：

• 精度: ログイット平均化手法は、すべての 9 つの（データセット × モデルサイズ）構成において、標準的な KL 正則化 GRPO に対して同等かそれ以上の回答精度を達成します。
• アブレーション: SFT アンカーをベースの Instruct モデルに置き換えると、性能が著しく低下し、 gains がフォーマット認識型の参照との混合に由来することを確認しました。
• ログイット対確率平均化: ログイット平均化（分布の幾何平均）は、確率平均化（算術平均）を上回ります。ログイット平均化は、より小さな混合重みでより高い解決率を達成し、$\alpha=0.5$ でピークに達するのに対し、確率平均化はより高い重みを必要とし、$\alpha=0.6$ でピークに達します。

理論的関連性:
本論文は、ログイット平均化が数学的に**専門家積（POE）**と等価であることを確立しています。その結果生じる方策は、RL 方策と参照方策への KL 発散の加重和を最小化します。この構造により、方策は両方の専門家が支持するトークンに確率を集中させることができ、SFT の代数的誤りを効果的に修正しつつ、そのフォーマット制約を維持します。

4. 意義と主張

著者らは、この手法が「フォーマットアンカー」を「勾配制約」から分離することにより、標準的な KL 正則化レシピに対して明確な利点を提供すると主張します。

• 限界の克服: KL 正則化が $\pi_\theta$ を $\pi_{\text{sft}}$ に近づけさせ（その推論上の欠陥を継承させる可能性がある）のに対し、ログイット平均化は、混合方策の構造的結合を通じてフォーマットの利点を維持しながら、$\pi_\theta$ が参照の限界を克服することを可能にします。
• 暗黙的な信頼領域: 混合は暗黙的な信頼領域として機能し、大規模な純粋 RL でしばしば問題となるフォーマットスタイルからの分布ドリフトを防ぎます。
• 密な事前分布: 最終報酬トークンからのみ信号を伝播する RL と異なり、ログイット平均化は、各デコードステップにおいて SFT アンカーからの密なトークンレベルの事前分布を組み込みます。
• スケーラビリティ: このアプローチは、小規模言語モデル（SLM）のアンサンブルに自然に拡張でき、システムが（コード、安全性、またはドメイン用語などのための）専門家の専門知識を、巨大な単一の方策をトレーニングするコストなしに活用することを可能にします。

本論文は、ログイット平均化が KL 正則化に対する堅牢な代替案を提供し、SFT と RL の相補性を効果的に活用して、フォーマットの品質を犠牲にすることなく推論精度を向上させることを結論付けています。