An Empirical Analysis of Compute-Optimal Inference for Problem-Solving with Language Models

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 03:14:08

Abstract: The optimal training configurations of large language models (LLMs) with respect to model sizes and compute budgets have been extensively studied. But how to optimally configure LLMs during inference has not been explored in sufficient depth. We study compute-optimal inference: designing models and inference strategies that optimally trade off additional inference-time compute for improved performance. As a first step towards understanding and designing compute-optimal inference methods, we assessed the effectiveness and computational efficiency of multiple inference strategies such as Greedy Search, Majority Voting, Best-of-N, Weighted Voting, and their variants on two different Tree Search algorithms, involving different model sizes (e.g., 7B and 34B) and computational budgets. We found that a smaller language model with a novel tree search algorithm typically achieves a Pareto-optimal trade-off. These results highlight the potential benefits of deploying smaller models equipped with more sophisticated decoding algorithms in end-devices to enhance problem-solving accuracy. For instance, we show that the Llemma-7B model can achieve competitive accuracy to a Llemma-34B model on MATH500 while using 2× less FLOPs. Our findings could potentially apply to any generation task with a well-defined measure of success.

Main paper: 
 claims: 
 Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: In Abstract and Introduction, we claim that we investigate the compute-optimal inference: designing models and inference strategies that optimally trade off additional inference-time compute for improved performance.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper.• The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.• The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.• It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
