Thought-Level Beam Search for Reasoning
Lijie Yang, Hongyin Luo, Jiawei Zhao, Tri Dao, Ravi Netravali
cs.AI
2026-08-08
把并行采样 256 条推理轨迹变成固定容量的束搜索:每 200 token 给轨迹打分,砍掉最差的 16 条、从最优前缀分叉补位,同硬件下 HMMT-24 比 prune 基线高 6.7%,token 总量比标准并行采样最多省 68.5%。
大推理模型的准确率靠测试时算力堆出来:一道题采样几百条推理轨迹,再多数投票。问题在于效率。一块 275GB 的 NVIDIA B300 跑完一道 AIME 题的 512 条轨迹要几个小时,而其中大多数轨迹走向错误答案,票是白投的。
现有改进分两派,各有各的死法。并行采样(self-consistency)把轨迹当独立试验,轨迹一多 KV cache 就爆,排队把延迟拉高三倍。打分剪枝派(STEP、DeepConf)提前掐掉没希望的轨迹,但省下的算力就闲置在那儿,并发数只会单调下降——过滤器能淘汰坏轨迹,却造不出好轨迹。这篇把问题重新表述成:算力总量固定时,该把它实时挪到哪里。
Gambit 在轨迹层面执行束搜索,核心是一个固定容量的「锦标赛」:
「思路级」指轨迹被换行符切成离散思考步,分叉发生在步的边界上,新轨迹从某个完整的中间思路继续往下想,而不是从半个句子里续写。
两个工程设计决定了它能不能落地。一是打分器只做轻量探测,分叉决策与物理显存管理解耦——vLLM 因显存压力逐出轨迹时,锦标赛仍按逻辑池容量做满员换血,否则会退化成反复从同一条最优轨迹分叉的贪心坍缩。二是最终聚合用分数加权投票,附位置权重惩罚。整个锦标赛机制的开销只占总运行时间的 0.97%。
三档模型(Qwen3-4B、DeepSeek-R1-8B、Phi-4-14B)、五个基准,每题固定预算 256 条完整轨迹,单卡 B300:
| 设置 | 对照 | Gambit | 提升 |
| 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 token | SC 5.56M | 1.75M | -68.5% |
| Qwen3-4B,AIME-26 吞吐 | STEP 0.098 | 0.216 | 2.2× |
由于 Gambit 与 STEP 用的是同一个打分器,准确率差异全部来自搜索拓扑本身。换上作者自训的历史感知打分器后差距进一步拉开:8B 模型 HMMT-24 上比 STEP 高 7.7%。一条能说明直觉的数据:在最难的 AIME 题上,从排名第一的前缀分叉 64 条续写,pass@1 达 87.5%,独立采样只有 6.2%。
多数投票的准确率上限取决于模型自己能采出多少条正确轨迹,Gambit 的分叉机制直接改写了采样分布,把「过滤坏轨迹」升级成「复制好前缀」,这解释了它为什么能突破投票天花板。对从业者更实际的是成本账:token 省 40%68%,延迟比并行采样快 2 倍,而准确率反而更高,三样通常不可能同时拿到。代码开源于 Dao-AILab,基于 vLLM 实现,工程上是可复用的。四条锦标赛超参在所有模型和基准上保持不变,调参负担比看起来小。
论文没有独立的局限一节,以下部分来自正文散落的自述。热身阈值、换血规模这些超参虽然做了消融,但实验全部在单卡 B300 上跑,更大规模多卡下的行为没有验证。基准以数学竞赛为主,GPQA 只有一项,代码、长上下文、agent 类任务完全缺席,而打分器对数学轨迹好坏的判断未必迁移。得分最高分叉的机制天然偏向收敛,如果 MLP 打分器在某个中间步骤系统性误判,错误前缀会被指数放大,论文只展示了 87.5% 的成功案例,没有系统分析这种失败模式的频率。作者与 Tri Dao 相关的开源实现质量一向可靠,但本篇的复现门槛在于需要拿到隐藏状态来跑打分器,这对 API-only 的使用者不可行。