Inference Scaling Laws: An Empirical Analysis of Compute-Optimal Inference for LLM Problem-Solving

Authors: yangzhen  wu
zhiqing  sun
shanda  li
sean  welleck
yiming  yang
Venue: NeurIPS_cc_2024_Workshop_MATH-AI
Published: 2024-09-19 06:20:28

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 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 indicate smaller models (e.g., Llemma-7B) can outperform larger models given the same computation budgets, and that smaller models paired with advanced inference algorithms yield Pareto-optimal cost-performance trade-offs. For instance, the Llemma-7B model, equipped with our novel tree search algorithm, consistently outperforms Llemma-34B with standard majority voting on the MATH benchmark across all FLOPs budgets. We hope these findings contribute to a broader understanding of inference scaling laws for LLMs.

Main paper: 
 introduction: 
 Scaling laws of neural networks [Hestness et al., 2017 , Rosenfeld et al., 2019] have been established across a range of domains, including language modeling [Kaplan et al., 2020 , Hoffmann et al., 2022 , OpenAI, 2023] , image modeling [Henighan et al., 2020 , Yu et al., 2022 , Peebles and Xie, 2023] , video modeling [Brooks et al., 2024] , reward modeling [Gao et al., 2023] , and board games [Jones, 2021] . These studies have demonstrated how model performance is influenced by both the size of the model and the amount of training computation. However, there is limited knowledge on how varying the compute during inference affects model performance after the model has been trained.To improve the task performance of large language models (LLMs), inference techniques typically involve additional computation as a performance maximization step at inference time [Nye et al., 2021 , Wei et al., 2022 , Wang et al., 2022b , Yao et al., 2023 , Chen et al., 2024b] . The computational cost of these techniques must be taken into account for compute-optimal inference. For example, a Monte Carlo Tree Search (MCTS) method [Jones, 2021] may improve task performance, but potentially require much more compute than simply sampling solutions multiple times. Generally speaking, we need a comprehensive understanding of how various inference-time methods (e.g., best-of-n, majority voting [Wang et al., 2022a] ) trade off between performance and cost. To improve our understanding, this paper presents a thorough empirical evaluation with careful analysis over various configurations of representative LLMs and inference algorithms.Specifically, we explore how to select an optimal size for the language model and an effective inference strategy (e.g., greedy search, majority voting, best-of-n, weighted voting, and their treesearch variants) to maximize performance (i.e., accuracy) with a given compute budget. We control the inference computation (FLOPs) of a fixed model by generating more tokens through the language model3 , sampling further candidate solutions, and ranking them with a reward model. We analyze the performance of fine-tuned models of various sizes given different inference FLOPs on mathematical reasoning benchmarks (e.g., GSM8K test set [Cobbe et al., 2021a] and MATH500 test set [Hendrycks et al., 2021 , Lightman et al., 2023b] ). Our experiments cover several model families, including general-purpose LLMs (e.g., Pythia [Biderman et al., 2023 ] & Mistral [Jiang et al., 2023] ) as well as math-specialized ones (e.g., Llemma [Azerbayev et al., 2023] ).Our results on Pythia (Fig. 1 ) illustrate how performance scales with increased inference compute across various model sizes. Typically, increasing the compute budget leads to higher accuracy until the accuracy reaches saturation. As the compute budget increases, smaller models initially perform better than larger ones, but once the accuracy of the smaller models saturates, the larger models have favorable performance. The right panel of Figure 1 demonstrates that the optimal model size for inference varies with different levels of computation. However, in real-world deployment, the available computation is typically much lower than the point where the accuracy of relatively small models saturates and larger models begin to show their advantage (as shown in Figure 2 , where the 7B model outperforms the 34B model until 128 Llemma 7B solutions are sampled). This indicates that relatively smaller models could be more compute-optimal for inference.We have also found that the commonly-used MCTS method does not perform well with weighted voting, as it often yields many unfinished solutions, hence having less effective votes. To address this issue, we propose a novel tree search algorithm, REward BAlanced SEarch (REBASE), which pairs well with weighted voting and achieves a Pareto-optimal trade-off between accuracy and inference compute. The key idea of REBASE is to use a node-quality reward to control node expansion, which eliminates the need for explicit rollouts while ensuring enough candidate solutions for voting.In our experiments, REBASE consistently outperforms sampling and MCTS methods across all settings, models, and tasks. Importantly, we find that REBASE with a smaller language model typically achieves a Pareto-optimal trade-off. For instance, we show that the Llemma-7B model can achieve competitive accuracy to a Llemma-34B model while using 2× less FLOPs when evaluating on MATH500 (Fig. 2 ) or GSM8K (Fig. 3 ). Moreover, Llemma-7B with REBASE outperforms Llemma-34B with standard majority voting across all compute budgets. Our results show the value of using smaller models with advanced inference-time algorithms, and the benefits of new algorithms for achieving better returns on inference-time compute.We explore the following question: Given a fixed FLOPs budget, how should one select an optimal model size for the policy model, and an effective inference strategy to maximize performance (i.e., accuracy)? .To address this, we represent the problem-solving error rate E(N, T ; S) as a function of the number of model parameters N , the number of generated tokens T and the inference strategy S. The computational budget C is a deterministic function FLOPs(N, T ; S), based on N and T . Our goal is to minimize E under the test-time compute constraint FLOPs(N, T, S) = C:(N opt (C), T opt (C); S) = arg min (N,T,S) s.t. FLOPs(N,T,S)=C E(N, T ; S)where N opt (C) and T opt (C) denote the optimal allocation of a computational budget C.Here, the inference computation (FLOPs) for a fixed model can be modulated by generating more tokens with the policy model, e.g., by sampling additional candidate solutions and subsequently ranking them using a reward model. We primarily consider sampling and tree-search approaches with reranking or Majority Voting as the means to consume more tokens. We examine sampling-based and tree-search inference strategies.In sampling methods, majority voting (a.k.a. self-consistency [Wang et al., 2022a] ) samples multiple reasoning paths, selecting the most frequent answer. When using a reward model, best-of-N selects the highest-scoring path, while weighted majority voting selects the answer with the highest weighted score, based on reward model values.Tree-search methods, recently adapted for LLMs [Yao et al., 2023 , Zhang et al., 2023 , Zhou et al., 2023 , Liu et al., 2024 , Choi et al., 2023 , Chen et al., 2024a , Tian et al., 2024 , Chen et al., 2024a] , often pair with value models to guide exploration. Monte Carlo Tree Search (MCTS) is common (Appendix C reviews MCTS), but it tends to be resource-heavy, requiring many more generated tokens than simpler methods. A more efficient inference strategy is needed, and comparisons of tree-search and sampling methods on computational cost are essential. We present our novel tree search method REBASE in 3.1. 
