Thought-Level Beam Search for Reasoning

2026-08-08

AI Summary

This work addresses the challenge of efficiently allocating limited computational resources during inference, a problem that often leads to memory bottlenecks or suboptimal hardware utilization. The authors formulate inference as a constrained partial trajectory computation allocation problem and introduce Gambit, an algorithm that enables beam search at the level of reasoning steps for the first time. Gambit employs a lightweight hidden-state scorer to dynamically evaluate trajectory quality, pruning unpromising paths and branching from high-quality prefixes to actively steer computation toward the most promising reasoning directions. Under identical hardware constraints, Gambit improves accuracy by 6.7% on HMMT-24 and 3.3% on AIME-25, achieves over twice the throughput of baseline methods, and reduces token consumption by up to 68.5%.

Abstract

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 \emph{how much} compute to spend, to \emph{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 \emph{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 $>2\times$ higher throughput on trace completion, and reduces total token consumption by up to 68.5\% relative to standard parallel sampling.

Problem

• test-time compute allocation
• reasoning trajectories
• hardware utilization
• compute efficiency
• large reasoning models

Innovation

• thought-level beam search
• test-time compute allocation
• reasoning efficiency
• dynamic trajectory pruning
• hardware-aware inference