技術要約: CATPO (Critique-Augmented Tree Policy Optimization)

問題提起

検証可能な報酬（Verifiable Rewards: RLVR）を用いた強化学習は、大規模言語モデル（LLM）の推論能力を向上させるための支配的なパラダイムとなっています。TREERPOのような木構造ベースの手法は、木構造のロールアウトを通じて、ステップレベルの密な報酬信号を生成することで、フラットな軌跡サンプリングよりも優れた成果を上げていますが、決定的な非効率性に直面しています。それは、「情報の乏しい木（uninformative trees）」による計算資源の浪費です。

現在の木構造ベースのアプローチは、訓練中にすべてのサンプリングされた木を等しく扱います。しかし、多くの木は有用な勾配信号をほとんど、あるいは全く提供しません：

• Dead-correct trees（正解で停滞した木）: すべてのリーフが成功している状態。モデルはその問題をすでに習得しており、改善のための対照的な信号を提供しません。
• Dead-wrong trees（不正解で停滞した木）: すべてのリーフが失敗している状態。ポジティブな分岐が存在しないため、グループ相対的なアドバンテージがゼロに収束し、ロールアウトが勾配更新に全く役に立たなくなります。
• Stale trees（陳腐化した木）: ポリシーの確信度がすでに報酬分布を完璧に追跡しており、収益が逓減している状態。

DAPOやGRESOといった既存の手法は、フラットなロールアウトにおける同様の問題に対処するためにプロンプトのフィルタリングを行っていますが、これらはプロンプトレベルでの操作に留まっています。これらの手法は、特定の分岐が驚きや誤りを含んでいる一方で、他の分岐はそうではないといった、木構造のロールアウト内部の構造を診断したり修復したりすることはできません。

手法

CATPO（Critique-Augmented Tree Policy Optimization）は、木のレベルでロールアウトを診断、修復、および再重み付けを行う4段階のパイプラインを通じて、この浪費に対処します。

1. 木の情報量スコア ($F(T)$)

CATPOは、ポリシー更新に対する木 $T$ の有用性を定量化する指標を導入します。このスコアは、サンプリング中に既に生成されているデータを用いて、追加コストなしで算出されます。これは2つの信号を組み合わせたものです：

• リーフの結果の多様性: リーフの通過率（$\hat{p}$）のベルヌーイ分散 $\hat{p}(1-\hat{p})$ によって測定されます。この項は、結果が混在しているとき（正解と不正解が混ざっている状態）に最大となり、すべてのリーフが一致するときにゼロになります。
• ポリシーと報酬の非相関性: ノードの対数確率と伝播された報酬の間のピアソン相関係数 $\rho_{\pi,r}$ を用いた $1 - \rho^2_{\pi,r}$ で測定されます。これは、ポリシーの確信度が報酬構造を線形に説明できていない、残留する不確実性を捉えます。

最終的なスコアは $F(T) = \hat{p}(1-\hat{p}) \cdot (1 - \rho^2_{\pi,r})$ です。木は閾値に基づいて以下の4つのレジームに分類されます：

• Dead-correct/Dead-wrong: 低い多様性（ほぼ全員一致の結果）。
• Stale: 中程度の信号。
• Informative: 高い多様性と非相関性。

2. クリティークによるヒーリング（修正）

Dead-wrong trees（すべての分岐が失敗している木）に対して、CATPOはロールアウトを破棄しません。代わりに、木を「ヒーリング（修復）」しようと試みます：

1. 失敗の特定: 幅優先探索を用いて、すべての子供が失敗（報酬 < $\epsilon$）している最も浅いノード $v^*$ を特定します。
2. クリティーク（批評）の生成: ポリシーモデル自体に対し、ノード $v^*$ におけるエラーの自然言語によるクリティークを生成するよう促し、論理的または数学的な間違いを特定させます。
3. グラフト（接ぎ木）: モデルに対し、問題、 $v^*$ までの部分的な解、および生成されたクリティークを入力として与え、$k$ 個の洗練された継続（continuation）を生成させます。
4. 再伝播: これらの新しい分岐を $v^*$ にグラフトし、再スコアリングを行い、報酬を下から上へと伝播させます。これにより、「死んだ」木をポジティブな報酬の多様性を持つ木へと変換します。

3. 情報量による重み付け損失

最終的な訓練目的関数は、各木のポリシー勾配への寄与を、正規化された情報量スコア $w(T) = F(T) / \bar{F}$ によってスケーリングします。

• Informative な木は、パラメータ更新を駆動するためにフルウェイトを受け取ります。
• Dead-correct や Stale な木は、ウェイトが下げられます。
• ヒーリングされた木は、ヒーリングプロセスが成功した場合、Informative として扱われます。

目的関数は、アドバンテージが $w(T)$ によって変調されるクリップされたサロゲート損失（PPO/TREERPOと同様）を使用し、全体の勾配の大きさを維持しながら、最も有用なデータに更新を集中させます。

主な貢献

1. 木レベルの診断: CATPOは、プロンプトレベルのフィルタリングを超えて、木構造のロールアウト内の情報を明示的に診断し、介入する初めての手法です。
2. ゼロコスト・スコアリング: 情報量スコア $F(T)$ は、補助的なモデルや追加のフォワードパスを必要とせず、既存の木サンプリング統計（リーフの結果と対数確率）を活用します。
3. クリティークによるヒーリング: 「死んだ」木から信号を回復するために、ポリシーモデル自身にエラーを批判させ、修正された分岐をグラフトするという、独自のメカニズムを提供します。これにより、無駄になった計算を有用な訓練データへと転換します。
4. 重み付け最適化: 高価値な木を動的に優先させる、情報量による重み付け損失を実現しました。

実験結果

実験は、MATHデータセットで訓練された Qwen2.5-Math-1.5B を用い、AIME24、MATH-500、OlympiadBench、MinervaMathの4つのベンチマークで実施されました。

• パフォーマンス: CATPOは 37.5% のマクロ精度を達成し、ベースラインであるTREERPO（35.6%）を 1.9%、フラットな軌跡のGRPO（32.7%）を 4.8% 上回りました。
• 難易度の高いタスク: 改善は、ベースラインがより多くの dead-wrong な木を生成する傾向にある困難なベンチマーク（AIME24およびOlympiadBench）で最も顕著でした。これは、CATPOが信号を回復する機会が多いことを示しています。
• 汎化性能: 学習分布外のMinervaMathにおいても大幅な向上が見られ、本手法が単なる過学習ではなく、転移可能な推論戦略の学習を助けていることが示唆されました。
• 仮説の検証:
  • H1: 情報量スコア $F(T)$ は、木ごとの勾配ノルムと強い正の相関（Spearman $\rho = 0.62$）を示し、訓練の有用性の信頼できるプロキシであることが検証されました。
  • H2: クリティークによるヒーリングは、少なくとも1つの正しい分岐を導入することによって、dead-wrong な木の 88.1% の情報量を向上させることに成功しました。

意義と主張

論文は、CATPOが、木ベースのサンプリングに内在する「浪費」に直接対処することで、LLMの推論におけるRLVRのより効率的かつ効果的なアプローチを提供すると主張しています。CATPOは、情報の乏しい木を区別し、後者を能動的に修復することで、計算資源を学習を促進するデータへと集中させます。著者らは、このアプローチが別途の報酬モデルや価値ネットワークを必要とせず、GRPO/TREERPOのシンプルさを維持しながら、複雑な推論タスクにおける性能を大幅に向上させることを強調しています。本手法は、従来の木ベースのRLを補完する進歩であり、特にロールアウトの内部構造を無視してきた従来の手法を標的としています。