inference scaling laws: 
 In order to compare the compute budgets of different models, we plot the figures with the number of FLOPs used per question during inference. We compute the inference FLOPs based on the standard formula from [Kaplan et al., 2020] .Scaling law of compute-optimal inference for model size. Figure 1 shows the relationship between inference compute and error rate for different model sizes. The error rate first decreases steadily and then starts to saturate. Initially, sampling many times from smaller models is computeoptimal. At larger compute budgets the larger models are preferable, since the performance of small models has saturated. As highlighted in the right panel of Figure 1 , the optimal model size varies based on the inference budget. We performed a regression analysis on inference FLOPs C and model sizes N to establish a relationship between a given computational budget and its optimal model size.The resulting equation, log 10 (C) = 1.19 log 10 (N ) + 2.03, lets us estimate the optimal inference model size for a specific compute budget.Llemma-7B model achieves competitive accuracy to Llemma-34B model with lower compute budget. Fig. 2 & 3 shows the relationship between error rate and inference FLOPs for Llemma 7B and Llemma 34B using different inference strategies. Llemma-7B requires around 2× less total FLOPs than Llemma-34B to achieve comparable accuracy. This held across inference strategies (sampling strategies, MCTS, REBASE) and tasks (MATH, GSM8K). This result suggests that, with the same training dataset and model family, generating more tokens with a suitable inference strategy using a smaller model can have more favorable cost-performance tradeoffs than using a larger model. 3 Compute-Optimal Inference 
