複数報酬の最適化効率を高める SA-MRPO 手法──AIME24 で最大 5% の精度向上を達成

論文概要

言語モデル（LLM）の事後学習における強化学習（RL）では、グループ相対アドバンテージを用いる手法がデファクトスタンダードとなっている。しかし、複数の報酬目的を同時に最適化する際、既存手法は固定の重み付き和で報酬ベクトルをスカラー化してからグループごとの標準化を行っていた。本論文では、この設計が「異なる報酬プロファイルを持つロールアウトが同一のアドバンテージを受け取る」「すべての目的が飽和度に関わらず固定の相対重みで最適化される」という 2 つの問題を引き起こすことを指摘している。その結果、すでに解決済みの目的に勾配リソースが浪費され続けるという課題に対し、報酬目的ごとに独立して標準化を行い、バッチ全体の飽和度推定量に基づいてその寄与を動的に割引く「Saturation Aware Advantage Reweighting for Multi-Reward Policy Optimization (SA-MRPO)」を提案する。

関連研究

LLM の報酬最適化においては、単一の報酬に基づくグループ相対方策最適化（GRPO など）や、複数報酬を扱う既存の複数報酬グループ相対方策最適化（GDPO）などが研究されてきた。しかし、これらは報酬のスカラー化や静的な重み付けに依存しており、最適化の進行に伴う目的ごとの学習進度の違い（飽和度）を動的に反映する仕組みを持っていなかった。本研究は、この静的な重み付けの限界を克服する初めてのアプローチである。

新規性と貢献

本研究の主要な貢献は、バッチ全体の推定に基づく飽和度を考慮した動的なアドバンテージ再重み付けメカニズムの導入である。これにより、単なるスケーリングの調整にとどまらず、更新の符号を反転させることも可能になり、未最適化の目的に対して効率的に勾配リソースを再配分できる点にある。数学的推論やコーディングタスクにおいて、既存手法を上回る実証的な成果を上げている。

提案手法の詳細

SA-MRPO は、複数の報酬目的のそれぞれに対して独立に標準化処理を適用する。さらに、バッチ内の各目的の達成度（飽和度）を動的に推定し、すでに十分な性能に達している目的への貢献度を適応的に割引く設計となっている。これにより、まだ達成されていない難しい目的に対して最適化の努力を集中させることが可能になる。また、必要に応じてパラメータ更新の方向（符号）自体を反転させることで、過剰最適化を防ぐ仕組みを内包している。

評価・考察

2つおよび3つの報酬目的を組み合わせた数学的推論の実験では、15のベンチマーク比較のうち 12 のケースで、GDPO と比較してより困難な正解目的の性能向上を達成した。具体的には AIME24 で最大 5% の精度向上を示している。さらに、適応的推論に関する 5 つのベンチマークすべてにおいて精度を改善し、AMC23 では平均 3.8%、最大 9.2% の向上を記録した。コーディングベンチマークにおいてもパス率が最大 2.3% 向上しつつ、すでに満たされている簡単な目的の性能を維持することに成功している。

応用例と今後の展望

医療文書の自動生成や金融データ分析など、複数の評価指標（正確性、安全性、フォーマット遵守など）を同時に満たす必要がある大規模言語モデルの実務的な事後学習プロセスにおいて、学習効率の劇的な改善が見込める。国内の大手テック企業や AI 研究機関が社内独自 LLM のマルチタスク・マルチリワード調整を行う際の計算コスト削減と精度向上の双方に直結するインパクトを持つ。今後は、より多様な報酬ドメインや大規模モデル群へのスケーラビリティ検証が求められる。

結論

SA-MRPO は、複数報酬を持つ強化学習において報酬の飽和度に応じた動的なアドバンテージ再重み付けを行うことで、未達成の目的に最適効率でリソースを集中させる画期的な手法である。各種ベンチマークでの顕著な改善により、複数目的を同時に最適化する際の従来手法の構造的欠陥を克服できることが示された。

注釈

• RL (Reinforcement Learning): 強化学習。試行錯誤を通じて環境から得られる報酬を最大化する方策を学ぶ機械学習の枠組み。
• GDPO (Group-Relative Policy Optimization): グループ内の相対的な報酬を用いて方策を最適化する既存の複数報酬対応手法。