PaperlayerView exact revision on arXiv

Source-linked AI summary

Thought-Level Beam Search for Reasoning

Lijie Yang, Hongyin Luo, Jiawei Zhao, Tri Dao, Ravi Netravali

arXiv:2608.08020v2cs.AI

TL;DR

Existing test-time scaling wastes compute through independent sampling or leaves hardware underused through subtractive pruning, motivating allocation over promising partial trajectories. Gambit addresses this with thought-level beam search that prunes weak traces, branches from high-quality prefixes, and maintains a fixed active pool. Across benchmarks and models, it improves accuracy over pruning baselines while increasing throughput and reducing token consumption relative to parallel sampling.

  • Problem

    Current test-time scaling is inefficient: parallel sampling wastes compute and strains memory, while pruning fails to reallocate capacity toward promising trajectories.

  • Method

    Gambit performs thought-level beam search using hidden-state scoring, periodic prune-and-branch operations, and zero-sum allocation under fixed hardware constraints.

  • Results

    +6.7% accuracy on HMMT-24 and +3.3% on AIME-25 over pruning methods, with over 2× higher productive trace throughput and up to 68.5% lower token consumption than parallel sampling.

  • Takeaways & Limitations

    Active branching shifts compute toward promising reasoning regions while preserving hardware utilization, producing consistent efficiency and accuracy gains across evaluated benchmarks and architectures.

  • Takeaways & Limitations

    The hardware budget is modeled primarily through limits on concurrently active traces and total KV-cache memory footprint.

Abstract

from arXiv · show

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.

1 Introduction

Test-time scaling improves reasoning but current allocation strategies waste compute under strict hardware limits. Gambit reframes allocation over partial trajectories and combines scoring, branching, and zero-sum resource management.

  • Test-time compute scaling drives reasoning gains, but the key challenge is deciding where to allocate compute rather than only how much to spend.
  • Parallel sampling independently generates many traces, causing memory pressure and wasting compute on paths that often yield incorrect answers.
  • Existing strategies must checkpoint promising partial trajectories, amplify them with additional compute, and keep hardware fully utilized.
  • Subtractive pruning removes weak traces but leaves freed capacity idle and does not sufficiently shift the output distribution toward correct answers.
  • Gambit performs thought-level beam search by scoring and checkpointing prefixes, branching from high-quality trajectories, and maintaining a constant active-trace pool.
  • +6.7% accuracy on HMMT-24 and +3.3% on AIME-25 over aggressive pruning accompany up to 68.5% lower token consumption than standard parallel sampling.

2 Related Work

Prior work scales inference through independent sampling, structured search, or subtractive pruning, but each approach faces efficiency or serving-system limitations.

  • Self-Consistency aggregates independently sampled traces, but repeated exploration of redundant pathways makes this scaling strategy inefficient.
  • Tree-of-Thoughts and Monte Carlo Tree Search guide structured reasoning with value functions or process reward models.
  • MCTS’s asymmetric expansion and asynchronous rollouts conflict with the synchronous, large-batch operation of modern language-model inference.
  • DeepConf, STEP, and Slim-SC reduce memory pressure through termination or deduplication, but subtractive methods do not reallocate freed capacity.

3 Motivation

The motivation is to concentrate computation on promising prefixes without sacrificing concurrency. Empirical observations show both the accuracy potential of branching and the hardware costs of unmanaged or subtractive allocation.

  • Effective test-time reasoning must allocate compute toward promising partial trajectories while maintaining high hardware utilization.
  • High-quality prefixes can yield large accuracy gains with roughly half the token budget, but committing to one prefix can produce confident errors.
  • Branching 64 continuations from a top-ranked prefix reaches 87.5% pass@1 versus a 6.2% baseline while sharing the parent KV-cache.
  • Controlled reallocation therefore maintains a candidate beam while progressively concentrating computation on the most promising subset.
  • Parallel sampling saturates KV-cache capacity and inflates latency, whereas pruning-only approaches reduce concurrency and increasingly leave hardware idle.

4 Method: Gambit

Gambit formulates inference as constrained compute allocation over partial trajectories and implements thought-level beam search with periodic score-based swaps. Its logical and physical state management preserves a fixed-capacity active pool despite memory pressure.

  • 4 Method: Gambit: Gambit maximizes the probability of deriving the ground-truth solution under a strict hardware budget covering concurrent traces and KV-cache memory.
  • 4 Method: Gambit: The method uses hidden-state scoring to identify promising prefixes and instantiates inference as thought-level beam search.
  • 4 Method: Gambit: Periodic zero-sum reallocation preserves a constant memory footprint while decoupled memory management separates logical search decisions from physical execution state.
  • 4.1 Thought-Level Beam Search: At capacity, Gambit removes the K lowest-scoring traces and branches K children from the highest-scoring prefixes, preserving exactly C active traces through shared KV-cache blocks.
  • 4.1 Thought-Level Beam Search: Trajectories are sequences of discrete thoughts, and tournament rounds occur every Δ steps over a bounded pool of concurrent traces.
  • 4.1 Thought-Level Beam Search: Under capacity, it branches eligible high-scoring traces to fill vacant slots, while the scorer can consume last-layer hidden states at thought boundaries.
  • 4.1 Thought-Level Beam Search: The Decoupled View Architecture keeps evicted traces logically active as ghost traces, ensuring capacity checks trigger balanced swaps rather than repeated under-capacity branching.
  • 4.1 Thought-Level Beam Search: Completed answers are aggregated with a score-weighted majority vote after physical execution terminates.

