论文笔记：Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling

导语

当一个大模型面对给定的算力，是应该更多的用于预训练还是推理？本文尝试回答这个问题。

• 会议：Arxiv 2024
• 链接：[2408.03314] Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters

摘要

本研究探讨了使大语言模型（LLM）在测试阶段利用更多计算能力以改善输出的关键步骤，关注推理时计算的扩展，特别是研究在固定但非平凡的推理计算条件下，LLM在应对具有挑战性的提示时的性能提升幅度。现有研究对各种测试时推理方法的扩展行为知之甚少，且大部分工作给出了多种策略的负面结果。本文分析了扩展测试时计算的两种主要机制：

1. 针对验证奖励模型（Verifier）进行搜索；
2. 在测试时根据提示自适应更新模型的响应分布。

研究发现，不同的测试时计算扩展方法在提示难度变化时效果差异显著。这一观察促使提出一种“计算最优”扩展策略，旨在根据每个提示有效分配测试时计算。采用这种计算最优策略后，测试时计算的扩展效率比best-of-N基线提高了4倍以上。此外，在对比计算量相同的评估中，发现对于一些较小基础模型能够达到较高成功率的问题，测试时计算甚至可以超越一个大14倍的模型。

1 引言

人类在面对困难问题时，往往会花更多时间思考，以可靠地提高决策质量。能否将这种能力赋予今天的大语言模型（LLM）？更具体地说，面对具有挑战性的输入查询，我们能否让语言模型在测试阶段有效利用额外的计算，从而提高其响应的准确性？ 理论上，通过在测试阶段应用额外的计算，LLM应该能够超越其训练时的表现。此外，这种测试时能力还有可能为智能体和推理任务开启新的方向。例如，如果可以在推理期间用额外的计算来平衡预训练模型的大小，那么将使得在较小的设备模型替代数据中心规模的LLM时得到应用。通过使用额外的推理计算来自动生成改进的模型输出，也为实现一种能减少人类监督的通用自我改进算法提供了途径。

先前研究推理时计算的结果存在不一致。一方面，一些研究表明，当前的LLM能够利用测试时计算来改善输出，另一方面，也有研究显示这些方法在更复杂的任务（如数学推理）上的有效性极为有限，尽管推理问题通常需要对现有知识进行推断，而不是获取新知识。这些相互矛盾的发现促使本文对不同的测试时计算扩展方法进行系统分析。

本文关注于了解扩展测试时计算的好处。可以说，扩展测试时计算最简单且研究最为广泛的方法是best-of-N抽样：从基础LLM“并行”抽样N个输出，并根据学习到的验证器或奖励模型选择得分最高的一个。然而，这并不是改善LLM的唯一方式。通过修改获取响应的提议分布（例如，要求基础模型“顺序”修正其原始响应）或通过改变验证器的使用方式（例如，训练一个基于过程的密集验证器并在此基础上进行搜索），测试时计算的扩展能力可以得到显著提升。

为理解扩展测试时计算的好处，本文在具有挑战性的MATH基准上进行实验，使用特别微调的PaLM-2模型，以修正错误答案或使用基于过程的奖励模型（PRM）验证答案中的每个步骤。实验发现，特定测试时计算策略的有效性在很大程度上依赖于具体问题的性质和所使用的基础LLM。例如，在较简单的问题上，基础LLM已经能够生成合理的响应，允许模型通过预测一系列N个修正（即修改提议分布）来迭代优化其初始答案，可能比并行抽样N个独立响应更有效。另一方面，对于更复杂的问题，这些问题可能需要搜索多种高层次的解决方案，此时独立并行重新抽样新响应或对基于过程的奖励模型进行树搜索可能是更有效的测试时计算方式。这一发现强调了部署适应性“计算最优”策略的必要性，即根据提示选择利用测试时计算的具体方法，以最佳利用额外的计算资源。

本文还展示了，从基础LLM的角度来看，问题难度的概念可以用来预测测试时计算的有效性，这使得能够根据提示实际实施这种“计算最优”策略。通过这种方式合理分配测试时计算，本文能够显著改善测试时计算的扩展效率，超越best-of-N基线的表现，同时仅使用约4倍更少的计算资源。

使用改进的测试时计算扩展策略后，本文进一步探讨测试时计算在多大程度上可以有效替代额外的预训练。本文在一个较小模型与额外测试时计算的性能与一个14倍大模型的预训练进行FLOPs匹配比较。结果显示，在简单和中等难度的问题上，甚至在某些困难问题上（取决于预训练和推理工作负载的具体条件），额外的测试时计算通常比扩展预训练更可取。这一发现表明，在某些情况下，预训练较小的模型以较少的计算资源，然后应用测试时计算来改善模型输出，可能比单纯扩展预训练更有效。然而，在最具挑战性的问题上，作者观察到扩展测试时计算几乎没有好处。相反，对于这些问题，使用额外的预训练计算取得进展更为有效，这表明当前的测试时计算扩展方法可能并不能与扩展预训练一一对应。总体而言，即使采用相对简单的方法，扩展测试时计算已经显示出比扩展预训练更为可取的潜力，随着测试时策略的成熟，仍有更多改进可待实现。长期来看，这暗示着未来在预训练阶段花费更少的FLOPs，而在推理阶段花费更多的FLOPs。

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

