文章标题: 推理缩放定律:LLM问题求解中计算最优推理的实证分析
作者/机构:
Yangzhen Wu (清华大学交叉信息研究院)
Zhiqing Sun, Shanda Li, Sean Welleck, Yiming Yang (卡内基梅隆大学计算机科学学院)

A1 主要贡献

本文研究了大型语言模型(LLM)的推理缩放定律(或称测试时缩放定律)计算最优推理,重点探讨了模型大小与不同推理策略下生成额外token之间的权衡。

核心问题与研究目标:
尽管LLM的训练缩放定律已被广泛研究,但其最优推理配置仍未得到充分探索。现有研究表明,增加推理时的计算量(如使用MCTS、多样本投票等方法)可以提升模型性能,但这需要系统地理解各种推理方法在性能和成本之间的权衡。本文旨在回答一个核心问题:在固定的浮点运算(FLOPs)预算下,如何选择最优的模型大小和有效的推理策略以最大化性能(即准确率)?

主要创新点与贡献:
本文通过在数学推理基准(如GSM8K和MATH)上对不同规模的微调模型(Pythia, Mistral, Llemma)进行全面的实证评估,提出了以下主要贡献:
1. 探索了新的推理缩放定律和计算最优推理: 论文评估了固定推理策略下不同模型规模的性能表现。研究发现,在相同的计算预算下,通过增加样本数量,较小的模型可以胜过较大的模型。例如,在推理计算预算较低时,小模型表现更优,而随着预算增加,大模型逐渐显示优势(如图1所示)。这表明最优推理模型的大小随计算预算而变化。
2. 对投票方法的缩放行为进行了新的理论分析: 论文为基于采样和投票的推理策略提供了收敛上界和收敛速度的理论分析。分析表明,即使有无限样本,这些简单策略的性能也会饱和,其上限由模型自身的概率分布决定,从而导致收益递减。这凸显了开发更复杂推理算法的必要性。
3. 提出了新颖的树搜索算法REBASE: 针对现有MCTS方法在加权投票中表现不佳(产生大量未完成解)的问题,论文提出了一种名为REward BAlanced SEarch (REBASE)的新型树搜索算法。REBASE利用节点质量奖励来控制节点扩展,无需显式的“rollout”过程,从而在保证足够候选解的同时,实现了计算效率和性能的帕累托最优权衡。实验证明,在所有设置、模型和任务中,REBASE与较小模型(如Llemma-7B)的组合始终优于其他方法,甚至超过了使用标准推理策略的大模型(如Llemma-34B)。


图 1: Pythia模型在GSM8K测试集上展现的推理缩放定律。我们评估了不同模型大小和加权多数投票中不同采样解决方案数量下的错误率(越低越好)。左图:随着推理计算量的增加,每个模型大小的错误率稳步下降,并最终收敛。右图:最优模型大小(星号表示)随推理时计算预算的变化而变化。例如,在$2^{41}$和$2^{44}$ FLOPs时,较小的模型是计算最优的。两个坐标轴均为对数尺度。

A3 背景知识

