Everything K-culture — comebacks to K-beauty, straight to your inboxGet it in your inbox

METAL MEDIA

Thought-Level Beam Search for Reasoning

arXiv:2608.080202026-08-10

Gambit dynamically redirects a reasoning model's computation toward its most promising answer attempts while they're still being generated

Large reasoning models boost accuracy by generating many parallel answer attempts and majority-voting (Self-Consistency), but most attempts end up wrong, wasting compute. Gambit periodically scores in-progress reasoning traces using their hidden states, cuts off the weakest ones, and immediately branches new continuations from the strongest prefixes, keeping GPU utilization high throughout. Under identical hardware budgets, it beats pruning-only baselines by up to +6.7 points accuracy on HMMT-24 and +3.3 points on AIME-25, more than doubles trace-completion throughput, and cuts total token usage by up to 68.5% versus standard parallel sampling.

METAL MEDIA explanatory visual

How Gambit actively reallocates compute

Evidence statusMeasured results reported

  1. Start with a fixed poolC parallel reasoning traces begin generating for one problem, filling the fixed hardware capacity
  2. Periodic scoringEvery Δ steps, a lightweight scorer reads each trace's hidden states to rank all active traces
  3. Zero-sum prune-and-branchThe K lowest-scoring traces are terminated while K new branches are spawned from the K highest-scoring prefixes via KV-cache inheritance, keeping the pool size fixed at C
  4. Completion and weighted voteFinished traces contribute their answers to a majority vote weighted by their own score
  5. Measured outcomesUp to +6.7 points accuracy over STEP, up to 68.5% less token usage than parallel sampling, and over 2x higher throughput
An explanatory diagram made by METAL MEDIA, not a figure supplied by the paper's authors.

What they did

  1. Reframes test-time compute scaling for large reasoning models (LRMs) as a compute-allocation problem: not how much to compute, but where to direct it among partial reasoning trajectories.
  2. Identifies a dichotomy in existing approaches: parallel sampling (Self-Consistency) treats each trace independently and overloads GPU memory via the KV-cache, while pruning-only methods free memory but leave it idle, starving hardware without reshaping the output distribution.
  3. Every Δ steps, Gambit scores all active traces using a lightweight model probing their hidden states, prunes the K lowest-scoring traces, and immediately branches K new continuations from the K highest-scoring prefixes, keeping the active pool size fixed (zero-sum reallocation).
  4. New branches inherit the parent's KV-cache via prefix caching, so they continue generation without recomputing earlier reasoning, sharply cutting token and memory cost.
  5. Measured across Qwen3-4B, DeepSeek-R1-8B, and Phi-4-14B on AIME, HMMT, and GPQA-Diamond benchmarks, Gambit consistently outperformed baselines on accuracy, throughput, and token efficiency at once.
