AI-Papers
- LLM強化学習で同じ誤りを繰り返す「サンプリング多様性崩壊」に対し、過去ロールアウトの中間表現を記憶・密度ベースクラスタリングで分析して動的ペナルティを付与するMEDSを提案
- pass@1で最大4.13ポイント、pass@128で最大4.37ポイントの向上を5データセット・3ベースモデルで実証し、高い再現性を確認
- エントロピー正則化とは異なる「履歴的行動信号の直接活用」により、GRPO系後訓練の多様性問題を実践的に解消するアプローチ
研究の背景
近年、大規模言語モデル（LLM）の推論能力を高める手法として、強化学習（RL）を用いた後訓練が注目を集めています。GRPOをはじめとするRL系手法は、モデルが自ら生成した解答に報酬信号を与えることで、段階的な推論能力の向上を実現しています。
しかしこうした手法には、訓練が進むにつれてポリシーが類似した誤答を繰り返す「サンプリング多様性崩壊」という課題が存在します。同じパターンの失敗が何度も生成されることで、モデルは新しい解法空間を探索できなくなり、訓練の効率と最終的な性能が低下してしまうのです。
多様性崩壊の何が問題か
RL後訓練では、モデルが複数の解答候補（ロールアウト）を生成し、それぞれに報酬を割り当てることで学習が進む仕組みです。理想的には、試行錯誤を通じて多様なアプローチが探索されるはずです。
ところが実際の訓練では、ポリシーが同一の誤ったアプローチに収束してしまう現象が起きやすい性質があります。同じ誤りが繰り返されても報酬は変わらないため、モデルはその失敗パターンから脱け出す動機を持ちにくくなっています。結果として、より広い解法空間の探索が阻まれ、訓練データの利用効率も著しく低下するのです。
従来のアプローチとしてはエントロピー正則化（出力の確率分布を分散させる手法）が知られていますが、出力分布の多様性を間接的に促すにとどまり、具体的にどのエラーパターンが繰り返されているかを直接把握する仕組みは持ちません。
MEDSの提案手法
本研究が提案するMEDS（Memory-Enhanced Dynamic reward Shaping）は、過去のロールアウトを「記憶」として活用する点が最大の特徴です。この手法は3つの段階で構成されています。
第1段階は「記憶の蓄積」です。モデルが生成した過去のロールアウトについて、モデル内部の中間表現（隠れ状態）を保存していきます。テキスト表面ではなく中間表現を用いることで、意味的に類似したエラーパターンを精度よく捉えられる仕組みになっています。
第2段階は「反復エラーの特定」です。蓄積された中間表現に対して密度ベースのクラスタリングを適用し、頻繁に出現するエラーのパターンを識別するのが、この段階の役割です。クラスタの密度が高いほど、同じ失敗が多く繰り返されていることを示します。
第3段階は「動的ペナルティの付与」です。現在生成中のロールアウトが高密度エラークラスタに属すると判定された場合、そのロールアウトに対してより重いペナルティが与えられます。これにより、モデルは特定の失敗パターンへの固着から離れ、未探索の解法空間へと誘導される仕組みです。
エントロピー正則化と比較したときの本質的な違いは、「履歴的な行動信号を報酬設計に直接組み込む」点です。過去に何が失敗したかを明示的に記憶しながら報酬を形成することで、より精度の高い探索促進が可能となります。
実験結果
MEDSの有効性は、5つのベンチマークデータセットと3種類のベースモデルを用いた広範な実験で検証されました。
コード生成タスクにおいては、pass@1（1回の生成で正解する割合）で最大4.13ポイント、pass@128（128回の生成のうち少なくとも1回正解する割合）で最大4.37ポイントの向上が確認されています。複数のデータセットとモデルにわたって一貫した改善が得られており、手法の汎化性能の高さを示す結果です。
また、LLMベースの評価と定量的多様性メトリクスを組み合わせた分析からは、MEDSが実際にサンプリング中の行動多様性を増加させていることも確認されています。スコアの向上が多様性の改善と連動している点は、手法の動作原理を直接裏付ける重要な知見です。
さらに、RAGEN-2がテンプレート崩壊を相互情報量で診断したアプローチと比較すると、RL後訓練における多様性劣化という共通課題に対して、記憶ベースと情報理論ベースという異なる方向からのアプローチが提案されつつある状況がわかります。
まとめと今後の展望
MEDSは、LLMの強化学習後訓練が抱える「サンプリング多様性崩壊」という実践的な課題に対し、過去ロールアウトの記憶と密度ベースクラスタリングを組み合わせたシンプルかつ効果的な解法です。
従来のエントロピー正則化が出力分布の間接的な制御にとどまるのに対し、MEDSは「何が失敗し、それがどれほど繰り返されているか」を直接把握して報酬設計に反映させます。この履歴的行動信号の活用という発想は、今後のRL後訓練フレームワークへの応用が期待されるアプローチです。
課題としては、中間表現の蓄積と密度クラスタリングに伴う計算オーバーヘッドがあります。大規模モデルへの適用や長期訓練での効率性については、引き続き検証が必要です。それでも5データセット・3ベースモデルという広範な実証実験の結果は、GRPO系手法を採用する開発者・研究者にとって参照価値の高い知見といえます。
