現在、強化学習(RL)は大規模言語モデル(LLM)の推論能力向上において大きな可能性を示しています。DeepSeek R1、Kimi K1.5、Qwen 3などのモデルは、LLMの複雑な推論能力を強化する上でのRLの有効性を十分に証明しています。
しかし、効果的な強化学習を実現するためには、根本的な課題、すなわち信用割り当て問題(credit assignment)を解決する必要があります。大規模言語モデルのシナリオでは、シーケンス全体(LLMの応答)の最終評価結果を、シーケンス内の具体的な決定行動(トークン)にどのように帰属させるかという問題です。
この問題の難しさは、報酬信号が非常に希薄であることにあります。成功または失敗の明確なフィードバックは、シーケンスの最後にしか得られません。
現在の主要な方法
強化学習では、通常、信用割り当て問題を解決するためにアドバンテージ推定(advantage estimation)の方法が採用されます。現在、大規模言語モデルに対する強化学習方法は主に2つのカテゴリに分けられ、それらの違いはアドバンテージ推定の粒度にあります。
DeepSeek R1が使用するGRPOのような粗い粒度の軌跡レベル(trajectory-level)の方法は、最終的な報酬のみに基づいてシーケンス全体のアドバンテージ値を計算します。この方法は効率的ですが、フィードバック信号が粗すぎ、LLMは間違った回答の中の正しい部分を報酬とすることができず、また正しい回答の中の冗長な部分を罰することもできません。
もう一つの極端な例は、古典的なPPOのような細かい粒度のトークンレベル(token-level)の方法です。これらの方法は各トークンのアドバンテージ値を推定するため、各トークンの状態価値(V値)を予測するための追加のクリティックモデルに依存する必要があります。しかし、大規模言語モデルの強化学習タスクでは、異なるプロンプトに対応する軌跡分布が大きく異なり、また学習中に各プロンプトからサンプリングされるモデルの応答数が非常に限られているため、クリティックモデルを適切に学習させることが難しく、トークンレベルのアドバンテージ推定に大きな誤差が生じます。
新しいSPOフレームワーク
このボトルネックを打破するため、中国科学院ソフトウェア研究所と香港城市大学の研究チームは、セグメント方策最適化(Segment Policy Optimization, SPO)フレームワークを革新的に提案しました。
論文タイトル:Segment Policy Optimization: Effective Segment-Level Credit Assignment in RL for Large Language Models
著者:Yiran Guo, Lijie Xu, Jie Liu, Dan Ye, Shuang Qiu
リンク:https://arxiv.org/abs/2505.23564
コードリンク:https://github.com/AIFrameResearch/SPO
SPOは、中程度の粒度であるセグメントレベル(segment-level)のアドバンテージ推定方式を採用しています。軌跡レベルの方法のように最後のステップでのみアドバンテージを計算するのではなく、またトークンレベルの方法のように各ステップでアドバンテージを計算するのではなく、生成されたシーケンスをいくつかの連続するセグメントに分割し、各セグメントのアドバンテージ値を計算します。
このセグメントレベルのアドバンテージ推定方式には、いくつかの明確な利点があります。
(1) より優れた信用割り当て:軌跡レベルの方法と比較して、セグメントレベルの方法はより局所的なアドバンテージフィードバックを提供し、モデルが誤った回答の中の価値ある部分を報酬と認識できるだけでなく、正しい回答の中の冗長で無効な部分を罰することも可能にします。
(2) より正確なアドバンテージ推定:トークンレベルの方法と比較して、セグメントレベルの方法は必要な推定点の数が少ないため、モンテカルロ(Monte Carlo, MC)サンプリングを効果的に利用して、追加の不安定なクリティックモデルに依存することなく、より正確で偏りのないアドバンテージ推定を得ることができます。
(3) より柔軟で調整可能:セグメントの分割方法は任意に定義でき、意味的な完全性を要求しないため、トークンレベルと軌跡レベルの間で粒度を柔軟に調整でき、異なるタスクや応用シナリオに適応できます。
SPOフレームワークは主に3つの核となる部分から構成されます。(1) 柔軟なセグメント分割戦略、(2) モンテカルロサンプリングに基づくセグメントレベルのアドバンテージ推定、(3) セグメントレベルのアドバンテージを利用した方策最適化。
このモジュール化された設計により、フレームワークは高度な柔軟性を備え、異なる部分が異なる実装戦略を持つことで、様々な応用シナリオに適応できます。
さらに、チームは異なる推論シナリオに対応するSPOフレームワークの2つの具体的なインスタンスを提案しました。短い思考連鎖(chain-of-thought, CoT)シナリオ向けにはSPO-chainを提案し、この方法はカットポイントベースのセグメント分割と連鎖型アドバンテージ推定を使用します。長いCoTシナリオ向けには、MCサンプリング効率を大幅に向上させるツリー構造のアドバンテージ推定方法を提案しました。
また、チームはトークン確率マスク(token probability-mask)方策最適化方法も提案しました。これにより、セグメント内のすべてのトークンではなく、低確率(重要)トークンに選択的に損失を計算します。著者らは、これらのトークンがモデルの推論軌跡が分岐する可能性のある場所であり、セグメントレベルのアドバンテージが生じる主要な原因であると考えています。この方法はSPO-chainとSPO-treeの両方で使用でき、信用割り当てをさらに強化することができます。
フレームワークと核となる技術
SPOフレームワークは、主に以下の3つの挑戦的な問題を中心に設計されています。(1) 生成されたシーケンスをどのように複数のセグメントに分割するか?(2) 各セグメントに対応するアドバンテージ値をどのように正確かつ効率的に推定するか?(3) セグメントレベルのアドバンテージをどのように利用して方策を更新するか?SPOの3つの核となるモジュールは、上記の3つの質問にそれぞれ答えるもので、各モジュールには異なるシナリオに適応するための複数のオプション戦略が含まれています:
1. セグメント分割 (Segment Partition):
a) カットポイントベースのセグメント分割 (Cutpoint-based Partition): 短い思考連鎖シナリオ向けに設計されており、セグメント分割点を状態価値(V値)が変化する可能性が高い場所に配置します。トークン確率に基づいてセグメント境界を動的に決定し、モデルが「躊躇する」または推論パスを変更する可能性のある重要な点(カットポイント)で優先的に分割することで、信用割り当てをより正確にします。例えば、下図の例では、赤でマークされたトークンがキーポイントであり、青でマークされた縦棒がセグメント分割結果を示しています。
b) 固定トークン数セグメント分割 (Fixed Token Count Partition): シーケンスを固定長のセグメントに分割し、ツリー構造の組織化とアドバンテージ推定を容易にするもので、SPO-tree向けに設計されています。
2. セグメントレベルアドバンテージ推定(Segment Advantage Estimation):
a) 連鎖型アドバンテージ推定 (Chain-based) 方法:短い思考連鎖シナリオでは、MCサンプリングのコストが高くないため、チームは直接的なセグメントレベルのアドバンテージ推定方式を採用し、各セグメント境界の状態価値(V値)を個別に推定し、その後セグメントレベルのアドバンテージ値を計算します。以下の式は連鎖型アドバンテージの推定方法を示しています。
b) ツリー型アドバンテージ推定 (Tree-based): 長い思考連鎖シナリオでは、MC推定のコストが非常に高いため、チームは効率的なツリー型推定方法を提案しました。サンプリングされた軌跡をツリー構造に編成し、ボトムアップの報酬集約を通じて状態価値(V値)を計算します。同じ親ノードの子ノードはグループを形成し、そのグループ内で各セグメントのアドバンテージ値を計算します。この方法では、V値推定に使用されるサンプルを方策最適化にも同時に利用することで、サンプル効率を大幅に向上させます。以下の式はツリー型アドバンテージ推定方法を示しています。
3. セグメントレベルアドバンテージに基づくトークン確率マスク方策最適化(Policy Optimization Using Segment Advantages with Token Probability-mask):
セグメントレベルアドバンテージ値を取得した後、信用割り当てをさらに向上させるために、チームは革新的にトークン確率マスク方策最適化方法を提案しました。方策更新時、セグメントレベルアドバンテージをセグメント内のすべてのトークンではなく、低確率(重要)トークンにのみ割り当てます。この方法は、重要な決定点により正確に報酬/罰を割り当てることができ、学習効率と効果を高めます。以下にSPO-chainとSPO-treeの最適化目標をそれぞれ示します。
a) SPO-chain最適化目標:
b) SPO-tree最適化目標:
ベースライン手法との比較
下図に示すように、短い思考連鎖シナリオにおいて、RhoMath1.1BをベースモデルとしてGSM8K訓練セットで訓練を行った結果、様々な訓練アルゴリズムと比較して、SPOで訓練されたモデルのテストセットの正解率がより高いことが示されました。
長い思考連鎖シナリオでは、下図に示すように、DeepSeek-R1-Distill-Qwen-1.5BをベースモデルとしてMATHデータセットで訓練を行った結果、同じ訓練時間でGRPOよりもテストセットの正解率が高いことが示されました。
以下の表は、長い思考連鎖シナリオにおけるさらなる比較結果を示しています。同じベースモデル(DeepSeek-R1-Distill-Qwen-1.5B)を使用し、GRPO手法で同時に訓練されたモデル(DeepScaleR、STILL-3)と比較して、SPOはMATHデータセットのみを使用し、最大コンテキスト長4Kでしか訓練されていないにもかかわらず、SPO-treeは各コンテキスト長評価で優れた性能を示しました。特筆すべきは、DeepScaleRが32Kのコンテキスト長評価で最高の性能を示したにもかかわらず、短いコンテキスト長(2Kと4K)では最も性能が悪く、元のベースモデルを下回った点です。これは、GRPO訓練方法がモデルのトークン効率を効果的に最適化できておらず、冗長な出力が多く、その結果、コンテキスト長が限られた状況で正解率が低下する問題を引き起こした可能性を示唆しています。
セグメント粒度の影響
実験の結果、非常に細かい粒度(int2、2つのカットポイントごとにセグメント化)は、中程度の粒度(int5)と比較してわずかな改善しか示しませんでしたが、粗すぎる粒度(int100)は、中程度の粒度(int5)と比較して正解率が大幅に低下しました。これは、SPOが中程度の粒度のアドバンテージ値を採用することの有効性を証明しています。
セグメント分割方法の影響
実験によると、短い思考連鎖シナリオでは、提案されたカットポイントベースのセグメント分割方法が最も優れた性能を示し、改行による分割(VinePPO)や固定トークン数分割(Fixed-token-count)よりも優れていることが分かりました。
トークン確率マスクのアブレーション
実験の結果、トークン確率マスクを削除するとSPO-chainの正解率が低下することが示されました。さらに注目すべきは、トークン確率マスクをGRPOに適用すると、その正解率が著しく向上することです。
異なるツリー構造の影響
実験の結果、より小さなツリー構造は初期段階でより高い正解率を示し、これはより多くのデータサンプルをより速く処理できるためと考えられます。しかし、学習が進むにつれて、より大きなツリー構造の方がより良い正解率を示し、これはより大きなツリー構造がセグメントレベルのアドバンテージ推定においてより正確であるためです。
まとめ
本研究では、中間粒度のセグメントレベルアドバンテージに基づくRL訓練フレームワークSPOを提案しました。これは、トークンレベルと軌跡レベルの間でより良いバランスを取り、軌跡レベルよりも優れた信用割り当てを持ち、かつ少数のアドバンテージ推定点しか必要としないため、追加のクリティックモデルなしで、効果的で偏りのないMC方式での推定が可能です。
本論文では、短い思考連鎖シナリオ向けに設計されたSPO-chainと、長い思考連鎖シナリオ向けに設計されたSPO-treeの2つのSPOのインスタンスも提案し、実験を通じてSPOフレームワークと2つのインスタンスの有効性を証明しました。
© THE END
転載は本公式アカウントに連絡し、許可を得てください。
寄稿または取材依頼:liyazhou@jiqizhixin.com