跳转至

Provable Test-Time Scaling for Beam Search in LLM Reasoning

会议: NeurIPS 2026
arXiv: 2609.38672
领域: LLM 推理
关键词: 测试时计算、束搜索、置信过滤、覆盖系数、前缀竞争性

一句话总结

本文在只能采样下一 token、不能读取完整 logits 的设定下,用低频 token 过滤改进束搜索,并在前缀似然与正确性对齐的条件下,将搜索所需样本对最难一步覆盖系数的依赖从最坏情形二次降至近线性;真实 LLM 矩阵乘法实验的准确率由 37.2% 提升至 38.8%。

研究背景与动机

测试时增加计算不一定要训练更大的模型,也可以让固定模型生成更多候选,再投票或使用奖励模型选择答案。Best-of-N 和 Best-of-Majority 的困难在于:它们通常先生成完整推理链,之后才判断哪条值得保留。若每一步正确的概率小于 1,完整正确路径的概率会随长度相乘;大量预算可能花在早就走偏的路径上。束搜索则边扩展边剪枝,让计算集中在少量仍有希望的前缀上。

但提前剪枝也引入了不可逆的错误。本文不假设能够访问准确的下一 token 概率,而是从每个前缀采样若干次,用出现频率估计概率,再累加对数概率排序。在大词表中,低概率 token 偶尔出现一次,其估计频率就变成了“至少一次除以样本数”;错误前缀之后若接上高概率延续,可能把真正正确但每步概率不高的路径挤出束。最终奖励模型再准确,也无法选择已经被剪掉的答案。

本文因此区分两件事:正确路径能否活到结尾,以及活到结尾之后能否被正确选中。前者由局部采样、前缀排序和过滤控制,后者仍受到奖励模型误差影响。核心 idea:在对数似然剪枝之前过滤缺乏重复采样支持的低频 token,抑制长尾估计噪声,并把局部搜索失败与最终奖励选择误差分开分析。

方法详解

整体框架

输入是提示、固定基础模型、束宽 \(b\)、最大长度 \(L\)、每个前缀的采样次数 \(N\),以及过滤阈值。算法从空前缀开始,每轮对当前束中的每个前缀独立采样 \(N\) 个下一 token,把相同 token 合并为频数统计;通过阈值的不同 token 才成为扩展候选。候选分数等于父前缀分数加上该 token 的经验对数概率,然后全局保留分数最高的至多 \(b\) 个前缀。

经过 \(L\) 轮后,标准 CF-Beam 用外部奖励模型从最终束中选一个完整回答。自一致版本则不使用奖励模型,直接选累计经验对数似然最高的完整路径。两者都不使用中间步骤奖励,也不更新基础模型参数;不能把它们理解为带过程奖励监督的新训练框架。

本文的主要贡献是这个简单过滤步骤的理论分析,而非新增网络模块,因此这里用前缀状态与定理条件解释算法,不把证明章节画成架构图。空束时算法返回预先固定的回答,分数并列使用固定规则;这些边界情况也纳入搜索失败事件。

关键设计

1. 采样式前缀评分:把全路径困难拆到每一步,但保留正确性条件

每个下一 token 的经验概率就是其出现次数除以 \(N\)。只有出现过的不同 token 会进入候选集,而不是让同一个 token 的多次出现占据多个束位置。前缀分数是沿途经验概率的对数之和,因而比较的是整个前缀的估计似然,而不是只比较最后一步。不同前缀使用独立采样随机性;共享祖先的路径仍共享早期分数,不能把所有路径得分当作独立变量。

为了描述难度,作者指定一条奖励最优路径,定义其第 \(t\) 步正确 token 的逆概率为 token 级覆盖系数。最难一步由最大值概括,完整路径难度则由这些逆概率的乘积概括:

\[ C_t^{\star}(x)=\frac{1}{\pi_{\mathrm{ref}}(a_t^{\star}\mid s_{t-1}^{\star},x)},\qquad C_{\max}^{\star}(x)=\max_t C_t^{\star}(x),\qquad \overline C(x)=\prod_{t=1}^{L}C_t^{\star}(x). \]

例如正确 token 在每一步都有 0.8 的概率,则最难一步的覆盖系数是 1.25,但完整路径的逆概率是 \(1.25^L\)。束搜索的希望在于反复保留正确前缀,不必等待一次独立完整采样恰好走完所有正确步骤。不过,这只是采样上的优势,不自动证明“高似然就是正确”。

定理 1 给出了普通束搜索的一个困难实例:正确路径前两步各有概率 \(1/C^{\star}\),第一步大量错误 token 的真实概率很小、采样时却各出现一次;这些错误前缀的下一步又分成两个各有概率 1/2 的延续。真实正确前缀依然更有竞争力,但有限样本下错误路径的估计乘积可以更大。该构造解释了为何仅增加束宽不一定消除长尾候选挤占,且二次样本下界是存在性最坏情形结论,不是每个任务的必需成本。

