技術的サマリー：モード崩壊を超えて：多様な推論のための分布整合

1. 問題定義

本論文は、オンポリシー強化学習（RL）手法、特に検証可能な報酬を用いた強化学習（RLVR）を通じて大規模推論モデル（LRM）を訓練する際に用いられる手法における根本的な限界として「モード崩壊」を特定している。グループ相対方策最適化（GRPO）などの手法は数学的問題解決やコード生成において成功を収めているが、一度高報酬の軌道が発見されると、確率質量を単一の軌道に集中させる傾向を示す。この振る舞いは、代替の解決戦略を抑制し、探索を停止させ、一般化に失敗する脆い推論へと至る。

著者らは、この現象が単なる実装上のアーティファクトではなく、最適化目的関数の直接的な帰結であると主張する。標準的なオンポリシー RL は、暗黙的に「逆 Kullback-Leibler（KL）ダイバージェンス」（$D_{KL}(\pi_\theta \| p^*)$）を最小化する。逆 KL は本質的に「モード探索的」であり、複数の多様な高報酬解の分布を維持するのではなく、最初に遭遇した成功軌道を増幅する。対照的に、「前方 KL ダイバージェンス」（$D_{KL}(p^* \| \pi_\theta)$）は「モード網羅的」な振る舞いを示し、方策がターゲット分布の全サポートに一致することを促すが、グローバルな報酬重み付き分布からのサンプリングを必要とするため、通常は計算的に困難である。

2. 手法：DMPO

計算的に困難なグローバルサンプリングを必要とせずにモード崩壊に対処するため、著者らは「分布整合方策最適化（DMPO）」を提案する。

核心的な洞察

DMPO は、グローバルではなく「グループレベル」で前方 KL の最小化を近似する。グローバルなターゲット分布からのサンプリングの代わりに、このアルゴリズムは、与えられたプロンプトに対してサンプリングされた軌道の特定のグループに対するターゲット分布を構築する。

アルゴリズム的定式化

1. グループレベルのターゲット分布: 現在の方策 $\pi_{\theta_{old}}$ からサンプリングされた $G$ 個の軌道 $\mathcal{O} = \{o_1, \dots, o_G\}$ のグループに対して、報酬の指数に比例するターゲット分布 $p(o_i | \mathcal{O})$（ボルツマン分布）を構築する。
   $$p(o_i | \mathcal{O}) = \frac{\exp(r(o_i)/\alpha)}{\sum_{j=1}^G \exp(r(o_j)/\alpha)}$$
   これはグループレベルの分配関数のみを必要とするため、計算的に実行可能である。

2. グループレベルの方策分布: このターゲットと方策を整合させるため、軌道の長さバイアスを考慮するために長さ正規化された対数尤度に基づく方策分布 $q_\theta(o_i | \mathcal{O})$ を定義する。
   $$q_\theta(o_i | \mathcal{O}) = \frac{\exp(\varphi(o_i))}{\sum_{j=1}^G \exp(\varphi(o_j))}$$
   ここで、$\varphi(o_i)$ はトークンあたりの平均対数確率である。

3. 分布整合損失: 方策をターゲットに整合させるため、$q_\theta$ と $p$ の間のダイバージェンスを最小化する。数値的安定性を確保し、方策の確率がゼロに近づく際の勾配爆発を防ぐため、KL ダイバージェンスの代わりに「平均二乗誤差（MSE）」を利用する。
   $$\mathcal{L}_{DM}(\theta) = \frac{1}{G} \sum_{i=1}^G (p(o_i | \mathcal{O}) - q_\theta(o_i | \mathcal{O}))^2$$

4. 統合目的関数: 最終的な目的関数は、標準的な GRPO の報酬最大化項と分布整合正則化項を組み合わせる。
   $$\mathcal{L}_{DMPO}(\theta) = \mathcal{L}_{GRPO}(\theta) + \lambda \mathcal{L}_{DM}(\theta)$$
   ここで、$\mathcal{L}_{GRPO}$ は利用（報酬最大化）を駆動し、$\mathcal{L}_{DM}$ は高報酬領域内で方策が単一のモードに崩壊するのを防ぐことで探索を強制する。