相关工作

  • 缩放定律。 近期研究已确立,模型性能与参数数量、训练数据集大小及可用计算资源之间存在可预测的幂律关系【21. Hestness et al., Deep learning scaling is predictable, empirically, arXiv 2017; 39. Rosenfeld et al., A constructive prediction of the generalization error across scales, ICLR 2020】。Kaplan等人【25. Kaplan et al., Scaling laws for neural language models, arXiv 2020】的开创性工作表明,语言模型的测试损失作为模型大小和数据量的函数,以一种高度规律的方式衰减。后续研究对这些初步观察进行了完善,并将其扩展到更多样化的情境中【22. Hoffmann et al., Training compute-optimal large language models, arXiv 2022; 2. Alabdulmohsin et al., Revisiting neural scaling laws in language and vision, NeurIPS 2022; 33. Muennighoff et al., Scaling data-constrained language models, NeurIPS 2023; 29. Lin et al., Selecting large language model to fine-tune via rectified scaling law, ICML 2024; 16. Goyal et al., Scaling laws for data filtering–data curation cannot be compute agnostic, CVPR 2024】。然而,这些现有工作大多主要关注训练阶段。Sardana等人【40. Sardana et al., Beyond chinchilla-optimal: Accounting for inference in language model scaling laws, ICML 2024】研究了同时考虑训练和推理的缩放定律,但仅考虑了固定的推理算法。相比之下,本文系统地展示了推理缩放定律,即LLM的问题解决性能会随着推理时计算预算的增加而提高,并提出和研究了计算最优推理。

  • LLM问题解决中的推理策略和推理时计算利用。 研究人员已经开发了多种推理策略来用训练好的模型生成序列【51. Welleck et al., From decoding to meta-generation: Inference-time algorithms for large language models, TMLR 2024】。确定性方法如贪心解码和束搜索(beam search)【45. Teller, Speech and language processing, 2000; 17. Graves, Sequence transduction with recurrent neural networks, arXiv 2012】能找到高概率序列,这些序列通常质量不错但缺乏多样性。采样算法(如温度采样【1. Ackley et al., A learning algorithm for boltzmann machines, Cognitive science 1985】)可以产生多样化的结果,然后通过聚合(如自洽性方法【49. Wang et al., Self-consistency improves chain of thought reasoning in language models, ICLR 2023】)来提高准确率。近期的方法将搜索算法与LLM结合,包括广度优先或深度优先搜索【53. Yao et al., Tree of thoughts: Deliberate problem solving with large language models, NeurIPS 2023】、蒙特卡洛树搜索(MCTS)【56. Zhang et al., Planning with large language models for code generation, ICLR 2023; 57. Zhou et al., Language agent tree search unifies reasoning, acting, and planning in language models, ICML 2024; 31. Liu et al., Don’t throw away your value model!, CoLM 2024; 10. Choi et al., KCTS: Knowledge-constrained tree search decoding with token-level hallucination detection, EMNLP 2023】和引导束搜索【52. Xie et al., Selfevaluation guided beam search for reasoning, NeurIPS 2023】。一些先前的研究还发现,在推理时输出“虚拟”token可以提高LLM的问题解决性能【15. Goyal et al., Think before you speak: Training language models with pause tokens, ICLR 2024; 38. Pfau et al., Let’s think dot by dot: Hidden computation in transformer language models, CoLM 2024】。所有这些方法都表明,在推理时使用搜索可以带来性能提升,但代价是增加推理时计算量,它们并未系统地描述成本-性能的权衡。本文是首个提出并研究计算最优推理的,分析了计算预算和问题解决性能之间的权衡,并提出了经验上帕累托最优的REBASE方法。同期,Snell等人【43. Snell et al., Scaling test-time compute optimally can be more effective than scaling LLM parameters, ICLR 2025】也研究了如何最优地扩展推理计算,并提供了互补的见解,因为我们考虑了更多的模型家族和尺寸,而他们研究了几种不同的推理策略。

  • LLM的数学推理。 近年来,大型语言模型取得了显著进展,并展现出强大的推理能力【7. Brown et al., Language models are few-shot learners, NeurIPS 2020; 22. Hoffmann et al., Training compute-optimal large language models, arXiv 2022; 26. Lewkowycz et al., Solving quantitative reasoning problems with language models, NeurIPS 2022; 11. Chowdhery et al., Palm: Scaling language modeling with pathways, JMLR 2023】。数学问题解决是衡量LLM推理能力的关键任务【12. Cobbe et al., Training verifiers to solve math word problems, arXiv 2021; 19. Hendrycks et al., Measuring mathematical problem solving with the MATH dataset, NeurIPS Datasets and Benchmarks Track 2021】。Ling等人【30. Ling et al., Program induction by rationale generation: Learning to solve and explain algebraic word problems, ACL 2017】首次开发了生成逐步解题过程直至最终答案的方法。随后,Cobbe等人【12. Cobbe et al., Training verifiers to solve math word problems, arXiv 2021】通过训练一个验证器来评估和排序解决方案,扩展了这项工作。后续研究显示了推理时技术(如多数投票和加权多数投票)的性能优势【26. Lewkowycz et al., Solving quantitative reasoning problems with language models, NeurIPS 2022; 27. Li et al., Making language models better reasoners with step-aware verifier, ACL 2023】。我们选择数学问题解决作为研究计算最优策略的任务,因为它能准确评估问题解决能力。


