Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters

文章标题：优化扩展LLM测试时计算比扩展模型参数更有效
作者/机构：Charlie Snell♦, 1, Jaehoon Lee2, Kelvin Xu♣, 2 and Aviral Kumar♣, 2 (♣共同指导, 1加州大学伯克利分校, 2Google DeepMind, ♦工作于Google DeepMind实习期间完成)

A1 主要贡献

本文旨在探讨如何通过增加测试时（test-time）的计算量来提升大型语言模型（LLM）的性能，特别是针对具有挑战性的任务。研究的核心问题是：在给定固定的、非平凡的推理计算预算下，一个LLM能在多大程度上提升其在复杂提示（prompt）上的表现？这一问题的答案不仅关系到LLM可达到的性能极限，也对未来LLM的预训练方向以及如何在推理时计算和预训练计算之间进行权衡具有重要意义。

核心问题与研究目标：
尽管现有研究表明LLM可以利用测试时计算来改进输出，但在数学推理等复杂任务上，这些方法的效果仍然有限。这种矛盾的结果凸显了对不同测试时计算扩展方法进行系统性分析的必要性。本文旨在系统地理解扩展测试时计算的益处，并探讨其与扩展模型预训练计算之间的权衡关系。

研究方法与创新点：
本文分析了两种扩展测试时计算的主要机制：
1. 基于密集、过程化验证器（process-based verifier）奖励模型的搜索：训练一个能够评估解决方案中每一步正确性的过程奖励模型（Process Reward Model, PRM），并利用树搜索等算法在该模型上进行搜索，以找到最佳答案。
2. 在测试时自适应地更新模型对响应的分布：通过让模型迭代地修正（revise）自己的答案，从而在测试时动态地改进其自身的提议分布（proposal distribution）。

核心发现：
1. 有效性依赖于问题难度：研究发现，无论是采用修正还是搜索策略，其有效性都严重依赖于具体问题的难度和所使用的基础LLM。例如，对于简单问题，迭代修正可能更有效；而对于难题，并行采样或树搜索等探索性更强的方法可能更优。
2. 提出“计算最优”扩展策略：基于上述发现，本文提出了一种“计算最优”（compute-optimal）的扩展策略，该策略能根据每个提示的难度自适应地分配测试时计算资源。通过使用一个可预测问题难度的指标，该策略能够为每个问题选择最有效的计算方式。
3. 显著提升计算效率：实验证明，采用计算最优策略，相比于传统的best-of-N基线，可以在修正和搜索两种场景下将测试时计算的效率提升超过4倍。
4. 测试时计算可替代预训练计算：在一个FLOPs匹配的评估中，本文发现，对于那些较小基础模型能取得一定成功率的问题，利用测试时计算可以使其性能超越一个参数量大14倍的模型。这表明，在某些场景下，将计算资源用于测试时优化而非单纯地预训练更大的模型可能更为有效。然而，对于最困难的问题，扩展预训练计算仍然是更有效的方法，说明两者并非完全可以互相替代。

A3 背景知识与设计原则

2. 测试时计算的统一视角：提议器与验证器

• 测试时计算的统一框架：本文将使用额外测试时计算的方法统一看作是在测试时根据给定提示自适应地修改模型预测分布的过程。理想情况下，测试时计算应修改分布以生成比直接从LLM采样更好的输出。修改LLM分布通常有两种方式：(1) 在输入层面，通过在给定提示中增加额外的token来让LLM进行条件化，从而获得修改后的分布；(2) 在输出层面，通过从标准语言模型中采样多个候选答案，并对这些候选答案进行处理。换言之，我们可以修改LLM自身的提议分布（proposal distribution），使其优于仅基于原始提示的分布，或者使用某种后验的验证器（verifier）或评分器来执行输出修改。这个过程类似于马尔可夫链蒙特卡洛（MCMC）【索引2，An introduction to mcmc for machine learning，2003】采样，即通过结合一个简单的提议分布和一个评分函数来从一个复杂的目标分布中采样。修改提议分布和使用验证器构成了本文研究的两个独立维度。
• 修改提议分布：一种改进提议分布的方法是通过强化学习（RL）风格的微调方法（如STaR或ReSTEM【索引35，Beyond human data: Scaling self-training for problem-solving with language models，2024；索引50，Star: Bootstrapping reasoning with reasoning，2022】）直接为特定推理任务优化模型。这些技术不使用额外的输入token，而是专门微调模型以引导出一个改进的提议分布。相对地，像自评（self-critique）【索引4，Constitutional ai: Harmlessness from ai feedback，2022；索引8，Improving factuality and reasoning in language models through multiagent debate，2023；索引23，Selfrefine: Iterative refinement with self-feedback，2023；索引30，Self-critiquing models for assisting human evaluators，2022】这样的技术，则通过指示模型在测试时迭代地批判和修正自己的输出来改进其提议分布。由于对现成模型进行提示在实现有效修正方面效果不佳，本文专门微调模型，使其能在复杂的推理场景中迭代修正答案。为此，本文采用了在策略（on-policy）数据上进行微调的方法，并利用Best-of-N引导模型响应的改进【索引28，Recursive introspection: Teaching foundation models how to self-improve，2024】。
• 优化验证器：在提议分布和验证器的抽象框架中，验证器用于从提议分布中聚合或选择最佳答案。最典型的方法是应用best-of-N采样，即采样N个完整解决方案，然后根据验证器选出最好的一个【索引7，Training verifiers to solve math word problems，2021】。然而，该方法可以通过训练一个基于过程的验证器（process-based verifier），或称过程奖励模型（PRM）【索引22，Let’s verify step by step，2023】，来进一步改进。PRM能够预测解决方案中每个中间步骤的正确性，而不仅仅是最终答案。然后，我们可以利用这些每步的预测在解决方案空间中执行树搜索，这提供了一种可能比朴素的best-of-N更高效、更有效的验证器搜索方式【索引6，Alphamath almost zero: process supervision without process，2024；索引10，Alphazero-like tree-search can guide large language model decoding and training，2024；索引48，Tree of thoughts: Deliberate problem solving with large language models，2023】。