reward balanced search (rebase): 
 The REBASE tree search method inherits the exploitation and pruning properties of tree search, while using the reward model alone to estimate the nodes' qualities without additional computation for estimating values by sampling children. The details are provided below:Notations. We view the fine-tuned LLM as a policy π θ which generates the solution step by step. Given a question x and the first k steps of a solution r 1 • • • r k , the (k + 1)-th step is sampled from π θ (•|xr 1 • • • r k ). REBASE) := R(qr 1 • • • r k ).Initialization. Given the question x, balance temperature T b > 0, and sampling number of solutions N , we sample N instances of the first step for the question, yielding all the nodes of depth 1 in the search tree. We let the sampling budget of depth 0, B 0 , to N at initialization.Reward modeling and update. In the i-th iteration, the PRM assigns the rewards to all the nodes at depth i. After that, the algorithm examines whether the solutions up to depth i are complete. Supposing there are C i completed solutions, we update the sampling budget usingB i ← B i-1 -C i .If B i = 0, the process ends, and we obtain N solutions.Exploration balancing and expansion. For all of the nodes n j with reward R(n j ) in the depth i of the tree, we calculate the expansion width of the n j as:W j = Round B i exp (R(n j )/T b ) k exp (R(n k )/T b ).Then we sample W j children for n j for all the nodes in depth i, and start the next iteration.Intuitively, when the balance temperature T b is small, this method encourages more exploitation which put much more compute budget on the nodes with high score, when T b is large, it encourages exploration where nodes with high score and low score are exlopred equally. In our experiment, we have found T b in the range of (0.1, 0.3) works well for our process reward model.REBASE is Pareto-optimal. REBASE consistently achieves the best cost-performance tradeoffs, outperforming the sampling-based methods in all settings when fixing the model and the evaluation task (Fig. 2 , 3, 4, and 5). For example, in Figure 2 , REBASE is the compute-optimal strategy at all inference compute budgets, with 7B typically the optimal model size. On the other hand, MCTS underperforms the sampling-based methods at each compute budget, likely due to its costly rollouts (Figure 2 ) compared to the efficient use of the reward model in REBASE.Table 1 shows that REBASE achieves better accuracy with a lower compute budget compared to sampling-based weighted voting. With the 7B model, REBASE achieves higher accuracy with 7 times less compute. This finding is novel, and differs from previous tree search methods that typically improve the performance at the cost of higher computational expense compared to sampling-based voting [Chen et al., 2024a , Xie et al., 2023] .Weaker models gain more from tree search. For example, our proposed REBASE leads to 5.3%, 3.3%, and 2.6% performance gains on MATH for Mistral-7B, Llemma-7B, Llemma-34B, respectively. The order of accuracy increase is inversely related to the model's corresponding greedy search accuracy on those datasets. This suggests that weaker models, as indicated by their lower greedy search accuracy, benefit more from tree search methods like REBASE.REBASE saturates later than sampling with higher accuray. From Fig. 2 and Fig. 3 , we observe that REBASE saturates later than the sampling methods, with lower final error rates. This is the evidence that REBASE improves the reasoning paths, due to the selection and pruning mechanism in the intermediate steps, REBASE discards the low quality partial paths and exploring more on good ones, results in a higher probability of sampling high-quality reasoning paths. 