图 2: 训练和推理中计算最优缩放定律的图示。Chinchilla缩放定律【22. Hoffmann et al., Training compute-optimal large language models, arXiv 2022】展示了在给定的训练计算预算下如何选择模型大小和训练token数量,而我们的工作则展示了在给定的推理计算预算下如何选择模型大小和推理策略。

A2 方法细节

计算最优的问题求解推理

  • 问题定义。 本文旨在探索一个核心问题:在固定的FLOPs预算下,应如何选择一个最优的策略模型大小和一种有效的推理策略,以最大化性能(即准确率)?这是首次对此问题进行明确的阐述和研究,与以往关注训练阶段的缩放定律研究有所区别(如图2所示)。

  • 数学形式化。 为了解决这个问题,我们将问题解决的错误率$E(N, T ; S)$表示为模型参数数量$N$、生成token数量$T$和推理策略$S$的函数。计算预算$C$是基于$N$和$T$的确定性函数$FLOPs(N, T ; S)$。我们的目标是在测试时计算约束$FLOPs(N, T, S) = C$下最小化$E$:


    其中$N_{opt}(C)$和$T_{opt}(C)$表示计算预算$C$的最优分配。这里的推理计算量(FLOPs)可以通过策略模型和推理策略生成更多token来调节,例如,采样额外的候选解并随后使用奖励模型对它们进行排序。作为推理策略,我们主要考虑采样和树搜索方法,并与重排序或多数投票相结合,具体包括贪心搜索、多数投票、best-of-n、加权投票及其树搜索变体。

3.1 推理策略

  • 策略分类。 本文考虑了实践中常用的以下几种推理策略:

    • 贪心搜索(Greedy search): 该策略每次生成一个token,选择每一步中概率最高的token。它计算效率高,但在多样性方面通常是次优的。
    • best-of-n: 该策略也称为拒绝采样,它生成一组候选解,并选择由奖励模型给出最高分的那个。
    • 多数投票(Majority voting): 在该策略中,生成一组候选解,问题的最终答案由所有输出中出现最频繁的答案决定。
    • 加权多数投票(Weighted majority voting): 该策略是多数投票的一个变体,其中候选解根据奖励模型给出的分数进行加权。
      我们将使用标准自回归采样算法(如温度采样)生成候选集的策略称为基于采样的策略(贪心搜索是独立的,因为它只有一个确定性候选解)。树搜索变体则使用树搜索来生成候选集。在讨论树搜索方法之前,我们先对基于采样的投票进行分析。
  • 基于采样的投票方法的理论分析。 本文在定理1和定理2中提出了关于投票策略在无限计算资源下渐进行为的理论结果。非正式地讲,研究表明标准/加权多数投票的准确率会随着样本无限增多而收敛,其极限仅取决于语言模型(和奖励模型)所建模的分布。这一理论发现与我们在4.2节中展示的经验发现一致,即在高采样预算下性能会饱和。证明过程见附录A。

  • 符号和假设定义。 设$V$为一个有限词汇表,$V^*$为其克莱尼闭包,即所有字符串的集合。给定一个问题$x$,如果模型输出$rey$,其中$r \in V^*$可以是任何“推理路径”,$e \in V$表示推理结束的特殊标记,我们就说语言模型对该问题回答为$y$。我们进一步假设答案字符串总是短于$L$个token,即$|y| \le L$,其中$L \in \mathbb{N}^*$是一个固定的常数,$|y|$表示$y$的长度。对于语言模型$\pi$,用$\pi(v|w)$表示给定输入(提示)$w$生成$v$的概率。对于奖励模型$\rho$,用$\rho(v)$表示它赋给字符串$v$的分数。我们用$I$表示指示函数。

  • 定理1(标准多数投票收敛性)。 考虑一个数据集$D = \{(x_i, y_i)\}_{i=1}^m$,其中$x_i$和$y_i$分别表示输入和真实答案。对于语言模型$\pi$,用$\text{acc}_{MV_n}(D; \pi)$表示使用$n$个样本的多数投票在$D$上的准确率。根据上述符号和假设,我们有:


    其中常数$c > 1$。

  • 定理2(加权多数投票收敛性)。 考虑一个数据集$D = \{(x_i, y_i)\}_{i=1}^m$。对于语言模型$\pi$和奖励模型$\rho$,用$\text{acc}_{WV_n}(D; \pi, \rho)$表示使用$n$个样本的加权多数投票在$D$上的准确率。根据上述符号和假设,我们有:


    其中常数$c > 1$。

  • 理论分析的启示。 定理1和2指出了随着样本数量增加,准确率会收敛,这表明对于任何固定模型,使用更多样本带来的性能增益将会饱和。极限由通过所有可能的推理路径生成正确答案的似然性决定(对于加权多数投票,该似然性应视为加权和)。这促使我们考虑能够搜索“好”的推理路径的推理算法,例如3.1.1和3.1.2节中详述的基于树搜索的变体。定理1和2还为比较标准多数投票与加权多数投票提供了见解。非正式地讲,只要奖励模型“优于随机”,即平均而言为正确的解决方案分配更高的奖励,加权多数投票的准确率极限就高于多数投票。在我们的实验中,我们一致发现加权多数投票优于多数投票。因此,我们在正文中主要关注best-of-n和加权多数投票,并将多数投票的结果放在附录D。