related works: 
 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., 2021b] . [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 finetuning the pre-trained language model on a large dataset to solve math word problems, a verifier is trained for evaluating solutions and ranking solutions. Nye et al. [2021] train models to use a scratchpad and improve their performance on algorithmic tasks. Wei et al. [2022] demonstrate that the reasoning ability of a language model can be elicited through the prompting. Subsequent research [Kojima et al., 2022 , Lewkowycz et al., 2022 , Zhou et al., 2022] in reasoning tasks has also highlighted the efficacy of rationale augmentation. We choose problem solving in mathematics as the task to study the compute-optimal strategy since it allows us to accurately evaluate the problem solving ability of LLMs. Inference Strategies of LLM Problem Solving. A variety of inference (also called decoding) 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 modern 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 Self-evaluation 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 trade-off between compute budget and LLM inference performance 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 the 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.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: The discussion is in the last section of the main paper.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.• The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.• The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.• The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.• The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.• If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.• While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
an empirical analysis of compute-optimal inference for problem-solving: 
 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 ) as a function of the number of model parameters N and the number of generated tokens T . The computational budget C is a deterministic function FLOPs(N, T ), based on N and T . Our goal is to minimize E under the test-time compute constraint FLOPs(N, T ) = C:EQUATIONwhere 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, including Greedy Search, Majority Voting, Best-of-N, Weighted Voting, and their variants on tree search methods.Greedy Search. This strategy generates tokens one at a time by selecting the highest probability token at each step, without considering future steps. It is computationally efficient but often suboptimal in terms of diversity.Best-of-n. This strategy, also known as rejection sampling, samples multiple solutions and chooses the one with the highest score given by the reward model.Majority Voting. In this strategy, multiple model outputs are generated, and the final answer to the problem is determined by the most frequently occurring answer in all the outputs.Weighted Majority Voting. This strategy is a variant of majority voting in which the votes are weighted based on the score given by the reward model.Monte Carlo Tree Search (MCTS) has proven effective in domains such as board games where strategic decision-making is required [Silver et al., 2016 , 2017 , Jones, 2021] . Recent work has shown that adapting MCTS to the context of LLMs can enhance the text generation process [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] . In this context, MCTS is often paired with a value model to score and guide the exploration steps. For additional background, we provide a review of MCTS in Appendix B.Recent work in MCTS or its variants (e.g., Tree of Thoughts [Yao et al., 2023] ) mainly focus on improving the performance (e.g., accuracy) on the studied tasks. However, generic comparisons of MCTS with conventional methods like Best-of-n and Majority Voting in terms of computational budget, measured in generated tokens or processing time, are either scarce or indicating inferencetime issues. For example, MCTS consumes substantially more resources, often requiring dozens of times more generated tokens than simpler methods. Specifically, a significant portion of the paths in the search tree are used to estimate and select nodes, and these paths do not necessarily become a part of the final candidate solution, although MCTS ensures that the sampled solutions comprise high-quality intermediate steps. In contrast, sampling methods generate multiple solutions in parallel and independently, and all the generated sequences are included in the candidate solutions. However, the intermediate steps in these sequences are not guaranteed to be of high quality, as there is no mechanism for pruning poor steps or exploiting promising ones.This highlights the need for developing a new tree search method that can achieve a comparable (or better) performance as MCTS, and that is computationally less costly, just like weighted majority voting and best-of-n. This need leads to the development of our new method named Reward Balanced SEarch (REBASE), as introduced next.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 efficiency is achieved by constraining the total expansion width of the tree at a certain depth. REBASE balances the expansion width among the nodes at the same depth based on the rewards given by the Process Reward Model (PRM). The details are provided below:Notations. We consider the fine-tuned LLM as a policy π θ . Given a question q and the first k steps of a solution x 1 , • • • , x k , the (k + 1)-th step is produced by π θ (x k+1 |q, x 1 • • • x k ). When generating solutions using tree search, the root of the tree corresponds to the question q. The node corresponding to x k+1 is the child of the node corresponding tox k if it is sampled from π θ (•|q, x 1 • • • , x k ). The reward of a node n(x k ) is determined by the PRM as R(n(x k )) = R(q, x 1 , • • • , x k ).Initialization. Given the question q, balance temperature T b , 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 set the sampling budget of depth 0 B 0 = N as 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 ).(2)Then we sample W j children for n j for all the nodes in depth i, and start the next iteration.Before empirically studying the scaling effects of increasing the inference-time compute budget, we present two theorems which will help us understand the experimental results later. These two theorems give an upper bound on the performance of sampling when fixing the LLM generator.We assume the vocabulary is limited and the sequence length is constrained, thus the number of possible solutions and answers are finite. The proofs are provided in the Appendix A. limEQUATION)where π(x|d) denotes the probability that the LLM answers x given input d and I denotes the indicator function.Theorem 2 shows that as long as the reward model assigns higher rewards than the policy for correct answers versus other answers in expectation, the upper bound of Weighted Majority Voting will be higher Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification:We only used open-source models in this work. The code will be released.Guidelines:• The answer NA means that paper does not include experiments requiring code.• Please see the NeurIPS code and data submission guidelines (https://nips.cc/ public/guides/CodeSubmissionPolicy) for more details.• While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).• The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https://nips.cc/public/guides/CodeSubmissionPolicy) for more details.• The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.• The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.• At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).• Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
setup: 
 Datasets. We conduct experiments on two mathematical problem-solving datasets to investigate the scaling effects of computation and our REBASE method for both challenging and simpler problems. Specifically, MATH [Hendrycks et al., 2021a] and GSM8K [Cobbe et al., 2021b] are datasets containing high school mathematics competition-level problems and grade-school level mathematical reasoning problems, respectively. Following [Lightman et al., 2023b , Wang et al., 2024 , Sun et al., 2024] , we use the MATH500 subset as our test set.Generators. We use Llemma-7B and Llemma-34B [Azerbayev et al., 2024] as our base models and finetune them on the MetaMath dataset [Yu et al., 2024] using full parameter supervised fine-tuning (Full-SFT), The detailed finetuning configuration is given in the Appendix. Additionally, we test the Mistral-7B model to expand our findings across different models.Reward Model. All of the experiments use the same Llemma-34B reward model, which we finetuned on the synthetic process reward modeling dataset, Math-Shepherd [Wang et al., 2024] . We added a reward head to make the model, enabling it to output a scalar reward at the end of each step. Each configuration is run multiple times to calculate the mean and variance, thereby mitigating the randomness and ensuring the reliability of our conclusions.In order to compare the compute budgets of 7B and 34B 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] . All inference configurations will saturate eventually. This is expected as Theorem 1 and Theorem 2 show. Also illustrated in Figures 5 and 6 , the slope of the erro rate curves start large, then decreases and the curves finally become nearly flat as the number of samples scales, showing the effect of saturation.LlemmaScaling law of compute-optimal inference. The findings in our experiments are consistent with the Theorem 1 and 2, After a threshold the accruacy of sampling more solutions saturate, we should scale the model size. We interpolate the smoothed test error rate curve in Figure 1 and Figure 4 , and fit power laws to estimate the optimal model size N and number of generated tokens T for any given amount of compute. We obtained a relationship N opt ∝ C a and T opt ∝ C b , where a = 1.0 and b = 0.0 for both sampling-based weighted voting and our tree-search method REBASE. Our fitted curves indicate that the optimal inference strategy is invariant to the amount of compute (e.g., re-ranking with 32 sampled solutions or REBASE tree search with a compute budget of 64 for MATH), and the optimal model size grows linearly with the increased compute budget.REBASE is Pareto-optimal. While MCTS undeperforms Sampling (Fig. 1 ), from Fig. fixing the model and the evaluation task. Table 1 shows that REBASE can achieve competitive accuracy with even a lower compute budget than the sampling method. This finding is novel, and differs from previous tree search works which typically improve the performance at the cost of higher computational expense compared to sampling [Chen et al., 2024a , Xie et al., 2023] . Table 2 shows that given the same compute budget (sampling 32 solutions for the 7B model and 8 solutions for 34B model), using REBASE yields higher accuray than sampling.Weaker models gain more from Tree Search. From Fig. 2 , we saw that compared with sampling, Mistral-7B, Llemma-7B, Llemma-34B increase 5.3%, 3.3%, 2.6% in MATH and 0.7%, 1.9%, 0.9% in GSM8K. The order of accuracy increase is inversely related to the model's corresponding greedy search 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 Figure 5 and Figure 6 , we observe that both sampling and REBASE saturate early in GSM8K and relatively late in MATH, which we attribute to the difference of the difficulty level. This can be explained through the LLM may assign high probability to the true answer in easy problems than those of harder problem, as suggested by Theorem 1 and 2 with their proofs A. On MATH (Figure 5 ), we see that REBASE finally saturates with a higher accuracy than sampling. We hypothesize the reason is that REBASE samples the truth answer with higher probability than sampling. And as demonstrated by Theorem 1 and 2, the upper bound becomes higher.  
