NeurIPS 2024 ｜ REBASE，比MCTS更高效的Inference Scaling算法

还在为大模型推理速度慢、成本高而烦恼吗？这篇来自 CMU 和清华大学的论文给你带来了一个颠覆性的解决方案！ 他们发现，在推理大型语言模型时，“大力出奇迹”不再是真理！ 通过深入研究推理过程中的计算与性能关系，他们提出了“推理缩放定律”，并革命性地推出了一种名为 REBASE 的全新算法。

REBASE 就像一位聪明的向导，能够巧妙地利用奖励信号，指引模型在推理的迷宫中高效探索，避免了传统方法中耗时费力的盲目搜索。 实验结果令人惊叹：在数学推理任务上，REBASE 搭配小模型，不仅速度更快，而且精度更高，轻松击败了更大的模型！

这篇论文的意义不仅仅在于提出了一种新算法，更在于它颠覆了我们对大模型推理的传统认知，为我们指明了一条通往更高效、更经济的 AI 之路。 想要了解如何让你的大模型推理效率飙升？快来阅读这篇论文，一起探索 LLM 推理的奥秘吧！

一、概述

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

• URL: https://arxiv.org/abs/2408.00724

• Authors: Yangzhen Wu, Zhiqing Sun, Shanda Li, Sean Welleck, Yiming Yang

• Code: https://thu-wyz.github.io/inference-scaling/

1.Motivation

• 目前关于大型语言模型 (LLM) 的研究主要集中在训练阶段的缩放定律上，而对于推理阶段的最佳配置研究较少。

• 现有的推理技术通常需要在推理时增加额外的计算，以最大化性能，但这些技术的计算成本必须考虑在内，以实现计算最优推理。

• 传统的采样算法会有性能上限（论文Figure 1），需要更好的推理算法来取得更好的结果。

2.Methods

文章评估了不同模型大小和推理算法在数学推理基准测试上的性能，设计了一个新的树搜索算法REBASE，该算法在计算效率上优于常用的采样和MCTS方法。

2.1 省流版总结:

核心思想： 研究在给定计算预算下，如何选择最优的语言模型大小和有效的推理策略，以最大化问题解决的准确性。

创新点： 提出了一个新的树搜索算法REBASE（Reward Balanced Search），该算法与加权投票结合使用，能够在准确性和推理计算量之间实现帕累托最优的权衡。

帕累托最优的权衡意味着：

• REBASE 算法在给定的计算预算下，能够获得最佳的性能（例如，最高的数学问题解决准确率）。

• 如果你想进一步提高准确率，你就必须增加计算预算（例如，生成更多的 tokens，进行更多的搜索）。

• 如果你想减少计算预算，你就必须接受较低的准确率。

2.2 REBASE详细方法和步骤:

目标： 在给定固定计算预算的情况下，如何选择最优的模型大小和推理策略，以最大化模型在特定任务上的性能。

2.2.1 REBASE 算法是干什么的的？

为了解决 MCTS 的计算成本高问题，论文提出了新的树搜索算法 REBASE。REBASE 通过奖励模型来控制节点的扩展，从而避免了 MCTS 中显式的 rollout 步骤。

REBASE 算法避免 MCTS 中显式 rollout 步骤的关键在于它巧妙地利用了奖励模型 (Reward Model) 和 平衡扩展 (Balanced Expansion) 策略。

2.2.2 MCTS 的 Rollout 步骤有什么问题？

在传统的 MCTS 中，为了评估一个节点的价值，需要进行 rollout（模拟）。Rollout 是指从当前节点开始，根据某种策略（例如随机策略或基于模型的策略）模拟生成多个完整的解决方案，然后根据这些解决方案的结果来估计该节点的价值。

Rollout 的问题在于：

• 计算成本高: Rollout 需要生成多个完整的解决方案，这在计算上非常昂贵，特别是当解决方案很长时。
• 估计偏差: Rollout 使用的策略可能与实际的策略不同，这会导致价值估计存在偏差。

2.2.3 REBASE 如何利用奖励模型控制节点扩展？

REBASE 摒弃了显式的 rollout 步骤，而是直接使用过程奖励模型 (PRM) 来评估节点的质量。具体来说，REBASE 在每次迭代中执行以下步骤：

• 奖励分配 (Reward Assignment):

对于当前深度 i 的每个节点 nj，REBASE 使用 PRM 来计算该节点的奖励 R(nj)。这个奖励反映了从根节点到节点 nj 的部分解决方案的质量。注意，这里直接利用了 PRM 的输出，而不需要进行 rollout。

