SoftmaxGRPO: Learning to Reason using Softmax Advantage Group Estimation

ロボットにパズルを解く方法を教えていると想像してみてください。あなたはロボットに問題の束を与えます。ロボットが問題を解くたびに、単純な「はい」または「いいえ」という成績が与えられます。正解すればハイタッチ、間違えれば「もう一度やってみて」と優しく促されます。これが**強化学習（Reinforcement Learning）**の世界であり、AIは試行錯誤を通じて学習します。しかし、ここでの難しい問題は、どの試行が最も学習に重要だったのかを、どうやってロボットに伝えるかということです。

かつて、研究者たちはGRPO（Group Relative Policy Optimization）と呼ばれる手法を使用していました。GRPOを、10人の生徒の回答を見る教師だと考えてみてください。もし9人が正解して1人が間違えたら、教師はその間違えた生徒に重点を置きます。しかし、もし10人全員が正解した場合（「簡単な」問題の場合）、教師は混乱してしまいます。なぜなら、GRPOが使用する数学は「平均」の差を見つけようとするため、意図せず、生徒がすでに解ける問題を解いている時に最も大きな声で叫んでしまい、難しい問題に対してはささやくようになってしまうからです。これは、スタープレイヤーが簡単なレイアップを外したことにコーチが怒鳴り、初心者がドリブルに苦戦しているのを無視しているようなものです。これでは、ロボットの脳の力を既知のことに浪費させ、新しいことを学ぶべき時に行き詰まってしまいます。

この論文は、SoftmaxGRPOと呼ばれる新しい手法を紹介しています。この「混乱する平均」の数学を使う代わりに、著者らは「温度スケール付きソフトマックス（temperature-scaled softmax）」アプローチを提案しています。これは、ロボットの注意が自然に最も興味深い試みに引き寄せられる「ヒートマップ」のようなものです。もし問題が簡単でロボットが正解した場合、この手法は「よくやった、でもこれ以上熱心に勉強する必要はない」と伝えます。もし問題が難しくてロボットが苦戦している場合、この手法は「これこそが集中すべきことだ！」と伝えます。これは、ロボットの注意を簡単な問題から外し、実際に学習できる問題へと自動的にシフトさせるスマートなフィルターとして機能します。研究者たちはこれを数学の問題、クリエイティブ・ライティング、会議の要約でテストし、たとえ「成績」が完璧なスコアではなく大まかな推定値であっても、彼らの新手法が従来の方法よりも一貫して速く、より良くロボットを学習させることを発見しました。

---

問題点：スタープレイヤーに怒鳴る教師

AIとその教師の物語を掘り下げてみましょう。AIの世界では、**グループベース強化学習（Group-Based Reinforcement Learning）**と呼ばれるテクニックをよく使います。例えば、AIに数学の問題を解くよう頼むとします。単に一度試すのではなく、10通りの異なる方法（これらは「ロールアウト」と呼ばれます）で試行させます。そして、その10個の回答をまとめて確認します。

旧来の手法であるGRPOは、グループの「平均」のパフォーマンスを計算する教師のように機能します。AIが質問に正解すれば、GRPSOは報酬を与えます。間違えれば、より低い報酬を与えます。問題は簡単な質問において発生します。もしAIがある特定の種類の数学問題にすでに非常に長けている場合、10回の試行のほとんどが正解になります。従来のGRPOの数学では、これが奇妙な状況を生み出します。「平均」が非常に高いため、「完璧な」回答と「ほぼ完璧な」回答の間のわずかな差が増幅されてしまうのです。数学の結果として、AIがすでにマスターしている行動に対して、その行動を変えるように激しく叫んでしまうことになります。これは、プロのバスケットボール選手がフリースローをわずかに外しただけでコーチが怒鳴り、ボールの持ち方すら分からない初心者には目もくれないようなものです。AIは壊れていないものを「直そう」としてエネルギーを浪費し、その結果、難しいことを学ぶためのエネルギーが残らなくなってしまうのです。

解決策：よりスマートなヒートマップ

Rice大学の著者たちは、SoftmaxGRPOと呼ばれる修正案を提案しました。彼らは、「zスコア」（数値が平均からどれだけ離れているかを測るもの）を使用する代わりに、**ソフトマックス（softmax）**関数を使用すべきだと考えました。

ソフトマックスを、注意の「ヒートマップ」と考えてください。これは報酬（成績）を取り込み、「温度（$\tau$）」設定を含む特別な公式を用いて重みに変換します。

• 高温（High Temperature）: ヒートマップは平坦です。すべての試行がほぼ同じ注意を引きます。これは、AIがゆっくりとランダムに学習する古いREINFORCE法のようなものです。
• 低温（Low Temperature）: ヒートマップは非常に鋭くなります。AIは最高の試行に集中的に注目し、それ以外を無視します。これは、単一の最良の答えを見つけるのには優れていますが、不安定になりやすいMaxRL法のようなものです。

