[Paper Review] 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 Techniques4 citations

TL;DR

This paper investigates compute-optimal inference for language models in mathematical reasoning, proposing a novel tree search algorithm, REBASE, that improves accuracy while reducing FLOPs. It demonstrates that smaller models like Llemma-7B can match larger models (e.g., Llemma-34B) in performance using half the FLOPs when paired with REBASE, especially under constrained compute budgets.

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.

Motivation & Objective

• To understand how inference-time compute budget affects model performance across different model sizes.
• To identify the optimal trade-off between inference compute (FLOPs) and accuracy in problem-solving with LLMs.
• To evaluate and compare the efficiency of various inference strategies—Greedy Search, Majority Voting, Best-of-N, Weighted Voting, and tree search variants—under fixed compute budgets.
• To design and validate a new tree search algorithm, REBASE, that enhances performance while minimizing computational cost.
• To demonstrate that smaller models with sophisticated inference strategies can outperform larger models in compute-constrained settings.

Proposed method

• Empirically evaluate multiple inference strategies (Greedy, Majority Voting, Best-of-N, Weighted Voting, and tree search variants) on mathematical reasoning benchmarks (GSM8K and MATH500).
• Control inference compute via token generation count and use a reward model to score and rank candidate solutions.
• Propose REBASE, a novel tree search algorithm that uses node-quality-based rewards to balance exploration and pruning.
• Integrate REBASE with weighted voting to improve solution quality and reduce reliance on incomplete or low-quality candidates.
• Measure performance in terms of accuracy and FLOPs to establish Pareto-optimal trade-offs across model sizes and inference strategies.
• Compare REBASE against standard sampling-based methods and MCTS, focusing on saturation points and computational efficiency.

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

Experimental results

Research questions

• RQ1 How does the optimal model size for inference vary with different compute budgets in mathematical reasoning tasks?
• RQ2 Which inference strategy—Greedy, Majority Voting, Best-of-N, Weighted Voting, or tree search—achieves the best accuracy-to-FLOPs ratio?
• RQ3 Can a smaller language model with a sophisticated decoding strategy like REBASE outperform a larger model under the same compute budget?
• RQ4 How does REBASE compare to standard sampling and MCTS in terms of accuracy and computational efficiency?
• RQ5 Does the performance gain from advanced inference strategies scale differently across model sizes, particularly for weaker models?

Key findings

• Smaller models like Llemma-7B achieve competitive accuracy to larger models (e.g., Llemma-34B) while using approximately half the FLOPs under the same inference strategy.
• On the MATH500 benchmark, Llemma-7B with REBASE achieves 46.8% accuracy using $148.0 \times 10^{12}$ FLOPs, matching Llemma-34B’s 43.6% accuracy with $176.8 \times 10^{12}$ FLOPs.
• REBASE achieves higher accuracy than sampling-based methods and MCTS, with a 5.3% performance gain on MATH for Mistral-7B and 2.6% for Llemma-34B.
• REBASE saturates later than sampling methods, indicating sustained performance gains under increasing compute, especially on harder tasks like MATH.
• Weaker models (lower greedy accuracy) benefit more from tree search: performance gains from REBASE are inversely proportional to their baseline greedy accuracy.
• The optimal model size for inference varies with compute budget—smaller models outperform larger ones at low to moderate FLOP levels, highlighting the importance of compute-optimal configuration.

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