Thought-Level Beam Search for Reasoning
Lijie Yang, Hongyin Luo, Jiawei Zhao, Tri Dao, Ravi Netravali
cs.AI
2026-08-08
Gambit reframes parallel trace sampling as a fixed-capacity beam search: every 200 tokens it scores partial trajectories, prunes the worst 16, and branches from top-scoring prefixes, gaining up to +6.7% accuracy on HMMT-24 over pruning baselines while cutting total token consumption by up to 68.5% versus standard parallel sampling.
Accuracy in large reasoning models is bought with test-time compute: sample hundreds of reasoning traces per problem, then take a majority vote. The problem is efficiency. A single 275GB NVIDIA B300 needs several hours to finish 512 traces for one AIME problem on vLLM with Qwen3-8B, and most of those traces walk into wrong answers, so the votes are wasted.
Existing fixes split into two camps, each with its own failure mode. Parallel sampling (self-consistency) treats traces as independent trials; a few hundred concurrent traces saturate the KV cache, and queueing inflates latency roughly 3×. Score-based pruning (STEP, DeepConf) terminates unpromising traces early, but the freed capacity just sits idle while concurrency decays. A filter can kill bad traces, yet it cannot manufacture good ones. The paper reframes the question: with a fixed hardware budget, where should compute go, in real time?
Gambit runs a beam search at the trajectory level via a fixed-capacity tournament:
"Thought-level" means traces are segmented into discrete reasoning steps by double newlines inside the thinking block; branching happens at step boundaries, continuing from a complete intermediate thought rather than mid-sentence.
Two engineering choices decide whether this survives contact with a real serving stack. First, the scorer is a lightweight probe, and search decisions are decoupled from physical memory management: when vLLM evicts traces under memory pressure, the tournament still sees the logical pool as full and runs a balanced swap. Without this, the system degenerates into repeatedly branching from a single top trace, a greedy collapse. Second, final answers are aggregated by score-weighted voting with a position-weighted penalty. The whole tournament mechanism costs 0.97% of total wall-clock time.
Three models (Qwen3-4B, DeepSeek-R1-8B, Phi-4-14B), five benchmarks, a fixed budget of 256 complete traces per question, one B300 GPU:
| Setting | Baseline | Gambit | Gain |
| Qwen3-4B, HMMT-24 | STEP 61.7 | 65.0 | +3.3 |
| Qwen3-4B, HMMT-24 | DeepConf 58.3 | 65.0 | +6.7 |
| Qwen3-4B, HMMT-24 | SC@256 50.8 | 65.0 | +14.2 |
| Qwen3-4B, AIME-25 | STEP 86.7 | 90.0 | +3.3 |
| Phi-4, HMMT-25 tokens | SC 5.56M | 1.75M | −68.5% |
| Qwen3-4B, AIME-26 throughput | STEP 0.098 | 0.216 | 2.2× |
Because Gambit uses the exact scorer STEP uses, the accuracy deltas are attributable to the search topology alone. Swapping in the authors' own history-aware scorer widens the gap further: +7.7% over STEP on HMMT-24 with the 8B model. One number captures the intuition: on the hardest AIME problem, branching 64 continuations from the top-ranked prefix reaches 87.5% pass@1 versus 6.2% for independent sampling.
Majority voting is capped by how many correct traces the model can produce organically. Branching rewrites the sampling distribution directly, upgrading "filter out bad traces" into "replicate good prefixes", which is why the voting ceiling breaks. For practitioners the economics matter more: 40–68% fewer tokens, 2× lower latency than parallel sampling, and higher accuracy at the same time, a trio that normally cannot coexist. The code is open-sourced under Dao-AILab on top of vLLM. The four tournament hyperparameters stay fixed across all models and benchmarks, so the tuning burden is smaller than it looks.
The paper has no dedicated limitations section; the points below come from remarks scattered in the text, plus what the evaluation leaves out. Everything runs on a single B300, and multi-GPU behavior at larger scale is untested. The benchmark suite is dominated by competition math, with GPQA as the lone science set; coding, long-context, and agent tasks are absent, and an MLP scorer that reads trajectory quality on math does not obviously transfer. The branch-from-the-top mechanism is inherently convergence-biased: if the probe systematically misjudges some intermediate step, a wrong prefix gets amplified exponentially. The paper shows the 87.5% success case without a systematic accounting of how often this failure mode occurs. Reproduction also requires access to hidden states for the scorer, which rules out API-only users.