3. 如何最优地扩展测试时计算

• 问题设定：给定一个提示和一个测试时计算预算，我们希望解决这个问题。在上述抽象框架下，有多种利用测试时计算的方式，每种方式的有效性可能因具体问题而异。本文旨在回答两个问题：如何为给定提示确定最有效的测试时计算使用方式？以及这种方式与简单地使用一个更大的预训练模型相比表现如何？
• 超参数选择的挑战：无论是优化提议分布还是通过验证器进行搜索，都存在多个超参数可以调整，以决定如何分配测试时计算预算。例如，当使用经过修正微调的模型作为提议分布，并使用结果奖励模型（ORM）作为验证器时，我们可以将全部计算预算用于并行生成N个独立样本并应用best-of-N，也可以用修正模型串行生成N个修正版本，然后用ORM选出最佳答案，或者在两者之间寻求平衡。直观上，“简单”问题可能更受益于修正，因为模型的初始样本可能已经方向正确，只需进一步完善。而对于难题，可能需要探索更多不同的高层解题策略，因此并行独立采样多次可能更可取。
• 不同搜索算法的选择：在使用验证器时，我们还可以选择不同的搜索算法（如波束搜索、前瞻搜索、best-of-N），每种算法的性能可能取决于验证器和提议分布的质量。对于更难的问题，更复杂的搜索过程可能比简单的best-of-N或多数投票基线更有用。
• 测试时计算最优扩展策略的定义：本文的目标是为给定问题选择测试时计算预算的最优分配方式。为此，对于任何一种利用测试时计算的方法（本文中为修正和验证器搜索），定义“测试时计算最优扩展策略”为：在给定测试时，为特定提示选择能带来最大性能收益的超参数。形式上，定义Target(x, θ, C)为给定提示x、使用测试时计算超参数θ和计算预算C时，模型引导出的自然语言输出token的分布。我们希望选择能最大化目标分布准确性的超参数θ。其形式化表达如下：
  

  

  
  其中，$y^*(x)$表示$x$的真实正确答案，$P^*_{\theta, C}(x)$代表问题$x$在计算预算$C$下的测试时计算最优扩展策略。
• 通过估计问题难度实现计算最优扩展：为了有效分析第2节中讨论的不同机制（如提议分布和验证器）的测试时扩展特性，本文将最优策略$P_{\theta, C}^*(x)$近似为一个关于给定提示统计量的函数。这个统计量估计了给定提示的“难度”。计算最优策略被定义为该提示难度的函数。尽管这只是对公式1所示问题的近似解，但我们发现它仍然能比采用临时或均匀采样方式分配推理时计算的基线策略带来显著的性能提升。
• 具体实现方法：我们估计问题难度的方法是将一个给定问题分配到五个难度级别之一。然后，我们可以利用这个离散的难度分类，在验证集上为给定的测试时计算预算估计$P_{\theta, C}^*(x)$，再将这些计算最优策略应用于测试集。具体来说，我们为每个难度分箱独立选择表现最好的测试时计算策略。通过这种方式，问题难度在设计计算最优策略时充当了一个问题的充分统计量。
• 问题难度的定义：本文遵循Lightman等人【索引22，Let’s verify step by step，2023】的方法，将问题难度定义为与给定基础LLM相关的函数。具体来说，我们将在测试集中的每个问题上，通过2048个样本估计模型的pass@1率，并将其分为五个分位数，每个分位数对应不断增加的难度级别。我们发现，这种特定于模型的难度分箱比MATH数据集中手动标记的难度分箱更能预测测试时计算的有效性。
• 模型预测的难度与实际部署：然而，上述难度评估方式需要预先知道正确答案，这在实际部署中是不可行的。为了在实践中可行，一个基于难度的计算最优扩展策略需要先评估难度，然后利用正确的扩展策略解决问题。因此，我们通过一个模型预测的难度概念来近似问题的难度，即对每个问题的2048个样本，使用一个学习到的验证器（而非真实答案检查）的平均最终答案得分进行同样的分箱操作。我们将此设置称为模型预测难度（model-predicted difficulty），而依赖于真实正确性的设置称为神谕难度（oracle difficulty）。
• 难度评估的计算成本：虽然模型预测的难度无需知道真实标签，但这种方式估计难度本身仍会产生额外的推理计算成本。不过，这种一次性的推理成本可以被包含在实际运行推理时策略的成本之内（例如，使用验证器时，可以利用相同的推理计算来运行搜索）。这类似于强化学习中的探索-利用权衡：在实际部署中，我们必须平衡用于评估难度的计算与应用计算最优方法的计算。这是一个重要的未来工作方向（见第8节），我们的实验为简化起见，并未计入这部分成本，因为我们的目标是首次展示有效分配测试时计算所能达到的效果。
• 交叉验证方法：为避免在计算难度分箱和选择计算最优策略时使用相同测试集带来的混淆，我们对测试集中的每个难度分箱使用两折交叉验证。我们在其中一折上根据性能选择最佳策略，然后在另一折上使用该策略测量性能，反之亦然，最后对两折测试的结果取平均值。