3. 主要な貢献

A. DMPO アルゴリズム

本論文は、モード崩壊に対する原理的かつ実用的な解決策として DMPO を提示する。グループレベルで前方 KL を近似することで、学習された分配関数やグローバルサンプリングを必要とせずにモード網羅的な振る舞いを達成する。この手法は既存の GRPO フレームワークと互換性があり、単一の正則化項の追加のみで済む。

B. MM-NP-Bench

著者らは、推論モデルにおける探索を厳密に評価するために設計されたマルチモーダルベンチマーク「MM-NP-Bench」を導入する。

• 範囲: テキストベースの NP-Bench を拡張し、10 種類の NP 困難な組み合わせ最適化タスク（TSP、グラフ彩色、頂点被覆など）の視覚的表現を含める。
• 二重指標評価: モード崩壊を可視化するため、ベンチマークは 2 つの指標を使用する。
  • 成功率（SR）: 制約満足度（実行可能性）を測定する。
  • 品質比（QR）: ヒューリスティックな基準に対する最適化の質を測定する。
  • 重要性: モード崩壊は、高い SR だが低い QR（有効だが最適ではない解の発見）として現れる。
• インフラ: 準最適基準を提供するためのパラメトリック生成器、ルールベースの検証器、ヒューリスティックソルバーを含む。

C. 実証的検証

本論文は、DMPO が複数のドメインにおいて強力なベースライン（GRPO、GSPO、GPG、FlowRL など）を上回ることを示す広範な実験を提供する。

4. 結果

最適化タスクにおける性能

• MM-NP-Bench（マルチモーダル）: DMPO は 43.1% の品質比を達成し、GRPO（38.4%）に対して 12% の相対改善を示した。また、最高成功率（61.9% 対 55.7%）も達成した。
• NP-Bench（テキストのみ）: DMPO は 43.9% の品質比を達成し、GRPO（40.1%）に対して 9% の相対改善を示した。
• アブレーション: この手法はハイパーパラメータの選択（$\lambda \in [0.3, 0.7]$ および $\alpha \in [0.8, 1.2]$）に対して堅牢であった。

一般化と転移

• 数学的推論: 最適化タスクにおける DMPO による訓練は、GRPO と比較して数学的推論ベンチマーク（AIME、AMC など）の性能を 2.0% 向上させた。
• ドメイン外タスク: DMPO はドメイン外の視覚的推論ベンチマーク（LogicVista、MathVista など）で 2.3% の改善を示し、多様性を維持する訓練が一般的な推論能力を向上させることを示唆している。

訓練ダイナミクス

訓練ログの分析により、GRPO が「長さ崩壊」（短く trivial な応答への退化）と早期の報酬プラトーに悩まされるのに対し、DMPO は継続的な改善を維持し、堅牢な推論チェーン（約 400 トークン）を保持することが明らかになった。DMPO は GRPO よりも低いエントロピーを維持しており、確率的な混乱ではなく「構造化された多様性」（異なる論理的戦略の自信ある実行）を示している。

5. 重要性と主張

本論文は、オンポリシー RL におけるモード崩壊を防ぐための実用的かつ原理的なアプローチとして「分布整合」を確立すると主張する。著者らは以下を論じる。

1. 根本原因の特定: RLVR におけるモード崩壊は、本質的に逆 KL 最小化のモード探索的な性質によって引き起こされる。
2. 実現可能な解決策: 従来、計算的に困難とされてきた前方 KL 最小化は、グループレベルで効果的に近似でき、解の多様性を維持できる。
3. 一般的な利益: 特定の最適化タスクの訓練中に多様性を維持することは、一般的な推論能力の向上に転移し、大規模推論モデルにおける堅牢な推論にとって探索が重要な要素であることを示唆している。

この研究は、DMPO を単なる多様性正則化項としてではなく、モード探索的な目的関数が見逃す高品質な解の発見を可能にするメカニズムとして位置づけており、これにより最適化性能とマルチモーダルにわたる推論の堅牢性の両方を向上させている。