AAPO: 優位性マージンを用いた LLM 推論能力の強化

技術的サマリー（日本語）

本論文「AAPO: Enhancing the Reasoning Capabilities of LLMs with Advantage Margin」は、大規模言語モデル（LLM）の推論能力を強化するための新しい強化学習（RL）アルゴリズム「Advantage-Augmented Policy Optimization（AAPO）」を提案するものです。特に、既存のグループ相対優位性推定（Group Relative Advantage Estimation）手法が抱える「推定された優位性がゼロに近づくことによる学習効率の低下」という課題を解決し、数学的推論タスクにおいて最先端（SOTA）の性能を達成しています。

以下に、問題定義、手法、主要な貢献、実験結果、および意義について詳細をまとめます。

1. 背景と問題定義

近年、LLM の推論能力向上には、教師あり微調整（SFT）よりも強化学習（RL）が有効であることが示されています。特に、PPO（Proximal Policy Optimization）に代わる手法として、GRPO（Group Relative Policy Optimization） が注目されています。GRPO は、価値モデル（Value Model）を不要とし、グループ内で生成された回答間の相対的な報酬に基づいて「優位性（Advantage）」を推定することで、計算コストを大幅に削減します。

しかし、既存のグループ相対優位性推定手法には以下の重大な限界があります：

• ゼロ優位性（Zero Advantage）の問題: グループ内の回答の報酬が均一である場合（すべて高品質、またはすべて低品質）、グループ内の平均に対する偏差が小さくなり、推定された優位性がゼロに近づきます。
• 勾配消失: 優位性がゼロに近づくと、ポリシー更新のための勾配もゼロに近づき、学習が停止または極めて非効率になります。
• 学習後半の停滞: 図 1 に示されるように、学習が進むにつれて生成される回答の質が均質化し、ゼロ優位性の割合が急増し、最終的には 100% に近づく現象が観察されました。

2. 提案手法：AAPO (Advantage-Augmented Policy Optimization)

AAPO は、上記の「ゼロ優位性」問題を解決するために、優位性マージン（Advantage Margin） を導入した新しい RL アルゴリズムです。

核心的な仕組み

AAPO は、既存のグループ内比較だけでなく、ポリシーモデル（$\pi_\theta$）と参照モデル（$\pi_{ref}$）の比較を組み合わせて優位性を再定義します。

1. サンプリング: ポリシーモデルと参照モデルからそれぞれグループ $G$ の回答をサンプリングします。

2. 報酬計算: 両モデルの回答に対して報酬を計算します。

3. 拡張された優位性（$\hat{A}^*$）の計算:
   従来のグループ相対優位性（$\hat{A}_{GRPO}$）に、ポリシーモデルの報酬と参照モデルの報酬の差（優位性マージン）を加算します。

   $$\hat{A}^*_{i,t} = \underbrace{\frac{r_{\theta i} - \text{mean}(r_\theta)}{\text{std}(r_\theta)}}_{\text{グループ相対優位性}} + \underbrace{\text{clip}(r_{\theta i} - r_{ref i}, \delta_{low}, \delta_{high})}_{\text{優位性マージン}}$$

   • $r_{\theta i}$: ポリシーモデルの回答報酬
   • $r_{ref i}$: 参照モデルの回答報酬
   • $\text{clip}$: 安定性のためのクリッピング操作

理論的利点

• ゼロ勾配の回避: 学習が進みグループ内の回答が均質化しても、参照モデル（学習中凍結）との比較により、ポリシーモデルが参照モデルよりも優れている限り、優位性マージンが正の値を維持します。これにより、ゼロ優位性による勾配消失を防ぎ、学習を継続可能にします。
• 安定性の保証: 論文では、AAPO の損失関数が有界であり、パラメータ更新が安定していること、および標準的な条件のもとで局所最適解に収束することを理論的に証明しています（定理 1, 2）。

3. 主要な貢献

1. 既存手法の限界の分析: GRPO やその派生手法（DAPO, GPG など）において、グループ内の回答が同質化すると優位性がゼロになり、学習効率が極端に低下する現象を定量的・理論的に分析しました。
2. AAPO アルゴリズムの提案: 参照モデルとの比較に基づく「優位性マージン」を導入し、クロスエントロピー損失を最適化する新しいアルゴリズムを提案しました。
3. 広範な実験的検証: 複数の数学的推論ベンチマーク（AIME24, MATH-500, AMC23 など）および異なるモデルシリーズ（Llama, Qwen, DeepSeek-R1 派生）において、AAPO が既存手法（GRPO, GPG）を上回る性能を示すことを実証しました。

4. 実験結果

実験は、1.5B、3B、7B パラメータのモデルを用いて行われました。

• 性能向上:
  • Qwen 1.5B モデル: AAPO-1.5B は、同じデータセットでトレーニングされた GRPO-1.5B や GPG-1.5B を上回り、MATH-500 や Minerva などで SOTA 性能を達成しました。特に、より多くのデータでトレーニングされた「Still-3-1.5B-Preview」を上回る結果も示されました。
  • Qwen 7B モデル: AAPO-7B は、SimpleRL-Zero-7B や GPG-7B に対して、平均スコアで 10% 以上の大幅な改善を示しました。
  • Llama シリーズ: Llama 1B/3Bにおいても、GRPO/GPG に対して 2.5%〜4.6% の絶対的な改善が見られました。
• 学習ダイナミクス:
  • 学習中の「ゼロ優位性の割合」を測定したところ、従来の手法では学習後半に 100% に近づくのに対し、AAPO ではゼロ優位性が発生せず、常に意味のある勾配が得られていることが確認されました（図 2）。
• 計算効率:
  • 参照モデルのサンプリングにより若干の時間オーバーヘッド（1.5B モデルでステップあたり 826 秒 vs GRPO 567 秒）が発生しますが、GPU メモリ使用量は GRPO と同等（1.5B で 52.2GiB）であり、追加のメモリコストは生じていません。

5. 意義と結論

AAPO は、LLM の推論能力を強化する RL 学習において、「価値モデル不要（GRPO の利点）」と「ゼロ勾配問題の回避（PPO の安定性）」を両立させた画期的な手法です。

• 実用性: 参照モデルとの比較というシンプルな仕組みにより、複雑な価値モデルの学習や追加の計算リソースを必要とせず、既存の RL 基盤に容易に統合可能です。
• 汎用性: 数学的推論タスクだけでなく、MMLU（一般知識）や LiveCodeBench（コーディング）などのドメイン外タスクにおいても、ベースモデルの能力を維持・向上させることが示され、汎化性能の低下がないことが確認されました。
• 将来展望: 学習時間の短縮（高速推論エンジン vLLM の活用など）や、参照モデルの更新戦略のさらなる最適化など、今後の研究課題が示唆されています。

総じて、AAPO は、LLM の推論能力を限界まで引き出すための RL 学習パラダイムにおいて、重要な進展をもたらす手法と言えます。