Figure 1: Reasoning topologies for test-time compute allocation. (Left) Parallel sampling treats traces as independent trials, leading to wrong answers in majority. (Center) Score-based pruning terminates unpromising paths but leaves the freed capacity idle, failing to actively shift the sampling distribution. (Right) Gambit actively reallocates compute: when low-scoring traces are pruned, it immediately branches from high-quality prefixes to amplify correct reasoning and saturate the compute budget.
Figure 1: Reasoning topologies for test-time compute allocation. (Left) Parallel sampling treats traces as independent trials, leading to wrong answers in majority. (Center) Score-based pruning terminates unpromising paths but leaves the freed capacity idle, failing to actively shift the sampling distribution. (Right) Gambit actively reallocates compute: when low-scoring traces are pruned, it immediately branches from high-quality prefixes to amplify correct reasoning and saturate the compute budget.
Table 1: End-to-end task accuracy and token consumption per question (×106) at N=256. Δ: relative token reduction (%) vs. SC@256. Bold: best accuracy / lowest tokens per dataset.
MethodAIME-25AIME-26HMMT-24HMMT-25GPQA
Tok. (Δ)AccTok. (Δ)AccTok. (Δ)AccTok. (Δ)AccTok. (Δ)Acc
Qwen3-4B-Thinking
SC@2566.0286.75.8286.77.6250.86.8665.02.2868.2
Slim-SC3.93 (-34.7)86.73.83 (-34.2)85.04.01 (-47.4)51.73.72 (-45.8)65.81.66 (-27.2)65.6
DeepConf3.27 (-45.7)90.03.19 (-45.2)86.74.28 (-43.8)58.34.15 (-39.5)66.71.52 (-33.3)67.6
STEP3.96 (-34.2)86.73.51 (-39.7)87.54.02 (-47.2)61.73.88 (-43.4)66.72.09 (-8.3)66.9
Gambit (Ours)3.07 (-49.0)90.03.45 (-40.7)88.33.00 (-60.6)65.03.49 (-49.1)67.51.83 (-19.7)70.2
DeepSeek-R1-8B
SC@2566.7683.36.8583.38.4755.87.6570.02.9267.1
Slim-SC5.87 (-13.2)83.36.08 (-11.2)83.37.32 (-13.6)55.86.93 (-9.4)70.72.26 (-22.6)67.1
DeepConf3.75 (-44.5)81.73.67 (-46.4)84.24.34 (-48.8)60.03.97 (-48.1)71.71.68 (-42.5)68.7
STEP3.71 (-45.1)83.33.66 (-46.6)83.34.15 (-51.0)63.34.08 (-46.7)75.82.44 (-16.4)68.2
Gambit (Ours)4.21 (-37.7)85.84.09 (-40.3)85.03.05 (-64.0)65.63.81 (-50.2)75.82.21 (-25.3)68.2
Microsoft-Phi
SC@2564.2486.74.2088.55.4356.75.5673.33.0576.3
Slim-SC3.43 (-19.1)85.03.50 (-16.7)86.74.72 (-13.1)55.74.48 (-19.4)74.22.24 (-26.6)72.3
DeepConf2.56 (-39.6)85.82.75 (-34.5)86.72.86 (-47.3)57.53.02 (-45.7)74.21.61 (-47.2)74.8
STEP2.39 (-43.6)87.52.44 (-41.9)90.02.75 (-49.4)56.72.78 (-50.0)75.02.10 (-31.1)76.7
Gambit (Ours)1.76 (-58.5)88.31.86 (-55.7)90.01.72 (-68.3)58.31.75 (-68.5)75.81.59 (-47.9)77.1
Figure 2: Branching from a high-quality prefix versus independent sampling on a hard AIME 2025 problem (Q27). Branching 64 continuations from the top-ranked prefix achieves 87.5% pass@1 (14× over 6.2% baseline) while sharing the parent’s KV-cache (green region) halves memory consumption.
Figure 2: Branching from a high-quality prefix versus independent sampling on a hard AIME 2025 problem (Q27). Branching 64 continuations from the top-ranked prefix achieves 87.5% pass@1 (14× over 6.2% baseline) while sharing the parent’s KV-cache (green region) halves memory consumption.
Table 2: Impact of scorer architecture on end-to-end downstream accuracy. Gambit consistently outperforms the pruning-only baseline (STEP) across different reward models and base LLMs. By actively reallocating compute to promising prefixes, Gambit effectively unlocks the potential of the guidance signal, whereas purely subtractive methods remain fundamentally bottlenecked by their inability to generate new traces.
Base ModelSystemAIME-25AIME-26HMMT-24HMMT-25GPQA
DeepSeek-R1-0528 -Qwen3-8BSTEP + SeqScorer83.384.261.775.866.7
Gambit + SeqScorer88.5 (+5.2)86.7 (+2.5)69.4 (+7.7)76.7 (+0.9)67.7 (+1.0)
Qwen3-4B -Thinking-2507STEP + SeqScorer86.786.760.065.067.1
Gambit + SeqScorer90.0 (+3.3)87.5 (+0.8)61.7 (+1.7)65.8 (+0.8)68.2 (+1.1)
Figure 3: GPU memory utilization and latency profiles during reasoning with batch_size=256 traces on HMMT-25 Q7. Parallel sampling exhausts KV-cache capacity, inflating latency by ∼3×. Pruning-only STEP (Liang et al. 2026) eliminates queueing but steadily loses concurrency. Gambit maintains high utilization throughout at ∼1.1× latency.
Figure 3: GPU memory utilization and latency profiles during reasoning with batch_size=256 traces on HMMT-25 Q7. Parallel sampling exhausts KV-cache capacity, inflating latency by ∼3×. Pruning-only STEP (Liang et al. 2026) eliminates queueing but steadily loses concurrency. Gambit maintains high utilization throughout at ∼1.1× latency.
Table 3: Latency decomposition for Gambit with Qwen3-4B on AIME-26 with a batch size of N=256. Algorithmic overhead accounts for less than 1% of total execution time, confirming that the beam search mechanism operates efficiently without introducing system bottlenecks.
Execution ComponentTime (s)% Wall Clock
GPU Compute (Forward Pass & Sampling)37,851.9799.03%
Beam-search Overhead (Scoring & Tree Mgmt)370.760.97%
Total Wall Clock38,222.73100.00%
Figure 4: End-to-end pipeline of Gambit on an AIME problem (capacity C=5, swap size K=2). During warmup, C=5 parallel traces are evaluated (color intensity denotes running score s¯∈[0,1]). Every Δ steps, a tournament ranks the active traces: the K=2 lowest-scoring traces are pruned (×) and replaced by new branches spawned from the highest-scoring prefixes via prefix caching (green arrows). This zero-sum reallocation maintains exactly C active traces throughout generation. Upon completion, answers are aggregated via a score-weighted majority vote, correctly selecting answer 29 (∑s¯=1.78 vs. 1.43).
Figure 4: End-to-end pipeline of Gambit on an AIME problem (capacity C=5, swap size K=2). During warmup, C=5 parallel traces are evaluated (color intensity denotes running score s¯∈[0,1]). Every Δ steps, a tournament ranks the active traces: the K=2 lowest-scoring traces are pruned (×) and replaced by new branches spawned from the highest-scoring prefixes via prefix caching (green arrows). This zero-sum reallocation maintains exactly C active traces throughout generation. Upon completion, answers are aggregated via a score-weighted majority vote, correctly selecting answer 29 (∑s¯=1.78 vs. 1.43).
Table 4: Systems-level efficiency at n=256 on a single B300-275 GB. Token consumption is reported in thousands (K); latency in wall-clock seconds per question. Bold indicates the best result per column within each model–metric block.
AIME-25AIME-26HMMT-24HMMT-25GPQA
Qwen3-4B-Thinking-2507
Tokens (K)SC6,0235,8177,6226,8572,276
Slim-SC3,9303,8294,0103,7231,659
DeepConf3,2663,1924,2784,1481,516
STEP3,9603,5064,0233,8842,089
Gambit (Ours)3,0713,4493,0023,4931,826
Latency (s)SC3,4653,3385,3584,313593
Slim-SC1,9051,8412,4121,995555
DeepConf2,0952,1883,2312,948605
STEP1,3901,2322,0161,507786
Gambit (Ours)1,3501,1771,9122,179527
DeepSeek-R1-0528-Qwen3-8B
Tokens (K)SC6,7646,8498,4657,6522,919
Slim-SC5,8746,0807,3236,9332,256
DeepConf3,7533,6664,3373,9721,679
STEP3,7113,6614,1474,0772,444
Gambit (Ours)4,2114,0893,0533,8062,211
Latency (s)SC4,0784,1075,6404,873880
Slim-SC3,4623,5645,0714,342771
DeepConf2,6922,5823,2562,800744
STEP1,5191,4691,9021,715730
Gambit (Ours)1,5711,8422,0142,169702
Phi-4
Tokens (K)SC4,2434,1985,4265,5593,050
Slim-SC3,4263,5034,7164,4822,242
DeepConf2,5572,7492,8603,0221,608
STEP2,3902,4352,7542,7802,097
Gambit (Ours)1,7561,8551,7171,7521,590
Latency (s)SC2,7602,6363,8484,0461,730
Slim-SC2,1232,1663,0692,7041,048
Figure 5: Efficiency vs. Accuracy trade-offs over all benchmarks. (a) Average Latency vs. Accuracy. (b) Average Token Consumption vs. Accuracy. Gambit (Ours) dominates the efficiency frontier, achieving the highest accuracy with the lowest overall token footprint and highly competitive wall-clock latency.
Figure 5: Efficiency vs. Accuracy trade-offs over all benchmarks. (a) Average Latency vs. Accuracy. (b) Average Token Consumption vs. Accuracy. Gambit (Ours) dominates the efficiency frontier, achieving the highest accuracy with the lowest overall token footprint and highly competitive wall-clock latency.