conclusion & limitations: 
 In this work, we have conducted a comprehensive empirical analysis of compute-optimal inference for problem-solving with language models. Our study has revealed several key findings. First, with an optimal inference configuration, a small language model can achieve competitive accuracy to a 4× larger model while using approximately 2× less total FLOPs under the same inference method (Sampling, MCTS, REBASE) and task (MATH, GSM8K), suggesting that training and inference with smaller models could be more favorable in terms of compute budget when combined with multiple sampling or search strategies. Second, our new REBASE tree-search method consistently outperforms sampling (and MCTS) across all settings, models, and tasks, achieving competitive accuracy with even lower compute budget compared to sampling. Our findings highlight the potential of deploying smaller models equipped with more sophisticated inference strategies like REBASE to enhance problem-solving accuracy while maintaining computational efficiency.Limitations First, our experiments focused specifically on mathematical problem-solving tasks, and the generalizability of our findings to other domains remains to be explored. Second, we only investigated a limited range of model scales, primarily focusing on 7B and 34B models. Future research could extend our analysis to a wider range of model sizes to gain a more comprehensive understanding of the scaling laws for compute-optimal inference. Finally, our experiments mainly utilized the MetaMath dataset for training the models. It would be valuable to explore the impact of different training datasets on the performance and efficiency of compute-optimal inference strategies for mathematical problem-solving. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: We used the standard training and test splits or MATH and GSM8K and report the hyperparameters in the appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.• The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes] Justification: The error bars are included in our figures.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)• The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean.• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.• For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).• If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: All the experiments are conducted on 8× H100 GPUs.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.• The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.• The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: Yes, we conform with the NeurIPS Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.• The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: We do not find significant positive societal impacts and negative societal impacts of our work.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.• The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.• If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification:Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.• Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.• We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We use the proper citations.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.• The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. 
