AI・機械学習
上級

推論時探索(Test-Time Search)(スイロンジタンサク)

LLMの推論時にツリー探索やモンテカルロ木探索(MCTS)などの体系的な探索アルゴリズムを適用し、最適な推論パスを発見する手法。AlphaGoの探索技術をLLM推論に応用したアプローチで、数学やコード生成で高い性能を達成している。

0 回閲覧
0 いいね

推論時探索とは

推論時探索(Test-Time Search)とは、LLMが推論を行う際にツリー探索アルゴリズムを活用して、可能な推論パスの空間を体系的に探索する技術である。単純な逐次生成やランダムサンプリングとは異なり、AlphaGoで実績のあるモンテカルロ木探索(MCTS)やビームサーチを推論に応用することで、最適な推論経路を効率的に発見する。

探索アルゴリズムの比較

アルゴリズム探索方式計算量精度適用場面
Greedy Decoding最尤トークン逐次選択O(1)基準低レイテンシ必須
Beam Search幅優先・固定ビーム幅O(B×L)中翻訳・要約
Best-of-NN回独立サンプリングO(N)中〜高単純な検証向き
MCTS木構造の確率的探索O(N×D)高数学・論理推論
A* Searchヒューリスティック誘導O(B×L)高制約付き探索
Iterative Refinement逐次改善ループO(K)中〜高文章改善

MCTSのLLM適用

モンテカルロ木探索をLLM推論に適用する際の4つのフェーズを説明する。

  • 選択(Selection): ルートノードから、UCB1スコア(exploitationとexplorationのバランス)に基づいてリーフノードまで下降する。各ノードはテキストの部分生成に対応
  • 展開(Expansion): リーフノードで新しい推論ステップ(トークン列)を生成し、子ノードとして追加。LLMの条件付き生成確率でサンプリング
  • シミュレーション(Simulation/Rollout): 展開したノードからランダムまたはポリシーに基づいて推論を完了まで進め、最終結果の品質を評価。PRMをロールアウト評価に使用するケースが多い
  • 逆伝播(Backpropagation): シミュレーション結果のスコアをルートまで逆伝播し、各ノードの訪問回数と累積報酬を更新

rStar-Math: MCTSの最前線

Microsoft Research が2025年に発表したrStar-Mathは、MCTSをLLM数学推論に適用した代表的なシステムである。

  • アーキテクチャ: 7BパラメータのSLM(Small Language Model)+ MCTS + 段階的検証PRM
  • 性能: AIME 2024で90.0%の正答率を達成(OpenAI o1の83.3%を上回る)
  • 特徴: 事前学習の追加なしに、テスト時の探索のみで性能を大幅に向上
  • 計算コスト: 1問あたり約500〜2000回のMCTSシミュレーションを実行。A100 GPUで1問あたり30秒〜5分
  • 自己進化: MCTS探索で得られた高品質な推論トレースを蒸留データとして収集し、ベースモデルを段階的に改善する自己進化ループ

ビームサーチの高度化

伝統的なビームサーチも、テスト時計算の文脈で大幅に改良されている。

  • PRMガイドビームサーチ: 各ステップでPRMスコアを参照し、高スコアのビームのみを残す。標準ビームサーチ(対数確率のみ)と比較して数学タスクで15〜25%の精度向上
  • 多様性ビームサーチ: ビーム間の多様性を明示的に促進し、類似した推論パスに収束することを防ぐ。ハミング距離やコサイン距離でビーム間の多様性を測定
  • 適応的ビーム幅: 問題の難易度に応じてビーム幅を動的に調整。初期段階で信頼度が高ければビーム幅を4に縮小し、不確実性が高ければ32に拡大
  • 段階的ビームサーチ: 粗い探索(ビーム幅64・浅い深度)で有望なパスを特定し、精密探索(ビーム幅4・深い深度)で最良解を洗練する2段階方式

探索空間の設計