5 Experiments

Gambit is evaluated against sampling and pruning baselines across models and reasoning benchmarks under fixed hardware and trace budgets. It improves accuracy while reducing token consumption and sustaining high throughput through active trace replenishment.

  • Experimental Setup: Gambit is evaluated on AIME, HMMT, and GPQA-Diamond across Qwen3-4B, DeepSeek-R1-8B, and Phi-4-14B models.The evaluation uses a single NVIDIA B300 GPU, a strict budget of N = 256 complete traces, and vLLM implementations.
  • Experimental Setup: The comparison includes Self-Consistency, Slim-SC, DeepConf, and STEP, spanning unmanaged sampling, deduplication, confidence-based stopping, and pruning.Gambit uses the same baseline MLP scorer as STEP, isolating search topology as the source of differences in the comparison.
  • Accuracy: Universal accuracy gains occur across benchmarks and model architectures from 4B to 14B when threshold-based pruning is replaced by active beam search.For the 4B model, Gambit matches DeepConf at 90.0% on AIME-25 and exceeds it by +6.7% on HMMT-24 and +2.6% on GPQA.
  • Accuracy: +3.3% accuracy on AIME-25 and HMMT-24 is achieved over STEP for the 4B model, with additional gains on larger models and Phi-4.The reported STEP improvements are +2.5% on AIME-25 for the 8B model and +1.6% on HMMT-24 for Phi-4.
  • Efficiency: 68.5% lower total token consumption is achieved on HMMT-25 for Phi-4 versus SC, while Gambit remains more than 2× faster than parallel sampling.The Phi-4 comparison is 1.75M tokens versus SC’s 5.56M; Gambit also reports a 60.6% reduction on HMMT-24 with Qwen3-4B.
  • Efficiency: 0.216 versus 0.098 trace throughput is reported on Qwen3-4B, giving Gambit over 2× higher throughput than the compared subtractive baseline.Continuous replenishment of the trace pool prevents the active batch from decaying as terminated traces are removed.

6 Analysis

Gambit’s efficiency gains come from branching from shared high-quality prefixes, which reduces new-token generation while shifting completed-trace lengths toward deeper reasoning. This saves tokens without proportionally reducing latency because more promising traces continue running.

  • Prefix Reuse and Unique Token Minimization: Gambit reduces unique tokens per completed trace by reusing shared high-quality KV-cache prefixes when exploring alternative reasoning paths.On Phi-4, the median is 5.2K unique tokens for Gambit, versus 14.5K for SC and 7.7K for STEP.
  • Distribution Shift in Total Sequence Length: Gambit produces deeper completed-trace length distributions than subtractive methods because more high-quality trajectories continue generating tokens.On DeepSeek-R1-8B, the median total sequence length is 35.8K for Gambit versus 16.5K for STEP.
  • Distribution Shift in Total Sequence Length: Branching does not lengthen individual traces; it increases the frequency of traces that continue generating, shifting the overall length distribution rightward.The per-trace generation dynamics remain unchanged despite the distributional shift among surviving traces.
  • Implications for Decoding Latency: Token savings do not translate directly into proportional latency reductions because prefix sharing removes redundant work while compute reallocation leaves more long-running traces.This distributional effect explains why Gambit’s wall-clock latency remains competitive with aggressive pruning methods.
  • Implications for Decoding Latency: Gambit achieves a favorable efficiency trade-off by reducing total token consumption while maintaining competitive wall-clock latency through targeted compute reallocation.The method reallocates compute toward promising trajectories rather than uniformly truncating them.

7 Conclusion

Gambit formulates test-time reasoning as thought-level beam search under constrained compute and hardware budgets. By branching from high-quality prefixes and reallocating computation, it improves accuracy and efficiency over sampling and pruning approaches.

  • Conclusion: Gambit treats test-time reasoning as constrained compute allocation over partial trajectories using thought-level beam search.The framework dynamically reallocates computation toward promising intermediate states rather than increasing independent samples.
  • Conclusion: Gambit combines high-quality-prefix branching with zero-sum allocation to improve accuracy, reduce token consumption by up to 68.5%, and maintain high hardware utilization.The conclusion reports consistent gains across benchmarks and strict dominance over parallel sampling and pruning-based approaches.