: 
 and their Tree Search variants) to maximize performance (i.e., accuracy) within a given compute budget. We manipulate the inference computation (FLOPs) of a fixed model by generating additional tokens through the policy model, sampling further candidate solutions, and ranking them with a reward model. We analyze the performance of a family of math-specialized LLMs (i.e., Llemma-7B and Llemma-34B [Azerbayev et al., 2023] ) fine-tuned on the MetaMath dataset [Yu et al., 2023] and measure the error rate on the GSM8K test set [Cobbe et al., 2021a] and MATH500 test set [Hendrycks et al., 2021b , Lightman et al., 2023b] .Our analysis shows that voting-based methods generally outperform the strategy which selects the best solution (i.e., Best-of-N), and weighted voting has the most favorable results (Section 4.3, Figure 5 & 6). However, neither method shows a desirable behavior at high levels of compute. For instance, weighted voting saturates when sampling more than 128 solutions (Figure 1 ). 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 votes. To address this issue, we propose a novel tree search algorithm, REward BAlanced SEarch (REBASE), which pairs well with weighted voting and improves the Pareto-optimal trade-off between accuracy and inference compute. The key idea of REBASE is to use a node-quality based reward to control the exploitation and pruning properties of tree search, while ensuring enough candidate solutions for voting or selection.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 (Figure 1 ) or GSM8K (Figure 4 ). These findings underscore the advantages of using smaller models with advanced inference-time algorithms on end-devices to improve problem-solving. 
 
 Appendix: A.1 Proof of Theorem 1Proof. Suppose the possible answers of the LLM arex 1 , x 2 , x 3 , • • • , x |A| , with π(x 1 |d) > π(x 2 |d) > • • • > π(x |A| |d).After sampling N solutions from the LLM, we denote the occurence ofx i as f (x i ), the probability that x 1 is not the most frequent output isEQUATIONWith Union bound, we getEQUATIONWhere ( 11) is by Chernoff Bound,δ 1 = π(x1|d)-π(x2|d) 2π(x1|d)andEQUATION. As N → ∞, we haveEQUATIONWhere M (x|N ) denotes the probability that majority voting over N sampled solutions gives x. The proof of original theorem is automatically completed by (13).Proof. The proof is similar to the Theorem 1, We rankx 1 , x 2 , • • • , x |A| with R(x 1 )f (x 1 ) > • • • > R(x |A| f (x |A| ).Denotes w(x i ) as the the total weights of answer x i after sampling N solutions. AsN → ∞, w(x i ) → R(x i )f (x i ).Same as proof in theorem 1, we haveEQUATION)Where v(x) = R(x)π(x|d), the remaining proof completely follows Theorem 1.The MCTS process can be represented 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 Where Q(s) represents the quality score of node s, N (s) is the number of visits to node s, and C is a constant determining the level of exploration.EQUATION)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, the updates are made as follows:N (n) ← N (n) + 1, Q(n) ← (N (n) -1) Q(n) + V (s) N (n) .Finetuning We put all the hyperparameters of fine-tuned models in the table 3. We preprocess the MetaMath Dataset to make the solutions in a stepwise format.Inference For all the inference strategies, the temperature of the LLM is set to 1.0. Max tokens for the output is 1024 and max tokens for one step is 256. For REBASE, we chose the balance temperature (the softmax temperature in the REBASE algorithm) as T b = 0.1. For MCTS, we set C in the UCT value as 1 and we expand 4, 8, 16 children for the root, 2 children for other selected nodes with total 32, 64, 128 expansions respectively.In (Fig. 7 , 8 ,9, 10) in this appendix. The experiments show that although the gap between MajorityVoting and Weighted Majority Voting on sampling is huge. This gap becomes much smaller if we apply REBASE. This phenomenon can be caused by the selection ability of tree search like REBASE.Once REBASE already samples solutions with high rewards, conducing weighted majority voting gains less since the sampled solutions may all have relatively high and stable rewards compared with those of sampling.NeurIPS Paper Checklist 