图 3: 奖励平衡搜索(REBASE)一次迭代的图示。

3.1.1 蒙特卡洛树搜索(MCTS)

  • MCTS在LLM中的应用与局限。 蒙特卡洛树搜索(MCTS)在需要战略决策的领域,如棋盘游戏,已被证明是有效的【41. Silver et al., Mastering the game of Go with deep neural networks and tree search, Nature 2016; 42. Silver et al., Mastering the game of go without human knowledge, nature 2017; 24. Jones, Scaling scaling laws with board games, arXiv 2021】。最近的研究表明,将MCTS应用于LLM的上下文可以增强文本生成过程【56. Zhang et al., Planning with large language models for code generation, ICLR 2023; 57. Zhou et al., Language agent tree search unifies reasoning, acting, and planning in language models, ICML 2024; 31. Liu et al., Don’t throw away your value model!, CoLM 2024; 10. Choi et al., KCTS: Knowledge-constrained tree search decoding with token-level hallucination detection, EMNLP 2023; 8. Chen et al., Alphamath almost zero: Process supervision without process, NeurIPS 2024; 46. Tian et al., Toward self-improvement of LLMs via imagination, searching, and criticizing, NeurIPS 2024; 8. Chen et al., Alphamath almost zero: Process supervision without process, NeurIPS 2024】。在这种背景下,MCTS与一个价值模型配对,用于评分和引导探索步骤。关于MCTS的更多背景知识,我们在附录B中进行了回顾。

  • MCTS的成本-性能权衡问题。 MCTS及其变体的近期研究主要关注于在所研究任务上提高性能(例如,准确率)。然而,在计算预算(以生成的token数或处理时间衡量)方面,MCTS与传统方法(如best-of-n和多数投票)的通用比较很少,或表明其可能存在不利的成本-性能权衡。例如,MCTS消耗的资源要多得多,通常需要比简单方法多几十倍的生成token。具体来说,搜索树中的一大部分路径被用于估计和选择节点,而这些路径不一定成为最终候选解的一部分,尽管MCTS确保了采样到的解决方案包含高质量的中间步骤。相比之下,采样方法并行且独立地生成多个解决方案,所有生成的序列都包含在候选解中。然而,这些序列中的中间步骤不能保证是高质量的,因为没有机制来修剪差的步骤或利用有前景的步骤。

  • 新方法的需求。 这凸显了需要一种新的树搜索方法,它能达到与MCTS相当(或更好)的性能,并且计算成本更低,其成本类似于加权多数投票和best-of-n。这激发了我们的新方法,即奖励平衡搜索(REward BAlanced SEarch, REBASE)。

