Variational Proximal Policy Optimization

本論文は、Stein変分勾配降下法とMixture-of-Expertsアーキテクチャを統合することで、方策のモード崩壊と分布ドリフトを軽減し、推論ベンチマークにおける大幅な性能向上とトークン効率を実現する、粒子ベースの変分推論フレームワークであるVariational Proximal Policy Optimization ($\textsc{VP}_2\textsc{O}$) を提案する。

要約

あなたは、巨大で非常に賢いロボットに、コードを書いたり、数学の問題を解いたり、あるいは人間が本当に好むような方法で会話をしたりすることを教えようとしていると想像してください。標準的な方法（PPOやGRPOと呼ばれます）は、少し厳格すぎるコーチのようなものです。「前回うまくいったことを正確にやりなさい。ただし、あまり変えすぎると罰を与えるぞ」と言うようなものです。

これは機能しますが、この論文では3つの大きな問題があると言います。

1. 「単調（One-Note）」問題： ロボットが高いスコアを得た特定の数少ない行動に固執してしまい、問題を解決するための他の創造的な方法を見逃してしまいます。
2. 「脆さ（Brittle）」問題： ロボットが新しいアイデアを試そうとすると、「どの程度の変化が許されるか」というルールが硬直的で恣意的であるため、混乱したり壊れたりすることがよくあります。
3. 「ドリフト（Drift）」問題： ロボットは、自分がどのように振る舞うべきだったかを徐々に忘れ、実際に役に立つことよりも、高いスコアを得るためにシステムをハック（攻略）するようになります。

新しい解決策: VP2O (Variational Proximal Policy Optimization)

著者らは、VP2Oと呼ばれる新しい手法を提案しています。理解を助けるために、いくつかの比喩を使ってみましょう。

1. 「専門家チーム」対「ジェネラリスト」

一つの巨大な脳にすべてをやらせる代わりに、この論文ではMixture-of-Experts (MoE) モデルを使用しています。これは、20人の異なるスペシャリスト（専門家）が部屋に座っている会社を想像してみてください。

• 従来の方法： マネージャー（ルーター）が仕事のために一人のスペシャリストを選び、彼らは全員が「同じ完璧なスペシャリスト」になろうとします。結局、彼らは皆同じように考え始め、チームは創造性を失います。
• VP2O の方法： マネージャーは各タスクに対して少人数の専門家チームを選びます。VP2Oは、各スペシャリストを独自の「粒子」または個体として扱います。目標は、彼らが皆同じになることではなく、それぞれが異なりつつも、それぞれの特定の仕事において優秀であることです。

2. 「磁力のあるダンスフロア」 (Stein Variational Gradient Descent)

これがこの論文の核心となる魔法です。20人のスペシャリストがダンスフロア上のダンサーであると想像してください。

• 引き寄せ（磁力）： フロアには「高報酬」ゾーン（最高の答えがある場所）があります。ダンサーたちは磁力によってこのゾーンへと引き寄せられます。
• 反発（パーソナルスペース）： 従来の方法では、ダンサーたちが同じ場所に密集し、互いに足を引っ張り合ってしまいます（これは「モード崩壊」と呼ばれます）。VP2Oはルールを追加します。「もし誰かに近づきすぎたら、押し返しなさい。」
• 結果： ダンサーたちは高報酬ゾーン全体に広がります。彼らはより広い範囲をカバーし、問題を解決するための「多くの異なる方法」（コードを書くなど）を見つけ出します。単一の「完璧な」方法だけに留まることはありません。

3. 「スマートなコーチ」対「クリッピング・ルール」

従来の方法では、コーチは「クリッピング」ルールを使用します。「もし君のダンスの動きが10%以上変わったら、私は君を止める」というものです。これは鈍い道具です。

