추론을 위한 사고 수준 빔 탐색

초록

테스트-타임 컴퓨트 확장은 대형 추론 모델(LRMs)의 성능을 결정하는 주요 동력이지만, 현재 접근 방식은 극심한 비효율성에 의해 제약을 받으며, 핵심 질문이 "얼마나 많은 컴퓨트를 사용할 것인가"에서 "어디에 할당할 것인가"로 이동하고 있다. 우리는 테스트-타임 추론을 부분 궤적에 대한 제약적 컴퓨트 할당 문제로 정식화한다. 고정된 하드웨어 예산 하에서 기존 패러다임들은 컴퓨트를 가장 유망한 부분적 진전에 능동적으로 할당하지 못한다. 전통적 병렬 샘플링은 궤적들을 독립적으로 처리하여 심각한 메모리 병목을 유발하는 반면, 감산적 프루닝은 하드웨어 활용을 저하시키고 출력 분포를 능동적이고 충분히 전환하지 못한다. 이 이분법을 극복하기 위해, 우리는 사고-수준 빔 서치를 수행하는 추론 알고리즘인 Gambit을 제안한다. Gambit은 주기적으로 유망하지 않은 궤적을 프루닝하고 고품질 프리픽스에서 즉시 분기함으로써, 히든 스테이트를 탐침하는 경량 스코어러를 통해 컴퓨트를 가장 유망한 추론 궤적에 동적으로 집중시키면서 지속적인 높은 하드웨어 활용도를 유지한다. 여러 모델과 벤치마크에 걸친 광범위한 평가는 Gambit이 기존 베이스라인들을 엄격히 지배함을 보여준다. 동일한 하드웨어 제약 하에서, 우리의 방법은 프루닝 베이스라인 대비 HMMT-24에서 최대 +6.7%p의 절대 정확도 향상과 AIME-25에서 +3.3%p의 향상을 달성하고, 궤적 완성에서 2배 이상 높은 처리량을 제공하며, 표준 병렬 샘플링 대비 총 토큰 소비를 최대 68.5%까지 감소시킨다.

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.