A4 实验环境

• 数据集：实验主要使用MATH基准数据集【索引13，Measuring mathematical problem solving with the math dataset，22021】，该数据集包含高中竞赛水平的数学问题，难度范围广泛。所有实验均采用Lightman等人【索引22，Let’s verify step by step，2023】使用的数据集划分，包含12k个训练问题和500个测试问题。
• 模型架构：本文分析使用了PaLM 2-S* (Codey) 基础模型【索引3，Palm 2 technical report，2023】。该模型在MATH上取得了不错的性能但尚未饱和，适合作为测试平台。
• 硬件配置：论文未明确提及。
• 软件配置：
  • 模型实现基于PaLM 2-S*，并针对修正（revision）和过程验证（PRM）任务进行了微调。
  • PRM训练使用了AdamW优化器，学习率为3e-5，批量大小为128。
  • 修正模型训练也使用了AdamW优化器，学习率为1e-5，批量大小为128。
  • 基础LLM的输出采用4-shot prompt引导，以生成分步解答格式。

A2 方法细节

5. 通过验证器扩展测试时计算

本节分析如何通过优化验证器来尽可能有效地扩展测试时计算。为此，我们研究了使用过程验证器（PRM）进行测试时搜索的不同方法，并分析了这些方法的测试时计算扩展特性。

5.1. 训练适合搜索的验证器

• PRM训练：最初的PRM训练【索引22，Let’s verify step by step，2023；索引42，Solving math word problems with process- and outcome-based feedback，2022】使用人类众包标签。虽然Lightman等人【索引22】发布了他们的PRM训练数据（PRM800k数据集），但我们发现这些数据对我们基本无效。我们发现，在该数据集上训练的PRM很容易被即使是朴素的best-of-N采样等策略所利用。我们推测这可能是由于他们数据集中GPT-4生成的样本与我们的PaLM 2模型之间存在分布偏移。我们没有采用昂贵的为PaLM 2模型收集众包PRM标签的过程，而是应用了Wang等人【索引45，Math-shepherd: Verify and reinforce llms step-by-step without human annotations，2023】的方法，在没有人为标签的情况下监督PRM，通过从解决方案的每一步运行蒙特卡洛（Monte Carlo） rollout来获得每步正确性的估计值。因此，我们的PRM每步预测对应于基础模型采样策略的未来奖励价值估计（reward-to-go），这与近期工作类似【索引31，Rl on incorrect synthetic data scales the efficiency of llm math reasoning by eight-fold，2024；索引45】。我们还与一个ORM基线进行了比较（附录F），但发现我们的PRM始终优于ORM。因此，本节中所有的搜索实验都使用PRM模型。更多PRM训练细节见附录D。
• 答案聚合：在测试时，基于过程的验证器可用于对基础模型采样的一组解决方案中的每个单独步骤进行评分。为了使用PRM选择best-of-N的答案，我们需要一个函数来聚合每个答案的所有步骤分数，以确定正确答案的最佳候选者。为此，我们首先聚合每个单独答案的步骤分数以获得完整答案的最终分数（步骤聚合），然后跨答案进行聚合以确定最佳答案（答案间聚合）。具体方法如下：
  • 步骤聚合（Step-wise aggregation）：我们不采用将各步骤分数相乘或取最小值【索引22，Let’s verify step by step，2023；索引45】的方法，而是使用PRM在最后一步的预测作为完整答案的分数。我们发现这种方法在我们研究的所有聚合方法中表现最好（见附录E）。
  • 答案间聚合（Inter-answer aggregation）：我们遵循Li等人【索引21，Making large language models better reasoners with step-aware verifier，2023】的方法，应用“best-of-N加权”选择，而非标准的best-of-N。Best-of-N加权选择将具有相同最终答案的所有解决方案的验证器正确性分数进行边缘化，选择总和最大的最终答案。

5.2. 基于PRM的搜索方法

• 搜索方法概述：我们通过搜索方法在测试时优化PRM。我们研究了三种搜索方法，它们都从一个few-shot提示的基础LLM中采样输出（见附录G）。图2展示了这些方法的示意图。
• Best-of-N加权（Best-of-N weighted）：我们从基础LLM中独立采样N个答案，然后根据PRM的最终答案判断选择最佳答案。
• 波束搜索（Beam search）：波束搜索通过在其每步预测上进行搜索来优化PRM。我们的实现类似于BFS-V【索引10，Alphazero-like tree-search can guide large language model decoding and training，2024；索引48】。具体来说，我们考虑固定数量的波束$b$和波束宽度$w$。然后运行以下步骤：
  1. 为解决方案的第一步采样$w$个初始预测。
  2. 根据PRM预测的每步未来奖励估计对生成的步骤进行评分（在此稀疏奖励设置中，这也对应于前缀的总奖励）。
  3. 只保留得分最高的$w/b$个步骤。
  4. 现在从每个候选步骤中，为下一步采样$b$个提议，总共得到$w/b \times b = w$个候选前缀。然后重复步骤2-4。
     我们运行此算法直到解决方案结束或达到最大波束扩展轮数（本文为40轮）。搜索结束时会得到N个最终答案候选，我们对其应用上述的best-of-N加权选择来做出最终答案预测。