SoftmaxGRPOはその中間に位置します。温度設定を使用して滑らかな曲線を作り出します。AIが簡単な質問に正解すると、ヒートマップはクールな状態を保ち、「よくやった、次へ進め」と指示します。AIが難しい質問に間違えると、ヒートマップは温かい状態を保ち、「これは重要だ、これを勉強せよ！」と指示します。

魔法のような点は、この方法が**重みを限定された範囲内に収める（keeps the weights bounded）**ことです。質問がいかに簡単であっても、「叫び」が無限大になることはありません。これにより、AIが簡単なプロンプトに対して脳の力を浪費することを防ぎます。

分かったこと：真実の証明

著者たちは単に推測したのではなく、数学的な検証を行い、テストを実行しました。

1. 数学は堅実である（バイナリ報酬に対して）
「正解」か「不正解」かの単純な答え（バイナリ報酬）を持つ質問に対して、SoftmaxGRPOが完璧で滑らかな目的関数を作成することを証明しました。彼らは、温度が低くなるにつれて、この手法が自然にMaxRL（最良の結果に焦点を当てる手法）へと変化し、グループサイズが非常に大きくなると最尤推定（Maximum Likelihood）（学習のゴールドスタンダード）のように振る舞うことを示しました。決定的なのは、GRPOとは異なり、SoftmaxGRPOは簡単な質問に対して決して暴走しないことを証明した点です。

2. 魔法の限界
彼らは限界についても発見しました。報酬が単なる「正解/不正解」ではなく、多くのレベル（例えば1から100までのステップがあるスコア）を持つ場合、数学は複雑になります。彼らは、報酬レベルが3つ以上あるグループの場合、あらゆるグループサイズに対して機能する単一の完璧な「スカラー目的関数（単純な公式）」を必ずしも見つけられないことを示しました。これは、この手法が理論的に最も完璧なのは単純な「正解/不正解」のシナリオにおいてであるが、より複雑なスコアに対しても実用上はうまく機能することを意味します。

3. 実世界の成果
彼らは、15億パラメータを持つモデル（中規模のAI）を用いて、いくつかのタスクでテストを行いました。

• 数学 (GSM8K, Countdown, DeepMath): 完璧な「検証器（verifier）」報酬（コンピュータが答えを正確にチェックするもの）を使用した場合、SoftmaxGRPOはDeepMathで**51.8%の精度に達し、従来のGRPOを上回りました。Countdownでは58.1%**に達しました。
• クリエイティブ・ライティング (詩): ここで非常に興味深いことが起こりました。詩には「正解」がありません。あるのは「類似度スコア」（その詩がどれほど優れた例に似ているか）だけです。これらは「弱い」かつノイズの多い報酬です。従来のGRPOはここで苦戦しました。しかし、SoftmaxGRPOは、**35.0%からスタートしたモデルを、詩のタスクにおいて68.0%**まで押し上げました。これは劇的な跳躍であり、教師が完璧でない場合でもこの手法が機能することを証明しています。
• 要約 (MeetingBank): 要約スコアを35%から70%へと向上させました。

4. 注意がどこに向かっているか
最も雄弁な証拠は、AIが「勾配予算（gradient budget）」（学習エネルギー）をどこに費やしたかを見ることで得られました。

• 旧GRPO: 「ほぼ解決済み」のプロンプト（AIがすでに90%以上の確率で正解できる質問）に、エネルギーの**36.4%**を費やしていました。これは、簡単なことに時間を浪費していたことを意味します。
• SoftmaxGRPO: それらの簡単なプロンプトには、わずか**10.0%**しか費やしませんでした。そのエネルギーを、AIが苦戦しているより難しい質問（20%から90%の範囲）へとシフトさせたのです。

まとめ

この論文は、重要度の計算方法を標準的な平均から温度スケール付きソフトマックスに置き換えるだけで、AIの学習における重大な欠陥を修正できることを示唆しています。これにより、AIがすでに知っていることに執着するのを止め、自身を賢くするために本当に必要な課題に集中させることができます。

数学的に最も厳密なのは単純な「正解/不正解」の報酬に対してですが、実験によれば、詩の作成や会議の要約に使用されるような、曖昧で弱い報酬に対しても素晴らしい効果を発揮します。これは「ドロップイン（そのまま置き換え可能）」な手法であり、つまり、コード内の小さな変更だけで、AIの推論能力を向上させる大きな改善をもたらします。著者らは、この手法が学習信号を再配分するための堅牢な方法であり、AIが最も重要な場面で時間を過ごせるようにすることを結論付けています。