技術的概要：開始の仕方が推論の仕方を決定する：RLVR における Prefix-Tuned Prior による探索の促進

1. 問題定義

検証可能な報酬を用いた強化学習（RLVR）は、数学的問題解決や定理証明などのタスクにおいて、大規模言語モデル（LLM）の推論能力を向上させる強力なパラダイムとして登場しました。しかし、著者らは現在の RLVR アプローチにエントロピーの崩壊という重要な限界を特定しています。

RLVR はしばしば単一ロールアウトの精度（Pass@1）を向上させますが、成功する推論経路の網羅性を拡大させることには失敗することが多いです。これは出力の多様性の低下として現れ、モデルが狭い範囲の推論経路に収束する現象です。エントロピー正則化や内在的報酬などの既存の探索技術は、ベースモデルの分布を固定されたものとして扱い、RL 目的関数やサンプリング手順を変更しようとします。本論文は、これらの手法が不十分であると主張します。なぜなら、RL はベースモデルの事前分布によって強く形作られるからです。さらに、受動的な探索技術は生成品質を低下させ、効果的な推論に寄与しないノイズの多いロールアウトをもたらす傾向があります。

核心的な課題は、ポリシーが冗長または低品質な経路に収束することなく、またモデル全体のパラメータを更新する計算コストを伴うことなく、効果的な探索（多様で、タスクに関連し、高品質な推論行動の発見）を促すことです。

2. 手法：IMAX フレームワーク

著者らは、最適化対象を LLM のパラメータからソフトプロンプトチューニングへ移行させる軽量フレームワーク、**情報最大化拡張探索（Information-Maximizing Augmented eXploration: IMAX）**を提案します。固定された事前分布の上に RL プロセスを修正するのではなく、IMAX は推論経路に対するモデルの事前分布を能動的に再構築するソフトプレフィックスのプールを学習します。

2.1 プレフィックス条件付き推論

中核的なメカニズムは、$C$個のソフトプレフィックス（連続的な仮想トークン）$\mathcal{E} = \{e_1, \dots, e_C\}$ のプールを含みます。与えられた入力プロンプト $x$ に対して、プレフィックス $e_z$ が入力埋め込みシーケンスに先付けられます。バックボーン LLM は凍結されたまま維持され、ポリシー $\pi_E(y|x, z)$ はプロンプトと選択されたプレフィックスの識別子 $z$ の両方に条件付けられます。

• 初期化: プレフィックスは、外部 LLM によって生成された意味的に意味のあるテキストプロンプト（例：「変数を明示的に導入する」、「ケースに分割する」）のプールから初期化されます。これにより、プレフィックスは異なる推論戦略で開始されます。
• 制御: 各プレフィックスは「制御ノブ」として機能し、異なるロールアウト分布を誘発し、モデルを推論空間の異なる領域（例：代数的 vs 図形的 vs コードのような解）へと誘導します。

2.2 情報最大化（InfoMax）報酬

学習されたプレフィックスプールが冗長な制御（異なるプレフィックスが同一の経路を誘発する状態）に収束するのを防ぐため、著者らはInfoMax 報酬を導入します。目的は、プロンプト $X$ に条件付けられたプレフィックス識別子 $Z$ と生成された経路 $Y$ との間の相互情報量（MI）を最大化することです：
$$I(Z; Y | X) = H(Y | X) - H(Y | X, Z)$$
直感的には、これは同じ質問に対して異なるプレフィックスが識別可能で有効な推論行動を生み出すことを促します。

2.3 変分下限とトレーニング

相互情報量を直接計算することは困難です。著者らは、プロンプト $x$ が与えられたときに特定の応答 $y$ を生成したプレフィックスを予測するようにトレーニングされた軽量分類器である変分事後分布 $q_\phi(z | x, y)$ を用いて、変分下限を導出します。

• 逆形式: MI の逆分解を利用します：$I(Z; Y | X) = H(Z | X) - H(Z | X, Y)$。
• 報酬の定式化: InfoMax 報酬は $\log q_\phi(z | x, y)$ として定義されます。プレフィックスプールの全目的関数は以下のようになります：
  $$\max_E \mathbb{E} [R(x, y) + \beta \log q_\phi(z | x, y)]$$
  ここで、$R(x, y)$ は標準的な検証可能な報酬（正解性）であり、$\beta$ はハイパーパラメータです。

2.4 交互最適化

トレーニング手順は、交互最適化スキーム（座標上昇変分推論）に従います：