• 前瞻搜索（Lookahead search）：前瞻搜索修改了波束搜索评估单个步骤的方式。它使用前瞻rollout来提高搜索过程中每一步PRM价值估计的准确性。具体来说，在波束搜索的每一步，前瞻搜索不是使用当前步骤的PRM分数来选择顶部候选，而是执行一个模拟，向前rollout最多$k$步，如果达到解决方案末尾则提前停止。为了最小化模拟rollout中的方差，我们使用温度为0进行rollout。该rollout结束时PRM的预测被用来为当前波束搜索步骤评分。换句话说，我们可以将波束搜索视为$k=0$的前瞻搜索的特例。在PRM准确的情况下，增加$k$应该会以额外的计算为代价提高每步价值估计的准确性。还需注意，此版本的前瞻搜索是MCTS【索引38，Reinforcement learning: An introduction，2018】的一个特例，其中移除了为促进探索而设计的MCTS的随机元素，因为PRM已经训练好并被冻结。这些随机元素主要用于学习价值函数（我们已经用PRM学到了），但在测试时当我们想要利用而非探索时，它们的作用较小。因此，前瞻搜索在很大程度上代表了MCTS风格的方法在测试时的应用方式。

5.3. 分析结果：使用验证器进行搜索的测试时扩展

• 搜索算法比较：我们首先对各种搜索设置进行了扫描。除了标准的best-of-N方法，我们还扫描了区分不同树搜索方法的两个主要参数：波束宽度$w$和前瞻步数$k$。虽然无法详尽扫描所有配置，我们扫描了以下设置，最大预算为256次生成：
  1. 波束宽度设为$w$的波束搜索，其中$w$是生成预算。
  2. 固定波束宽度为4的波束搜索。
  3. 将$k=3$的前瞻搜索应用于上述两种波束搜索设置。
  4. 将$k=1$的前瞻搜索应用于波束搜索设置1。
• 计算成本的公平比较：为了公平地比较不同搜索方法作为生成预算函数的性能，我们建立了一个估算每种方法成本的协议。我们将一次“生成”定义为从基础LLM采样一个答案。对于波束搜索和best-of-N，生成预算分别对应于波束数量和$N$。然而，前瞻搜索利用了额外的计算：在搜索的每一步，我们额外向前采样$k$步。因此，我们将前瞻搜索的成本定义为$N \times (k+1)$个样本。
• 结果分析：如图3（左）所示，在较小的生成预算下，波束搜索显著优于best-of-N。然而，随着预算的增加，这些优势大大减弱，波束搜索的表现常常低于best-of-N基线。我们还看到，前瞻搜索在相同的生成预算下通常表现不如其他方法，这可能是由于模拟前瞻rollout所引入的额外计算。搜索带来的收益递减可能是由于对PRM预测的过度利用（exploitation）。例如，我们看到一些实例（如图29），其中搜索导致模型在解决方案的末尾生成信息量低的重复步骤。在其他情况下，我们发现过度优化搜索可能导致解决方案过短，仅包含1-2个步骤。这解释了为什么最强大的搜索方法（即前瞻搜索）表现最差。我们在附录M中包含了一些由搜索找到的此类示例。
• 搜索对不同难度问题的影响：为了理解如何计算最优地扩展搜索方法，我们现在进行难度分箱分析。具体来说，我们比较了波束搜索（$w=4$）和best-of-N。在图3（右）中我们看到，虽然在高生成预算下，波束搜索和best-of-N的总体性能相似，但按难度分箱评估其有效性时，却揭示了非常不同的趋势。在简单问题上（级别1和2），两种方法中更强的优化器——波束搜索，随着生成预算的增加性能下降，显示出利用PRM信号的迹象。相反，在较难的问题上（级别3和4），波束搜索持续优于best-of-N。在最难的问题上（级别5），没有任何方法取得有意义的进展。
• 结果的直观解释：这些发现符合直觉：我们可能预期在简单问题上，验证器对正确性的评估大多是正确的。因此，通过波束搜索进行进一步优化，只会放大验证器学到的任何虚假特征，导致性能下降。在更难的问题上，基础模型本身采样到正确答案的可能性要小得多，因此搜索可以帮助引导模型更频繁地产生正确答案。
• 计算最优搜索：鉴于上述结果，问题难度显然可以作为一个有用的统计量，来预测在给定计算预算下应使用的最优搜索策略。此外，最佳搜索策略的选择会随着这个难度统计量的变化而急剧变化。因此，我们在图4中可视化了“计算最优”扩展趋势，即在每个难度级别上表现最佳的搜索策略。我们看到，在低生成预算范围内，使用神谕难度和预测难度，计算最优扩展几乎能用少至4倍的测试时计算（例如16次生成对64次生成）超越best-of-N。而在高预算范围内，使用预测难度时，这些优势有所减弱，但使用神谕分箱我们仍然看到最优扩展测试时计算带来的持续改进。这个结果展示了通过自适应地分配测试时搜索计算所能获得的性能增益。

• 关于验证器计算最优扩展的要点：我们发现，任何给定的验证器搜索方法的有效性都严重依赖于计算预算和手头的问题。具体来说，波束搜索在较难的问题和较低的计算预算下更有效，而best-of-N在较容易的问题和较高的预算下更有效。此外，通过为给定的问题难度和测试时计算预算选择最佳的搜索设置，我们几乎可以用少至4倍的测试时计算超越best-of-N。

6. 优化提议分布