• VP2Oのアプローチ： 硬い停止を行う代わりに、VP2Oは幾何学を使用します。それはダンサーの動きの「形状」を見ます。そしてこう言います。「君が始まった時の位置に対して、この特定の幾何学的な形状の中に留まっている限り、好きなだけ動いてよい。」
• これにより、より自然で流動的な動きが可能になります。ロボットは、ルールが恣意的な数字に基づいているのではなく、学習プロセスの実際の形状に基づいているため、ルールを破ることなく新しいアイデアを探索できます。

4. 「直交（Orthogonal）」する目標

スペシャリストたちが互いにコピーし合わないように、VP2Oは**直交化（Orthogonalization）**と呼ばれるルールを追加します。

• 比喩： 二人の専門家に数学の問題を解くよう頼む場面を想像してください。もし二人が全く同じ方法を使うなら、それは非効率です。VP2Oは、彼らが（一人が代数を用い、もう一人が幾何学を用いるように）異なる方法を使うよう強制します。これにより、チームはあらゆる問題に対処するための幅広いツールを持つことができます。

これを試した結果はどうなったか？

著者らは、330億のパラメータを持つ大規模モデルと20の専門家を用いてテストを行いました。結果は以下の通りです。

• コーディング (Codeforces): これが最大の勝利でした。新しい手法は、ロボットのコーディングスコアを179ポイント向上させました（競技プログラミングにおける巨大な飛躍です）。ロボットは単に良くなっただけでなく、コード問題を解決するためのより多様な方法を見つけ出しました。
• 数学 (AIME): ロボットはより多くの数学問題を正しく解きました。興味深いことに、内部的な推論（思考）にはより多くの時間を費やしたにもかかわらず、最終的な答えを説明するためにより少ない言葉を使用しました。より効率的になったのです。
• 指示への追従性: ロボットは複雑な指示に従うことがはるかに上手くなりました。これは、おそらく「万能な一つの手法」に陥っていなかったためです。

まとめ

この論文は、AIの「脳」を、互いに異なるように促される（磁力による反発を用いる）多様なスペシャリストのチームとして扱うことで、AIが以下のようになることを主張しています。

1. より創造的になる（問題を解決するためのより多くの方法を見つける）。
2. より安定する（クラッシュしたり行き詰まったりしない）。
3. より効率的になる（仕事をこなすために必要なトークンを減らす）。

著者らは、AIが長文で複雑な回答（16,000トークンなど）を書かなければならない場合、多様な「専門家」のチームを持つことは、単一の硬直した戦略を持つことよりも価値が高くなるため、この手法が最も効果を発揮すると強調しています。

技術概要

技術要約: Variational Proximal Policy Optimization (VP2O)

1. 問題提起

Proximal Policy Optimization (PPO) およびその派生手法（例：GRPO）を利用した人間からのフィードバックによる強化学習 (RLHF) は、以下の3つの持続的な限界に直面している：

1. ポリシーのモード崩壊 (Policy Mode Collapse): ポリシーが高報酬の狭い行動セットに収束してしまい、人間の好みの全領域を捉えるために必要な多様性が犠牲になる。
2. 非効率な探索 (Inefficient Exploration): 探索が脆弱であり、特に疎な報酬やノイズの多い報酬ランドスケープにおいて顕著である。これはエントロピー・ボーナスのようなヒューリスティックに依存することが多い。
3. 分布のドリフトと不安定性 (Distributional Drift and Instability): ポリシーが誤設定された報酬モデルに過学習し、「報酬ハッキング」を引き起こす可能性がある。さらに、PPO/GRPOにおけるトークンレベルの重要度比（importance ratios）は、トレーニング中に高分散のノイズを導入し、長いシーケンスにおいて不安定さを引き起こすため、固定クリッピングやKLスケジュールの調整といったアドホックな修正を必要とする。

GRPOなどの現在のアプローチは、勾配ペナルティを通じて安定性を向上させているが、多様性を考慮した最適化や探索のための原理的なメカニズムを欠いている。

2. 手法: Variational Proximal Policy Optimization (VP2O)