LLMの推論を探索問題として定式化する際、探索空間の設計が性能に大きく影響する。

  • トークンレベル: 1トークンごとにノードを作成。探索空間が巨大(語彙サイズ^系列長)で非効率だが、最も細粒度の制御が可能
  • 文レベル: 1文(ピリオドや改行まで)ごとにノードを作成。探索空間が適度に圧縮され、実用的なバランス
  • ステップレベル: 論理的な推論ステップ(数式の1行、コードの1ブロック)ごとにノードを作成。PRMとの親和性が高く、最も広く採用されている
  • チャンクレベル: 意味的に自己完結したチャンク(50〜200トークン)ごとにノードを作成。固定長ではないため、LLMの判断で区切り位置を決定

実装例: vLLM + MCTS統合

推論パイプライン概要:
1. 問題文 → LLM(vLLM推論エンジン)
2. MCTS管理ノード → 選択フェーズ(UCB1)
3. vLLMバッチ推論 → 複数ノードの展開を一括処理
4. PRM検証 → 各ノードのスコア算出
5. 逆伝播 → ツリー更新
6. 反復終了条件判定(信頼度閾値 or 最大シミュレーション回数)
7. 最高累積スコアのパス → 最終回答

性能最適化のポイント:

  • vLLMのContinuous Batchingで複数のMCTS展開を並列処理し、GPUの利用効率を最大化
  • KVキャッシュの共有: 同一のプレフィックス(問題文+共通推論パス)のKVキャッシュを複数ノードで再利用。SGLangのRadix Attention技術で実装
  • 早期打ち切り: PRMスコアが閾値を大幅に下回るパスは途中で打ち切り、計算資源を有望なパスに再配分

2025〜2026年の最新動向

  • AlphaProof(DeepMind, 2025年): IMO 2024で金メダルレベルの数学証明を自動生成。MCTSとLean 4形式検証の統合で、探索結果の正確性を100%保証
  • 並列MCTS(2025年): GPU上で数千のMCTSシミュレーションを並列実行する技術。A100 8枚で1問あたりのMCTS時間を10分の1に短縮
  • 学習済み探索ポリシー(2026年): MCTSのUCB1に代わり、強化学習で学習された探索ポリシーを使用。問題のタイプに応じてexploration/exploitationバランスを自動調整
  • コード実行ガイド探索(2026年): コード生成タスクでMCTSの各ノードでテストケースを実行し、通過するパスに報酬を与える。SWE-Benchで人間エンジニアの80%の解決率を達成

よくある質問(FAQ)

Q1: MCTSは全てのLLMタスクで有効か? A: いいえ。MCTSは「正解を検証可能」かつ「推論ステップが分解可能」なタスクで最も効果的である。数学・コード生成・論理パズルが適しているが、創造的文章生成・対話・要約では効果が限定的またはむしろ品質が低下する場合がある。

Q2: MCTSの計算コストはどの程度か? A: 問題の難易度とシミュレーション回数に依存するが、rStar-Mathの報告では1問あたりA100 GPU 1枚で30秒〜5分。100問のバッチ処理ではA100 8枚で約1〜2時間。コストはGPUの時間単価で$3〜$10/100問程度。

Q3: 自前でMCTSベースの推論システムを構築するには何が必要か? A: 最低限必要なのは、(1) 推論用LLM(DeepSeek R1-Distill-7Bなどオープンソース)、(2) PRM(Math-ShepherdをLLaMA 7Bでファインチューニング)、(3) MCTS制御ロジック(Python 500行程度)、(4) vLLMまたはSGLangの推論サーバー。A100 1枚で7Bモデルの場合、1問あたり1〜5分の推論時間。

まとめ

  • 推論時探索はAlphaGoのMCTS技術をLLM推論に応用した手法
  • rStar-Mathは7BモデルでAIME 90.0%を達成し、探索の威力を実証
  • ビームサーチ+PRMが実用的なバランスで最も広く採用
  • 探索空間はステップレベルの設計がPRMとの親和性で最適
  • 2026年現在、並列MCTSと学習済み探索ポリシーにより計算効率が大幅に改善