Findings

  • Under identical hardware constraints, Gambit improved accuracy over the pruning-only baseline STEP by up to +6.7 points on HMMT-24 and +3.3 points on AIME-25.
  • Relative to standard parallel sampling (Self-Consistency), Gambit reduced total token consumption by up to 68.5% (Phi-4 on HMMT-25: 1.75M vs 5.56M tokens) and delivered more than 2x higher trace-completion throughput (0.216 vs 0.098 on Qwen3-4B).
  • On the 4B model, Gambit matched the tightly calibrated DeepConf baseline on AIME-25 (90.0%) while outperforming it on HMMT-24 (+6.7 points) and GPQA (+2.6 points).
  • The algorithmic overhead of the beam search mechanism itself was measured at under 1% (0.97%) of total wall-clock execution time, confirming it does not create a systems bottleneck.
  • In a case study on one AIME 2025 problem, doubling the sample budget with plain parallel sampling (n=512) still failed to find the correct answer, while Gambit reached it using far fewer generated tokens.
Figure 6: Trace throughput on AIME-26. Gambit effectively doubles the completion rate.
Figure 6: Trace throughput on AIME-26. Gambit effectively doubles the completion rate.

Where it can be used

  • Could be applied as an inference-time compute allocation strategy for competition-level math or graduate-level science problems where correct traces are rare and reasoning chains are long.
  • Relevant for operators of LLM serving engines like vLLM who want to reduce both KV-cache memory pressure and idle GPU time simultaneously.
  • Systems already using lightweight internal scoring signals like STEP or DeepConf could consider replacing pure pruning with active branch reallocation based on this idea.
