← 返回 2026-08-14

面向推理的思维级束搜索 Thought-Level Beam Search for Reasoning

Lijie Yang, Hongyin Luo, Jiawei Zhao, Tri Dao, Ravi Netravali 📅 2026-08-11 👍 15 2026-08-19 18:30
LLM推理加速 推理系统 束搜索 测试时计算扩展 解码策略

把测试时推理重构为思维级束搜索:零和剪枝并从高分前缀分支,同硬件下精度与效率双升。

前置知识

测试时计算扩展

指在推理阶段投入额外计算来换取更高答案质量的技术路线,典型做法是让大推理模型(LRM)并行生成多条长思维链再聚合答案,而非只生成一条。其收益类似缩放定律:算力越多答案越准,但边际收益递减且成本急剧上升,因此关键问题已从“花多少算力”转向“算力投在哪里”。

本文的全部分析都建立在这一范式之上:Gambit 本质上就是在固定硬件预算内优化计算分配策略,不理解测试时扩展就无法理解其动机与贡献。

自洽性多数投票(Self-Consistency)

Wang et al. 2023 提出的经典解码策略:对同一问题独立并行采样大量推理轨迹(如 256 或 512 条),各自抽取最终答案,再通过多数投票选出出现次数最多的答案。它简单有效,是当前测试时扩展的标准基线,但轨迹之间完全独立,大量算力被浪费在重复探索相同的错误路径上。

Gambit 的对照基线 SC@256 就是它,且本文的核心批评正是其“独立采样”的低效,理解 SC 才能理解改进的出发点。

束搜索

经典序列解码算法:每一步只保留评分最高的 $B$ 个候选序列(束宽),其余全部剪掉,在指数级搜索空间中做有界启发式搜索。经典版本作用于 token 级别、以序列似然为评分。本文将其提升到“思维”粒度,评分函数改为预测“该前缀能否引出正确答案”的代理值,且束宽对应硬件并发容量。

Gambit 的整个算法框架就是束搜索在轨迹层面的重述,理解束宽、剪枝、分支这些概念是读懂算法的前提。

KV-Cache 与前缀缓存

自回归模型为避免重复计算,会把每层注意力用过的 Key/Value 缓存在显存中;KV-Cache 随序列长度线性增长,是长推理批处理的主要显存瓶颈。前缀缓存允许多个请求共享相同前缀的 KV 块:新分支只需增量生成后续 token,无需重算公共前缀,显存与计算都省一半以上。

Gambit 的分支操作能“低成本”展开新轨迹、并把总 token 消耗降低最高 68.5%,全靠子轨迹继承父前缀的 KV-Cache,这是其系统效率的物理基础。

隐藏状态探针评分器

不训练昂贵的独立过程奖励模型(PRM),而是直接在语言模型最后一层隐状态上接一个轻量分类头(如 STEP 使用的两层 MLP),在每步输出后判断当前推理前缀“通向正确答案”的概率。这类内部信号近乎免费,但推理早期阶段信号噪声大、不可靠,需要谨慎设定启用时机。

Gambit 的锦标赛排序完全依赖这类评分器给出的轨迹平均分 $\bar{s}_\tau$,并为此设计了 12000 token 的 warmup 阈值来规避早期噪声,这是算法可信度的关键。

研究动机

现有测试时扩展的两条主流路线在固定硬件预算下都有系统性缺陷。并行采样把每条轨迹当作独立试验:在 vLLM 上用 Qwen3-8B 完成 512 条轨迹解一道 AIME-2025 题,一块顶级 NVIDIA B300 需要数小时,而绝大多数轨迹以错误答案收尾;更糟的是,大批量长上下文很快耗尽 KV-Cache,推理引擎被迫排队,端到端延迟膨胀约 3 倍(Figure 3 中平均显存占用 88.5% 但吞吐仅 1.5K tokens/s)。纯剪枝路线(STEP、DeepConf、Slim-SC)用内部信号提前终止低分轨迹,虽缓解显存压力,却引入“硬件饥饿”:被剪掉的轨迹不会得到补充,并发数随生成单调下降,GPU 后期大量空闲。Figure 2 量化了问题的另一面:在 AIME-2025 Q27 上独立采样 pass@1 仅 6.2%,而只从排名第一的高质量前缀分支 64 条续写就达到 87.5%,AIME-25 Q12 上也有 1.6%→39.1%,说明算力被极度错配,且纯减法策略不改变输出分布、无法突破多数投票的天花板。

