推論時探索(Test-Time Search)(スイロンジタンサク)
LLMの推論時にツリー探索やモンテカルロ木探索(MCTS)などの体系的な探索アルゴリズムを適用し、最適な推論パスを発見する手法。AlphaGoの探索技術をLLM推論に応用したアプローチで、数学やコード生成で高い性能を達成している。
推論時探索とは
推論時探索(Test-Time Search)とは、LLMが推論を行う際にツリー探索アルゴリズムを活用して、可能な推論パスの空間を体系的に探索する技術である。単純な逐次生成やランダムサンプリングとは異なり、AlphaGoで実績のあるモンテカルロ木探索(MCTS)やビームサーチを推論に応用することで、最適な推論経路を効率的に発見する。
探索アルゴリズムの比較
| アルゴリズム | 探索方式 | 計算量 | 精度 | 適用場面 |
|---|---|---|---|---|
| Greedy Decoding | 最尤トークン逐次選択 | O(1) | 基準 | 低レイテンシ必須 |
| Beam Search | 幅優先・固定ビーム幅 | O(B×L) | 中 | 翻訳・要約 |
| Best-of-N | N回独立サンプリング | 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と学習済み探索ポリシーにより計算効率が大幅に改善