技術要約：テスト時スケーリングにおける固定トークン予算とツリー探索ポリシーの整合化

問題提起

大規模言語モデル（LLM）は、モデルのパラメータを更新することなく推論時の計算量を割り当てるテスト時スケーリング（Test-Time Scaling）を通じて、回答の品質を向上させることができます。モンテカルロ木探索（MCTS）などのツリー探索デコーディングは、効果的なテスト時スケーリングの一形態ですが、現実世界のデプロイメントでは、多くの場合クエリごとの固定トークン予算が課されます。

既存のツリー探索ポリシーは、その大部分が**予算に依存しない（budget-agnostic）**ものです。これらは通常、トークン予算を単なる終了条件として扱っており、探索行動を導く動的な制約としては扱っていません。この不一致は、主に以下の2つの失敗モードを引き起こします。

1. 終盤の過剰な分岐（Late-stage over-branching）: 探索プロセスが終盤に差し掛かっても浅いノードの拡張を続けてしまい、有望な候補を洗練したり完成させたりする前に予算を使い果たしてしまう。
2. 早期終了（Premature termination）: 探索が早期に停止してしまい、割り当てられた予算の一部が未使用のまま残ってしまう。

既存の早期停止バリアントは、トークンの節約には焦点を当てていますが、残りの予算が減少するにつれて探索戦略を「広範な探索」から「洗練」へと動的にシフトさせるポリシーは欠いています。

手法：予算誘導型MCTS (BG-MCTS)

著者らは、残りのトークン予算に応じて探索ポリシーを明示的に条件付けるツリー探索デコーディングアルゴリズム、予算誘導型MCTS（BG-MCTS）を提案しています。コアとなるメカニズムは、以下の式で定義される予算充足比 $\rho \in [0, 1]$ に基づいています。

$$\rho = 1 - \frac{C_{used}}{B}$$

ここで、$C_{used}$ は生成された累積トークン数、$B$ は総予算です。$\rho$ が減少する（0に近づく）につれて、ポリシーは探索（Exploration）から活用（Exploitation）および完了（Completion）へとシフトします。

BG-MCTSは、2つの予算条件付き制御メカニズムを導入しています。

1. 予算誘導型選択 (BG-PUCT)

標準的なPUCT選択は、探索ボーナスをアニーリングし、完了バイアスを導入するように修正されています。

$$\text{BG-PUCT}(p, s, \rho) = \frac{\tilde{W}(s, \rho)}{m_s} + \rho \cdot c \cdot P(s|p) \sqrt{\frac{\ln(m_p)}{m_s}}$$

• 探索のアニーリング: 探索項に $\rho$ が掛けられます。予算が枯渇するにつれ（$\rho \to 0$）、新しい枝を探索するインセンティブが減少します。
• 完了バイアス: 累積値 $\tilde{W}(s, \rho)$ は、ノードの深さ $d(x)$ と $(1-\rho)$ に比例する項によって補正されます。これにより、既に回答が完了しているノードを不当に押し上げることなく、最終的な回答の生成に近い深いノードへの選択を促します。

2. 予算誘導型ツリー・ワイドニング（Tree Widening）

無駄な終盤の分岐を防ぐため、BG-MCTSは「ワイドニング（新しい子ノードの生成）」を、既存の子ノードと並列の選択可能なアクションとして扱います。仮想生成子ノード（virtual generative child） $s_{gen}(p)$ が選択セットに追加されます。そのスコアは以下の通りです。

$$E_{gen}(p, \rho) = \mu(p) + \lambda \cdot \rho \cdot \sigma^2(p)$$

• $\mu(p)$ と $\sigma^2(p)$ は、既存の子ノードの平均および分散を表します。
• この項は $\rho$ によってスケールされます。予算が豊富なときは、ワイドニング（特に子ノード間の意見が食い違っている場合）が推奨されます。予算が少なくなると、ワイドニングへのインセンティブが抑制され、既存の有望な軌跡を深化させるよう強制されます。

統一アルゴリズム

アルゴリズムは標準的なMCTSループに従いますが、選択とワイドニングの両方の決定を $\rho$ に依存させます。これにより、「広から深へ（wide-to-deep）」のスケジュールが誘発されます。すなわち、探索の初期段階では広範な探索を行い、予算が尽きるにつれて集中的な洗練と回答の完了へと移行します。

主な貢献

1. 予算を考慮したポリシー設計: 本論文は、予算を単なる静的な停止信号として扱うのではなく、残りのトークン予算に合わせて探索行動（探索 vs 活用 vs ワイドニング）を動的に調整する、初のツリー探索デコーディングフレームワークを導入しました。
2. 新規メカニズム: 予算条件付きの探索（BG-PUCT：探索アニーリングと深さバイアス付き値補正を含む）および、予算を考慮したワイドニングのための仮想生成子ノードを提案しました。
3. 実証的検証: オープンウェイトLLM（Llama-3.1-8B, Qwen2.5-7B, Qwen3-32B）を用い、数学的推論ベンチマーク（MATH500, AIME24/25）および物理推論ベンチマーク（UGPhysics）における包括的な評価を行いました。

実験結果

著者らは、固定予算（10k, 20k, 30kトークン）の下で、BG-MCTSを予算非依存のベースライン（Standard MCTS, AB-MCTS-M, LiteSearch）およびスケーリングのベースライン（Repeated Sampling, Sequential Refinement）と比較評価しました。

• パフォーマンス: BG-MCTSは、すべてのモデルおよび予算において、一貫してすべてのベースラインを上回りました。例えば、Qwen2.5-7Bを用いたMATH500において、BG-MCTSは30kトークンで**71.1%の精度を達成しましたが、次に優れたツリー探索ベースラインであるMCTSは65.9%**でした。
• 予算利用率: 早期終了することが多いLiteSearchとは異なり、BG-MCTSは予算の枯渇に近づくにつれて精度を向上させ続けます。
• 定性的挙動: 探索木の分析により、BG-MCTSは探索の初期段階では広い幅を維持するが、後半には特定の枝を深化させる方向にシフトすることが明らかになりました。対照的に、標準的なMCTSは予算が尽きるまで浅いノードのワイドニングを続けてしまい、有望なパスを完成させることに失敗します。
• 物理推論: この手法の有効性はUGPhysicsにも及び、予算に適合した「広から深へ」の挙動が、純粋な数学的推論を超えても有益であることを示しました。

意義と主張

本論文は、ツリー探索の決定を、予算を単なる停止条件として扱うよりも、残りの予算に基づいて条件付ける方が効果的であると主張しています。

• 予測可能なスケーリング: BG-MCTSは、ユーザーがトークン予算を指定すれば、ポリシーがそのコスト範囲内で探索から完了へのスケジュールを自動的に適応させる、予測可能なテスト時スケーリングのためのシンプルなインターフェースを提供します。
• データ合成への示唆: 著者らは、BG-MCTSは問題インスタンスあたりの回答済みノード数は少なくなる傾向にあるものの、精度（正確性）が大幅に高まることを指摘しています。この「より少ないインスタンスで、より良い回答を得る」というレジームは、高品質な推論トレースが広範なカバー率よりも重要となるデータ合成において、特に価値が高いと強調されています。
• 堅牢性: 本手法は、弱すぎる、あるいは飽和した信号を提供する不完全なプロセス報酬モデル（PRM）を使用している場合でも効果的であり、予算を考慮した構造的変更が、ノイズの多い評価信号を補完できることを示唆しています。

結論として、本研究は、固定されたトークン制約に探索ポリシーを適合させることは、リソース制約のある現実世界の環境でLLMをデプロイするための極めて重要なステップであると述べています。