本文的目标是本文要把测试时推理显式形式化为“部分轨迹集合上的受限计算分配问题”:给定问题提示 $P$,在峰值算力与显存硬约束 $\Omega(\pi)\le B$ 下,寻找分配策略 $\pi$ 最大化聚合答案命中真值的概率 $\max_\pi \Pr[\mathcal{A}_\pi(P)=y^*]$。具体目标有三:(1) 及时“检查点化”值得继续的高质量前缀,避免一个下游错误就丢弃此前全部有效进展;(2) 把剪枝腾出的算力主动再分配给优质前缀,通过分支放大它们在最终投票中的权重,真正把输出分布推向正确区域;(3) 全程维持恒定数量的活跃轨迹,让 GPU 始终饱和运行,既不排队也不挨饿。实验目标是:在与基线完全相同的硬件预算(单卡 B300、每题 256 条完整轨迹)下,同时在准确率、总 token 消耗、轨迹完成吞吐三个维度严格优于并行采样与纯剪枝方法。

与已有工作不同的是,本文的独特之处在于把“分配”而非“终止”作为第一性问题,并从算法与系统两层同时切入。算法上它与三类主流工作都不同:不同于自洽性的完全独立采样,它维护一个逻辑相关的候选池;不同于 Tree-of-Thoughts / MCTS 的异步非对称扩展——那与 vLLM 这类同步大批量服务引擎天然冲突——它把搜索重构成周期性、同步、零和的 prune-and-branch 锦标赛;也不同于 STEP/DeepConf/Slim-SC 的纯减法剪枝,它在剪掉最低分轨迹的同一瞬间从最高分前缀派生新轨迹,活跃数恒为容量 $C$,输出分布被持续推向高分区。系统上,论文首次指出“逻辑搜索决策与物理执行状态耦合”这一失败模式(显存驱逐会扭曲搜索策略),并提出调度器视图/树视图解耦的幽灵轨迹机制,这在同类工作中是全新的视角。

核心方法

直觉上,Gambit 把“一次发 N 条独立采样”的被动模式换成一场持续的“锦标赛”:系统始终维护恰好 $C$ 条(实验中 256 条)活跃轨迹,轻量评分器持续为每条轨迹打分;每隔 $\Delta$ 个 token 触发一轮排序,把分数最低的 $K$ 条处决,同时立刻从分数最高的 $K$ 条前缀经前缀缓存复制出 $K$ 条新分支补位——一减一支、总数不变,算力与显存始终满载。技术路线分三块:(1) 思维级评分:轨迹以 '\n\n' 切分为离散步骤,评分器 $f_\theta$ 作用于每步末的最后一层隐状态 $h_{i,j}$,轨迹分数为累计平均 $\bar{s}_\tau=\frac{1}{n}\sum_{j=1}^n f_\theta(h_{i,j})$;(2) 束搜索主循环,含 warmup 门槛与欠容量/满容量两个路由分支;(3) 调度器/树视图解耦的内存管理,防止物理驱逐扭曲搜索策略。结束后用分数加权多数投票聚合答案。主实验 $C=256$、$K=16$、$\Delta=200$、$w=12000$,实现于 vLLM,开销 <1%。

