Thought-Level Beam Search for Reasoning

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.

• The paper introduces Gambit, a thought-level beam search that scores partial reasoning traces, prunes weak prefixes, and branches from promising ones while maintaining a fixed active capacity.
• Across three open-weight reasoning models and multiple benchmarks, Gambit improves accuracy over pruning baselines, cuts generated tokens by up to 68.5% versus self-consistency, and achieves more than twice the productive trace throughput.
• The method shows how KV-cache prefix sharing and scorer-guided compute allocation can correct errors through targeted continuations, although performance remains dependent on scorer calibration, branch diversity, and benchmark generalization.

Thought-Level Beam Search for Reasoning

Research problem and central thesis

“Thought-Level Beam Search for Reasoning” (2608.08020) addresses a systems and algorithmic limitation in test-time compute scaling for large reasoning models (LRMs). Conventional self-consistency (SC) increases inference-time computation by generating many independent chains of thought and aggregating their final answers. Although effective in some settings, this strategy allocates equal computational resources to trajectories that differ substantially in quality. Long prefixes of incorrect traces may consume most of the inference budget, while promising partial solutions receive no additional computational investment.

The paper’s central claim is that test-time scaling should be treated as a constrained compute-allocation problem over partial reasoning trajectories rather than as a fixed-size collection of independent samples. Under a fixed hardware and KV-cache budget, an effective inference policy must identify valuable prefixes, preserve them, allocate additional continuations to them, and maintain high GPU utilization throughout decoding. The proposed method, Gambit, implements this policy as a thought-level beam search.

The distinction from conventional beam search is important. Gambit does not optimize token-level log probability, which is generally an unsuitable objective for long-form mathematical reasoning. Instead, it periodically ranks partial trajectories using hidden-state-based quality estimates, prunes low-scoring traces, and branches from high-scoring prefixes. The resulting search topology combines selective exploration, prefix reuse, and fixed-capacity scheduling.

The paper frames existing approaches as occupying two unsatisfactory extremes:

• Parallel sampling preserves a large active batch but wastes computation on redundant or low-quality trajectories and rapidly exhausts KV-cache capacity.
• Subtractive pruning reduces memory pressure by terminating weak traces but does not replenish the active population, causing declining concurrency and leaving freed computation unused.

Gambit is designed to resolve this dichotomy through active, zero-sum compute reallocation.

Figure 1: Gambit reallocates freed compute from pruned trajectories to continuations of high-quality prefixes while maintaining a fixed active-trace population.

Empirical motivation for prefix-based allocation

The method rests on the empirical premise that successful and unsuccessful reasoning traces often share substantial intermediate structure before diverging. Consequently, discarding an entire trajectory after a late error wastes the computation invested in its valid prefix. The paper argues that intermediate hidden states contain enough information to estimate continuation quality before a trace has completed.

A motivating experiment pauses a set of traces midway through a difficult AIME 2025 problem, ranks their prefixes with a lightweight hidden-state scorer, and branches from the highest-ranked prefix. For one reported problem, generating 64 continuations from the selected prefix achieves an 87.5% pass@1 rate, compared with 6.2% for independent sampling, a reported 14-fold improvement. Prefix sharing also approximately halves memory consumption because the child traces inherit the parent’s KV-cache rather than recomputing the shared prefix.

This result supports two distinct conclusions. First, the distribution of continuation quality is highly heterogeneous across prefixes. Second, prefix reuse makes it computationally feasible to amplify a promising state without paying the full cost of independently regenerating its history. However, the experiment also exposes a risk: concentrating all computation on a single high-scoring prefix can amplify a confidently incorrect trajectory. Gambit therefore retains a beam of candidates and uses repeated rank-based replacement rather than deterministic Top-1 expansion.

Figure 2: Branching from a high-quality prefix substantially improves pass@1 while reusing the parent’s KV-cache to reduce memory consumption.

The systems motivation concerns the interaction between long contexts and large batches. In standard parallel sampling, every trace continues independently, and the aggregate KV-cache footprint grows until the serving engine must queue or preempt requests. On the reported HMMT-2025 example, this behavior inflates latency by approximately threefold. Pruning-only systems eliminate the queueing pressure but suffer from the opposite problem: once traces terminate, the active batch progressively shrinks.

Gambit maintains the active population through paired pruning and branching. Its reported latency profile remains close to the initial high-utilization regime, at approximately 1.1 times the favorable baseline in the motivating system comparison, while avoiding the concurrency decay of pruning-only methods.

Figure 3: Parallel sampling incurs KV-cache-induced latency inflation, pruning-only inference loses concurrency, and Gambit sustains high utilization through replacement.