3.1.2 奖励平衡搜索(REBASE)

  • REBASE的核心思想。 REBASE树搜索方法如图3所示,它继承了树搜索的利用(exploitation)和剪枝(pruning)特性,同时仅使用奖励模型来估计中间节点的质量。与MCTS等方法相比,这节省了推理计算量,因为它不涉及使用显式的“rollout”来估计节点质量。简而言之,其基本思想是使用一个过程奖励模型来决定在每个深度上每个节点应扩展多少。即,REBASE根据节点的softmax归一化奖励分数来扩展给定深度的节点,同时受限于总扩展预算。下面我们更详细地描述这个过程。

  • 符号定义。 我们将微调后的LLM视为一个策略$\pi_\theta$,它逐步生成解决方案。给定一个问题$x$和解决方案的前$k$个步骤$r_1 \cdots r_k$,第$(k+1)$个步骤从$\pi_\theta(\cdot|xr_1 \cdots r_k)$中采样。REBASE在推理过程中生成一个解决方案树,其中根节点是问题$x$,其他节点对应于解决方案的步骤。在生成解决方案树时,我们通过从$\pi_\theta(\cdot|xr_1 \cdots r_k)$中采样来生成$r_k$的子节点。我们用相应的解决方案步骤来表示一个节点。节点$r_k$的奖励由过程奖励模型(PRM)生成:$R(r_k) := R(xr_1 \cdots r_k)$。


图 4: MATH数据集上不同推理策略和模型大小的推理缩放情况(越低越好)。详细的MCTS配置可在附录B中找到。左/右图显示了基于加权多数投票/best-of-n的MATH错误率。在所有预算下,REBASE都是计算最优策略,其中7B模型通常是最佳模型大小。

  • 初始化。 给定问题$x$、平衡温度$T_b > 0$和目标生成的解决方案数量$N$,我们为问题采样$N$个第一步的实例,从而得到搜索树中深度为1的所有节点。我们在初始化时将深度0的采样预算$B_0$设为$N$。

  • 奖励分配与更新。 在第$i$次迭代中,PRM为深度$i$的所有节点分配奖励。之后,算法检查截至深度$i$的解决方案是否已完成。假设有$C_i$个已完成的解决方案,我们使用$B_i \leftarrow B_{i-1} - C_i$来更新采样预算。如果$B_i=0$,过程结束,我们得到$N$个解决方案。

  • 探索平衡与扩展。 对于树中深度$i$的所有节点$n_j$及其奖励$R(n_j)$,我们计算$n_j$的扩展宽度为:


    然后,我们为深度$i$中的所有节点$n_j$采样$W_j$个子节点,并开始下一次迭代。

A4 实验环境

数据集:
* MATH 【19. Hendrycks et al., Measuring mathematical problem solving with the MATH dataset, NeurIPS Datasets and Benchmarks Track 2021】:包含高中数学竞赛级别的问题。实验遵循【28. Lightman et al., Let’s verify step by step, ICLR 2024; 47. Wang et al., Math-shepherd: Verify and reinforce LLMs step-by-step without human annotations, ACL 2024; 44. Sun et al., Easy-to-hard generalization: Scalable alignment beyond human supervision, NeurIPS 2024】的做法,使用MATH500子集作为测试集。
* GSM8K 【12. Cobbe et al., Training verifiers to solve math word problems, arXiv 2021】:包含小学水平的数学推理问题。

模型架构:
* 策略模型(解决方案生成器):
* Pythia 【5. Biderman et al., Pythia: A suite for analyzing large language models across training and scaling, ICML 2023】:用于研究模型大小对性能缩放的影响,因其提供了多种尺寸的模型。
* Llemma 【4. Azerbayev et al., Llemma: An open language model for mathematics, ICLR 2024】:数学专用模型,用于研究不同推理策略下的缩放规律。
* Mistral-7B 【23. Jiang et al., Mistral 7b, arXiv 2023】:用于扩展研究发现到不同的模型和架构。
* 奖励模型:
* 所有实验均使用同一个Llemma-34B奖励模型。该模型在合成的过程奖励建模数据集Math-Shepherd 【47. Wang et al., Math-shepherd: Verify and reinforce LLMs step-by-step without human annotations, ACL 2024】上进行微调,并增加了一个奖励头,使其能在每一步结束时输出一个标量奖励值。

软件配置:
* 微调: 所有策略模型均在MetaMath数据集【55. Yu et al., Metamath: Bootstrap your own mathematical questions for large language models, ICLR 2024】上使用全参数监督微调(Full-SFT)进行训练。详细的微调超参数见附录。
* 推理配置: 使用采样和树搜索方法生成多个候选解,并通过best-of-n、多数投票或加权投票选择答案。每个配置都运行多次以计算均值和方差,以减少随机性影响。除非另有说明,图中每个数据点对应$2^i$个样本,其中$i$是从0开始的整数。