核心创新是“零和剪枝-分支”与“解耦内存管理”这对组合,与已有方法的本质区别在于主动改变输出分布。STEP、DeepConf 只做减法:终止低分轨迹后释放的算力闲置,剩余候选上的分布基本不变,因此无法突破多数投票的准确率天花板。Gambit 把每次剪枝强制配对一次分支:剪掉分数排名第 $N-k+1$ 名的轨迹 $\tau_{(N-k+1)}$ 的同时,从排名第 $k$ 名的前缀 $\tau_{(k)}$ 派生子轨迹;子轨迹经前缀缓存继承父前缀 KV-Cache、只需增量生成,边际成本极低,还可施加温度乘数促进多样性,从而把概率质量持续移向高分区域。第二个关键发现是:若搜索算法直接感知引擎在显存压力下的抢占/驱逐,活跃数会跌破 $C$,触发反复针对 Top-1/2 轨迹的欠容量分支,形成把算力反复砸向少数高分局部的贪婪塌缩、搜索分布退化。Gambit 让被驱逐的轨迹变成不占显存、不生成 token、但逻辑上仍活跃的“幽灵轨迹”,容量检查始终按树视图的 $N=C$ 判定,永远走均衡的满容量交换,避免策略被系统调度扭曲。

方法步骤详情

流程如下。输入:提示 $P$、容量 $C$、交换尺寸 $K$、间隔 $\Delta$、warmup $w$、评分器 $f_\theta$。第 1 步:从 $P$ 启动 $C$ 条初始轨迹,登记到调度器视图与树视图。第 2 步:轨迹每次推进一个 thought,每完成一步就用 $\bar{s}_\tau \leftarrow \bar{s}_\tau+\frac{f_\theta(h_n)-\bar{s}_\tau}{n}$ 增量更新平均分。第 3 步:显存逼近饱和时,把排名最低的运行中轨迹驱逐为幽灵轨迹(不生成、不占显存、逻辑仍活跃)。第 4 步:每 $\Delta$ 个 token 触发锦标赛:按 $\bar{s}$ 降序排列,生成 token 数超 $w$ 的轨迹才是合法分支父本;若活跃数 $N<C$,从最佳 $\min(C-N,|B|)$ 条各分支一条补齐;若 $N=C$,剪掉最差 $K$ 条、从最佳 $K$ 条前缀分支 $K$ 条新轨迹。第 5 步:完成轨迹移入 $F$,运行数归零或 $|F|=C$ 时终止。第 6 步:输出加权投票答案 $a^*=\arg\max_a \sum_{\text{ans}(\tau)=a}\bar{s}_\tau$,位置加权偏向后期置信度。

技术新颖性

技术新颖性体现在四个层面。其一,形式化创新:论文首次把测试时扩展写成“固定硬件预算下的部分轨迹分配”这一约束优化问题 $\max_\pi \Pr[\mathcal{A}_\pi(P)=y^*]\ \text{s.t.}\ \Omega(\pi)\le B$,把束宽解释为硬件并发容量约束,使算法设计与推理引擎的物理限制显式对齐。其二,粒度创新:经典 token 级束搜索优化序列似然,本文把束操作提升到 thought 级、评分目标改为“前缀续写价值”的代理信号,避免了对单 token 概率的病态偏好,并以 warmup 门槛规避评分器早期噪声。其三,机制创新:零和 prune-and-branch 在恒定显存占用下实现输出分布的主动迁移,这是纯剪枝方法(STEP/DeepConf/Slim-SC)在原理上做不到的——它们只有减法没有加法。其四,系统创新:调度器/树双视图与幽灵轨迹机制解决了“内存驱逐扭曲搜索策略”这一此前未被指出的耦合失败模式,保证策略在真实引擎(vLLM)中按设计执行,且整套机制在单卡 B300 上以 <1% 的开销运行,实现持续硬件饱和与 2 倍以上的轨迹完成吞吐。

End-to-end pipeline of Gambit on an AIME problem (capacity C=5, swap size K=2)
Figure 4: End-to-end pipeline of Gambit on an AIME problem (capacity C=5, swap size K=2)
Thought-Level Beam Search in Gambit
Algorithm 1: Thought-Level Beam Search in Gambit

实验结果