VP2Oは、RLHFのポリシー最適化を変分推論 (variational inference) 問題として再定義する。プロキシマル制御を単なるスカラーのクリッピング・ルールとして扱うのではなく、最適化プロセスを Mixture-of-Experts (MoE) アーキテクチャ内での Stein Variational Gradient Descent (SVGD) へとマッピングする。

コア・フレームワーク

• 変分再定式化 (Variational Reformulation): 本論文では、PPO/GRPOの目的関数を、現在のポリシー $\pi_\theta$ と最適なポリシー分布 $p^*$ の間のカルバック・ライブラー・ダイバージェンス $D_{KL}(\pi_\theta \parallel p^*)$ を最小化することとして再構成している。
• 粒子ベースの最適化 (Particle-Based Optimization): 最適な分布 $p^*$ は、単一のポリシーではなく、一連の「粒子」によって近似される。VP2Oにおいて、これらの粒子は、疎なMoEレイヤー内の個々のエキスパートとして実現される。
• Stein 輸送場 (Stein Transport Fields): 最適化には、粒子の更新にSVGDを利用する。この更新則は2つの力を組み合わせている：
  1. 駆動力を伴う力 (Driving Force): 粒子（エキスパート）を $p^*$ の高報酬領域へと移動させる。
  2. 反発力 (Repulsive Force): 粒子が単一のモードに崩壊するのを防ぎ、多様性を維持する。

主要なアーキテクチャ構成要素

• 変分アンサンブルとしてのMoE: MoEレイヤー内の各エキスパート $i$ は、個別のポリシー成分 $\pi_{\theta_i}$ として機能する。ルーター $\phi(\cdot)$ は、各トークンに対してエキスパートの疎なサブセット（Top-K）を選択する。
• プロトタイプ上の関数カーネル: 高次元のパラメータ空間を扱うため、VP2Oはパラメータ空間ではなく、出力空間におけるカーネル $K$ を定義する。各エキスパートの出力投影行列の主固有ベクトルから導出される単位ノルムのプロトタイプ $p_i$ を保持する。カーネルは、これらのプロトタイプの角度的な類似性を測定する。
• デカップルされた輸送場 (Decoupled Transport Field): Stein 更新は、ルーティングの活性化に基づいてデカップルされる：
  • 吸引 (Attraction): 共活性化されたエキスパート（ルーターによって選択されたもの）は、カーネル重み付きの勾配平均を介して情報を共有する。
  • 反発 (Repulsion): 非活性または稀に共活性化されるエキスパートは、エキスパート間の専門化を促すために、カーネル勾配項を介して互いに押し離される。
• エキスパートの直交化 (Expert Orthogonalization): さらなる崩壊を防ぐため、補助的な損失関数がエキスパート表現の直交性を促進し、同一のTop-Kグループ内における異なるエキスパート出力間の投影を最小化する。
• 幾何学的信頼領域 (Geometric Trust Regions): VP2Oは、固定クリッピングや静的なKLペナルティを、2つの幾何学ベースの制御に置き換える：
  1. アンカー・プロトタイプ予算 (Anchor Prototype Budget): 「アンカー」ポリシーのスナップショットに対する低次元プロトタイプ空間におけるステップサイズを制限する。
  2. オンポリシー・ビヘイビア予算 (On-Policy Behavior Budget): ドリフト診断（KLダイバージェンスおよび有効サンプルサイズ）に基づくイベント駆動型の同期を使用し、固定スケジュールではなく必要な場合にのみアクター・ポリシーを更新する。

3. 主な貢献