Decision: The paper explores how best to balance LLM model size and the subsequent inference (by RL) to improve accuracy. The method “Reward Balanced Search (REBASE) is proposed In the context of MCTS balances the expansion width among nodes at the same depth based on the rewards given by Process reward model (PRM).

Reviewers were overall  mixed but the majority tended to reject.. The discussion did convinced one reviewer to raise its rank but another reviewer remained highly critical. There was no further discussion beyond the discussion with the authors. The issues raised had to do with how novel is the approach and whether or not the empirical results are impressive enough, if there is comparison to state of the art solvers, or earlier solvers and the scope of the work (focused on Math problems).  Reading the paper myself I think it lacked a clear problem statement and context of what is the underlying problem that is attempted to be solved here.  Is the context solving RL instances? MDPsd? In that sense the paper was not self-contained which made it hard to evaluate.

Reviews: review: paper_summary: paper_summary: This paper investigates the optimal training configurations of large language models (LLMs) during inference. The proposed inference strategy, REward BAlanced SEarch (REBASE), combines the strengths of Monte Carlo Tree Search (MCTS) with reduced inference costs, resulting in improved performance on math-domain tasks.
strength_weakness: strength_weakness: ### Major 

1. Did you take into account the inference cost of the reward model (RM) in your analysis? As the REBASE frequently uses RM to judge the quality of immediate solutions than other sampling strategies, such as, weighted major voting, It's crucial to consider this aspect to provide a holistic view of the efficiency and practicality of your proposed strategy.

2. The base model with post-training techniques such as SFT and RLHF inherently limits the upper bound of performance. It seems that adding more tricks during inference could improve performance, but the marginal effect may result in diminished returns when using models already tuned by the RLHF process. Could you compare the performance gains of REBASE between the base model, the SFT model, and the Chat model? Is the performance gain only significant in models that have not been tuned?

3. In Section 4.2, the observation in "Scaling law of compute-optimal inference" indicates that the optimal inference strategy is invariant to the amount of compute but depends on the model size, i.e., the model's inherent capacity. This raises a concern: does the inference strategy significantly improve the model's performance, or does it only take effect in certain scenarios, such as with base models that have not been aligned?

4. The paper focuses solely on the math domain. To strengthen your claims, a more comprehensive evaluation across general domains using widely adopted benchmarks, such as MMLU, SuperGLUE, HumanEval, etc,  is necessary. 

5. There appears to be no significant improvement in the GSM8K datasets than MATH500 dataset. 

### Minor

1. Figures. 2 and 3 are not referenced in the main manuscript. 