2. 置信过滤:先去掉偶然出现的长尾 token,再做硬剪枝

CF-Beam 相比 Vanilla-Beam 只增加一步:经验概率低于阈值的下一 token 不参与扩展,其余仍沿用累计对数似然和 top-\(b\) 选择。这里的“置信”通过经验频率阈值实现,并不是逐 token 计算校准后的置信区间。过滤发生在扩展阶段,不是在完整回答生成后去重或投票。

主定理的阈值取正确 token 真实概率的 \(1-1/L\) 倍,即 \(\beta_t=(1-1/L)/C_t^{\star}(x)\)。这可以压住长尾 token,又为正确 token 的采样波动留下余量。但它使用未知的正确 token 概率,因此是 oracle 阈值,而不是部署时能直接计算的规则。无奖励版本改用 \(\beta_t=1/(2C_t^{\star}(x))\),同样是 oracle 设定。

实际版本在每个正在访问的前缀上,将阈值设成最大经验频率的固定比例:

\[ \beta_t(s)=\gamma\max_{a\in\mathcal A}\hat\pi(a\mid s,x),\qquad \gamma\in(0,1). \]

实验默认 \(\gamma=0.3\)。这条规则不需要知道正确答案或覆盖系数,但合理性依赖于正确延续在局部也具有竞争力,使其概率与局部最大概率同量级。全局前缀竞争性本身不能直接替代这项局部代理条件;论文没有为该经验阈值建立与 oracle 完全相同的定理。阈值过高会滤掉有用路径,阈值过低则逐渐退回普通束搜索。

3. 前缀竞争性与遗憾分解:说明什么时候剪枝可靠、什么时候增加预算无效

作者要求最优路径在每个深度都具有累计平均对数似然优势,而且要战胜该深度的所有不同前缀。它不是最终答案更常出现,也不是每一个正确 token 都逐步贪心最优,而是整段前缀的概率排序条件:

\[ \frac{1}{t}\log\frac{\pi_{\mathrm{ref}}(s_t^{\star}\mid x)}{\pi_{\mathrm{ref}}(s\mid x)}\geq\kappa,\qquad s\neq s_t^{\star},\quad t\in[L],\quad \kappa\in(0,1]. \]

这个间隔 \(\kappa\) 让经验排序在样本增加后能够接近正确的前缀排序。若没有此条件,命题 1 用一个单步反例说明问题:正确 token 的概率为 1/8,错误 token 的概率为 7/8,束宽为 1;任意采样次数下,正确答案被剪掉的概率至少为 3/4。更准确地估计一个错误的似然排序并不能修正它。

证明把失败分成“正确 token 没通过过滤”和“通过了过滤、前缀却被 top-\(b\) 剪掉”。前一种用二项分布下尾和跨深度的联合界处理;后一种统计有多少经过过滤的竞争前缀得分不低于正确前缀。过滤后的二项矩及逆矩控制比分波动,再利用真实前缀概率总质量求和,避免直接对指数数量的候选逐个付出代价。

为了处理剪枝后的采样依赖,附录在分析中预先为每个可能前缀生成独立频数直方图,只有算法访问时才揭示。这个耦合让被提前剪掉的正确路径也有可定义的经验频率;它不要求真实算法预采样整棵树,也不增加实际查询预算。共享前缀造成的相关性通过 Cauchy–Schwarz 控制,而不是假设竞争路径彼此独立。

在奖励位于 \([0,1]\)、最优奖励为 1、参考分布下奖励均方误差至多 \(\epsilon_{\mathrm{RM}}^2\)、最优点误差至多 \(\epsilon_{\mathrm{opt}}\) 的条件下,定理 2 得到:

\[ \mathrm{Reg}(x)\leq L\exp\!\left(-\frac{N}{2C_{\max}^{\star}L^2}\right) +\frac{2C_{\max}^{\star}}{b}\exp\!\left(-\frac{\kappa^2N}{48C_{\max}^{\star}}\right) +\epsilon_{\mathrm{opt}}(x) +2\sqrt{\overline C(x)\epsilon_{\mathrm{RM}}^2(x)}. \]

遗憾是最优真实奖励减去算法输出的期望真实奖励。这里还要求 \(L\geq2\),并满足 \(N\geq(48C_{\max}^{\star}/\kappa^2)\max\{2,\log(2C_{\max}^{\star})\}\)。前两项随采样增多下降,后两项不会因本界中的 \(N\) 或 \(b\) 增大而下降;不能把“搜索项摆脱完整路径覆盖系数”写成“整个遗憾都摆脱指数难度”。