• 平衡扩展 (Balanced Expansion):

Wj = Round ( Bi * exp(R(nj) / Tb) / Σk exp(R(nk) / Tb) )

其中：

直观理解： 这个公式的含义是，每个节点的扩展宽度与该节点的奖励的指数成正比。奖励越高的节点，其指数值越大，因此获得的扩展宽度也越大。通过这种方式，REBASE 将更多的计算资源分配给更有希望的节点。

Bi 是深度 i 的可用采样预算（即还可以生成的总节点数）。

R(nj) 是节点 nj 的奖励。

Tb 是一个平衡温度参数，用于控制扩展的“平滑”程度。较高的 Tb 会导致更均匀的扩展，而较低的 Tb 会导致更集中于高奖励节点的扩展。

Round() 函数将结果四舍五入到最接近的整数。

分母 Σk exp(R(nk) / Tb) 对所有节点 nk 的奖励进行 softmax 归一化，确保所有节点的扩展宽度之和等于 Bi。

REBASE 的目标是根据节点的奖励来“平衡”地扩展节点。这意味着奖励越高的节点，将被分配越多的计算资源（即扩展出更多的子节点）。

具体来说，REBASE 计算每个节点 nj 的扩展宽度 Wj，公式如下：

• 节点扩展 (Node Expansion):

根据计算出的扩展宽度 Wj，REBASE 从策略模型 πθ 中采样 Wj 个子节点，作为节点 nj 的子节点。

2.2.4 总结一下REBASE方法的改进点

REBASE 通过以下方式避免了 MCTS 的 rollout 步骤：

• 直接使用奖励模型评估节点质量: 无需进行耗时的 rollout 模拟。

• 基于奖励模型进行平衡扩展: 根据节点的奖励动态分配计算资源，将更多资源分配给更有希望的节点，从而更有效地利用计算预算。

通过这种方式，REBASE 在保持树搜索的“探索”和“利用”特性的同时，显著降低了计算成本，使其在准确性和计算效率之间取得了更好的平衡，从而实现了帕累托最优的权衡。

3.Conclusion

• 较小的模型（例如 Llemma-7B）在相同的计算预算下可以胜过较大的模型，并且较小的模型与先进的推理算法相结合可以产生帕累托最优的成本性能权衡。

• 论文提出的 REBASE 算法在所有测试的计算预算中都实现了最佳性能，优于常用的加权多数投票（Majority Voting）和 MCTS 方法。

• 通过采样更多的样本，基于采样的多数投票策略不可避免地会饱和到一个取决于底层生成策略的分布。因此，通过设计替代推理策略来改变采样分布是有意义的。

4.Limitation

• 该分析主要针对数学问题解决，推理缩放规律和计算优化推理在其他任务中是否适用还有待进一步研究。

• 在GSM8K和MATH500数据集上的评估可能不足以泛化到其他任务和数据集。

• 模型大小和推理策略的探索仍然存在局限性。

二、详细情况

1.实验设置

其他推理策略对比： 论文主要考虑了以下几种推理策略：

• 贪婪搜索（Greedy Search）：每次选择概率最高的token。

• Best-of-n：生成多个候选答案，选择得分最高的答案。

• 多数投票（Majority Voting）：生成多个候选答案，选择出现频率最高的答案。

• 加权多数投票（Weighted Majority Voting）：根据奖励模型给出的分数对候选答案进行加权，然后进行投票。

• 树搜索算法（Tree Search）：包括 MCTS 和作者提出的 REBASE。

模型和数据集：

• 使用了 Pythia 和 Llemma 系列的多种大小的模型，并在数学推理基准测试集 GSM8K 和 MATH500 上进行了实验。

• 使用 MetaMath 数据集对模型进行微调。

理论分析：

• 对基于采样的投票方法进行了理论分析，表明随着采样次数的增加，准确率会收敛到一个极限，这个极限取决于语言模型和奖励模型的分布。

实验评估：

• 通过实验，论文比较了不同模型大小和推理策略在不同计算预算下的性能。结果表明，较小的模型与高级的推理算法相结合可以实现帕累托最优的成本效益。

2.最优错误率和模型大小以及预算的关系

2.1 错误率和计算量以及模型大小的关系（左）

• 随着计算量的增加，错误率下降： 所有的曲线都显示，随着推理计算量的增加 (即横轴向右移动)，错误率逐渐下降，说明模型性能在提升。这是因为增加计算量通常意味着模型可以进行更多的采样，从而找到更好的解。