2. Figures. 2 and 3 appear to be in draft form and are somewhat vague.
questions: questions: See Weakness.
limitations: limitations: See Weakness.
paper_summary: paper_summary: The paper presents an approach to select an optimal inference strategy for LLMs and empirical analysis on Math problem solving tasks. The main idea is to select an inference strategy based on a computational budget (FLOPs). The underlying policy model samples solutions by generating tokens based on the budget and a ranking model consumes these tokens. A new reward model is developed  to explore the solution space more effectively. The reward acts as a weighted majority function over the solutions.
Experiments are performed on Math problem solving benchmarks. Some of the key insights from the experiments is that a smaller LLM can outperform the larger LLM in terms of using a smaller computational budget while maintaining similar accuracy. They also show that the proposed approach with a smaller budget has comparable accuracy than sampling with a larger budget.
strength_weakness: strength_weakness: - In terms of the method itself, I was not sure if it is very novel. It seems to be a smaller variation on the tree search methods that search for solutions in the generated space
- In terms of comparisons, I was not sure about the significance of the benchmark, i.e., are there some properties that make the proposed reward reranking more optimal in Llema model specifically (due to the structure of math problems, etc.). In general, since the main contribution of the paper is empirical, I think there should be experiments or discussions different LLMs to make the contribution more significant. 
-Overall, the empirical conclusions seem very tied to the specific benchmarks, so I was a little unsure regarding the significance of the conclusions.
questions: questions: - Is the comparison based on state of the art inference strategies for compute-optimal inference? Specifically, the other methods are all agnostic of the computational limits, so I was wondering if there are other approaches that do take computational limits into account (the related works do not mention any so it is possible there are not)?
limitations: limitations: Limitations regarding the datasets are mentioned.
paper_summary: paper_summary: This paper explores compute-optimal inference for large language
models (LLMs), focusing on designing models and strategies that
balance additional inference-time computation with improved
performance. The study evaluates the effectiveness and efficiency of
various inference strategies, including Greedy Search, Majority
Voting, Best-of-N, and Weighted Voting, across different model sizes
(e.g., 7B and 34B) and computational budgets. Experimental results
indicate that smaller models with advanced tree search algorithms can
achieve a Pareto-optimal trade-off, offering significant benefits for
end-device deployment. For example, the Llemma-7B model matches the
accuracy of the Llemma-34B model on the MATH500 dataset while using
half the FLOPs. These findings suggest that smaller models with
sophisticated decoding algorithms can enhance problem-solving accuracy
across various generation tasks.
strength_weakness: strength_weakness: - Although the paper offers quite thorough experimental analysis, it
  does not look deep in terms of theoretical ideas (although there are
  2 theorems), which may be a problem for a flagship venue like
  NeurIPS.
- Overall findings on the possibility to train an equally accurate
  model with fewer computational resources do not look surprising.
- The paper would benefit from additional proof-reading as there are a
  large number of typos present.
questions: questions: N/A
limitations: limitations: The paper concentrates on mathematical problem-solving tasks using 7B
and 34B models, with findings potentially not applicable to other
domains. Future research should explore a broader range of model sizes
and different training datasets to better understand compute-optimal
inference in mathematical problem-solving.

I should also say that these limitations have been explicitly
discussed by the authors themselves (so not a criticism).

Comments: comment: Thank you for your review. We would appreciate it if you could confirm whether our response has addressed your additional concerns. We have clarified the practical implications of inference scaling, the distinction between PRM in inference and RM in RLHF, and the significance of inference-time algorithms in the last thread. Please let us know if you have any further questions, we hope you might consider increasing the score, thanks!
Thank you for increasing our score. Besides the math tasks, we also added benchmark MBPP in code generation, please see the detailed results in our general response.
Thank you for your response and for taking the time to reevaluate our work. While the rebase method does improve the performances on benchmarks, we would like to clarify that our paper is not primarily focused on proposing new methods. Instead, our key contributions lie in the development of inference scaling laws and the formulation of the inference compute-optimal problem. To illustrate our novel findings, we use the Rebase method, demonstrating that employing smaller models with advanced inference strategies is, in fact, compute-optimal.
Thanks for the detailed response. Based on the response, I think the proposed approach seems to yield a state of the art inference strategy. I still think since the main contribution focuses on engineering a better inference strategy and empirical analysis of LLMs, the type of datasets could be broader as suggested by other reviewers as well. At the same time, the proposed approach could improve the usability of LLMs in specific domains (e.g. MATH understanding). I will increase my score based on the discussions.
Thank you for the feedback. We respond to your follow-up concerns below.

> uncertain about the practical implications of the inference scaling law

Existing research shows that either scaling up model sizes or applying advanced inference strategies improves LLMs’ task performance. However, given an inference compute budget, it is not clear whether using large models with simple inference strategies or small models with sophisticated inference strategies is more favorable. We believe that identifying the right model size and inference strategy is as important as, if not more important than, the research on training scaling law [1] which studies the trade-off between model sizes and training number of tokens. 

Our empirical analysis reveals that using smaller models with advanced inference strategy (like Rebase) is compute-optimal under budget constraints. We note that this is a novel and useful finding which provides important guidance to model deployment on end devices.