到目前为止，我们研究了针对验证器进行搜索的测试时计算扩展特性。现在我们转向研究修改提议分布（第2节）的扩展特性。具体来说，我们让模型能够迭代地修正自己的答案，从而在测试时动态地改进其自身的分布。简单地提示现有LLM纠正自己的错误，在推理问题上通常无法有效提升性能【索引15，Large language models cannot self-correct reasoning yet，2023】。因此，我们基于Qu等人【索引28，Recursive introspection: Teaching foundation models how to self-improve，2024】提出的方法，并针对我们的设置进行修改，微调语言模型以迭代地修正自己的答案。我们首先描述如何训练和使用通过顺序地依赖于自身先前尝试来优化自身提议分布的模型。然后，我们分析修正模型的推理时扩展特性。

6.1. 设置：训练和使用修正模型

• 修正模型微调流程：我们微调修正模型的流程与【索引28】相似，但我们引入了一些关键差异。为了微调，我们需要由一系列不正确答案后跟一个正确答案组成的轨迹，然后我们可以在这些轨迹上运行监督式微调（SFT）。理想情况下，我们希望正确答案与上下文中提供的不正确答案相关联，以便有效地教导模型隐式识别上下文中示例的错误，并通过进行编辑来纠正这些错误，而不是完全忽略上下文中的示例并从头再来。
• 生成修正数据：Qu等人【索引28】获取多个多轮rollout的在策略（on-policy）方法被证明是有效的，但在我们的基础设施中，由于运行多轮rollout的计算成本，这并不完全可行。因此，我们在较高温度下并行采样64个响应，并从这些独立样本中事后构建多轮rollout。具体来说，我们遵循【索引1，Training revision models with synthetic data，2024】的方法，将每个正确答案与一系列不正确的答案配对作为上下文，以构建多轮微调数据。我们在上下文中最多包含四个不正确的答案，上下文中解决方案的具体数量是从0到4的类别中均匀随机抽样的。我们使用字符编辑距离度量来优先选择与最终正确答案相关的不正确答案（见附录H）。请注意，token编辑距离并不是一个完美的关联度量，但我们发现这个启发式方法足以将上下文中的不正确答案与正确的目标答案关联起来，以促进训练一个有意义的修正模型，而不是随机地将不相关的不正确和正确响应配对。
• 在推理时使用修正：给定一个微调好的修正模型，我们可以在测试时从模型中采样一系列修正。虽然我们的修正模型只在上下文中最多有四个先前答案的情况下进行训练，但我们可以通过将上下文截断为最近的四个修正响应来采样更长的链。在图6（左）中，我们看到随着我们从修正模型中采样更长的链，模型在每一步的pass@1逐渐提高，表明我们能够有效地教导模型从上下文中先前答案所犯的错误中学习。
• 推理时的分布偏移问题：尽管如此，在推理时存在分布偏移：模型只在上下文中包含不正确答案的序列上进行训练，但在测试时，模型可能会采样到正确的答案并将其包含在上下文中。在这种情况下，模型可能会在下一次修正步骤中偶然地将正确答案变成不正确答案。我们发现，确实，与Qu等人【索引28】类似，使用我们的修正模型采用朴素方法时，大约38%的正确答案会被转换回不正确的答案。因此，我们采用一种基于顺序多数投票或基于验证器的选择机制，从模型做出的一系列修正中选择最正确的答案（见图5）来产生最佳答案。
• 比较串行与并行采样：为了测试通过修正来修改提议分布的有效性，我们对顺序采样N个修正和并行采样N次尝试的性能进行了公平比较。我们在图6（右）中看到，无论是使用基于验证器的选择机制还是基于多数投票的选择机制，顺序采样解决方案的性能都优于并行采样。

6.2. 分析结果：使用修正进行测试时扩展

• 串行与并行计算的权衡：我们之前看到，顺序提出答案优于并行提出。然而，我们可能预期顺序和并行采样具有不同的特性。并行采样可能更像一个全局搜索过程，原则上可以覆盖许多完全不同的解题方法，例如，不同的候选方案可能采用完全不同的高层方法。而顺序采样则更像一个局部优化过程，修正那些已经大致正确的响应。由于这些互补的优势，我们应该通过将一部分推理时间预算分配给并行采样（例如$m$）和其余部分分配给顺序修正（例如$k$）来在这两个极端之间取得平衡。我们现在将展示存在一个计算最优的顺序与并行采样比率，并根据给定提示的难度来理解它们的相对优缺点。
• 权衡顺序与并行测试时计算：为了理解如何最优地分配顺序和并行计算，我们对多种不同的比率进行了扫描。我们在图7（左）中看到，确实，在给定的生成预算下，存在一个理想的顺序与并行比率，可以达到最高的准确率。我们还在图7（右）中看到，理想的顺序与并行比率取决于给定问题的难度。特别是，简单问题更受益于顺序修正，而对于难题，在顺序和并行计算之间取得平衡是最佳选择。这一发现支持了这样一个假设：顺序修正（即改变提议分布）和并行采样（即使用验证器进行搜索）是扩展测试时计算的两个互补轴，它们可能在每个提示的基础上更有效。我们在附录L中包含了我们模型的生成示例。更多结果见附录B。
• 计算最优修正：鉴于顺序和并行采样的有效性取决于问题难度，我们可以根据难度分箱选择理想的顺序与并行计算比率。在图8中，我们绘制了使用我们的神谕和预测难度概念时，采用这种计算最优扩展策略的结果。在这两种情况下，我们都能通过修正来改进提议分布，从而显著改善测试时计算的扩展性。特别是，我们看到在较高的生成预算下，并行采样的性能似乎趋于平稳，而计算最优扩展则表现出持续的改进。对于神谕和预测难度分箱，我们都看到计算最优扩展可以用少至4倍的测试时计算（例如64个样本对256个）超越best-of-N。总的来说，这些结果展示了通过在每个提示的基础上调整提议分布来改善测试时计算扩展的潜力。