硬件配置:
* 论文未明确提供具体的硬件信息(如GPU型号、数量等)。

A4 实验结果

计算最优的模型大小

  • 模型大小的计算最优推理缩放定律。 图1展示了不同模型大小下推理计算量与错误率的关系。错误率先是稳步下降,然后开始饱和。在计算预算较小时,从小模型中多次采样是计算最优的。当计算预算较大时,大模型则更受青睐,因为小模型的性能已经饱和。如图1右侧面板所示,最优模型大小随推理预算而变化。通过对推理FLOPs $C$和模型大小$N$进行回归分析,得到了一个关系式:$\log_{10}(C) = 1.19 \log_{10}(N) + 2.03$,该式可用于估算特定计算预算下的最优推理模型大小。

  • Llemma-7B以更少计算量达到与Llemma-34B相当的准确率。 图4和图5展示了Llemma 7B和Llemma 34B在使用不同推理策略时的错误率与推理FLOPs的关系。Llemma-7B达到与Llemma-34B相当的准确率所需的总FLOPs大约少2倍。这一发现在不同的推理策略(采样策略、MCTS、REBASE)和任务(MATH、GSM8K)中都成立。这一结果表明,在相同的训练数据集和模型家族下,使用较小模型并配合合适的推理策略生成更多token,可能比使用较大模型具有更有利的成本-性能权衡。


图 5: GSM8k上不同推理策略和模型大小的推理缩放情况(越低越好)。左/右图显示了基于加权多数投票/best-of-n的GSM8K问题解决错误率。由于MCTS的计算-准确率权衡较差,未包含在比较中。REBASE是计算最优的推理策略,最优模型大小有所不同。

计算最优的推理策略

  • REBASE是帕累托最优的。 在固定模型和评估任务的情况下,REBASE在所有设置中始终实现了最佳的成本-性能权衡,优于基于采样的方法(见图4、5、6和7)。例如,在图4中,REBASE在所有推理计算预算下都是计算最优策略,其中7B模型通常是最佳模型大小。另一方面,MCTS在每个计算预算下都表现不如基于采样的方法,这可能是由于其昂贵的rollout过程(图4),相比之下REBASE能高效地利用奖励模型。

  • REBASE的计算效率优势。 表1显示,与基于采样的加权投票相比,REBASE能用更低的计算预算达到更高的准确率。对于7B模型,REBASE以7倍少的计算量实现了更高的准确率。这是一个新颖的发现,与以往的树搜索方法不同,后者通常以比基于采样的投票更高的计算开销为代价来提高性能【8. Chen et al., Alphamath almost zero: Process supervision without process, NeurIPS 2024; 52. Xie et al., Selfevaluation guided beam search for reasoning, NeurIPS 2023】。

  • REBASE在难题上增益更大。 MATH数据集为每个问题分配了1到5的难度等级。我们将MATH测试集分为MATH-easy(1-2级)和MATH-hard(3-5级),并在这两个子集上比较REBASE与采样的性能。结果如图8所示,在简单问题上,两者的性能相当。然而,在较难的问题上,REBASE显示出显著优势。这表明像树搜索这样的高级推理策略对于解决难题尤其有效。

  • REBASE比采样饱和得更晚且准确率更高。 从图6和图7中,我们观察到采样和REBASE在GSM8K上都较早饱和,而在MATH上则相对较晚。这归因于GSM8K和MATH的难度差异。在难题上,模型可能需要聚合更多的解路径才能收敛。在MATH上(图6),REBASE最终以比采样更高的准确率饱和。我们推测,从REBASE中抽样相当于从一个为真实答案分配高概率的策略中抽样,这提高了理论上的准确率上限。


图 6: MATH上不同推理策略和模型的推理缩放情况(越低越好)。测试模型为Llemma-7B(左)、Llemma-34B(中)和Mistral-7B(右)。图例中,W.M.和BoN分别指加权多数投票和best-of-n。


图 7: GSM8K上不同推理策略和模型的推理缩放情况(越低越好)。测试模型为Llemma-7B(左)、Llemma-34B(中)和Mistral-7B(右)。图例中,W.M.和BoN分别指加权多数投票和best-of-n。


