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

We thank you for taking the time to review our paper and are very glad to see you like our paper. We appreciate that you find our paper to be insightful and our experiments to be comprehensive. We have added additional experiments on model Llama-3-Instruct and on coding task MBPP to address your concerns. We include the discussions, figures and results in appendix D and will move them to the main part in later revisions.

Regarding generalization to the tasks beyond mathematical reasoning. Thank you for your constructive comment. We additionally conduct experiments on the code generation task MBPP using the Llama3-Instruct 8B model. The results are shown below:

Llama3-Instruct 8B on MBPP:

The above results are consistent with our findings on compute-optimal inference: Well-designed inference algorithms like REBASE surpasses widely used sampling methods, indicating the large potential of advanced inference algorithms for better inference scaling on diverse tasks.

Regarding the sensitivity to model architectures:

Thank you for the question. Our original experiments have tested three different model families with different architectures, i.e., Pythia, Mistral, and Llemma. We further extend our study to Llama3 and our additional experimental results are shown below:

Llama3-Instruct 8B MATH:

Llama3-Instruct 8B GSM8K:

We find that experimental results on Llama3 share similar trends with our original observation. In particular, REBASE still offers non-trivial performance gains on the Llama3 model, especially on the challenging MATH dataset. We are grateful for your suggestion, which helps us better demonstrate the robustness of our findings and keep the research up-to-date.

Regarding the presentation of the figures. Thank you for the feedback. We will update those figures for better clarity.

We sincerely hope that our responses address your concerns and you reevaluate our work based on the responses. Thank you again for your time!

Why is test error reported instead of task accuracy? Please report the test accuracy for the sake of comparison to other work.

We report test error following the GPT4 technical report [1] and existing works on training scaling laws [2,3], showing the error (inference) / loss (training) decreases with more computation. We agree that reporting test accuracy can facilitate comparisons to other works. For this sake, we add a table reporting the test accuracy in the appendix in our updated manuscript. We hope this addresses your concern.

[1] Achiam, Josh, et al. "Gpt-4 technical report." arXiv preprint arXiv:2303.08774 (2023).

[2]Hoffmann, Jordan, et al. "Training compute-optimal large language models." arXiv preprint arXiv:2203.15556 (2022).

[3]Kaplan, Jared, et al. "Scaling laws for neural language models." arXiv preprint arXiv:2001.08361 (2020).

Can you report the number of samples you were able to draw from each model / inference method corresponding to the number of FLOPs used?

In all the figures, the datapoint represents $2^i$ samples, where integer i starts from 0. For 7b models, our results cover from 1 sample to 128 or 256 samples. For the 34B model we have covered results of 1 to 64 samples. We have also added the clarification in the revision.

We sincerely hope that our responses address your concerns and you reevaluate our work based on the responses. Thank you again for your time!

Some results seem insignificant. Figure 7 shows a lot of overlap amongst the methods tested. This is especially the case for Mistral-7B (which is perhaps the strongest model tested). This overlap exists in Figure 6 as well; albeit, less so. Again, this makes me concerned about how much the results extend to better models. And, in Figures 4 and 5, ReBASE with the larger models doesn’t seem to be significantly better than sampling.

Overlaps typically happen in the leftmost parts of the figures. In each scaling curve, the result begins from only one sample. Thus, these leftmost parts indeed correspond to comparing different methods using only 1-4 sample(s). In such cases, tree search approaches almost degenerate into naive sampling, and overlapped curves are expected.

We agree with you that REBASE offers larger performance gains on MATH (Fig. 6) than GSM8K (Fig. 7). We believe that REBASE can be more advantageous on more challenging datasets. To elaborate, the average length of the model-generated answers on GSM8K is $2\times$ shorter than that on MATH. This indicates that solving problems on MATH requires longer reasoning chains, i.e., more steps, compared with GSM8K. Our results show that REBASE is more advantageous under scenarios where more reasoning steps are required, demonstrating the effectiveness of the search strategy. This finding also suggests that the notion of compute-optimal inference can depend on the data property.

In Figs. 4 & 5, we believe that REBASE consistently shows decent performance gains. 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. Our additional experiments on Llama3 above also exhibit the same trend, demonstrating the robustness of the finding.

What is the difference between ReBASE and beam search with sampling turned on? They appear to be quite similar.

Beam search expands to an equal number of next nodes given current partial solutions and retains a fixed number of highest solutions. Each partial solution is explored equally. In contrast, REBASE is encouraged to explore the partial solutions with higher rewards and prune less promising ones.

