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

Concern 1: 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.

Our main focus is on formalizing the compute-optimal inference problem, designing a new method, and empirical analysis. To further aid in reasoning about inference strategies in our problem setting, we provided theoretical analysis. We respectfully disagree that these contributions are together insufficient for a venue like NeurIPS. We also follow your suggestion and additionally prove the asymptotic convergence bounds for majority voting and weighted majority voting to understand the performance saturation of the voting methods.

Notations and assumptions. Let $\mathcal{V}$ be a finite vocabulary and $\mathcal{V}^*$ its Kleene closure, i.e., the set of all strings. Given a problem $x$, we say a language model answers $y$ to this problem if the model outputs $r\mathrm{e}y$ where $r\in\mathcal{V}^*$ can be any ``reasoning path'' and $\mathrm{e}\in\mathcal{V}$ denotes a special token that marks the end of reasoning. We further assume that the answer string is always shorter than $L$ tokens, i.e., $|y|\leq L$ for some fixed $L\in\mathbb{N}^*$ where $|y|$ denotes the length of $y$. For a language model $\pi$, denote by $\pi(v|w)$ the probability of generating $v$ given input (prompt) $w$. For a reward model $\rho$, denote by $\rho(v)$ the score it assigns to the string $v$. We use $\mathbb{I}$ to denote the indicator function.

Theorem. Consider a dataset $\mathcal{D}=\\{(x_i, y_i)\\}^m_{i=1}$ where $x_i$ and $y_i$ denote input and true answer, respectively. For a language model $\pi$, denote by $\mathrm{acc}^{\mathrm{MV}}_n (\mathcal{D}; \pi)$ the accuracy on $\mathcal{D}$ using Majority Voting with $n$ samples. Following the notations and assumptions defined above, we have:

$\mathbb{E}\left[\mathrm{acc}_ n^{\mathrm{MV}}(\mathcal{D}; \pi)\right] = \frac{1}{m}\sum_{i=1}^m \mathbb{I}\left[y_i = \arg\max_{|y|\leq L} \sum_{r\in\mathcal{V}^*}\pi(r\mathrm{e}y|x_i)\right] - \mathcal{O}(c^{-n})$ for some constant $c>1$.

For weighted voting, similar convergence results hold, i.e., the accuracy of weighted voting also saturates at the speed of $\mathcal{O}(c^{-n})$ where $c>1$. We will include these results and proofs in the paper revision.

Concern 2: Overall findings on the possibility to train an equally accurate model with fewer computational resources do not look surprising

There might be some misunderstanding, since we study the compute-optimal inference strategy in this paper (rather than training). Our main findings are: 1. Generally more computation at inference time leads to higher performance, but finally saturates at some point. 2. Using the same training dataset, smaller models can be more compute-optimal than larger models during inference. They typically achieve comparable performance with less computation. 3. While typical tree search algorithms like MCTS are not compute-optimal (in that they may improve performance, but by using much more computation), we found that our proposed new tree search algorithm REBASE can get higher performance with less computation than sampling. This gives the compute-optimal inference strategy: using a smaller model with a sophisticated tree search algorithm (REBASE).

Concern 3:The paper would benefit from additional proof-reading as there are a large number of typos present.

Thanks for pointing it out! We have fixed some typos and will do more proof-reading for the final version.

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

Concern 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.

We didn’t take into account the inference cost of the reward model, because when measuring the inference computation, the cost of the RM is negligible compared to the policy model. While the policy model generates a sequence of tokens in an autoregressive way, the reward model only runs one forward pass to get logits for calculating the score. The dominant part of computation is the decoding process, hence we can only take that into consideration.

Concern 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?

All of the models in our current paper have been tuned on the GSM8k and MATH training set (i.e., the MetaMath dataset). In order to see how scaling laws of inference vary based on the post-training methodology (or lack thereof), we conducted additional experiments using Llama3-base and Llama3-Instruct on the MBPP benchmark, with the results shown below.

Llama-8B-Base

Llama3-8B-Instruct