1. RLHFのためのSVGD: 本論文は、KL正則化された報酬最大化を、Stein Variational Gradient Descent を用いて $D_{KL}(\pi_\theta \parallel p^*)$ を最小化することとして再解釈している。これにより、PPOのクリッピングを、報酬と多様性を同時に最適化するカーネル重み付きの更新へと置き換えている。
2. 共同エキスパート専門化 (Joint Expert Specialization): 直交損失とルーティング多様化損失を組み合わせた特化した目的関数を導入している。これにより、標準的なPPO/GRPOで見られるエキスパートの崩壊を抑制し、各エキスパートが機能的に異なる挙動を発展させることを保証する。
3. 統一されたフレームワーク: ポステリア・サンプリングとポリシー制約をVP2Oへと統合し、アドホックな信頼領域に頼ることなく、多様で不確実性を考慮し、かつ人間の好みに適合したポリシーを可能にする。

4. 実験結果

著者らは、VP2Oを 33B/4B 疎なMoEモデル（総パラメータ数33B、トークンあたりのアクティブパラメータ数4B）を用いて評価し、同一条件下でGRPOで訓練されたベースラインと比較した。

パフォーマンス・ベンチマーク

• 数学的推論 (AIME): VP2Oは一貫した向上を示した。AIME 2024において、8Kコンテキストで+2.6%、16Kコンテキストで+1.6%の改善を達成した。特筆すべきは、VP2OはAIME 2024においてベースラインよりも約2,000ステップ早く収束したことである。
• 科学的推論 (GPQA): 8Kコンテキストではほぼ同等であったが、16Kコンテキストでは明確な+1.8%の優位性を示し、長い生成予算の下での複雑な多段階タスクにおける多様性の恩恵を示唆した。
• コード生成 (Codeforces): 最も顕著な利得は16Kコンテキストで見られ、VP2Oはベースラインを +179 ELO および +3.6 Pass@1 ポイント上回った。著者らは、これをエキスパートを構造的に異なる解決戦略へと押し出す反発力の効果であるとしている。
• 指示遂行 (IFBench/IFEval): VP2Oは、指示遂行の全指標において最も一貫した向上を実現しており、コンテキスト長や指標に応じて+3.6%から+5.7%の改善が見られた。

効率性とトークン使用量

• トークン効率: VP2Oは解決の効率性を示した。AIME 2025 (8Kコンテキスト) において、より高い精度を達成しながら、32%少ないトークン（130トークン少ない）を使用した。
• 推論パターン: 「思考トークン (thought tokens)」と「解決トークン (solution tokens)」の分析により、VP2Oは（特にコーディングや数学のタスクにおいて）「思考」は多い（中間的な推論トークンを多く生成する）が、最終的な回答はより簡潔に書く傾向があることが明らかになった。

学習ダイナミクス

• 安定性: VP2Oは学習の初期段階から安定した優位性を確立し、ベースラインで見られた学習終盤の劣化（報酬の過剰最適化）を回避した。
• 収束: フレームワークは、特に16Kコンテキストの設定において、より速い収束を示した。

5. 意義と主張

本論文は、VP2Oが、現在RLHFを支配しているヒューリスティックな制約（クリッピング、固定KLスケジュール）に代わる原理的な選択肢を提供すると主張している。エキスパートの集団に対する変分推論の観点からポリシー最適化を捉えることで、VP2Oは以下を実現する：

• ヒューリスティックへの依存を軽減: アドホックなクリッピングを、データから導出された幾何学的プロキシマル制御に置き換える。
• 多様性を強化: Steinの反発力と直交性制約を通じて、モード崩壊とエキスパートの重複を明示的に防止する。
• 長文生成の向上: 利点はより長い生成コンテキスト（16Kトークン）で最も顕著であり、これは、機能的な多様性を維持することが、長期的な強化学習の軌道を安定させるために極めて重要であることを示唆している。

著者らは、結果が単一のモデルファミリー（33B/4B MoE）において有望であるものの、より大きなスケール（例：70B以上）や異なるモデルアーキテクチャにおけるさらなる検証は未解決の課題であると述べている。現在の研究は、RLHFの訓練を安定化および多様化させるための変分フレームワークの有効性に焦点を当てている。