Besides, to generate $N$ samples with beam search, in each iteration, $NK$ nodes will be expanded where $K$ denotes the beam size. In contrast, REBASE will only explore at most $N$ nodes, keeping the tree width constant by careful pruning.

Therefore, REBASE can be more computationally efficient and more accurate compared with beam search.

Not a question, but I think that you may want to discuss the “Dot-by-dot” paper (https://arxiv.org/abs/2404.15758) since it also concludes that generating more tokens can be helpful; albeit their setup is much different from yours.

Thank you for the reference! It is indeed a nice paper which studies an interesting approach of scaling up inference compute and makes nice connections to complexity theory. We believe the “dot-by-dot” method advances scientific understanding of language model inference. However, from the angle of compute-optimal inference, it may not be the most favorable approach, because standard chain-of-thought reasoning is as powerful as, if not more powerful than, the “dot-by-dot” method. We have added this paper in our references.

In the expression denoting the reward of a node r_k (last line of page 6), what was q?

Thank you for pointing out the typo. Here, $q$ should have $x$ (the input question) defined in line 320. We have already fixed this in our updated manuscript.

In Figure 5, what happens if you continue sampling from the 7B model with ReBASE until the maximum number of FLOPs considered (i.e., if we extend the orange line all the way to the right, what does it look like)?

This is a good question. We have additionally tested the 7B model in this setting and updated the figures in our manuscript. We find that the model accuracy can continually improve with more FLOPs. However, the best configuration in this setting is the 34B model with the REBASE algorithm.

We thank you for the time for reviewing and the constructive feedback! We are glad to see you like our paper and we appreciate that you have found the topics in our paper attractive and important. We address your concerns below.

Mathematics Tasks Only. The experiments only study two datasets, both of which are related to mathematical reasoning. As such, I’m concerned about the generalizability to new domains and new tasks. I think this paper could be strengthened with experiments on additional tasks like question answering and/or its variants.

We add results on a code task, MBPP, to show that REBASE is also compute optimal compared to sampling on the code benchmark MBPP.

Llama3-Instruct 8B on MBPP:

Limited models tested. The set of models tested leaves something to be desired. In my understanding and experience, Pythia models tend to be weak in comparison to newer open-source models, such as the Llama3.1 family of models. Lemma is a special purpose model for the mathematical domain; Mistral-7B is generally understood to be a strong model. As such, I’m a bit worried that the conclusions would not extend to better models. As is reported, the proposed method, ReBASE, tends to help the weakest models and might not provide the same advantages for a model like Llama3.1-7b-instruct.

We additionally conduct experiments using Llama-3-8B-Instruct, and report the accuracy and update the figures in the paper. The results are averaged over multiple independent experiments. We also add the figures and discussions in the appendix D and will move them to the main part in later revisions.

Llama3-Instruct 8B MATH:

Llama3-Instruct 8B GSM8K:

We find that experimental results on Llama3 share similar trends with our original observation. In particular, REBASE still offers non-trivial performance gains on the Llama3 model, especially on the challenging MATH dataset. We are grateful for your suggestion, which helps us better demonstrate the robustness of our findings and keep the research up-to-date.

Dear Reviewer,

Thank you again for your time and effort in reviewing our paper. If you could spare a moment to review our response and share any further comments, it would be greatly appreciated!

Warm regards,

Authors of Paper 6864

Thank you for your thorough responses and providing additional clarity on the various points/questions I raised and making corresponding updates to your manuscript. I sincerely appreciate your running of additional experiments!

As you mention, in your updated experimentations, the results seem most significant for MATH (and MBPP) and less so for GSM8k. The results seem to hint at length/complexity of the task affecting the magnitude of the improvements you're seeing. I believe these results will help to strengthen your paper and alleviate some doubts related to stronger models.

The experiments in Section 4.3 show inconsistent performance trends of REBASE compared with sampling methods on the MATH and GSM8K datasets. The authors attribute this inconsistency to differences in dataset difficulty, which seems unconvincing. The authors should provide more credible explanations.

We find that the average length of the model-generated answers on GSM8K is $2\times$ shorter than that on MATH. This indicates that solving problems on MATH requires longer reasoning chains, i.e., more steps, compared with GSM8K. Our results show that REBASE is more advantageous under scenarios where more reasoning steps are required, demonstrating the effectiveness of the search strategy.

To further justify this claim, we additionally conduct evaluations on the hard problems (levels 3-5) in MATH based on weighted majority voting. The results are shown below:

The results show that REBASE has stronger performance particularly on hard problems. For example, the sampling accuracy of Llemma 34B does not improve when increasing the sample size from 32 to 64, while REBASE shows better scaling trends. We hope the additional results make our claim more convincing.

In the scenarios where REBASE is purportedly advantageous (i.e., less compute-intensive situations), the performance on the MATH dataset with Llemma-34B and Mistral-7B models is either weaker than or merely comparable to the Sampling method. This indicates that REBASE does not consistently achieve improvements. An explanation for this should be provided.

We believe that you are referring to the leftmost parts of the figures. In each scaling curve, the result begins from only one sample. Thus, these leftmost parts indeed correspond to comparing different methods using only 1-4 sample(s). In such cases, tree search approaches almost degenerate into naive sampling. The FLOPs budget is so limited that there is little room to compare different algorithms.

More meaningful comparisons can be observed in the middle ranges in these figures, where the budget is still limited but offers decent room for applying advanced algorithms. In this regime, we see that REBASE performs consistently better than sampling, e.g., in Fig. 6.

The paper only uses two baselines: sampling and MCTS, without elaborating on their implementation details. It would be beneficial to include the latest baselines and provide detailed introductions and comparisons of the baselines and the proposed REBASE method.

Thank you for the comment. As has been made clear in the keywords and abstract, the primary focus of our paper is inference scaling laws and compute-optimal inference, rather than designing the world’s best mathematical reasoning algorithms. As the initial step in this direction, we begin with systematic theoretical and empirical evaluations of simple methods. We are also aware of many recent works on new methods for mathematical problem solving with LLMs. Studying them in the context of inference scaling laws and compute-optimal inference can be interesting directions for future research.

For the baselines in our work, we have documented them in Sec. 3 and the Appendix. We plan to release our full code upon paper acceptance. We are also happy to include more details in the Appendix if the reviewer finds it helpful.

We sincerely hope that our responses address your concerns and you reevaluate our work based on the responses. Thank you again for your time!

While the paper employs a PRM as the reward model within REBASE, it fails to detail the training and inference processes of the reward model. Moreover, it does not address how to ensure the reliability of the rewards generated by this model. An unstable or suboptimal reward model could significantly impact the overall framework, and the authors should consider this aspect.

Thank you for the comment. We report the training hyper-parameters in Appendix C. More specifically, we finetune Llemma 34B on the Math-shepherd to obtain the PRM [1]. Math-shepherd is a popular dataset with step-wise labels on MATH and GSM8K train split. Following [1], the PRM is trained by predicting the step-wise reward through the logits of the step tag tokens.

Our PRM is trained in a standard way following prior works. While we agree that the quality of reward models can be critical, our work does not aim at optimizing reward model training methods since it is orthogonal to our focus on inference scaling laws and compute-optimal inference. Our results show that building upon existing datasets and reward model training techniques already demonstrates better inference scaling behavior than naively scaling up model sizes, which validates our main claim.

[1] Wang, Peiyi, et al. "Math-shepherd: Verify and reinforce llms step-by-step without human annotations." Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2024.

The paper provides theoretical analyses for four inference strategies: greedy search, majority voting, best-of-n, and weighted voting, but does not offer comparative experimental results. A detailed experimental comparison involving different model sizes and inference strategies should be included.

All the figures in our experiment sections are comparative studies from various angles. We note that there are multiple varying factors in our study: FLOPs, model sizes, and inference strategies, resulting in hundreds of possible configurations in total. To conduct meaningful comparisons, we present controlled experiments, presenting both comparisons among different inference strategies with fixed model sizes and comparisons among different model sizes with fixed inference strategies.

In our paper revision, we follow your suggestion and add a table in the appendix, including the results with varying model sizes, inference strategies on both GSM8K and MATH datasets. We hope this addresses your concern.

Similar to traditional MCTS methods, REBASE necessitates extensive intermediate interactions with LLMs before obtaining the final generated results. Additionally, it introduces the extra PRM component. This could lead to significant resource consumption. The authors should provide data analysis on resource usage and runtime to demonstrate whether REBASE is more efficient than conventional MCTS methods.

In all our presented results, we evaluate the models on the same devices (A6000 GPUs with 48GB memory), and analyze the model test error given FLOPs consumption. Therefore, all the comparisons have taken both model quality and resource usage into account. For example, Fig. 4 clearly demonstrates REBASE shows better cost-performance trade-off than MCTS. We also point out that reward generation only requires a forward pass on the reward model, which is much more efficient than autoregressive generation and does not contribute to the computational bottleneck.

According to the experiments in Section 4.2, both small and large LLMs show high error rates under limited computation scenarios, suggesting that both models are not good enough in these situations. Claiming that smaller models perform better in such conditions is not sufficiently convincing.

Our claim is that within a fixed compute budget, using a small LLM can potentially outperform the Larger ones since it can generate more samples. We believe our experiments justify this claim. First, in the experiment, the overall trend is consistent: Smaller models outperform larger ones under the same compute-budget when the budget is relatively limited. Second, our experiments in Section 4.2 demonstrate an accuracy of at least 25% on MATH and 68% on GSM8K, with peak performance reaching approximately 50% on MATH and 90% on GSM8K—results that are non-trivial achievements for language models with simple methods and limited compute. Therefore we believe our comparisons are meaningful and the results are well justified.

Thank you for taking the time to review our paper. We are glad to see that you find our work studying the different inference strategies and model sizes with various compute budgets is meaningful. We answer each of your questions below.

The paper lacks technical contribution. The proposed REBASE framework merely integrates process reward models (PRM) into the existing Monte Carlo Tree Search (MCTS) framework. Considering these components are well-established and the simplicity of the REBASE framework, there are concerns regarding the novelty of the work.

We feel that there are some misunderstandings on the scope and contributions of our paper. Besides the REBASE algorithm, we would like to emphasize that our main contributions also include rigorous formulation, theoretical characterization, and empirical evaluation of inference scaling laws and compute-optimal inference, which have not been formally studied in any prior works and contributed to the main novel technical contributions. To reiterate, our novel insight leads to the following observations:

• Smaller models can outperform larger ones under the same compute budget with more samples, and the compute-optimal inference exhibits a scaling law in model size.
• Sampling and voting methods have performance limits and diminishing returns (which we justified theoretically in Sec. 3.1 and empirically in Sec. 4).
• Well-designed inference algorithms like REBASE further supasses widely-used sampling and MCTS methods, indicating the large potential of equipping smaller models with advanced inference algorithms for better inference scaling.

We also clarify that REBASE is not an integration of process reward models (PRM) and Monte Carlo Tree Search (MCTS) - it is indeed drastically different from MCTS and significantly better than MCTS given compute budgets.

Finally, we point out that the goal of designing REBASE is not presenting a very complicated tree search method that surpasses any other algorithm. Rather, it is positioned as a part of investigation on compute-optimal inference, demonstrating that applying well-designed inference strategies is more favorable than scaling up the model size.

Dear Reviewer,

Thank you again for your time and effort in reviewing our paper. If you could spare a moment to review our response and share any further comments, it would be greatly appreciated!

Warm regards,

Authors of Paper 6864

Thank you for the time for reviewing, we are glad to see you like our paper and find our observations on the inference scaling laws significant. We conduct further experiments to address your concerns and answer your questions below.

The authors analyzed inference scaling laws using Pythia and Llemma but did not include REBASE experiments for Pythia, opting instead for Llemma and Mistral-7B. Including REBASE experiments on Pythia could strengthen the findings.

We use Pythia models to demonstrate the inference scaling law: Smaller models could outperform larger ones by generating more samples since Pythia models include many different sizes and are relatively small, which could be analyzed given our compute resources. For studies on compute-optimal inference, we used Llemma and Mistral models, which are stronger and more suitable for tree search since they perform well on both MATH and GSM8K. In contrast, Pythia is too weak for the MATH dataset. We also added the Llama-3 8B Instruct model to strengthen our conclusions.

The paper focuses on mathematical datasets, yet the title "LLM Problem-Solving" suggests a broader scope beyond math. The methods and analysis in Section 3 do not specifically mention "Math", indicating potential applicability to other areas. Exploring additional problem-solving domains, like coding or logical reasoning, could enhance the study's relevance and impact.

Thank you for your valuable suggestion. We follow your suggestion and additionally experiment on a code generation task, MBPP, with Llama3-Instruct-8B models. The results show the potential of using compute-optimal inference methods in other domains, e.g., coding.

Llama3-Instruct 8B on MBPP:

While the models Pythia, Mistral-7B, and Llemma have been around for some time, recent advancements with models like Llama3 show significant progress. Analyzing these newer models could provide valuable insights and keep the research up-to-date.

Thank you for your comment. We have added the Llama3-Instruct 8B results in the paper appendix D and also show the results here for your reference.

Llama3-Instruct 8B MATH:

Llama3-Instruct 8B GSM8K:

We find that experimental results on Llama3 share similar trends with our original observation. We are grateful for your suggestion, which helps us better demonstrate the robustness of our findings and keep the research up-to-date.

We sincerely hope that our responses address your concerns and you reevaluate our work based on the responses. Thank you again for your time!

Dear Reviewer,

Thank you again for your time and effort in reviewing our paper. If you could spare a moment to review our response and share any further comments, it would be greatly appreciated!

Warm regards,

Authors of Paper 6864