• 通过修正来优化提议分布的计算最优扩展要点：我们发现，顺序（例如修正）和并行（例如标准的best-of-N）测试时计算之间存在权衡，而理想的顺序与并行测试时计算比率严重依赖于计算预算和具体问题。具体来说，较容易的问题受益于纯粹的顺序测试时计算，而较难的问题通常在某个理想的顺序与并行计算比率下表现最佳。此外，通过为给定的问题难度和测试时计算预算最优地选择最佳设置，我们可以用少至4倍的测试时计算超越并行的best-of-N基线。

A4 实验结果

本文通过一系列实验，系统地评估了不同测试时计算扩展策略的有效性，并将其与扩展预训练计算进行了比较。

• 实验一：基于验证器的搜索策略评估

  • 实验内容：比较了三种基于过程奖励模型（PRM）的搜索方法：Best-of-N加权、波束搜索和前瞻搜索。分析了它们在不同计算预算和问题难度下的性能。
  • 实验结果：在低计算预算下，波束搜索优于Best-of-N；但在高预算下，其优势减弱甚至表现更差，尤其是在简单问题上，这可能是由于对PRM的过度利用。Best-of-N在简单问题和高预算下更稳健。前瞻搜索由于计算成本高，整体表现不佳（图3）。
  • 分析结论：搜索策略的有效性与问题难度和计算预算密切相关。基于此，提出的“计算最优”搜索策略（根据问题难度自适应选择最佳搜索方法）能以少4倍的计算量达到甚至超越Best-of-N基线的性能（图4）。

• 实验二：基于提议分布优化的修正策略评估

  • 实验内容：训练了一个能迭代修正自身答案的模型，并比较了串行修正采样与并行独立采样的性能。同时，探索了在不同比例下结合这两种采样方式的效果。
  • 实验结果：纯串行修正（局部优化）在简单问题上表现最好，而对于难题，结合一定比例的并行采样（全局探索）效果更优。存在一个依赖于问题难度的最优串行/并行计算比率（图7）。
  • 分析结论：“计算最优”修正策略（根据问题难度选择最优的串行/并行比率）同样能以少4倍的计算量超越纯并行采样的Best-of-N基线，显示出动态调整提议分布的巨大潜力（图8）。

• 实验三：预训练计算与测试时计算的权衡

  • 实验内容：在一个FLOPs匹配的设定下，比较了两种提升性能的路径：1) 使用一个参数量大14倍的模型进行预训练；2) 使用较小的基础模型并投入等效的额外FLOPs进行测试时计算（使用前述的计算最优策略）。
  • 实验结果：结果显示，两者的“可交换性”取决于问题难度和推理负载（推理token与预训练token的比值$R$）。
    • 在简单和中等难度的问题上，或者当推理负载较低（$R \ll 1$，如自提升流水线）时，投入测试时计算通常比扩展预训练更有效。
    • 在非常困难的问题上，或者当推理负载很高（$R \gg 1$，如大规模生产环境）时，扩展预训练是更有效的途径（图9）。
  • 分析结论：测试时计算和预训练计算并非可以一对一等价交换。对于模型能力范围内的问题，测试时计算可以有效弥补预训练的不足。但对于超出模型能力范围的难题，预训练带来的基础能力提升更为关键。

A5 结论

• 论文结论：本文对扩展测试时计算以提升LLM在数学推理任务中的性能进行了深入分析。研究发现，无论是改进验证器搜索还是优化LLM的提议分布，其策略的有效性都与基础LLM视角下的问题难度高度相关。基于此，本文引入了“计算最优”扩展的概念，即根据提示自适应地分配测试时计算资源，从而将计算效率提升了2到4倍。在一个FLOPs匹配的比较中，本文首次证明，即使是使用相对简单的测试时计算方法（如修正和搜索），在特定类型的提示上也能很好地扩展，其收益超过了将同等FLOPs用于预训练。
• 未来工作展望：
  1. 进一步改进测试时计算扩展：未来工作应探索将多种方法（如修正与PRM树搜索、批判与修正等）结合起来，以进一步提升扩展效果。同时，需要开发新的方法来克服当前策略在难题上收益有限的瓶ভিন。
  2. 快速评估问题难度：本文使用的难度评估方法计算成本较高。未来可以研究更高效的方法，如预训练模型直接预测问题难度，或动态地在评估难度和尝试解决问题之间切换。
  3. 融合测试时计算与训练时计算：本文主要关注测试时计算扩展。未来，一个重要的方向是将测试时计算的输出蒸馏回基础LLM，形成一个在开放式自然语言上运行的迭代式自我改进循环，从而将测试时计算的成果用于提升模型本身。

A6 附录

A. 相关工作