[1] Hoffmann, Jordan, et al. "Training compute-optimal large language models." 
 
> It seems more appropriate to apply the RM during the post-training stage rather than the inference stage

We feel that there may be some misunderstandings regarding the RMs. The RM in the RLHF scores the whole LLM output by learning from human feedback. In contrast, Rebase uses the process reward model (PRM) [2-3] which scores each reasoning step and navigates for a better path of reasoning. Thus, **both the training methods and working mechanisms of the two “RM”s in Rebase and RLHF are different.** We will clarify this in our paper.

What’s more, recent research has shown that the RL-aligned (with RM) model with vanilla inference strategy underperforms the model using inference strategies with RMs. [4] shows that RL-aligned Llemma 7b mode with the reward model has 34.0% accuracy (greedy decoding), which underperforms the weight majority voting of corresponding non-RL model(42.7%). [5] shows that after alignment against a reward model, the policy model still benefits from reward-model-based tree search, with the performance improved from 61.84% to 67.04% on MATH dataset. Thus, we believe that **using RM in the post-training / inference stage can be two orthogonal choices** to improve performance.

[2] Lightman, Hunter, et al. "Let's verify step by step." 

[3] Ma, Qianli, et al. "Let's reward step by step: Step-Level reward model as the Navigators for Reasoning." 

[4] Sun, Zhiqing, et al. "Easy-to-hard generalization: Scalable alignment beyond human supervision." 

[5] Chen, Guoxin, et al. "AlphaMath Almost Zero: process Supervision without process." 

> The Inference strategy doesn’t enhance the model’s inherent capability.

Besides enhancing the model’s inherent capability, inference-time improvements are also important since our ultimate goal is to achieve better task performances. Generally learning and inference are complementary methods [6] for improving LLMs task performances. There are many works focused on inference improvement. For example, [7] introduces search during inference, [8] refine the LLMs output to get better results and [9] integrates the LLM’s output with programs to enhance the reasoning capability.

[6] Jones, Andy L. "Scaling scaling laws with board games."

[7] Yao, Shunyu, et al. "Tree of thoughts: Deliberate problem solving with large language models." NeurIPS 2024.

[8]  Madaan, Aman, et al. "Self-refine: Iterative refinement with self-feedback." NeurIPS 2024.

[9] Gao, Luyu, et al. "Pal: Program-aided language models." International Conference on Machine Learning. PMLR, 2023.

> The performance improvements don’t seem significantly better than the vanilla inference strategy.

In our general response, one can see that Rebase leads to huge improvements on the accuracy rate. Specifically, Rebase improves the baseline performance by 5.2%-8% (Llama3-Base) and 2.4%-6.6%(Llama3-Instruct) under different sampling numbers.

In our paper, we compare the Rebase with weighted majority voting in the plots and paper, where both rebase and weighted majority voting have used the reward model. To compare the performance of Rebase and vallina inference strategy (majority voting or greedy decoding), we present the results here:

**Mistral 7B on MATH**

|   |greedy decoding|Maj@256|Rebase@256|
|---|-------|------|------|
|Accuracy|28.6%  |  37.8% |  46.8% |

**Llemma 7b on MATH dataset**
|   |greedy decoding|Maj@256|Rebase@256|
|---|-------|------|------|
|Accuracy|30%  |  41.8% |  50.6% |

**Llemma 34b on MATH dataset**
|   |greedy decoding|Maj@64|Rebase@64|
|---|------|-------|-------|
|Accuracy|33%  |  45% |  51.4% |

Rebase improves over greedy decoding by ~20 points, and naive majority voting by 8 points on average. We believe those performance gains are significant.

We hope our responses address your concerns and you reconsider the score. Thank you for your time!
Thank you for your detailed responses. However, I'm still uncertain about the practical implications of the inference scaling law. It seems more appropriate to apply the RM during the post-training stage rather than the inference stage, as it doesn't enhance the model's inherent capacity. Additionally, the performance improvements don't seem significantly better than the vanilla inference strategy. Thus, I  choose to retain my score.
Thank you for your response.