Gambit’s thought-level search algorithm

Gambit maintains a fixed-capacity set of logical trajectories. Let $C$ denote the maximum number of active traces, $K$ the number exchanged during each tournament, $\Delta$ the interval between tournaments, and $w$ the warmup threshold before a trajectory becomes eligible as a branching parent.

Each trajectory is segmented into discrete reasoning steps, operationalized through thought boundaries such as double-newline delimiters. At each boundary, a scorer evaluates the most recent hidden state. The running trajectory score is an average over the scores assigned to its observed reasoning steps. The implementation uses either the two-layer MLP scorer from STEP or a custom history-aware sequence scorer introduced in the appendix.

The algorithm proceeds in four stages. During warmup, all traces generate independently and are not eligible for branching. Once sufficient reasoning depth has been reached, Gambit repeatedly performs tournament rounds. At a round, it ranks active traces by cumulative score, removes the bottom $K$, and creates $K$ new child traces from high-ranked eligible parents. Branching is implemented through KV-cache prefix sharing. The active capacity therefore remains constant: every pruned trace is replaced by a new continuation.

If the system falls below nominal capacity because traces complete or are physically evicted, Gambit fills vacant slots by branching from eligible high-scoring prefixes. The warmup restriction prevents newly created children from being immediately branched again, limiting cascaded expansion from immature states.

The paper also introduces a decoupled scheduler/tree architecture. The scheduler view tracks traces that physically hold KV-cache blocks, whereas the tree view tracks the logical search population. A trace that is evicted from physical memory becomes a “ghost trace”: it produces no further tokens but remains logically active until selected for pruning. This design prevents accidental under-capacity branching, which could otherwise trigger repeated expansion from a small number of top-ranked trajectories and cause search collapse.

Figure 4: Gambit alternates generation, scoring, pruning, and prefix-cache branching while preserving a constant logical capacity.

At termination, completed traces are aggregated using a score-weighted majority vote. Rather than counting each answer equally, Gambit weights an answer by the cumulative score of the trajectory that produced it. This aggregation is consistent with the method’s broader objective: not merely to generate many answers, but to increase the representation of high-quality reasoning paths in the final ensemble.

Experimental design

The evaluation covers three open-weight reasoning models:

• Qwen3-4B-Thinking-2507
• DeepSeek-R1-0528-Qwen3-8B
• Phi-4-reasoning-plus-14B

The benchmarks include AIME 2025, AIME 2026, HMMT 2024, HMMT 2025, and GPQA-Diamond. All methods are implemented on vLLM and evaluated using a single 275 GB NVIDIA B300 GPU. The principal experiments allocate a budget of 256 complete traces per problem.

The comparisons include unweighted SC, Slim-SC, DeepConf, and STEP. A particularly important aspect of the design is that Gambit’s primary evaluation uses the same off-the-shelf MLP scorer as STEP. Thus, the principal comparison isolates search topology: Gambit actively reallocates computation, whereas STEP only terminates low-scoring traces.

The selected Gambit configuration uses $C=256$, $K=16$, $\Delta=200$ tokens, and a warmup threshold of 12,000 tokens. The paper reports smooth sensitivity across swap size, tournament interval, warmup duration, and memory utilization, with the selected configuration lying in a broad high-performing region rather than at a sharply tuned optimum.

Accuracy and token-efficiency results

Across model families and benchmarks, Gambit generally improves accuracy over both SC and subtractive pruning while reducing token consumption relative to independent sampling. The strongest reported gains over STEP include:

• +3.3 percentage points on AIME 2025 for Qwen3-4B.
• +3.3 percentage points on HMMT 2024 for Qwen3-4B.
• +2.5 percentage points on AIME 2025 for DeepSeek-R1-8B.
• +1.6 percentage points on HMMT 2024 for Phi-4.
• Up to +6.7 percentage points over pruning baselines on HMMT-24 in the paper’s headline summary.

For Qwen3-4B, Gambit reaches 65.0% on HMMT-24 compared with 61.7% for STEP, and 70.2% on GPQA compared with 66.9% for STEP. For DeepSeek-R1-8B, it reaches 65.6% on HMMT-24 versus 63.3% for STEP. On Phi-4, Gambit obtains 77.1% on GPQA, compared with 76.7% for STEP and 76.3% for SC.

The most substantial efficiency results concern generated tokens. Relative to SC, Gambit reduces token consumption by as much as 68.5%, reported on HMMT-25 with Phi-4. On that setting, token consumption falls from 5.56 million to 1.75 million tokens. The Qwen3-4B result on HMMT-24 shows a 60.6% reduction, while the DeepSeek-R1-8B result on the same benchmark shows a 64.0% reduction.

