[论文解读] Inference Scaling Laws: An Empirical Analysis of Compute-Optimal Inference for Problem-Solving with Language Models

Yangzhen Wu, Zhiqing Sun+3 more|arXiv (Cornell University)|Aug 1, 2024

Natural Language Processing Techniques被引用 4

一句话总结

本文研究了语言模型在数学推理任务中的计算最优推理，提出了一种新颖的树搜索算法 REBASE，该算法在提升准确率的同时降低了 FLOPs。结果表明，在计算资源受限的条件下，搭配 REBASE 后，较小的模型（如 Llemma-7B）可使用一半的 FLOPs 达到与更大模型（如 Llemma-34B）相当的性能，尤其在计算预算有限时表现更优。

ABSTRACT

While the scaling laws of large language models (LLMs) training have been extensively studied, optimal inference configurations of LLMs remain underexplored. We study inference scaling laws (aka test-time scaling laws) and compute-optimal inference, focusing on the trade-offs between model sizes and generating additional tokens with different inference strategies. As a first step towards understanding and designing compute-optimal inference methods, we studied cost-performance trade-offs for inference strategies such as greedy search, majority voting, best-of-$n$, weighted voting, and two different tree search algorithms, using different model sizes and compute budgets. Our findings suggest that scaling inference compute with inference strategies can be more computationally efficient than scaling model parameters. Additionally, smaller models combined with advanced inference algorithms offer Pareto-optimal trade-offs in cost and performance. For example, the Llemma-7B model, when paired with our novel tree search algorithm, consistently outperforms the Llemma-34B model across all tested inference strategies on the MATH benchmark. We hope these insights contribute to a deeper understanding of inference scaling laws (test-time scaling laws) for LLMs.

研究动机与目标

• 理解推理时计算预算如何影响不同模型尺寸下的模型性能。
• 识别在大语言模型解决问题时，推理计算量（FLOPs）与准确率之间的最佳权衡。
• 在固定计算预算下，评估并比较多种推理策略（贪婪搜索、多数投票、Best-of-N、加权投票及树搜索变体）的效率。
• 设计并验证一种新型树搜索算法 REBASE，以在最小化计算成本的同时提升性能。
• 证明在计算受限环境下，小型模型若搭配复杂推理策略，可超越大型模型。

提出的方法

• 在数学推理基准（GSM8K 和 MATH500）上，通过实验评估多种推理策略（贪婪、多数投票、Best-of-N、加权投票及树搜索变体）的性能。
• 通过生成的 token 数量控制推理计算量，并利用奖励模型对候选解进行评分和排序。
• 提出 REBASE，一种新颖的树搜索算法，利用基于节点质量的奖励来平衡探索与剪枝。
• 将 REBASE 与加权投票结合，以提升解的质量并减少对不完整或低质量候选解的依赖。
• 以准确率和 FLOPs 为指标衡量性能，建立不同模型尺寸与推理策略之间的帕累托最优权衡。
• 将 REBASE 与标准采样方法及 MCTS 进行对比，重点关注饱和点与计算效率。

Figure 1 : The inference computation scaling laws of Pythia exhibited in error rate on the GSM8K test set. We evaluate Pythia model using various sizes and various numbers of sampled solutions for majority voting. The left panel shows the error rate for each model size decreases steadily when the co

实验结果

研究问题

• RQ1 在数学推理任务中，不同计算预算下，推理的最优模型尺寸如何变化？
• RQ2 在贪婪搜索、多数投票、Best-of-N、加权投票或树搜索中，哪种推理策略能实现最佳的准确率-FLOPs 比率？
• RQ3 在相同计算预算下，能否通过高级解码策略（如 REBASE）使小型语言模型超越大型模型？
• RQ4 与标准采样方法和 MCTS 相比，REBASE 在准确率与计算效率方面表现如何？
• RQ5 高级推理策略带来的性能提升是否在不同模型尺寸间表现出不同规模效应，特别是对性能较弱的模型？

主要发现

• 在相同推理策略下，较小的模型（如 Llemma-7B）可实现与更大模型（如 Llemma-34B）相当的准确率，同时仅使用约一半的 FLOPs。
• 在 MATH500 基准上，Llemma-7B 搭配 REBASE 达到 46.8% 的准确率，仅消耗 $148.0 \times 10^{12}$ FLOPs，而 Llemma-34B 以 $176.8 \times 10^{12}$ FLOPs 达到 43.6% 的准确率。
• 与基于采样的方法及 MCTS 相比，REBASE 实现了更高的准确率，在 Mistral-7B 上提升 5.3%，在 Llemma-34B 上提升 2.6%。
• 与采样方法相比，REBASE 的性能饱和点更靠后，表明在计算量增加时仍能持续提升性能，尤其在 MATH 等更难的任务上表现更优。
• 性能较弱的模型（贪婪准确率较低）从树搜索中获益更多：REBASE 带来的性能增益与基线贪婪准确率成反比。
• 推理的最优模型尺寸随计算预算变化而变化——在低至中等 FLOP 水平下，小型模型优于大型模型，凸显计算最优配置的重要性。

Figure 2 : Illustration of compute-optimal scaling laws in training and inference. The Chinchilla scaling law (Hoffmann et al., 2022 ) shows how to choose a model size and number of training tokens under a training-compute budget, while our work shows how to choose a model size and an inference stra