conclusions and limitations: 
 Conclusions. In this work, we conducted a comprehensive empirical analysis of inference scaling law and compute-optimal inference for problem-solving with language models. We examined the scaling effect of computation during inference across different model sizes and found that while increased computation generally leads to higher performance, the optimal model size varies with the available compute budget. When the computation budget is limited, smaller models can be preferable. Additionally, we introduce our novel tree search method, REBASE, which is more compute-optimal than both sampling methods and Monte Carlo Tree Search (MCTS). REBASE typically achieves higher accuracy while using several times less computation than sampling methods. Our results underscore the potential of deploying smaller models equipped with sophisticated inference strategies like REBASE to enhance problem-solving accuracy while maintaining computational efficiency.Limitations. Our empirical analysis specifically targets mathematical problem-solving. Investigating the inference scaling laws and compute-optimal inference strategies for tasks beyond mathematical problem-solving would be a valuable direction for future research. Additionally, we mainly evaluate the proposed REBASE on the GSM8K and MATH500 datasets. We speculate that the REBASE algorithm, which assumes access only to a function that assigns scores to nodes, will be effective in tasks beyond those studied here.Figure 4 : MATH inference scaling across inference strategies and models (lower is better). The tested models are Llemma-7B (left), Llemma-34B (middle), & Mistral-7B (right). In the legend, W.M. and BoN refer to weighted majority and best-of-n, respectively.  
 
 Appendix: In this section, we present the results of Llemma models and Mistral model to see how the REBASE performs on different models. Fig. 4 and 5 shows the scaling behaviors of Llemma models and Mistral model on MATH and GSM8k datasets. Table 1 shows the advantage of REBASE in specific compute settings. Mathematical Reasoning with LLMs. Large language models have made significant progress in recent years, and have exhibited strong reasoning abilities [Brown et al., 2020 , Hoffmann et al., 2022 , Chowdhery et al., 2022 , Lewkowycz et al., 2022] . Mathematical problem solving is a key task for measuring LLM reasoning abilities [Cobbe et al., 2021a , Hendrycks et al., 2021] . Ling et al. [2017] first developed the method of producing step by step solutions that lead to the final answer. Later, [Cobbe et al., 2021b] extended the work by training a verifier for evaluating and ranking solutions. Subsequent research (e.g., Lewkowycz et al. [2022] ) has shown the performance benefits of inference-time techniques such as majority voting [Wang et al., 2022a] and weighted majority voting [Li et al., 2023] . We choose problem solving in mathematics as the task to study compute-optimal strategies since it allows us to accurately evaluate problem solving ability.Inference Strategies of LLM Problem Solving. A variety of inference strategies have been developed to generate sequences with a trained model. Deterministic methods such as greedy decoding and beam search [Teller, 2000 , Graves, 2012] find highly probable sequences, often yielding high quality results but without diversity. Sampling algorithms (e.g., temperature sampling [Ackley et al., 1985] ) can produce a diverse set of results which are then aggregated to achieve higher accuracy (e.g., via majority voting [Wang et al., 2022a] ). Recent methods combine search algorithms with LLMs, including breadth-first or depth-first search [Yao et al., 2023] , Monte-Carlo Tree Search (MCTS) [Zhang et al., 2023 , Zhou et al., 2023 , Liu et al., 2024 , Choi et al., 2023] , and guided beam search [Xie et al., 2023] . All of these methods show that using search at inference time can lead to performance gains in various tasks. However, the trade-off for the improved performance is the use of compute to perform the search. Analyzing the resulting cost-performance trade-offs remains understudied. In this paper, we systematically analyze the trade-off between compute budget and problem-solving performance, and propose a tree search method that is empirically Pareto-optimal.Process Reward Models. Process reward models (PRMs) have emerged as a technique to improve the reasoning and problem-solving capabilities of LLMs. These models assign rewards to the intermediate steps of the LLM generated sequences. PRMs have been shown effective in selecting reasoning traces with a low error rate, and for providing rewards in reinforcement learning-style algorithms [Uesato et al., 2022 , Polu and Sutskever, 2020 , Gudibande et al., 2023] . Ma et al. [2023] applies a PRM to give rewards on the intermediate steps and guide the multi-step reasoning process. The PRM can be either trained on human labeled data [Lightman et al., 2023a] or model-labeled synthetic data [Wang et al., 2023] . In our work, we use the PRM as the reward model for selecting generated solutions, and for selecting which partial solutions to explore in tree search.In this section, we present additional background on the Monte Carlo Tree Search (MCTS) algorithm.The MCTS process can be formulated as the following steps:Selection. The process begins at the root node. Here, the algorithm recursively selects the child node that offers the highest Upper Confidence Bound applied to Trees (UCT) value, continuing until a node is reached that has not been expanded. The UCT is calculated using the formula U CT (s) = Q(s) + C ln (N (Parent(s))) N (s) , where Q(s) denotes the quality score of node s, N (s) is the number of visits to node s, Parent(s) denotes the parent node of s, and C is a constant determining the level of exploration.Expansion and evaluation. Upon reaching a non-terminal node s, the node is expanded by generating multiple child nodes. Each child node c is then evaluated using a value function V (c), which predicts the potential quality of continuing the sequence from node c.Backpropagation. After evaluation, the algorithm updates the UCT values and the visit counts for all nodes along the path from the selected node back to the root. For any node n in this path, theFollowing Uesato et al. [2022], we refer to the main language model generating outputs as the policy model. It can be paired with a reward model, which scores outputs from the policy model to facilitate inference. 



