技術要約: RREDCoT – 推論モデルのためのセグメントレベル報酬再分配

問題定義

推論言語モデルの最近の進歩は、特にGroup Relative Policy Optimization (GRPO) アルゴリズムとReinforcement Learning with Verifiable Rewards (RLVR) パラダイムを組み合わせた強化学習（RL）のファインチューニングに大きく依存しています。この設定では、モデルは最終的な回答を生成する前に、Chain-of-Thought (CoT) トレースを生成します。しかし、ここには決定的なボトルネックが存在します。それは、報酬信号が通常、CoTトレース全体が完了し、最終的な回答が検証された後にのみ割り当てられるという点です。これは**遅延報酬問題（delayed reward problem）**を引き起こします。

標準的な手法であるGRPOは、RLにおけるモンテカルロ法に対応しており、報酬が全エピソードに対して一様に分配されるため、個々の推論ステップに対して直接的な監督を提供できず、高い分散（バリアンス）に悩まされます。Monte Carlo Sampling (MCS) は中間状態の値を推定できますが、あらゆるセグメントに対して複数の補完（completion）を生成する必要があり、長いコンテキストにおける高粒度なクレジット割り当て（credit assignment）においては、計算コストが極めて高くなります。既存の代替案として、ステップごとの分析を行うジャッジモデルの使用や、モデル統計から中間ユーティリティを抽出する方法がありますが、これらはノイズを導入したり、追加の計算リソースを必要としたりすることがよくあります。

手法: RREDCoT

著者らは、RUDDERアルゴリズムの原理をCoT生成の特定の構造に適応させた手法であるRREDCoT (Reward REDistribution for Chain of Thoughts) を提案しています。核心となる目的は、追加のモデルや追加の生成ステップを必要とせずに、より密な学習信号を提供するために、最終的な報酬をCoTトレースの中間セグメントへと再分配することです。

1. MDP定式化

CoT生成プロセスは、状態 $s_t$ が元のクエリと $t-1$ までの生成されたセグメントのシーケンスで構成されるマルコフ決定過程（MDP）としてモデル化されます。アクション空間は以下のように分割されます：

• $Y$ (Answer): 最終的な出力をトリガーするアクション。
• $U$ (Continuation): さらなるCoTセグメントを生成するアクション。
  思考を生成することに対する即時報酬はゼロであり、報酬は正しい最終回答が生成された時にのみ実現されます。

2. 最適な報酬再分配

RREDCoTは、連続する状態間の期待リターンの差に基づき、最適な報酬再分配 $\sigma$ を導出します。この手法は、補助的なLSTM（オリジナルのRUDDERで使用されているもの）や広範なサンプリングを必要とせず、生成言語モデル自体を利用して価値関数を近似します。

再分配は次のように計算されます：
$$\sigma_t \propto \hat{R}_{t+1} - \hat{R}_t + \hat{v}_{t+1} - \hat{v}_t$$
ここで：

• $\hat{R}_t$ は即時報酬の推定値（現在のトレースが与えられた時の参照回答の確率）。
• $\hat{v}_t$ は状態の推定価値（将来の補完の期待ユーティリティ）。

これらの値を効率的に推定するために、著者らは参照解のパス（$Y^\star, U^\star$）に対するインポータンスサンプリングを用いた確率報酬 (Probability Reward: PR) 推定器を採用しています。これにより、モデルは現在のセグメントから高ユーティリティな状態に到達する可能性を推定するためのダイナミクス関数（世界モデル）として機能することができます。

3. ハイブリッド・セグメンテーション戦略

クレジット割り当てを扱いやすくするために、RREDCoTはハイブリッド・キーワード・エントロピー・セグメンテーション戦略を導入しています。