Figure 7: Distribution of unique tokens generated per completed trace on AIME-2026. By branching from shared high-quality KV-cache prefixes, Gambit drastically reduces the median number of new tokens required to explore alternative reasoning pathways compared to SC and STEP.
Figure 7: Distribution of unique tokens generated per completed trace on AIME-2026. By branching from shared high-quality KV-cache prefixes, Gambit drastically reduces the median number of new tokens required to explore alternative reasoning pathways compared to SC and STEP.

Limits and open work

  • Results were measured only on math/science competition benchmarks (AIME, HMMT, GPQA-Diamond) with specific open-weight models (Qwen3-4B, DeepSeek-R1-8B, Phi-4-14B); generalization to other domains or much larger models is not reported.
  • Performance depends on the quality of the scoring function; the paper only compares STEP's off-the-shelf MLP scorer and its own trained sequence scorer.
  • Scoring and checkpointing are deliberately withheld until a warmup threshold (w=12K tokens) because early reasoning signals are noisy; whether this threshold generalizes to other problem types is not separately validated.
  • Results were measured on a single NVIDIA B300 GPU; reproduction across multi-GPU setups or other hardware generations is not reported in the paper.
Figure 9: Pairwise ranking accuracy vs. fraction of reasoning steps used (k%). Higher accuracy indicates better separation of correct and incorrect traces given only a partial prefix.
Figure 9: Pairwise ranking accuracy vs. fraction of reasoning steps used (k%). Higher accuracy indicates better separation of correct and incorrect traces given only a partial prefix.

Why it matters

The biggest cost of deploying reasoning-heavy AI models is the GPU time and memory spent generating many candidate answers to boost accuracy. Gambit shows a concrete way to get higher accuracy, faster completion, and lower token usage on the same hardware, which matters directly for anyone optimizing inference cost of reasoning models.

Figure 10: Hyperparameter sensitivity of Gambit on AIME-25 with DeepSeek-R1-0528-8B (batch size 256, C=256). Each panel varies one parameter while fixing the others to the selected configuration (K=16,Δ=200,w=12​K,r=0.9). The dashed line denotes the SC@256 baseline (83.3%). The selected configuration (red circle) lies within a broad high-performing regime, consistently outperforming the baseline.
Figure 10: Hyperparameter sensitivity of Gambit on AIME-25 with DeepSeek-R1-0528-8B (batch size 256, C=256). Each panel varies one parameter while fixing the others to the selected configuration (K=16,Δ=200,w=12​K,r=0.9). The dashed line denotes the SC@256 baseline (83.3%). The selected configuration (red circle) lies within a broad high-performing regime, consistently outperforming the baseline.

Terms in this paper

  • Self-Consistency · A technique that samples multiple independent reasoning traces for the same question and picks the final answer by majority vote
  • KV-cache · Memory that stores a transformer's previously computed key/value vectors so they can be reused when generating the next tokens
  • beam search · A search strategy that keeps only the top-scoring candidate paths at each step; here applied at the level of 'thoughts' rather than individual tokens
  • pruning · Terminating a reasoning trace early when it's judged unlikely to succeed, to save compute
  • hidden state · An internal representation vector produced by the model's layers, used here as a lightweight signal to score how promising a partial trace is

Original abstract (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.

Authors · Lijie Yang

Read on arXiv

Latest papers

All papers →

Latest from METAL MEDIA

Figures: Lijie Yang et al., arXiv:2608.08020, CC BY-SA 4.0