Reviews: review: title: title: This paper investigates inference scaling laws for large language models (LLMs) in problem-solving tasks. It introduces the concept of compute-optimal inference, exploring trade-offs between model size and inference strategies. The authors present REBASE, a tree search method that outperforms sampling and MCTS in terms of compute efficiency. Empirical results on mathematical reasoning benchmarks demonstrate that smaller models with sophisticated inference strategies can achieve competitive performance while maintaining computational efficiency.
main_review: main_review: **Quality**

Strengths:
- Comprehensive analysis of inference scaling laws across different model sizes and compute budgets.
- Introduction of REBASE, the tree search method with clear explanation and empirical validation.
- Well-structured mathematical foundation for REBASE, using a softmax-like formula for exploration balancing.
- Empirical evidence showing REBASE's superior performance and compute efficiency compared to MCTS and sampling methods.

Weaknesses:
- The paper lacks a detailed discussion of limitations or potential drawbacks of the proposed method.
- No theoretical analysis of REBASE's convergence properties or regret bounds.
- Limited discussion on hyperparameter sensitivity, particularly the balance temperature T_b.
- Absence of comparative analysis with tree search variants beyond MCTS (e.g., beam search, A* search).
- No explicit discussion on how REBASE handles reward normalization across different tree depths.
- Lack of analysis on REBASE's performance across problems with different structures (e.g., sparse rewards, deceptive rewards).
- No discussion on memory requirements for storing the search tree, which could be a limiting factor for very large trees.

**Clarity of writing**

Strengths:
- Clear problem formulation and explanation of the research question.
- Well-structured explanation of the REBASE algorithm.

Weaknesses:
- The introduction could benefit from a more comprehensive review of related work.
- Lack of explicit discussion on potential limitations of REBASE and areas for future work.

**Originality of the paper**

Strengths:
- Novel exploration of inference scaling laws for LLMs in problem-solving tasks.
- REBASE's unique approach of balancing exploitation and exploration without additional value estimation computations.

Weaknesses:
- While REBASE is novel, it builds on existing concepts (e.g., tree search, sampling methods). The paper could benefit from a more detailed comparison with other advanced tree search methods.
- The reward-based node evaluation in REBASE is similar to some existing reinforcement learning techniques, which could be acknowledged more explicitly.

**Significance of this work**

Strengths:
- Implications for deploying smaller models with sophisticated inference strategies to enhance problem-solving accuracy while maintaining computational efficiency.
- Demonstration that smaller models (e.g., Llemma-7B) can achieve competitive accuracy to larger models (e.g., Llemma-34B) with lower compute budgets.
- The paper opens up new research directions in compute-optimal inference for LLMs, which is crucial for practical applications of AI in resource-constrained environments.