A.1 Evaluating Search Generalization: A History-Aware Sequence Scorer

The study evaluates a history-aware sequence scorer that maps full hidden-state reasoning prefixes to quality scores for Gambit's search. Its causal architecture and last-step training objective are designed to use complete trajectory context during inference.

  • Scorer design: The sequence scorer fθ maps a prefix of hidden-state vectors to a scalar quality score in [0,1].It provides an alternative to off-the-shelf scorers for evaluating trajectories.
  • Scorer design: Hidden states are extracted at reasoning-step boundaries, yielding a per-trace sequence for scoring.The states are taken from the final model layer immediately before each boundary.
  • Scorer design: The scorer is implemented as a compact Transformer.
  • Scorer design: Causal attention conditions each prediction on all prior reasoning steps, helping identify globally inferior steps that appear locally plausible.The causal mask lets position t attend to steps 1,...,t.
  • Training objective: The model is trained with last-step binary crossentropy because inference queries the output at the final observed step.Only the final position contributes gradient signal, encouraging credit assignment over the full trajectory context.

A.2 Generalization Across Scoring Architectures

Gambit’s active search remains effective when the scoring architecture changes, consistently outperforming subtractive pruning across models and benchmarks. The comparison attributes this robustness to branching from promising prefixes rather than merely filtering traces.

  • Generalization across scorers: Gambit consistently outperforms STEP across different reward models and base LLMs by actively reallocating compute to promising prefixes.
  • Search topology: Pruning-only systems can terminate unpromising traces but cannot actively construct or discover correct ones.
  • Search topology: Gambit actively branches from high-scoring prefixes, structurally shifting sampling toward promising solution regions.
  • Results: +7.7% on HMMT-24 and +5.2% on AIME-25 are the 8B model’s accuracy gains over STEP.
  • Results: +3.3% on AIME-25 and +1.7% on HMMT-24 are the 4B model’s accuracy gains over STEP.

A.3 Latency Decomposition and System Overhead

Latency profiling shows that Gambit’s thought-level beam search adds negligible systems overhead. GPU forward passes and sampling dominate runtime, while search-management interventions occupy less than one percent of total execution time.

  • Motivation: The latency decomposition was performed to test whether search introduces CPU-bound bookkeeping or memory-management bottlenecks.
  • Latency breakdown: Less than 1% of total execution time is algorithmic overhead for the beam search mechanism.The profiling uses Qwen3-4B on AIME-26 with batch size N =256.
  • Latency breakdown: 99.03% of Gambit’s wall-clock time is consumed by autoregressive forward passes and GPU sampling routines.
  • Latency breakdown: 0.97% of execution time covers algorithmic interventions, with approximately 0.60% attributed to guidance synchronization.The synchronization consists of RPC communication for transmitting hidden states and retrieving trajectory scores.
  • Implication: The profiling concludes that dynamic thought-level beam search can run with virtually no systems penalty.The reported overhead does not prevent token reductions from translating into latency improvements.

A.4 Hyperparameter Ablations

Gambit remains robust across tested hyperparameter choices, with a moderate configuration providing strong performance. A runtime example shows how branching can correct a parent trace’s arithmetic error and steer the vote toward the correct answer.

  • Hyperparameter Sensitivity: K=16, Δ=200, w=12K, and r=0.9 lie in a broad high-performing regime that consistently outperforms the SC@256 baseline.The selected configuration uses K=16, Δ=200, w=12K, and r=0.9.
  • Hyperparameter Sensitivity: Performance varies smoothly across swap size, beam interval, warmup threshold, and GPU memory ratio, indicating robustness to hyperparameter choices.The tested values include K ∈ {4,8,16,32}, Δ ∈ {100,200,400}, and w ∈ {8K,12K,16K}.
  • Hyperparameter Sensitivity: Moderate settings balance accuracy and efficiency: small swaps limit trace refresh, large swaps disrupt decoding, and Δ=200 provides a stable update cadence.Early branching can harm reasoning quality, whereas excessive warmup delays selection; w=12K provides a moderate warmup.
  • Runtime Example of Gambit: 42 of 256 completed traces return the correct answer 204, and score-weighted voting selects 204 over the runner-up.The score-weighted vote totals 24.02 for 204 versus 21.81 for the runner-up.
  • Runtime Example of Gambit: Gambit uses an order of magnitude fewer generated tokens than unweighted plurality, which selects an incorrect answer with 90 votes.The plurality vote gives answer 3 90 votes, or 17.6%, while 204 ranks only third.
  • Runtime Example of Gambit: 3,649 new tokens from child τB repair the parent’s missing initial-region term, changing 203 to the correct result 204.The child performs a smaller-case sanity check and writes 1 + 203 = 204 after inheriting the parent prefix through caching.
Loading 2608.08020v2…