• 模型性能最终趋于饱和： 每条曲线最终都趋于平缓，说明模型性能有一个极限。增加更多的计算量后，错误率的下降幅度会变得越来越小，最终趋于饱和。

• 不同大小的模型在不同阶段表现不同： 在计算量较少的时候，较小的模型 (例如 410M, 1.4B) 下降速度可能更快。随着计算量增加，较大模型 (例如 6.9B, 12B) 会逐渐展现出更好的性能，但其饱和也需要更大的计算量。

2.2 最佳模型大小和预算的关系（右）

• 最佳模型大小随计算预算变化： 图中显示，在不同的推理计算量预算下，最佳的模型大小是不同的。

例如，在较小的计算量预算下 (241 和 244 FLOPs)，较小的模型可能是最优的。

随着计算预算的增加 (247 FLOPs)，较大的模型可能表现更优。

• 小模型在较低计算量下表现出色： 如图中所示，在 241 和 244 FLOPs 的计算量下，较小的模型是计算最优的。这意味着，如果计算资源有限，与其使用大模型，不如使用小模型搭配适当的推理策略可能效果更好。

3.REBASE 算法在计算效率上有明显优势（7B模型+REBASE算法可以超过MCTS等方法）

方法：在 MATH 数据集上，对比不同推理策略 (weighted majority voting 和 best-of-n) 和不同模型大小对模型性能的影响 。

结论：

• 在 MATH 数据集上，REBASE 算法在不同的推理策略（加权多数投票和最佳 n 选 1）下，均优于其他方法（包括传统的采样和 MCTS）。

• REBASE 算法的优越性不仅仅局限于特定的模型大小，其在使用 7B 模型时通常能够达到最佳的性能。

• MCTS 方法在数学推理任务中，性能不如其他方法。

4.REBASE 算法可以在更低的计算预算下实现更高的准确率

• 列: 包括 # SAMPLES (采样数量), FLOPS (推理计算量) 和 MATH500 (在 MATH500 数据集上的准确率) 三列。

• 行: 比较了使用采样方法和 REBASE 算法在不同模型 (Mistral-7B, Llemma-7B, Llemma-34B) 下的性能。

总结:

• REBASE 算法可以在更低的计算预算下实现更高的准确率: 对于所有模型，REBASE 算法在使用更少的样本和更低的计算量 (FLOPS) 的情况下，都比采样方法实现了更高的 MATH500 准确率。

• 验证了 REBASE 的高效性： 这个表格用具体数值验证了 REBASE 的高效性，即能够在降低计算成本的同时，实现更高的性能。

5.消融实验

实验说明：

• Sampling W.M.：使用采样方法和加权多数投票策略。

• Sampling BoN：使用采样方法和 best-of-n 策略。

• REBASE W.M.：使用 REBASE 算法和加权多数投票策略。

• REBASE BoN：使用 REBASE 算法和 best-of-n 策略。

实验结论：

• REBASE 算法性能普遍优于采样： 在所有模型 (Llemma-7B, Llemma-34B, Mistral-7B) 和所有计算预算下，使用 REBASE 算法 (无论是加权多数投票还是 best-of-n) 的曲线都低于使用采样方法的曲线，说明 REBASE 算法性能更佳。

• REBASE + 加权多数投票 策略表现最佳： REBASE 算法配合加权多数投票策略通常可以达到最优的性能。

• 不同模型之间在相同策略下，性能随计算预算增加有差异： 例如，在低预算下，Llemma-7B 在 REBASE 策略下表现可能更好；在高预算下，不同模型之间的差距可能缩小。

三、总结

结论1: 计算最优的推理策略依赖于模型大小和计算预算。 论文表明，在推理时使用更小的模型并进行多次采样，在计算成本方面可能比使用更大的模型更有效。此外，论文提出的REBASE算法在多个数据集上优于MCTS方法和传统的采样方法。 

结论2: 基于采样的投票策略具有性能上限，需要更复杂的推理算法。 论文的理论分析表明，基于采样的投票方法（如多数投票和加权多数投票）在计算预算增加时，性能会饱和，且达到一个基于底层生成策略的极限，这表明了需要更复杂的推理策略。 

结论3: REBASE算法在计算效率和准确率方面都有优势。 REBASE算法通过奖励模型来控制节点的扩展，从而避免了显式的rollout，从而在计算效率上优于MCTS。同时，其性能在所有测试的计算预算下优于采样方法。

本文转载自 ​​NLP PaperWeekly​​，作者： NLP PaperWeekly