主实验在单卡 B300、每题 256 条完整轨迹预算下覆盖 3 个模型×5 个基准。Qwen3-4B:AIME-25 达 90.0%(与 DeepConf 持平,较 SC 与 STEP 的 86.7% 提升 3.3 个百分点),HMMT-24 65.0%(较 STEP 61.7% 提升 3.3、较 SC 50.8% 提升 14.2 个百分点),GPQA 70.2%,token 较 SC 减少 19.7%–60.6%。DeepSeek-R1-8B:AIME-25 85.8%(较 STEP +2.5),HMMT-24 65.6%(较 SC 提升 9.8 个百分点,token 省 64.0%),HMMT-25 75.8% 追平 STEP。Phi-4:HMMT-25 token 从 5.56M 压到 1.75M(省 68.5%),准确率 75.8% 仍略升。效率上,Qwen3-4B 轨迹完成吞吐 0.216 条/秒,约为 SC 的 2.2 倍;Phi-4 每轨迹中位新增 token 仅 5.2K(SC 14.5K),但存活轨迹更长(35.8K vs STEP 16.5K),总 token 大降而延迟与 STEP 相当、较 SC 快逾 2 倍。因与 STEP 用同一评分器,增益严格归因于搜索拓扑。

End-to-end task accuracy and token consumption per question (×10^6) at N=256
Table 1: End-to-end task accuracy and token consumption per question (×10^6) at N=256
Efficiency vs. Accuracy trade-offs over all benchmarks
Figure 5: Efficiency vs. Accuracy trade-offs over all benchmarks
Trace throughput on AIME-26
Figure 6: Trace throughput on AIME-26
Distribution of unique tokens generated per completed trace on AIME-2026
Figure 7: Distribution of unique tokens generated per completed trace on AIME-2026
Distribution of total sequence length (including inherited prefixes) for completed traces on AIME-2026
Figure 8: Distribution of total sequence length (including inherited prefixes) for completed traces on AIME-2026
查看结构化数据
任务指标本文基线提升
AIME 2025(Qwen3-4B-Thinking) 准确率(%,N=256) 90.0 SC@256 86.7 / STEP 86.7 / DeepConf 90.0 较 STEP +3.3 个百分点,与最优 DeepConf 持平,token 消耗较 SC 省 49.0%
HMMT 2024(Qwen3-4B-Thinking) 准确率(%,N=256) 65.0 STEP 61.7 / DeepConf 58.3 / SC 50.8 较 STEP +3.3、较 DeepConf +6.7、较 SC +14.2 个百分点,token 较 SC 省 60.6%(7.62M→3.00M)
AIME 2025(DeepSeek-R1-0528-Qwen3-8B) 准确率(%,N=256) 85.8 STEP 83.3 / SC 83.3 / DeepConf 81.7 较 STEP +2.5 个百分点,token 较 SC 省 37.7%
HMMT 2024(DeepSeek-R1-0528-Qwen3-8B) 准确率(%,N=256) 65.6 STEP 63.3 / SC 55.8 较 STEP +2.3 个百分点,token 较 SC 省 64.0%
HMMT 2025(Phi-4-reasoning-plus) 每题 token 消耗(×10^6) 1.75(准确率 75.8%) SC@256 5.56(准确率 73.3%)/ STEP 2.78(75.0%) token 较 SC 省 68.5%,准确率 +2.5 个百分点,实现效率-精度前沿的严格支配
GPQA-Diamond(Qwen3-4B-Thinking) 准确率(%,N=256) 70.2 SC 68.2 / STEP 66.9 / DeepConf 67.6 较 SC +2.0 个百分点,验证对研究生级科学推理的泛化性
AIME 2026(Qwen3-4B-Thinking) 轨迹完成吞吐(条/秒) 0.216 SC@256 0.098 / STEP 0.147 约为 SC 的 2.2 倍,在三个模型上均实现 2 倍以上

局限与改进