We can see that when using REBASE as an inference algorithm, Llama3-8B-Base gets more improvement (7.4%, 7.6%, 8% and 5.2% compared to sampling 8, 16, 32, 64 respectively) than the Llama3-8B-Instruct (6.6%, 3%, 3.4%, 2.4%). This is similar to the “weaker models gain more from the tree search” observation in our paper (section 4.3). However, the performance of Llama3-8B-Instruct–which has undergone several post-training techniques (e.g., SFT, RL)--still significantly improves when using more samples, and the tree-search-based REBASE still outperforms vanilla sampling. Overall thank you for giving this suggestion, which encourages us to conduct extra experiments to strengthen our conclusion that weaker models tend to benefit more from the tree search method.

Concern 3: 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?

In the paper, all the models (Llemma and mistral) are SFT-aligned. To further address your concern, we provide new experiment results showing that the inference strategies also take effect in RL-tuned models like Llama3-Instruct-8B.

Concern 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.

We conduct additional experiments on MBPP to show our findings are also applicable to code generation. Rebase consistently outperforms the sampling method, the accuracy raises from 62.8% to 68% and 79% to 81.4% for Llama3-base and Llama3-Instruct when sampling 64 solutions.

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

This is because GSM8K is much easier than MATH500, so the error rate ranges of these two datasets are different. In Table 2, we can see that GSM8K’s error rate is typically around 10-15%, so 0.7-1.9% are already relatively significant improvements. For MATH500, the error rate is around 55-60%, so that we can achieve an absolute accuracy improvement of 2.6-5.3%.

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

Concern 1: 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

Our emphasis in this work is on formulating and studying a new setting of compute-optimal inference. As part of this, we design a new tree search algorithm that performs well in terms of accuracy and compute budget. This results in two kinds of novelty:

• First, we are the first to formulate the compute-optimal inference problem (at least in the context of LLM problem solving). Previous research on inference methods for problem solving mainly targets accuracy improvements and pays little attention to the additional increase in compute budget.
• Second, we propose the REBASE method which combines the merits of MCTS and sampling in a novel way. REBASE combines a highly parallel tree search with a new branch cut and exploration mechanism. This results in performance improvements from the new exploration mechanism, with a cost comparable to sampling. While on the surface it may appear as a variation of existing tree search methods, it is nontrivial to design such an algorithm and achieve a good accuracy-cost tradeoff. We are not aware of other methods that have done so, and/or used the same mechanisms that we propose.

Concern 2: 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.

MATH and GSM8K are the two most widely used benchmarks for evaluating LLM math abilities. Although it is common in recent papers to use only these two benchmarks, to address your concern we additionally provide results on a code generation task (the MBPP benchmark) with different LLMs. Please check the detailed results in the general response.

In our original paper, we discussed various models, including Lemma models and Mistral, which have entirely different architectures. For the new code generation task, in addition to presenting the results from Llama3-8B, we have included the results from CodeLlama-7B-Instruct to demonstrate the effectiveness of REBASE across different LLMs.

CodeLlama-7B-Instruct

Question: 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)?

We use state-of-the-art inference strategies (for instance, weighted majority voting and MCTS have been used to achieve state-of-the-art accuracy in problem solving tasks [1-4]), but a key contribution of our paper is pointing out that these methods may not be optimal when cost is taken into account. Since this compute-optimal setting is new when it comes to problem solving, we are not aware of a “state-of-the-art” for compute-optimal inference. Instead, we use strong, commonly used inference strategies, analyze their performance tradeoffs, and show that designing a tree search algorithm with cost in mind can lead to a better tradeoff. The comparison of our Rebase with weighted majority voting and MCTS is shown in figure 1. We also present an example comparison of the Llema 7B performance on the different inference strategies here:

[1] Wang, Xuezhi, et al. "Self-consistency improves chain of thought reasoning in language models." arXiv preprint arXiv:2203.11171 (2022).

[2] Zhang, Shun, et al. "Planning with large language models for code generation." arXiv preprint arXiv:2303.05510 (2023).

[3] Liu, Jiacheng, et al. "Making ppo even better: Value-guided monte-carlo tree search decoding." arXiv preprint arXiv:2309.15028 (2023).

[4] Tian, Ye, et al. "Toward Self-Improvement of LLMs via Imagination, Searching, and Criticizing." arXiv preprint arXiv:2404.12253 (2024).

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