「思考レベルビームサーチによる推論」

要旨

テスト時計算スケーリングは、大規模推論モデル（LRM）の性能を左右する主要な要因であるが、その極端な非効率性が現在の手法の限界を規定しており、重要な問いは「どれだけ計算を費やすか」から「どこに計算を配分するか」へと移行している。本稿では、テスト時推論を、部分的な軌道（トラジェクトリ）に対する制約付き計算配分問題として定式化する。固定ハードウェア予算の下では、既存のパラダイムは最も有望な部分進捗に計算を能動的に配分できない。従来の並列サンプリングはトレースを独立に扱い、深刻なメモリボトルネックを引き起こす一方、減算的枝刈りはハードウェアを遊休させ、出力分布を能動的かつ十分にシフトさせることができない。この二分法を克服するため、我々は思考レベルのビームサーチを実行する推論アルゴリズムGambitを提案する。Gambitは、有望でない軌道を定期的に枝刈りし、高品質なプレフィックスから即座に分岐することで、隠れ状態をプローブする軽量スコアラを通じて計算を最も有望な推論トレースに動的に集中させるとともに、継続的な高いハードウェア利用率を維持する。複数のモデルとベンチマークにわたる広範な評価により、Gambitが既存のベースラインを厳密に凌駕することが示された。同一のハードウェア制約の下で、本手法は枝刈りベースラインと比較して、HMMT-24では最大+6.7%ポイント、AIME-25では+3.3%ポイントの絶対精度向上をもたらし、トレース完了では2倍以上のスループットを達成し、標準的な並列サンプリングに対して総トークン消費を最大68.5%削減する。

English

Test-time compute scaling is a primary driver of performance in large reasoning models (LRMs), but extreme inefficiency bounds current approaches, shifting the critical question from how much compute to spend, to where to allocate it. We formalize test-time reasoning as a constrained compute allocation problem over partial trajectories. Under a fixed hardware budget, existing paradigms fail to actively allocate the compute to the most promising partial progress: traditional parallel sampling treats traces independently and induces severe memory bottlenecks, while subtractive pruning starves hardware and fails to actively and sufficiently shift the output distribution. To overcome this dichotomy, we introduce Gambit, an inference algorithm that executes thought-level beam search. By periodically pruning unpromising trajectories and immediately branching from high-quality prefixes, Gambit dynamically concentrates compute onto the most promising reasoning traces via a light-weight scorer probing hidden states while maintaining continuous high hardware utilization. Extensive evaluations across multiple models and benchmarks demonstrate that Gambit strictly dominates existing baselines. Under identical hardware constraints, our method yields up to a +6.7\% absolute accuracy gain on HMMT-24 and +3.3\% on AIME-25 over pruning baselines, delivers >2times higher throughput on trace completion, and reduces total token consumption by up to 68.5\% relative to standard parallel sampling.