本文先统一测试时计算的使用方法，然后分析一些代表性的方法。首先，从在给定提示条件下自适应修改模型预测分布的角度来看待额外测试时计算的使用。理想情况下，测试时计算应当修改分布，从而生成比直接从LLM采样更好的输出。通常，有两个方法可以引入对LLM分布的修改：

1. 在输入层面：通过在给定提示中增加一组额外的标记，使LLM在此基础上获得修改后的分布；
2. 在输出层面：通过从标准语言模型中采样多个候选输出并对这些候选进行处理。

换句话说，可以直接修改LLM自身诱导的提议分布，以改善简单条件下的输出，或者使用一些后验验证器或评分器来进行输出修改。这一过程类似于从复杂目标分布进行马尔科夫链蒙特卡洛（MCMC）采样，但通过结合简单的提议分布和评分函数。直接修改提议分布和使用验证器的过程构成了本研究的两个独立维度。

修改提议分布（Modifying the proposal distribution）

改善提议分布的一种方法是通过受强化学习启发的微调方法（如STaR或ReSTEM）直接优化模型以适应特定推理任务。这些技术不利用任何额外的输入标记，而是专门微调模型，以诱导改进的提议分布。另一方面，自我批判（self-critique）等技术使模型能够在测试时通过指示其批判和修正自身输出，以迭代方式改善其提议分布。由于直接提示现成模型在测试时有效性不足，本文特别微调模型，使其在复杂推理环境中能够迭代修正答案。为此，利用了通过Best-of-N引导的优化方法，在策略数据上进行微调，以改善模型响应。

优化验证器（Optimizing the verifier）

在对提议分布和验证器的抽象中，验证器用于汇总或选择提议分布中的最佳答案。最经典的使用验证器的方法是best-of-N抽样，即采样N个完整解决方案，然后根据验证器选择最佳答案。然而，通过训练基于过程的验证器或过程奖励模型（PRM），这一方法可以进一步改进，该验证器能够对每个解决方案中的中间步骤的正确性进行预测，而不仅仅是最终答案。可以利用这些逐步预测在解决方案空间内进行树搜索，从而提供比简单best-of-N方法更高效和有效的搜索方式。

3 如何最佳扩展测试时计算

在统一了各种方法之后，本文希望了解如何最有效地利用测试时计算，以提高语言模型在特定提示上的表现。具体来说，作者想回答以下问题：

在给定一个提示和一个测试时计算预算的情况下，如何有效利用测试时计算来解决问题？根据上述抽象，不同的方法可能在特定问题上效果不同。我们如何确定在给定提示上最有效的测试时计算利用方式？与简单使用更大预训练模型的效果相比，这种方法表现如何？

在提议分布优化或对验证器进行搜索时，有多个不同的超参数可以调整，以确定测试时计算预算的分配方式。例如，当使用经过微调的修正模型作为提议分布，并使用ORM作为验证器时，可以选择将全部测试时计算预算用于从模型中并行生成N个独立样本，然后应用best-of-N方法；或者，可以使用修正模型顺序抽样N个修正，然后使用ORM选择序列中最佳的答案，或者在这两者之间寻求平衡。直观上，可能认为“较简单”的问题会更受益于修正，因为模型的初始样本更有可能接近正确答案，只需进一步优化。另一方面，对于具有挑战性的问题，可能需要探索多种高层次的解决策略，因此在这种情况下，独立并行采样可能更为可取。

在验证器的情况下，也可以选择不同的搜索算法（如束搜索、前瞻搜索、best-of-N），每种算法可能根据所使用的验证器和提议分布的质量表现出不同的特性。对于更困难的问题，更复杂的搜索程序可能比简单的best-of-N或多数基线更有用。

3.1 测试时计算优化扩展策略

一般来说，人们希望为特定问题选择最佳的测试时计算预算分配。为此，对于任何使用测试时计算的方法（例如本文中的修正和对验证器的搜索，其他地方的各种方法），本文将“测试时计算优化扩展策略”定义为为特定提示在测试时实现最大性能收益而选择的超参数策略。形式上，定义Target(θ, 𝑁, 𝑞)为模型在给定提示$q$下产生的自然语言输出标记的分布，使用测试时计算超参数$\theta$和计算预算$N$。本文的目标是选择超参数$\theta$，以最大化给定问题的目标分布的准确性，将其形式化表示为：

$${\theta }_{q,a^{\ast }(q)}$$