奖励误差项来自搜索输出与参考分布之间的分布迁移:即使奖励模型在常见回答上均方误差很小,搜索仍可能偏向参考模型极少生成的回答。oracle 阈值让联合输出概率受到至多 \(4\overline C\) 倍参考概率的控制,随后转移奖励误差。定理 3 的小间隔困难实例说明该路径覆盖依赖在最坏情形下不能直接删除;这不等于固定正间隔下搜索永远无法只留下正确答案。

4. 自一致最终选择:不靠奖励模型,但必须坚持正确性与似然对齐

自一致 CF-Beam 最后选经验似然最高的完整路径,而不是奖励分数最高的路径。由于没有奖励误差迁移,可以把过滤阈值保持为正确 token 概率的一半,不再让阈值随 \(L\) 增长而逼近均值。推论 2 因此允许每个前缀的样本数为 \(N\gtrsim(C_{\max}^{\star}/\kappa^2)\log(eLC_{\max}^{\star}/\delta)\),使遗憾至多 \(\delta\);固定束宽、覆盖系数、间隔和目标精度时,总查询预算只随长度近线性增长。

不过,“正确路径留在束里”不足以保证这种最终选择成功,证明还必须排除结尾处任何得分不低于正确答案的竞争者。这是它相比奖励选择额外控制的事件。它也不是常见的“多个不同思维链按最终答案投票”:本文直接对完整路径的累计经验似然做最大化,正确性依赖前缀 top-1 间隔。

附录 B 给出两项条件扩展。奖励最优路径可以不唯一,但须指定一条满足间隔与最优点误差条件的见证路径,且它仍须胜过其他正确路径的不同前缀;理论并没有利用全部正确路径的总概率质量。奖励选择版本也可允许每个深度至多 \(k-1\) 个前缀违背间隔,只要 \(1\leq k\leq b\),将剪枝界中的 \(b\) 替换为 \(b-k+1\);无奖励的最终最大似然选择不能直接继承这个放宽。

损失函数 / 训练策略

本文没有新增训练损失或微调阶段,基础策略和最终奖励模型均视为给定。奖励模型误差是理论假设,不是 CF-Beam 自己训练出的保证;实验中的 AceMath-7B-RM 也只对完整回答评分。

查询成本按“从某个前缀抽一次下一 token”计数,实际最多为 \(N+(L-1)bN\leq LbN\)。每轮剪枝前最多处理 \(bN\) 个不同候选,过滤只减少候选,不引入额外分支。这些计数不等于神经网络前向次数、延迟或 KV 缓存大小:本地模型可复用一次前向的分布采样多次,API 与批处理的实际成本则不同。

实验关键数据

主实验

真实 LLM 实验使用 Qwen3-1.7B-Base,采样温度 1.3,在 500 道合成的 \(4\times4\) 整数矩阵乘法题上评估;需要奖励选择的方法使用 AceMath-7B-RM。两个束搜索方法均设置 \(N=12\)、\(b=2\),CF-Beam 使用经验阈值 \(\gamma=0.3\)。下表来自附录 H.2 表 2,准确率为比例,预算单位为千次 token 级策略查询。

方法 准确率 查询预算(千次)
Best-of-N 0.332 16.8
Majority Voting 0.304 16.8
Best-of-Majority 0.326 16.8
Vanilla-Beam 0.372 18.6
Empirical CF-Beam 0.388 16.2

CF-Beam 比 Vanilla-Beam 提高 1.6 个百分点,同时报告预算从 18.6 千次降为 16.2 千次;相对 Best-of-N 提高 5.6 个百分点。这里是相近、并非严格相同的实测预算,不能据此声称获得等 FLOPs、等延迟的收益。

消融实验

附录 H.2 表 3 固定 \(N=12\),考察过滤比例与束宽;\(\gamma=0\) 表示关闭过滤。下表摘取能显示“适度过滤有益、过强过滤有害”的配置,未列出全部阈值。

束宽 \(\gamma=0\) \(\gamma=0.15\) \(\gamma=0.3\) \(\gamma=0.6\) \(\gamma=0.9\)
\(b=2\) 0.372 0.378 0.388 0.318 0.290
\(b=4\) 0.482 0.512 0.472 0.374 0.302
\(b=8\) 0.536 0.550 0.480 0.364 0.300

最佳阈值并不固定:\(b=2\) 时为 0.3,而 \(b=4,8\) 时为 0.15。三个束宽下,相对无过滤的最佳准确率增益分别为 1.6、3.0、1.4 个百分点。加宽束提高了最佳准确率,但该表没有对应预算列,不是等总预算的束宽比较。

理论结果也可作为分析表阅读。下表中的近线性结论都需要保留条件,而不是直接从真实 LLM 准确率反推定理已获验证。

