Thought-Level Beam Search for Reasoning

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

arXiv:2608.08020 · 2026-08-16 공개 · arXiv · PDF

llm-inference token-efficiency reasoning beam-search hmmt-24 aime-25 thought-level gambit

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 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.

한국어 요약

한 줄 요약

Gambit는 토큰 소비를 최대 68.5% 절감하면서도 AIME-25에서 +3.3%의 정확도 향상을 달성한 사고 수준 빔 서치 알고리즘입니다.

핵심 기여도

핵심 아이디어

기존의 병렬 샘플링은 추적을 독립적으로 처리하여 메모리 병목 현상을 유발하고, 제거적 트리밍은 하드웨어 활용률을 낮추며 출력 분포를 충분히 변경하지 못한다. Gambit는 이 이분법을 해결하기 위해 사고 수준 빔 서치를 제안한다. 이는 고질량 프리픽스에서 즉시 분기하고, 유망하지 않은 추적을 주기적으로 제거함으로써, 유망한 추적 경로에 컴퓨팅 자원을 동적으로 집중한다. 이 과정에서 가벼운 스코어러가 히든 상태를 탐색하여 추적의 잠재력을 평가하며, 하드웨어 활용률을 지속적으로 유지한다.

기술적 접근법

주요 결과

의의 및 한계

Gambit는 테스트 시간 추론을 부분 추적에 대한 제약된 컴퓨팅 할당 문제로 정식화함으로써, 정확도와 효율성을 동시에 향상시키는 새로운 접근법을 제시한다. 기존 방식이 하드웨어 예산 내에서 최적의 추적을 선택하지 못하는 문제를 해결하며, 병렬 샘플링 및 제거적 트리밍 기반 접근법을 엄격히 우월하게 만든다. 그러나, 특정 유형의 추적에서 과도한 제거가 정확도에 부정적 영향을 줄 수 있으며, 이는 추적 평가 스코어러의 정확성에 크게 의존한다는 한계가 있다.

실용적 활용

Gambit는 수학 문제 해결, 과학적 추론, 복잡한 논리적 추론이 필요한 산업 및 연구 분야에서 유망한 추적 경로를 효율적으로 탐색하는 데 활용될 수 있다. 특히, 고비용의 병렬 샘플링을 대체하면서도 정확도를 유지하거나 향상시킬 수 있어, 대규모 추론 모델의 실용적 적용에 기여할 수 있다.