• 语言模型推理：近年来，语言模型在挑战性数学推理任务上的性能迅速提升，这主要归功于四个因素：1) 在大量数学语料上进行持续预训练；2) 改进LLM提议分布，如通过RL微调或让模型迭代修正；3) 通过微调验证器使LLM能从额外测试时计算中受益。本文的工作建立在后两个研究方向之上，分析了通过优化提议分布和验证器搜索来改善测试时计算扩展的程度。
• 分析测试时计算扩展：Jones【索引16，Scaling scaling laws with board games，2021】曾研究过在Hex棋盘游戏中使用蒙特卡洛树搜索时，训练时和测试时计算的权衡。本文则将分析重点放在了大规模语言模型的数学推理问题上。Villalobos和Atkinson【索引44，Trading off compute in training and inference，2023】的一篇综述分析了多个领域中训练和推理的权衡，但其语言模型分析多集中在已知真实答案的场景，而本文关注的是未知真实答案的场景。
• 用测试时计算增强LLM：除了验证器和修正，还有其他方法使语言模型能利用测试时计算进行推理。例如，Wang等人【索引46，Hypothesis search: Inductive reasoning with language models，2024】通过层次化假设搜索实现归纳推理。一些工作通过在测试时为语言模型配备工具来提升性能。还有一些工作提出了无监督学习“思考token”的方法。虽然本文只分析了两种主要机制，但文中的分析方法（如根据问题难度进行计算最优扩展）原则上也可应用于这些其他方法。

B. 额外的修正结果

• 多数投票选择的结果：我们在图10中绘制了使用我们的PaLM 2-S*修正模型进行多数选择的额外结果。使用多数选择时，我们观察到的趋势与图7中验证器选择的趋势大体相似。

C. 无监督难度分箱

• “预测难度”的计算方法：我们通过在每个问题上对2048个样本的PRM最终答案分数进行平均，来计算没有神谕真实正确性信息的难度分箱，从而获得对应于该问题的价值估计。然后，我们将测试集中每个问题的价值分箱到五个五分位数中（使用与神谕难度分箱相同的程序）。我们称之为“预测难度”而非“神谕难度”。技术上，这个过程成本极高，因为它需要生成大量样本。虽然我们在分析中没有考虑这个成本，但在实际生产环境中，这个成本会是个问题。一个更有效的方法是微调一个模型来直接根据问题预测其正确性。我们没有在工作中探索这个方向，但将这种更廉价的难度估计方法的探索留给未来的工作。
• 预测难度分箱的有效性：在图12中，我们绘制了使用我们的难度分箱的PRM搜索结果，在图11中，我们绘制了相应的修正结果。我们看到，在这两种设置中，这些预测分箱都显示出与神谕分箱相似的趋势。

D. PRM训练细节

• 模型与损失函数：我们将PRM微调为一个二元分类器，模型在解决方案的每一步预测一个0到1之间的值。我们使用从蒙特卡洛rollout中获得的软值来训练模型，采用二元交叉熵损失函数（即 $-(y \log(\hat{y}) + (1-y) \log(1-\hat{y}))$，其中$y$对应于软真实值，$\hat{y}$是模型的预测值）。
• 优化器与超参数：我们使用AdamW优化器微调基础模型，学习率为3e-5，批量大小为128，dropout为0.05，Adam betas为(0.9, 0.95)。我们进行早停，选择在随机留出的验证集（占原始PRM800k训练划分中问题的10%）上验证损失最低的检查点。
• 训练数据生成：我们对每个问题从相应的few-shot提示基础模型中采样16个样本来微调PRM。在每一步，我们使用16个蒙特卡洛rollout，使用相同的基础模型和提示，来估计步骤级的值。我们从训练数据中过滤掉所有未能输出有效、可解析最终答案的样本，因为我们发现在初步实验中这些样本会损害PRM性能。
• 数据格式：生成样本时，基础模型被提示以换行符分隔的逐步格式输出答案，如Lightman等人【索引22】所做。然后我们使用简单的换行符分割程序将每个答案分成步骤。我们在附录G中包含了关于我们提示的详细信息。

E. 比较PRM聚合策略

• 聚合策略比较：我们比较了聚合每步PRM分数以产生完整解决方案最终分数的不同方法。具体来说，我们比较了：1) 取所有步骤中的最低分，如Lightman等人【索引22】所做（即“min”）；2) 取所有步骤正确性概率的乘积（即“prod”）；以及3) 只取最后一步的预测（即“last”）。我们在图13中看到，取最后一步的表现优于其他两种方法。
• 结果差异的可能原因：先前的工作【索引22，45】发现“min”是最好的聚合器。我们认为这种差异是由于我们的验证器是用软MC回报标签训练的，这与二元正确性标签的呈现方式非常不同，因此其他聚合策略可能没有相同的效果。
• PRM作为ORM的思考：有趣的是，当使用最后一步聚合时，我们实际上是在像使用ORM一样使用PRM。然而，我们看到PRM的表现优于ORM，这表明在我们的情况下，每步PRM训练可能主要作为一种表示学习的形式有用，而不仅仅是作为推理时的工具。未来的工作应进一步探索这一思路。

F. PRM与ORM的比较

• 实验设置与结果：我们使用PaLM 2-S*基础语言模型训练了一个PRM和一个ORM模型。在图14中，我们看到PRM的表现优于ORM，并且随着使用样本数量的增加，PRM和ORM之间的差距越来越大。我们使用PRM的最后一步预测来对答案进行评分，如附录E所述。

G. 提示细节

• Prompt设计：为了使基础模型能够以PRM可以应用的逐步格式输出答案，我们使用了一个4-shot提示，该提示由从Lightman等人【索引22】发布的PRM800k数据中随机选择的正确答案示例组成。具体来说，我们使用了第一阶段训练划分中的答案。这些答案对应于GPT-4生成的正确答案示例，其中包括正确的逐步格式。在初步实验中，我们发现这种提示程序产生的结果与Lewkowycz等人【索引20，Solving quantitative reasoning problems with language models，2022】使用的提示相似。我们使用这个提示为PRM和修正模型生成训练数据。在测试集上对PRM进行搜索时，我们也使用这个提示。为了给这个提示预测的最终答案评分，我们使用了Lightman等人【索引22】发布的评分函数。