结果 条件与保证 适用边界
定理 1:普通束搜索下界 存在 $50\leq C^{\star}\ll \mathcal A
定理 2 / 推论 1:奖励选择 CF-Beam oracle 阈值、前缀间隔与奖励误差假设;固定间隔和精度时,总预算 \(\widetilde O(bL^3C_{\max}^{\star})\) 只将搜索失败压低;奖励误差仍带 \(\overline C\)
推论 2:自一致 CF-Beam oracle 半概率阈值、top-1 前缀间隔;总预算 \(\widetilde O(LbC_{\max}^{\star}/\kappa^2)\) 不需奖励误差假设,但最终最大似然必须对应最优奖励
命题 1:似然错位反例 \(b=1\),正确/错误概率为 1/8、7/8;任意 \(N\) 下失败概率至少 3/4 更多采样无法单独弥补似然与正确性的错位

关键发现

  • 模拟器每个实例运行 300 次独立试验,用 0.01 的对称标签翻转模拟奖励误差;完整正确路径概率取 0.01、0.05、0.3 时,过滤的优势在最难设置最大。
  • 长度实验考察 \(L\in[2,40]\),完整正确路径概率设为 \(0.8^L\)。作者报告序列级方法下降更快,但缓存未提供曲线的可读逐点数值,因此不补写具体准确率。
  • 经验版可能在固定有限预算下超过 oracle 版:oracle 阈值距正确概率仅有相对 \(1/L\) 的余量,长度增大时更易受向下采样波动影响;“oracle”不表示每个预算下都最优。
  • 原文图标题把难度实验标为图 2、长度实验标为图 3,但正文与 H.1 多次把两组都引用为图 3;本笔记按实验内容区分,不自行补改图号。算法交叉引用也出现 Algorithm 3 与已展示 Algorithm 1 不一致的情况。

亮点与洞察

  • 把“长尾 token 被偶然采到”与“错误前缀后续看起来很确定”连起来,解释了普通束搜索并非只受正确 token 是否出现限制。过滤是在错误路径占据束之前阻断这一链条。
  • 最有价值的区分是局部搜索覆盖与最终奖励的分布迁移。前者可以避免完整路径逆概率的乘积,后者仍可能因稀有路径上的奖励误差恶化。
  • 证明不把自适应剪枝后的候选当作独立样本,而是借助预采样耦合、截断矩与概率质量求和。这个分析思路可用于其他“先采样估计、再硬选择”的解码器。

局限与展望

  • 前缀竞争性很强:正确答案必须在所有深度都保持累计似然领先。现实中罕见但正确的推理步骤可能不满足它,不能将理论推广为任意数学推理的保证。
  • 部署所用经验阈值与主定理的 oracle 阈值不同;实验支持前者有效,但没有弥合二者的理论条件。值得研究基于可观测量的间隔诊断与自适应阈值。
  • 真实模型实验只有一个基础模型、一个矩阵乘法任务和 500 个问题,所列结果未附置信区间;提升尚不足以支持广泛任务、不同温度或多模型上的稳定收益。
  • 奖励选择版本的误差界仍受完整路径覆盖系数放大;改进奖励校准、引入中途验证,或在似然错位提示上切换到序列级方法,比无条件加大采样更有针对性。
  • 成本是抽样查询计数,不是实际运行性能。后续应同时报告前向次数、生成长度、批处理、延迟与内存,并利用多条正确路径的共享前缀获得更贴近实际的保证。

相关工作与启发

  • vs Best-of-N / Best-of-Majority:二者先生成完整回答再选择,覆盖难度通常含完整路径的逆概率;CF-Beam 提前维护前缀,在似然对齐时改善搜索项,但引入不可逆剪枝及间隔假设。
  • vs 自一致投票:传统方法按答案出现频率聚合不同完整推理链;本文无奖励变体按路径经验似然选择,不应将其近线性预算结论直接归给答案投票。
  • vs 过程奖励引导搜索:Lightman 等与过程验证器工作可以对中间步骤提供额外正确性信号;本文刻意只允许结尾奖励,因此揭示了纯似然剪枝的能力与边界。
  • 可迁移启发:首先判断候选排序信号是否与最终目标对齐,再决定应提前剪枝还是延迟选择;过滤比例也应随束宽和局部采样可靠性调整,而不是固定视为通用最优值。

评分

  • 新颖性: 4/5 — 过滤机制简单,但给出普通束搜索反例及覆盖依赖改进,理论问题明确。
  • 实验充分度: 3/5 — 有模拟器、真实 LLM 对比和阈值消融,模型与任务覆盖仍有限。
  • 写作质量: 3/5 — 搜索与奖励误差的分解清楚,但图号和算法交叉引用有不一致。
  • 价值: 4/5 — 有助于识别束搜索真正节省计算的条件,避免把增加预算等同于可靠推理。