1. 事後分布（$q_\phi$）の更新: プレフィックスプールを固定した状態で、分類器 $q_\phi$ を更新して、プレフィックス識別子をプロンプト - 応答ペアから予測します（情報漏洩を防ぐため、ソフトプレフィックスなしで生成されたペアを使用します）。
2. プレフィックスプール（$E$）の更新: $q_\phi$ を固定した状態で、InfoMax 報酬を強化したグループ相対方策最適化（GRPO）アルゴリズムを用いてソフトプレフィックスを更新します。
3. 層化サンプリング: トレーニング中、各プロンプトは異なるプレフィックス識別子とペアリングされ、各 GRPO グループ内でプレフィックス空間の均一なカバレッジが確保されます。

3. 主要な貢献

1. 軽量 RLVR パラダイム: 本論文は、フルモデルパラメータではなくソフトプロンプトのプールを最適化することで推論を改善する手法を導入します。これにより、フルファインチューニングの計算オーバーヘッドなしに、凍結されたモデルの条件付き生成分布を再構築します。
2. 多様性のための InfoMax 報酬: 著者らは、異なるプレフィックスが識別可能でタスクに関連する推論経路を誘発し、プレフィックスプールが冗長な制御に収束するのを防ぐ、応答レベルの InfoMax 報酬を提案します。
3. 実証的検証: 3 つのモデルスケール（Qwen2.5-1.5B、Qwen3-4B、Qwen3-8B）における数学的推論ベンチマーク（GSM8K、MATH-500、Minerva、AMC）での広範な実験により、IMAX が推論の品質と多様性の両方を向上させることが示されました。

4. 実験結果

著者らは、IMAX をフルポリシー GRPO、テキストプロンプト初期化、および他の多様性指向の RLVR 手法（DIVER、MERCI）などのベースラインと比較して評価しました。

• 性能向上: IMAX は、標準的なプレフィックスチューニング RLVR ベースラインを一貫して上回りました。
  • Qwen3-4Bにおいて、IMAX は特定のベンチマークにおいて、標準的なベースラインと比較して Pass@4 を最大11.60%、Avg@4 を**10.57%**向上させました。
  • Qwen3-8Bにおいて、DAPO+IMAX は GSM8K、MATH-500、Minerva において DAPO-prefix よりも顕著な改善を示しました。
• 多様性とカバレッジ:
  • t-SNE 分析: 応答埋め込みの可視化により、IMAX によって誘発されたプレフィックスは他の手法と比較してより分離されたクラスターを生み出し、グループ内での多様性が強いことが示されました。
  • 推論コンポーネント: ヒートマップは、異なるプレフィックスが異なる推論パターン（例：プレフィックス 1 は代数的操作を好む、プレフィックス 2 はケース分割と検証を好む）を成功裡に誘発したことを示しており、それらの意味的初期化と一致していました。
  • オラクル性能: 複数のプレフィックスの組み合わせ性能（オラクル）は、最良の単一プレフィックスを上回り、プレフィックスが問題空間の相補的な成功領域を網羅していることを確認しました。
• 効率性: IMAX は、軽量なプレフィックス埋め込みと分類器ヘッドのみを更新するため、フルパラメータ GRPO に比べてトレーニングコスト（GPU 時間）が大幅に少なくなります（3.0 倍〜6.0 倍削減）。

5. 意義と主張

本論文は、モデルが推論プロセスを開始する方法（事前分布）が、本質的にその推論の仕方を決定すると主張します。ソフトプロンプトを再利用可能な制御変数として扱うことで、IMAX は RLVR における探索問題に対するスケーラブルな解決策を提供します。

• アルゴリズム非依存: このフレームワークは、既存の RLVR パイプライン（例：GRPO、DAPO）にシームレスに統合されるように設計されています。
• 構造化された多様性: 単にランダム性を増大させる手法とは異なり、IMAX は構造化された推論多様性を促進し、探索が意味的に異なり高品質な経路につながることを保証します。
• 限定的な限界: 著者らは、この手法は主に推論タスク向けに設計されており、一般的な指示遂行タスクでは限られた利益しか示さないことを認めています。さらに、効果性は初期テキストプロンプトの品質に依存し、ドメイン横断的に手動設計または再チューニングが必要になる可能性があります。学習されたプレフィックスは、推論中に持続的ではあるが小さなオーバーヘッドをもたらします。

要約すると、IMAX は、ソフトプレフィックスを介して生成プロセスの「開始」を最適化することが、LLM における探索を促進し推論能力を向上させるための極めて効果的な戦略であることを示しており、フルポリシー強化学習に対する計算効率的な代替手段を提供します。