These results support the paper’s principal algorithmic claim: pruning is not sufficient because it only removes computation, whereas branching converts the removed capacity into additional samples conditioned on informative prefixes. The resulting distribution is not merely smaller; it is actively reshaped.

Figure 5: Gambit occupies a favorable accuracy–efficiency frontier, combining higher accuracy with lower token consumption and competitive latency.

The throughput results reinforce this interpretation. The paper reports more than twofold higher productive trace throughput than independent sampling in representative settings, including a completion rate of 0.216 versus 0.098 for Qwen3-4B. Gambit also reports less than 1% wall-clock overhead for scoring, tree management, synchronization, and scheduling, with GPU computation accounting for approximately 99.03% of execution time.

Figure 6: Prefix inheritance reduces unique token generation, replenishes the trace distribution, and introduces less than 1% system overhead.

Prefix reuse and the token–latency relationship

A technically important result is that token savings do not translate linearly into latency savings. Gambit reduces the number of newly generated tokens per completed trace by inheriting long prefixes, but it also preserves and amplifies long-running high-quality trajectories. Consequently, its completed traces can have longer total logical sequence lengths than those of pruning-only systems, even though their unique generated-token counts are much smaller.

For Phi-4 on AIME-2026, the paper reports a median of approximately 5.2K unique tokens per Gambit trace, compared with 14.5K for SC and 7.7K for STEP. The full logical sequence length includes inherited tokens, whereas unique-token accounting measures only tokens generated after the relevant branch point. This distinction is essential for interpreting the efficiency claims.

The apparent contradiction—large token reductions but latency comparable to pruning-only methods—follows from autoregressive decoding. Gambit maintains a larger population of deep continuations, so the surviving traces incur higher context-dependent attention costs. STEP can exhibit lower raw latency in some settings because it permanently eliminates traces and allows concurrency to decay. That lower latency is therefore partly purchased by reducing the effective vote population and discarding potentially useful computation.

Figure 7: KV-cache sharing substantially reduces the number of unique tokens required to explore alternative continuations.

Scorer dependence and guidance quality

Gambit is scorer-agnostic in its algorithmic formulation, but its behavior depends on the ranking quality of the hidden-state evaluator. The paper therefore trains a history-aware sequence scorer: a compact causal Transformer that attends over the sequence of step-level hidden states and predicts final-answer correctness from the full partial trajectory.

With this alternative scorer, Gambit continues to outperform STEP. On DeepSeek-R1-8B, it improves over STEP by +5.2 percentage points on AIME-25 and +7.7 percentage points on HMMT-24. On Qwen3-4B, the corresponding gains include +3.3 points on AIME-25 and +1.7 points on HMMT-24.

These experiments clarify the conceptual difference between the two search topologies. In STEP, the scorer functions as a negative filter: it can terminate trajectories but cannot create additional evidence in favor of a promising region. In Gambit, the same signal becomes an allocation mechanism. A high score causes the system to generate more continuations from the associated prefix, translating ranking information into a change in the sampling distribution.

The ranking analysis nevertheless points to an important limitation. Early partial prefixes are intrinsically ambiguous, and hidden-state scores may have limited discrimination before sufficient reasoning depth. This motivates the warmup threshold. More generally, Gambit’s gains are bounded by scorer calibration, temporal reliability, and the degree to which local prefix quality predicts final correctness.

Figure 8: Prefix-ranking accuracy improves the separation of correct and incorrect trajectories as more reasoning steps become available.

Runtime behavior and error correction

The appendix provides a detailed runtime example on AIME 2025 Q12. Gambit produces 1,595 logical nodes, including 256 initial traces, 1,339 spawned branches, 1,083 pruned traces, and 256 completed traces. Forty-two completed traces produce the correct answer, 204, and score-weighted aggregation selects it over the strongest incorrect alternative.

The example is particularly informative because it illustrates correction through branching rather than merely selection. A high-scoring parent produces the incorrect answer 203 after making an arithmetic omission. A later child inherits more than 31,000 tokens of the parent’s prefix but generates only 3,649 new tokens. During its continuation, the child performs a smaller-case sanity check, identifies the missing initial region, and produces the correct answer 204.

This behavior is a direct consequence of the search topology. The parent’s valid intermediate work is retained, while the child receives a new opportunity to inspect and repair the final inference. A pruning-only method could terminate the parent, but it could not reuse the parent’s partial derivation to generate a targeted corrective continuation.