Weaknesses:
- The paper could elaborate more on the broader implications of these findings for the field of AI and mathematical reasoning.
- Limited discussion on how REBASE might perform in domains beyond mathematical reasoning tasks.
title: title: This paper studies the scaling inference law for math problems solving with LLMs. They find that small model could have advantage when facing limited computation budget. To address the resource limit, they also propose a new algorithm REBASE and conduct numerical experiments to verify its performance.
main_review: main_review: This paper conducts a comprehensive empirical analysis of the scaling inference law and studies the inference strategies for solving math problems with Large Language Models. The authors investigate the selection of optimal model size and inference strategy under a fixed computational budget to maximize model performance. They discover that while increasing inference computation leads to better performance, the optimal model size varies depending on the available budget. Small models is likely to be more advantageous when the computation budget is limited.  To address the high cost of traditional MCTS, the authors propose a new Reward Balanced Search algorithm (REBASE). Numerical Experiments show that it is more computational efficient than MCTS and some other sampling methods.
Strengths:
1. By conducting extensive experiments across different model sizes and inference strategies, the authors reveal how inference computational budget affects model performance. The numerical experiment is convincing.
2. The intuition behind the new algorithm is simple but innovative by just replacing MCTS with some novel tree search method, but it remains good performance especially when per question Inference FLOPs is large
3. Also the study highlights the advantage of combining small models with REBASE when computational budget is limited, which sheds lights on how to deploying LLMs in resource-limited environments.
Weakness:
1. Lack of theoretical analysis and intuition. If the author could give us some mathematical intuition about the new algorithm design, then the paper could be more convincing.

2. I suggest the author comparing their method with more algorithms to better illustrate the power of the REBASE algorithm.

Questions: I am not an expert in LLM and empirical study, so please forgive me if my questions are naive.

1. Could you please provide some theoretical intuition about why REBASE works better? For example, why do you choose that exponential weighting:

$W_j=\mbox{Round}(B_i\frac{\exp(R(n_j)/T_b)}{\sum_k\exp(R(n_k)/T_b)})?$

2. The mathematical problem-solving tasks are very interesting. Are there any possibilities that your strategy could be applied to other types of problems such as natural language tasks? In other words, I’m curious that are there any special structures of mathematical problem-solving tasks that make your innovative method works well?
title: title: review
main_review: main_review: The paper investigates inference scaling laws of the task performances v.s. model parameters and inference FLOPs, and compare tree search variants to sampling-based methods. The paper finds that a smaller language model with a novel tree search algorithm typically achieves a Pareto-optimal tradeoff. The authors further propose REBASE, a tree search method that utilizes a separate reward model to determine the expansion of nodes. The empirical results show that REBASE performs better than other baselines (MCTS, sampling).

Pros:

- the paper provides solid empirical results to justify the advantages of the proposed REBASE method.

Cons:

- the setting of inference-time search often relies on a seperate reward model to aggregate the results of different trajectories (except for majority voting). The performance of REBASE would also likely to depend on the process reward model.
title: title: Well-Executed Evaluation Framework with Opportunities for Further Exploration
main_review: main_review: # Pros
## 1. Strong Paper Writing, with Minor Typos:
The paper demonstrates good writing quality, although some typos are present, such as "Mjority" which should be corrected to "Majority." Please review the document for any additional typos.

## 2. Comprehensive Evaluation Settings:
The evaluation settings are thorough and well-detailed, covering aspects such as model size and inference FLOPs per question, which provides valuable insights.

# Cons
## 1. Need for a Theoretical and Detailed Scaling Formula:
While the paper does describe the scaling trends, I would prefer a more theoretical and detailed scaling formula to deepen the understanding of the trends observed.

## 2. Inclusion of Larger-Sized Language Models:
The largest model tested is LLemma-34B. Including experiments with even larger models could provide additional findings and more comprehensive insights.

## 3. Additional Benchmarks, e.g., AIME2024:
Adding more benchmarks, such as AIME2024, could help to further confirm whether the scaling law holds across different tasks, increasing confidence in the generality of the results.