图 8: 在MATH-easy(1-2级)和MATH-hard(3-5级)问题上,使用加权多数投票的采样和REBASE的比较。测试模型为Llemma-7B和Llemma-34B。

Table 1: REBASE以较低的计算预算实现了比采样以较高计算预算更好的准确率。我们对采样和REBASE都使用加权投票来聚合候选解。

A5 结论

本文研究了任务性能与推理时计算消耗量之间的关系,涵盖了不同模型大小、模型家族和推理策略,从而形成了经验性的推理缩放定律。这些关系使我们能够探讨计算最优推理,即在给定计算预算下获得最佳性能的推理配置。

研究结果主要有三个 takeaways:
1. 小模型+多推理计算 > 大模型: 使用较小的模型并通过推理策略生成更多token,通常在固定的计算预算下优于使用较大的模型。这对于现实世界中部署的模型具有重要意义,因为推理计算受到各种限制。具体来说,部署较小的模型并采用更复杂的推理策略可能在成本-性能权衡上更有利。
2. 采样方法的局限性: 基于采样的多数投票策略在无限计算(通过抽取更多样本分配)的极限下,会不可避免地饱和到一个取决于底层生成策略的分布。因此,通过设计替代的推理策略来改变采样分布是值得研究的。
3. REBASE的优越性: 本文设计了一种新颖的树搜索推理策略——REBASE,并发现它是帕累托最优的,即在所有测试的计算预算下都取得了最佳性能。值得注意的是,它优于广受欢迎且被广泛使用的加权多数投票和MCTS方法。这一发现不仅展示了REBASE的强大,也表明通过推理时算法来提升语言模型性能仍有很大的改进空间。

A6 附录

A 省略的证明

  • 定理1的证明。

    • 思路概述。 证明的核心思路是利用概率论中的Hoeffding不等式和Borel-Cantelli引理。首先,将模型对一个问题$x$输出特定答案$y$的概率$\tilde{\pi}(y|x)$定义为所有能推导出该答案的“推理路径”的概率之和。
    • 关键步骤。 随后,定义$y^*$为模型认为最可能的答案。证明多数投票在$n$次采样中未能输出$y^*$的事件$E_n$发生的概率。$E_n$发生意味着存在某个其他答案$y''$的出现次数不小于$y^*$。通过联合界(union bound),可以将$P(E_n)$的上界表示为所有其他答案$y'$的出现次数不小于$y^*$的概率之和。
    • 应用不等式。 $y^*$和$y'$的出现次数之差是一个独立同分布随机变量的和,其期望为$\delta = \tilde{\pi}(y^*|x) - \tilde{\pi}(y'|x)$。应用Hoeffding不等式可以得到该差值小于等于0的概率呈指数级下降。
    • 收敛性证明。 最终,通过Borel-Cantelli引理证明,当$n \to \infty$时,$E_n$发生的概率为0,这意味着多数投票的答案几乎必然收敛到$y^*$。将此结论推广到整个数据集,即可完成定理的证明。
  • 定理2的证明。

    • 思路概述。 定理2的证明过程与定理1非常相似。主要区别在于需要重新定义每个答案的“概率”或权重。
    • 关键步骤。 对于加权多数投票,我们将每个答案$y$的“加权概率”定义为所有能推导出该答案的推理路径$r$的概率$\pi(rey|x)$与该路径奖励$\rho(rey)$的乘积之和。
    • 应用定理1的证明框架。 在这个新的定义下,定理1的证明技术可以被直接应用,从而完成定理2的证明。

B MCTS 细节

  • MCTS流程概述。 MCTS算法的过程可以概括为以下几个步骤:
    1. 选择(Selection): 从根节点开始,算法递归地选择具有最高“树的上置信界”(UCT)值的子节点,直到达到一个未被扩展的节点。UCT的计算公式为:

      其中,$Q(s)$是节点$s$的质量得分,$N(s)$是节点$s$的访问次数,$\text{Parent}(s)$是$s$的父节点,$C$是决定探索程度的常数。
    2. 扩展与评估(Expansion and evaluation): 到达一个非终止节点$s$后,通过生成多个子节点来扩展该节点。每个子节点$c$随后由一个价值函数$V(c)$进行评估,该函数预测从节点$c$继续序列的潜在质量。
    3. 反向传播(Backpropagation): 评估后,算法更新从所选节点回到根节点的路径上所有节点的UCT值和访问次数。对于该路径上的任何节点$n$,更新方式如下:

