INFERENCE SCALING LAWS: AN EMPIRICAL ANALYSIS OF COMPUTE-OPTIMAL INFERENCE FOR LLM PROBLEM-SOLVING

文章标题: 推理缩放定律: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 背景知识

相关工作


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

A2 方法细节

计算最优的问题求解推理

3.1 推理策略


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

3.1.1 蒙特卡洛树搜索(MCTS)

3.1.2 奖励平衡搜索(REBASE)


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

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 实验结果

计算最优的模型大小


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

B MCTS 细节

C 超参数

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

D 额外的实验结果


图 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和加权投票。