論文要約：SPS (Steering Probability Squeezing)

タイトル: Reinforcement Learning for Large Language Models におけるより良い探索のための Steering Probability Squeezing (SPS)

1. 背景と課題 (Problem)

大規模言語モデル（LLM）の推論能力向上において、強化学習（RL）は重要な役割を果たしています。特に、ルールベースの報酬信号を用いた RL 学習（RLVR）は、Pass@1（単一サンプルの正答率）を向上させるのに効果的です。しかし、従来の RL 学習には以下の重大な限界が存在します。

• 探索の不足: RL 学習は、単一の正解への収束を促進する一方で、多様な推論経路の探索（Exploration）が不十分です。その結果、Pass@k（k 個のサンプル中いずれかが正解である確率）の向上が Pass@1 に比べて限定的になります。
• 確率の圧縮効果 (Squeezing Effect): 著者らの予備分析により、RL 学習における「確率の圧縮効果」が特定されました。これは、低確率のトークンに対して負の勾配を適用した際、その確率質量が均等に再分配されるのではなく、逆にすでに支配的（高確率）なトークンへ偏って集中してしまう現象です。
  • Softmax 関数の正規化特性により、低確率トークンの確率が減ると、相対的に支配的なトークンの確率が増加します。
  • この結果、出力分布が尖り（Sharpening）、エントロピーが低下し、モデルが多様な推論経路を発見できなくなる「エントロピー崩壊」を引き起こします。

2. 提案手法 (Methodology: SPS)

この問題を解決するため、著者らは**「Steering Probability Squeezing (SPS)」**という新しいトレーニングパラダイムを提案しました。SPS は、従来の RL と逆強化学習（IRL）を交互に実行するハイブリッドなアプローチです。

核心的なアイデア

RL による「確率の圧縮」を、単に抑制するのではなく、**探索に有利な方向へ「誘導（Steer）」**します。具体的には、RL で生成された軌跡（rollouts）を「デモンストレーション」として扱い、IRL を用いて分布を明示的に再整形します。

具体的なアルゴリズム

SPS は以下の 2 つのフェーズを反復的に実行します（Algorithm 1）：

1. オンポリシー RL フェーズ:

   • 標準的な RL（GRPO など）を用いて、データセット上で探索を行い、複数の軌跡（rollouts）を生成します。
   • この段階で、正解と不正解の報酬信号に基づき、モデルが更新されます。

2. IRL フェーズ (確率再分配):

   • RL フェーズで生成された軌跡の一部（特に低確率だが潜在的に価値のある軌跡）を「専門家軌跡」として扱います。
   • 前方 KL 発散 (Forward-KL) を目的関数として使用し、現在のポリシー分布を RL で生成された軌跡分布に近づけます。
   • L2TE (Low-Likelihood Trajectory Emphasis): 低確率の軌跡を強調的にサンプリングする戦略を採用し、RL による過度な集中を相殺し、探索領域を広げます。
   • このプロセスでは外部の教師データや追加の報酬モデルは使用せず、モデル自身の生成物のみを利用します。

この RL と IRL の交互ループにより、モデルは単に高報酬の解に収束するのではなく、未探索の領域への確率質量の再分配を継続的に行い、多様な推論経路を発見できるようになります。

3. 主要な貢献 (Key Contributions)

1. RL 学習ダイナミクスの分析と「圧縮効果」の特定:

   • RL 学習における確率質量の再分配メカニズムを理論的に分析し、低確率トークンへの負の勾配が支配的トークンへの確率集中を引き起こす「圧縮効果」を明らかにしました。
   • これにより、Pass@k の向上には理論的な上限（経験的上限）が存在することを示唆しました。

2. SPS アプローチの提案:

   • 従来の RL に IRL を組み合わせた新しいトレーニング手法を提案し、分布の再整形を通じて探索能力を向上させました。
   • 反復的な RL-IRL ループ戦略を導入し、事前の収束を防ぎました。

3. 広範な実験的検証:

   • 5 つのオリンピックレベルの数学ベンチマーク（AIME, BRUMO, HMMT など）で評価を行いました。
   • 多様なモデルサイズ（1.5B, 7B）とデータ規模（3k, 5k, 10k）で、SPS が既存の RL 手法（GRPO, DAPO, GSPO）を凌駕する結果を示しました。

4. 実験結果 (Results)

• Pass@k の大幅な向上:
  • SPS は、Pass@128（128 個のサンプル中どれか一つが正解）において、既存の RL 手法と比較して一貫して大幅な改善を示しました。
  • 例として、Qwen2.5-Math-1.5B モデルにおいて、BRUMO ベンチマークで Pass@128 が 63.33 となり、従来の GRPO に対して +10.00 ポイント の向上を達成しました。
• モデルサイズによる効果の違い:
  • 小規模モデル（1.5B）では、従来の GRPO は分布の尖りを悪化させ探索を阻害しましたが、SPS はこれを解消し、データ制約下でも高い探索能力を発揮しました。
  • 大規模モデル（7B）でも、SPS は GRPO 以上の性能を維持・向上させました。
• 多様性の向上:
  • 生成された軌跡間の類似性を測定したところ、SPS はベースモデルや GRPO に比べて軌跡の多様性（類似度の低さ）が高く、より広範な推論経路を探索できていることが確認されました。
• 計算コスト:
  • IRL フェーズはトレーニング時間の約 3% しか占めず、計算コストの増加は最小限に抑えられています。

5. 意義と結論 (Significance)

本論文は、LLM の推論能力向上における RL 学習の限界を「確率の圧縮効果」という観点から解明し、それを克服するための実用的な解決策を提示しました。

• 理論的洞察: RL 学習がなぜ多様性を失うのかを、勾配更新と Softmax 正規化の相互作用という根本的なメカニズムから説明しました。
• 実用的な手法: 外部の教師データや複雑な報酬モデルを必要とせず、モデル自身の生成物を用いて探索を促進する SPS は、実用的でスケーラブルなアプローチです。
• 将来への示唆: 推論指向の LLM において、単なる RL 最適化ではなく、IRL との交互学習が、モデルの探索能力を拡張し、より複雑な問題解決を可能にする有効な道筋であることを示しました。

要約すると、SPS は「RL による分布の尖りを逆手に取り、IRL を用いて意図的に多様な軌跡へ確率を再分配する」ことで、大規模言語モデルの推論能力、特に多様な解の発見能力（Pass@k）を飛躍的に向上させる画期的な手法です。