スパース報酬を密なトークン監視に変換——LLM算術推論向け「TBRM-RLOO」フレームワーク

課題：中間ステップへの学習信号がない

大規模言語モデル（LLM）を多段階の算術推論タスクに適用する際、標準的な強化学習手法は大きな壁にぶつかる。報酬がシーケンスの末尾にしか与えられない「スパース終端報酬」の問題だ。例えば、与えられた数字と四則演算を組み合わせて目標数値を導くCountdownタスクでは、最終答えが正解かどうかだけが評価されるため、途中の推論ステップには意味ある学習信号が届かない。

これに加え、LLMのトークン生成という広大な行動空間は勾配分散を著しく増大させ、高品質な軌跡が稀なほど膨大なサンプルを必要とするサンプル非効率の問題も重なる。論文の著者らは「スパース終端報酬と密なステップワイズ監視の根本的なミスマッチ」こそが多段階推論学習の中核的ボトルネックだと指摘している。

提案手法：TBRMを核とした3要素フレームワーク

著者らはこれらの課題に対し、ベルマン整合性を保ちながらトークンレベルの密な監視信号を生成するBellman-consistent Trajectory RLフレームワークを提案した。フレームワークは以下の3つの核心コンポーネントから構成される。

1. 軌跡ベルマン残差最小化（TBRM）

TBRMはKL正則化付きの強化学習パラダイムの下で、軌跡全体を対象としたベルマン整合性を最適化する。従来の遷移単位のベルマン残差最小化（BRM）が「二重サンプリング問題」——非線形遷移下で不偏勾配を得るには同一の状態行動ペアから2つの独立サンプルが必要になる問題——に悩まされるのに対し、TBRMは評価を完全な軌跡スケールに移すことでこの問題を回避する。

さらにTBRMは、モデルの未正規化対数確率（ロジット）を直接利用してQ関数・V関数・方策関数を解析的に導出するため、独立した批評家（critic）ネットワークを必要としない。単一の1回のロールアウトだけで動作し、スパース報酬をトークンレベルの密な監視信号に変換できる点が大きな特徴だ。

2. Reinforce Leave-One-Out（RLOO）グルーピング戦略

ポリシー勾配の分散を削減するため、同一プロンプトからサンプリングした複数の出力をグループとして扱い、グループ内の相対的な優位性を基準にするRLOO戦略を採用した。著者らはRLOOグルーピングがソフトベルマン残差の分散を証明可能な形で削減することを示している。比較対象のGRPO（Group Relative Policy Optimization）よりも一貫して優れた性能を発揮するとしている。

3. ヒンドサイト目標再ラベリング（HGR）

目標条件付き強化学習で用いられるHGRをLLM推論に応用した手法だ。モデルが誤った最終答えを生成した場合、その中間の多段階導出を「実際に達成された数学的結果」として再ラベリングし、失敗したロールアウトを別問題に対する成功例として変換する。これにより失敗サンプルからも意味ある勾配信号を抽出でき、中間推論ステップの学習を加速させる。著者らは、構造化言語生成と複雑な推論タスクへのHGR応用は新規の拡張であると述べている。

Countdownタスクでの実験結果

Countdownタスクでの評価において、TBRM-RLOOフレームワークの性能は以下のように報告されている。

• ロールアウト精度：RLOO採用により0.61→0.73（+20%）
• Pass@1：0.55→0.56（RLOO単独）、さらに目標再ラベリング追加で0.64に到達
• Pass@16：0.74→0.84（+14%）、複数条件の組み合わせで最高84.0%

比較ベースラインとの比較では、SFT（教師あり微調整）比で+8.0ポイント、RLOOベースライン比で+10.0ポイントの絶対的改善を達成し、Pass@16で最高性能を記録したと著者らは報告している。特筆すべきは、既存手法が72.0%前後で停滞するプラトーを、テスト時スケーリングによって突破した点だ。また100ステップ以内という低計算コスト領域においても、Pass@1で56.1〜64.0%という優れた単サンプル効率を示すとしている。

議論と今後の展望

著者らはTBRMがスパース報酬を密なトークンレベル監視に変換する点で有効であり、RLOOがプロンプトや問題難易度に起因する最適化分散を緩和することで訓練を安定させ、HGRがデータを豊かにして中間ステップ学習を加速させると総括している。一方で、モデル規模の複雑性や、より広範なタスクへの汎化については制約があると認めており、今後は多様な多段階推論ドメインへのスケーラビリティ検証を予定しているとしている。

---

出典: Bellman-Consistent Trajectory RL for LLM Arithmetic Reasoning