1. 初期分割: トレースを一般的なキーワード（改行など）によって分割します。
2. 反復結合: 目標とするセグメント数に達するまで、結合されたエントロピーが最も低い連続するセグメントを結合していきます。
3. 根拠: 高エントロピーのトークンは、高い多様性（決定点）を表しており、中間的な寄与を推定するための自然な「出口ポイント」として機能します。この戦略により、セグメントを可能な限りデカップル（分離）させ、価値関数の推定におけるバイアスを最小限に抑えます。

4. RL目的関数への統合

再分配された報酬は、標準的なポリシーグラディエントの目的関数に統合されます。この手法は**リターン等価性（return-equivalence）**を保証しており、つまり再分配された報酬の合計が元のエピソードのリターンと一致することを意味します。これにより、最適ポリシーを維持しながら収束を加速させます。

主な貢献

1. RREDCoTアルゴリズム: 追加のモデルや追加の生成ステップを必要とせず、価値推定に生成モデル自体を活用する、CoTトレースのための扱いやすいクレジット割り当ておよび報酬再分配アルゴリズム。
2. エントロピーに基づくセグメンテーション: キーワードベースの分割とエントロピー駆動の結合をバランスさせ、中間的な寄与の推定を最適化するために設計された、RREDCoTに特化した新しいセグメンテーション戦略。
3. 理論的および経験的分析: 提案された推定器を、モンテカルロ・サンプリングや様々な属性付け（attribution）手法（Shapley値やLeave-One-Outなど）と比較した詳細な分析。これにより、正解に条件付けられていない属性付け手法は、報酬再分配とは異なる側面を捉えていることが示されました。

実験結果

著者らは、AIME、MATH、Minervaのデータセットを含むいくつかのベンチマークでRREDCoTを評価しました。

• 長文コンテキスト性能: 生成長が最大25kトークンに達するNumina-CoTデータセットにおいて、RREDCoTは標準的なGRPOと比較して優れた性能を示しました。例えば、Qwen3-4B Instructモデルにおいて、RREDCoTはAIME24のスコアでGRPOの0.850に対し、0.908を達成しました。
• 小規模な洗練: 1.5Bパラメータの小型推論モデルを用いたオンライン・リファインメントに適用した場合、RREDCoTは、ハイパーパラメータ調整を行ったGRPOでチューニングされたモデルと比較して、同等または優れた結果を示しました（特にAMC23およびMATH500ベンチマークにおいて）。
• 分散とバイアス: 研究の結果、切り捨てられたモンテカルロ・サンプリング（truncated Monte Carlo sampling）は有効な補完の著しい喪失（バイアスの導入）を招く一方で、RREDCoTの推定器は、極めて低い計算コストでMCサンプリングの結果とよく一致することが示されました。
• 相関分析: 本論文では、勾配ベースの属性付け手法とLeave-One-Out (LOO) 手法は互いに高い相関があるものの、RREDCoTとは相関が低いことが判明しました。これは、属性付け手法が報酬再分配とは異なる信号特性を捉えていることを示唆しており、トレーニングの重み付けのために属性付けのみを使用することは、冗長的または最適ではない可能性があることを示しています。

意義と主張

本論文は、RREDCoTが、推論モデルの強化学習における高分散および遅延報酬の問題に対処することを主張しています。モデル自身の予測能力に基づいて報酬を中間セグメントに再分配することで、モンテカルロ・サンプリングのような計算上の制約を受けることなく、収束を加速し、最終的な性能を向上させます。

著者らは、自身のアプローチが最適報酬再分配の扱いやすい近似であることを強調しています。この手法は、参照解の利用可能性と、ある程度のサブゴール構造（例：数学問題における一意の解）に依存していることを指摘しています。本手法は、素朴なGRPOと比較して1.5〜2倍の計算オーバーヘッドを導入しますが、フル・モンテカルロ・サンプリングと比較すると大幅に効率的です。本研究は、RREDCoTを、従来の属性付け技術とは異なるクレジット割り当ての新しい視点を提供する、学習信号の増幅器および分析ツールとして位置付けています。