作者承认的结构性局限是延迟:分支使存活轨迹普遍更长(DeepSeek-R1-8B 中位总长 35.8K vs STEP 16.5K),省下的 token 无法完全转化为墙钟加速,只能与 STEP 相当、个别设置下略慢。我的独立观察有四点。其一,评估集中于数学竞赛(AIME/HMMT)与单一 GPQA,而算法依赖 '\n\n' 作 thought 边界,在代码生成、多轮工具调用等无清晰分段结构的任务上未必成立。其二,锦标赛与加权投票都建立在隐藏状态探针上,warmup 虽缓解早期噪声,但评分器误判仍可能造成“赢者通吃”式错误集中,主表只用 STEP 现成两层 MLP,校准误差敏感度未充分分析。其三,超参($C=256$、$K=16$、$\Delta=200$、$w=12000$)跨基准固定,虽有消融,但对小 $C$ 的消费级算力场景缺乏适配证据。其四,DeepSeek-R1-8B 的 GPQA 上 Gambit(68.2)未超过 DeepConf(68.7),泛化收益并非无条件。

独立分析的弱点

第一,评分器单点依赖:整个锦标赛由一个两层 MLP 驱动,若探针在某类问题上系统性失准,beam 会把算力集中到错误前缀且难以恢复,分支越多错得越集中;改进方向是引入多评分器集成或不确定性感知的保守分支,对低置信区间延迟剪枝。第二,幽灵轨迹占用逻辑名额:被显存驱逐但尚未跌入 bottom-K 的幽灵会一直占据 $N=C$ 中的位置,极端情况下锦标赛可能长期在“少量物理活跃+大量幽灵”上进行,实际多样性下降;改进方向是给幽灵设置过期时间,或按树视图中幽灵占比动态调整交换尺寸 $K$。第三,'\n\n' 分段假设脆弱:对输出不含空行的模型或公式密集的推理,thought 粒度会退化,评分与锦标赛触发节奏失真;可改为基于语义边界或句法信号的自适应分段。第四,聚合公式对后期步骤置信度加位置权重惩罚,但权重形式未在正文消融,存在过拟合基准答案分布的风险;建议补充跨基准的聚合函数敏感性分析,并报告投票熵等诊断指标。

未来方向

作者在附录中提出的方向包括训练带历史感知的自定义序列评分器以替代 STEP 的现成 MLP,并在 A.4 中消融各超参的鲁棒性,这暗示评分器质量仍是方法上限所在。基于本文成果可自然延伸:其一,自适应容量与交换尺寸——按题目难度或运行时评分方差动态调整 $C$、$K$、$\Delta$,把预算向困难题倾斜;其二,把零和束分配推广到代码修复、智能体多步决策等有可验证中间状态的领域,thought 边界可由测试用例或工具调用结果定义;其三,与推测解码、KV 量化等正交的服务优化叠加,进一步压缩延迟短板;其四,在线学习式评分:用完成轨迹的最终对错作为弱标签,在推理过程中持续微调探针,缓解分布漂移;其五,理论层面为“分支自高质量前缀优于独立采样”给出更一般的样本复杂度刻画——Figure 2 中 6.2%→87.5%(14 倍)的极端现象值得形式化解释。

复现评估

复现条件较好但硬件门槛不低。代码已开源(github.com/Dao-AILab/Gambit),实现基于 vLLM,论文明确给出全部超参:容量 $C=256$、交换尺寸 $K=16$、检查间隔 $\Delta=200$ token、warmup $w=12000$ token、硬底 $\delta=0.1$,且声明在所有模型与基准上保持不变,大幅降低调参成本。基准(AIME 2025/2026、HMMT 2024/2025、GPQA-Diamond)全部公开,评测协议(每题 256 条完整轨迹、分数加权投票)描述清楚。主要障碍在算力:主实验在单块 275GB 的 NVIDIA B300 上进行,并要求引擎支持前缀缓存与批内驱逐/恢复,普通实验室难以原样复刻;不过方法的相对增益逻辑(对比 SC@256 与 STEP)应可在单张 80GB A100/H100 上以更小的 $C$ 近似验证。评分器直接复用 STEP 的两层 MLP,无需额外训练即可复现主表;若要复现附录的自定义评分器则需额外训练数据与流程。总体属于“开源扎实、代码完整、硬件敏感”的中等偏低复现难度。