H. 修正模型微调细节

• 数据生成流程：为了微调修正模型，我们遵循了第6.1节中概述的程序。我们首先为每个问题采样64个输出。然后我们过滤掉所有以无效解决方案结束的答案。对于每个正确答案，我们然后从0到4之间均匀采样一个数，表示在上下文中包含多少个不正确的答案用于训练。正确答案被用作轨迹中的最后一个答案（我们训练模型来生成它），而不正确的答案则包含在上下文中。如果采样的数大于0，我们然后根据字符级编辑距离度量找到最接近的不正确答案，作为轨迹中最后一个不正确的答案包含进来。这里的目标是选择一个与正确答案有些相关的不正确答案，以改善学习效果。其余的不正确答案，我们从可用答案集合中随机采样。在采样到的不正确答案少于4个的情况下，我们将均匀分布的最大值截断以匹配不正确样本的数量。我们使用这个程序为训练数据中的所有问题生成轨迹。
• 模型训练：然后，我们对这些生成的轨迹中的正确答案解决方案对基础语言模型进行微调。我们使用AdamW优化器，学习率为1e-5，批量大小为128，dropout为0.0，Adam betas为(0.9, 0.95)。
• 检查点选择：我们发现，通常在如上所述生成的轨迹组成的评估集上评估损失，并不能为早停提供一个好的信号。相反，我们发现评估损失开始增加很久之后的检查点在修正能力上要强得多。这可能是因为在微调修正模型后，评估集代表了离策略数据，这自然会与模型自身在策略下生成的轨迹存在分布差异。因此，我们选择的修正模型检查点略晚于我们在验证集上观察到过拟合的点。

I. 修正模型选择标准

• 选择标准的必要性：如第6.1节所述，为了有效使用我们的修正模型，我们需要部署一个标准来选择一个修正轨迹内部以及多个并行轨迹之间的最佳答案。我们使用两种方法：1) ORM验证器；和2) 多数投票。
• ORM验证器选择方法：对于ORM验证器，我们根据附录J中的程序，在修正模型的输出上训练一个ORM。在推理时，我们然后使用这个验证器来选择最佳答案。由于我们有两个聚合的维度（每个修正轨迹内部和多个轨迹之间），我们部署了一个分层策略，首先在每个修正轨迹内部选择最佳答案，然后将这些选定的答案在轨迹之间进行聚合。为了在每个轨迹内部选择最佳答案，我们执行best-of-N加权聚合，然后选择得分最高的解决方案。然后，为了在所有修正链中选择最终答案，我们使用每个修正链中的最佳答案进行另一轮best-of-N加权选择。这第二轮best-of-N加权后的答案代表了我们的最终答案预测。
• 多数投票选择方法：对于多数投票，我们发现当轨迹长度或轨迹数量太少时，分层聚合会产生问题。问题在于，没有足够的样本，多数投票无法有效地选择最佳选项。因此，对于多数投票，我们简单地将所有轨迹中的所有答案一次性全部拿来，并取它们的多数作为最终答案。我们发现这比分层方法产生了更平滑的扩展行为。

J. 修正模型验证器训练

• 训练新验证器的原因：我们发现，我们在PaLM 2-S基础模型输出上微调的PRM在应用于PaLM 2-S修正模型的输出时效果不佳（见图15(a)），这可能是由于与修正模型的分布偏移。因此，我们训练了一个单独的ORM验证器与我们的PaLM 2-S*修正模型一起使用。我们本可以也训练一个PRM，但由于生成每步PRM标签的成本很高，我们选择了ORM。
• 验证器训练的修改：我们对标准ORM在修正设置中做了一些修改，通过在上下文中包含先前的修正来微调ORM，这样验证器就可以访问与修正模型相同的上下文，从而允许验证器在评分当前答案时看到修正模型先前的答案尝试。所有其他实验细节与训练PRM时使用的相同。
• 实验结果：从经验上看，我们发现在上下文中包含修正历史略微提高了性能（见图15(b)）。此外，即使上下文中没有修正，我们看到顺序修正仍然略微优于并行，表明顺序采样的改进不仅仅是由于验证器的上下文。

K. ReSTEM修正模型实验

• 实验设置：我们通过使用一个简化的RL算法ReSTEM【索引35】来训练模型，尝试进一步优化我们的PaLM 2-S*修正模型。具体来说，我们为MATH训练集上的每个问题生成了64个最大长度为5的修正轨迹。我们在每个轨迹中第一个正确答案处停止修正模型。使用这些生成的数据，我们然后在正确答案数据上微调基础LM。为了帮助模型学习任务，我们明确地平衡了轨迹长度的分布。
• 实验结果与分析：在图16中，我们绘制了这个新修正模型在改变顺序与并行比率时的性能。我们看到，额外的顺序修正显著损害了这个新模型的性能。我们假设这种性能下降是由于从运行ReSTEM获得的在线数据加剧了修正数据中的虚假相关性，导致优化后的模型未能学会修正任务。我们认为，使用更离线的数据收集策略，如Qu等人【索引28】所做，可能更有效，并将进一步的探索留给未来的工作。

L. 修正模型输出示例

• 内容说明：在图17至图23中，我们包含了我们修正模型输出的一些选定示例。

略。

M. PRM波束搜索输出示例

• 内容说明：在图24至图29中，我们包含了PRM波束搜索的一些选定示例。我们在示例中包含了每一步的PRM分数，范围在0到1之间。

略。