The example also provides a strong comparison with SC. At a larger budget of 512 independent traces, SC consumes approximately 21 million tokens and fails to recover the correct answer because an incorrect answer receives the largest plurality. Gambit solves the same instance with roughly an order of magnitude fewer generated tokens under its active reallocation policy. This example should be interpreted as an illustrative case rather than a substitute for aggregate evaluation, but it concretely demonstrates how prefix reuse and branching can overcome a majority-vote failure.

Theoretical and practical implications

The paper contributes a useful reformulation of test-time inference. Compute scaling is usually described in terms of increasing sample count or decoding length. Gambit instead treats inference as a controlled allocation process over a dynamic population of partial trajectories. This perspective connects LLM reasoning with sequential decision-making, beam search, population-based search, and resource-constrained scheduling.

The theoretical implication is that the marginal value of inference computation is state-dependent. A token generated from a low-quality trajectory may have substantially lower expected utility than a token generated from a high-quality prefix. Consequently, the relevant optimization variable is not total FLOPs alone but the allocation policy governing where those FLOPs are spent.

Practically, the method is most relevant to inference regimes in which:

• reasoning traces are long enough to create severe KV-cache pressure;
• batch-level throughput is constrained by memory rather than arithmetic alone;
• hidden-state or reward-model signals are available at intermediate boundaries;
• prefix caching is supported by the serving stack;
• the cost of generating independent traces is substantial.

The method also suggests a path toward tighter integration between model architecture and serving infrastructure. Future inference engines could expose first-class support for logical search trees, shared KV-cache blocks, branch priorities, and scorer-driven scheduling. More advanced variants could use adaptive tournament intervals, uncertainty-aware branching, diversity regularization, or learned allocation policies that optimize accuracy, latency, and energy jointly.

Several limitations remain. The reported experiments focus heavily on mathematical and scientific reasoning benchmarks, where answer extraction and score calibration are relatively tractable. Generalization to open-ended generation, coding, theorem proving, and multimodal reasoning is not established. The fixed hyperparameters may also interact with model-specific thought segmentation and termination behavior. Moreover, score-weighted majority voting assumes that the scorer’s confidence is meaningfully calibrated across trajectories and models. Miscalibration could amplify a systematic but incorrect reasoning mode.

The strongest claim—that active branching “strictly dominates” existing baselines—should therefore be understood relative to the evaluated models, hardware configuration, scorer choices, and benchmark suite. Gambit does not eliminate the exploration–exploitation trade-off; it relocates that trade-off to scorer quality, beam width, swap size, warmup duration, and branch diversity. In settings where the scorer is weak or biased, active reallocation could amplify errors more efficiently than independent sampling.

Figure 9: Gambit remains effective across a broad range of swap sizes, update intervals, warmup thresholds, and memory-utilization settings.

Future developments in AI inference

The paper’s design points toward inference algorithms that are jointly optimized across three layers: model representations, search policy, and serving infrastructure. Hidden-state probes may evolve into multimodal value estimators or calibrated continuation-value models trained specifically for branch selection. Rather than assigning a scalar score to each trace, future systems could estimate uncertainty, expected answer diversity, and conditional improvement under additional computation.

Search policies may also become adaptive. A system could increase branching around prefixes with high expected value but high epistemic uncertainty, while allocating fewer continuations to prefixes that are either clearly poor or already sufficiently validated. Diversity-aware allocation could prevent collapse onto a single erroneous mode by penalizing redundant branches in hidden-state or semantic space.

At the systems level, dynamic KV-cache trees could support more general forms of speculative execution, including branch merging, partial-prefix deduplication, and asynchronous thought-level scheduling. Hardware-aware policies might allocate compute according to the marginal latency and memory cost of each branch rather than treating all continuations identically. Such systems would move beyond static batch inference toward search-native LLM serving.

Conclusion

“Thought-Level Beam Search for Reasoning” (2608.08020) argues that test-time compute scaling should allocate computation selectively across partial reasoning states. Gambit operationalizes this principle by scoring intermediate prefixes, pruning low-quality trajectories, branching from high-quality ones, and maintaining a constant logical capacity through decoupled memory management.

Across multiple reasoning models and benchmarks, the method reports accuracy improvements over pruning-only baselines, token reductions of up to 68.5% relative to SC, more than twofold higher productive trace throughput, and less than 1% search-management overhead. Its primary contribution is not simply a new decoding heuristic but a systems-aware formulation of reasoning inference as dynamic resource allocation. The results indicate that preserving and amplifying partial progress can be more effective than increasing the number of independent complete traces, provided that scoring reliability, branch diversity, and KV-cache management are controlled.