C 超参数

  • 微调。 模型微调的所有超参数见表2。在微调前,对MetaMath数据集进行了预处理,使解决方案呈现分步格式。

Table 2: 微调超参数:LR指学习率,BS指批量大小。Pythia、Llemma-7B和LLemma-34B是我们在实验中使用的生成器,RM是奖励模型的缩写。我们仅使用GSM8K中的问题来训练Pythia模型。

  • 推理。 对于所有推理策略,LLM token生成的温度设置为1.0。输出的最大token数为1024,单步最大token数为256。对于REBASE,平衡温度(即公式(1)中的参数$T_b$)设置为0.1。对于MCTS,UCT值中的$C$设置为1,我们分别为根节点扩展4、8、16个子节点,为其他选定节点扩展2个子节点,总扩展次数分别为32、64、128次。新扩展的节点将由PRM分配值,然后通过上一节描述的过程反向传播Q值。

D 额外的实验结果

  • 多数投票实验结果。 本节补充了多数投票方法的实验结果,并与加权多数投票进行了比较(图9, 10, 11, 12)。实验表明,尽管在采样方法中,多数投票和加权多数投票之间的差距巨大,但如果应用REBASE,这个差距会变得小得多。这种现象可能是由REBASE等树搜索的选择能力造成的。一旦REBASE已经采样了高奖励的解决方案,进行加权多数投票的收益就会减少,因为与采样相比,这些解决方案可能都具有相对较高且稳定的奖励。

  • 在Llama3模型上的额外实验。 我们使用Llama3-8B-Instruct【13. Dubey et al., The llama 3 herd of models, arXiv 2024】模型在MATH和GSM8K数据集上进行了额外实验,如图13所示。代码生成任务MBPP【3. Austin et al., Program synthesis with large language models, 2021】的结果呈现在表3中。这些实验表明,我们的结论可以推广到Llama3架构和编码任务,证实了增加计算投入会提高性能直至达到饱和,并且REBASE达到了最优的性能-计算权衡。

    • 在数学推理任务中,REBASE在不同的答案选择策略(包括best-of-n、多数投票和加权投票)下始终优于采样方法。在每个数据集上的最高性能都是使用REBASE实现的。
    • 对于代码生成任务MBPP,我们通过pass rate评估来分析缩放行为和计算最优推理。结果证实REBASE比采样更具计算效率。这归因于使用了一个评估部分代码解决方案的奖励模型。通过执行一次REBASE迭代,我们的方法剪除了次优的部分解决方案,同时鼓励探索有前景的方案,从而提高了计算效率和解决方案质量。
  • 不同模型下不同策略的比较。 表4展示了在特定计算预算下不同策略的准确率。


图 9: 不同模型在MATH测试集上问题解决错误率的推理缩放定律。测试模型为Llemma-7B(左)、Llemma-34B(中)和Mistral-7B(右)。图例中,M.V.指多数投票。


图 10: 不同模型在GSM8K测试集上问题解决错误率的推理缩放定律。测试模型为Llemma-7B(左)、Llemma-34B(中)和Mistral-7B(右)。图例中,M.V.指多数投票。


图 11: 不同模型在MATH测试集上问题解决错误率的推理缩放定律。测试模型为Llemma-7B(左)、Llemma-34B(中)和Mistral-7B(右)。图例中,M.V.和W.M.分别指多数投票和加权多数投票。


图 12: 不同模型在GSM8K测试集上问题解决错误率的推理缩放定律。测试模型为Llemma-7B(左)、Llemma-34B(中)和Mistral-7B(右)。图例中,M.V.和W.M.分别指多数投票和加权多数投票。

Table 3: 采样和REBASE在MBPP代码生成任务上的零样本pass rate。


图 13: GSM8K(左)和MATH(右)上不同推理策略和模型的推理缩放情况(越低越好)。测试模型为Llama3-instruct-8B。图例中,M.V.、W.M.和BoN分别指多数投票、加权多数投票和best-of-n。

Table 4: 特定计算预算下不同推理配置的准确率。MV、BoN和WV分别表示多数投票、best-of-n和加权投票。