Adaptive Inference-Time Compute: LLMs Can Predict if They Can Do Better, Even Mid-Generation

Q5: Could you consider extending the method to learn to detect whether a given generation is the best among N>2 generations?

A5: Absolutely! Preference datasets usually only have two responses per query but this is a great suggestion, we could definitely extend this method to consider the probability that N>2 number of generations will result in at least one better generation. This is an interesting way of augmenting the dataset, as well as producing potentially more useful predictions to more efficiently allocate samples to a given query. We can try to incorporate this for the final version of the paper and will add this to the discussion section of our paper for completion.

Q6: Can you provide ablation studies, such as applying temperature annealing to basic Best-of-N?

A6: Thank you for the suggestion. We find that our annealing schedule method (annealing based on samples generated so far) seems to provide gains even without adaptive sampling at small Best-of-Ns (+2% on GSM8K for Best-of-2). However, benefits taper off at larger Best-of-Ns. This makes sense as our scheduling strategy only applies significant annealing to the first few samples. We will introduce this and additional ablations for the final revision of the paper.

Q7: Could you explore a simple strategy that detects prompt complexity and allocates compute accordingly?

A7: This is a great suggestion. However, practically, we had found limited success with this approach as estimating the prompt complexity is quite difficult to use to then predict the required amount of compute given only the query. This is what motivated the experiments with mid-generation self-evaluation (early pruning), allowing the model to perform a limited amount of inference (e.g 64 or 128 tokens) and then prune according to the complexity of the problem. We will include this as a formal baseline for the final revision of the paper.

Q8: Could you clarify why you need the model's prior distribution over tokens to be similar after training to keep the underlying policy unchanged? Isn't there a risk of catastrophic forgetting when fine-tuning for self-evaluation?

A8: Thank you for bringing this up! This is a small but important detail in practice. We need the model’s distribution over tokens to be similar before training, as this will minimize the amount the model needs to change to minimize the loss during training. This could be thought of as reducing the risk of catastrophic forgetting, but more concretely, we do not want the model to overgeneralize and learn to only output the target tokens (“Yes” and “No”). We found that If the self-evaluation prompt elicits these tokens by default (this should be verified before training), the model does not overgeneralize and we effectively avoid this problem. One way to check if this was done correctly is to make sure that the loss is already reasonable before fine-tuning.

We thank the reviewer for their comments and for engaging with the paper. To address the concerns, we add a more complex reasoning domain, Math 500 and add additional clarity to the computational efficiency of our algorithms with the measure of the number of FLOPs. We are happy to clarify any questions you may have. Please let us know if your concerns are addressed and if so we would be grateful if you would be willing to raise your score. We would be happy to discuss if you have any concerns

Q1: Are the improvements over baselines small?

A1: Thank you for your question. Please see the sections “What are we improving?” and “How much are we improving it by?” in the overall response.

In brief, self-evaluations and early pruning reduce costs (computation, memory, and energy consumption) by a factor of 4 without introducing additional latency. Adaptive sampling further reduces costs by approximately a factor of 6 to 8, though it comes with the tradeoff of increased latency. Additional details on these measurements are included in the overall response.

Q2: Why doesn't the paper mention process-based reward models?

A2: Thank you for bringing this up and for the helpful reference! Note, we do not intend to improve the performance of the given inference-time method, which we fix in this work to be Best-of-N. We simply aim to make it more cost-effective by removing the need for an external reward model as well as adaptively reducing the amount of computation needed depending on the query.

In this work, we learn reward models with an on-policy preference dataset, utilizing prompts from a large open-source general dataset, LMSYS. In contrast, PRMs in the context of reasoning require the need for ground truth outcome supervision. For domains such as Math and Code, supervision can be provided by symbolically checking the final answer in the context of math or executing test cases in code. A standard paradigm explored in works such as Snell et al, 2024 and OmegaPRM (Luo et al 2024) is to learn a PRM through MC rollouts. These are two complementary approaches for training reward models, with different assumptions and different dataset compositions. However, the inference time optimizations we propose such as adaptive sampling + pruning, could be readily applied to any reward model, allowing for savings in FLOPs during inference. We will try to incorporate a PRM for the final version of the paper to showcase the efficacy of this approach. We will additionally add these works to the related works section of the paper and will map out how our method can be extended to other inference-time methods very naturally.

Q3: How does the need to generate sequentially affect real-time deployment, and can you address this limitation more thoroughly? Is the idea of exponentially increased batch size sufficient to mitigate this issue?

A3: Thank you for raising this point. Please refer to the sections “What are we improving?” and “How much are we improving it by?” in the overall response.

In summary, self-evaluations and early pruning reduce costs (computation, memory, and energy consumption) by a factor of 4 without additional latency. Adaptive sampling further reduces costs by a factor of 6 to 8 but introduces approximately 2x additional latency, which can be mitigated using exponentially increasing batch sizes. More details on these measurements are available in the overall response.

Q4: Can you discuss how your training method effectively trains the model to estimate the quantile of the generation according to the external reward model?

A4: If $\epsilon = 0$ in the construction of the preference dataset (i.e ties only if the response is identical), capability aware self evaluation reduces to estimating the quantile of the given response. However, practically this is difficult to model (the Best-of-N performance suffers), which is one of the reasons we introduced a non-zero epsilon as a relaxation of this interpretation. Intuitively, it is easier to predict with high confidence if the model can’t do significantly better rather than if the model can’t do better by any amount. This relaxation is fine as we do not strictly care about the quantile of a response, but instead if a new response could be significantly better than the current generated response. This is the reason our model can routinely make confident predictions such as P(Win or Tie) = 0.99, indicating that the probability of a significant improvement is very low.

We thank the reviewer for their continued engagement!

You require sequential generations while the standard BoN can be distributed

To clarify, we propose two approaches: early pruning and adaptive sampling. Early pruning does not require any sequential processing of samples, while adaptive sampling does. Specifically:

1. Early Pruning: This approach reduces FLOPs by 4x to match the accuracy of Best-of-16 with no additional latency. It reduces computational overhead by halting the generation of less promising samples during parallel generation. For example, after generating 128 tokens for 16 samples in parallel, the 12 least promising samples can be pruned, effectively halving the total number of tokens generated and achieving significant computational savings.

2. Adaptive Sampling: This approach achieves a larger reduction in FLOPs by 6-8x to match the accuracy of Best-of-16, albeit with additional latency due to sequential generations. However, the average latency can be as low as 1.1 sequential generations, thanks to further optimizations such as the exponentially annealed temperature schedule, that we propose, resulting in limited additional overhead over a batch of questions.

PRM as a baseline

We note that the Process Reward Model (PRM) can readily be replaced with the self-evaluation function we propose. The inference speedups we present (adaptive sampling + pruning) remain valid and transferable with this function.

The key difference between a PRM and our mid-generation self-evaluation approach lies in the dataset composition used to train these functions. PRMs (e.g., Snell et al., 2024; Wang et al., 2024, Math Shepherd) are trained on policy rollouts evaluated for correctness via outcome verification. In contrast, our self-evaluation frameworks collects the same on-policy rollouts but utilizes relative preference feedback instead of absolute feedback. This decision to train on general queries from preference data (e.g., LMSYS) enables our single model to generalize across a wide range of domains.

In many domains, learning an accurate PRM is challenging due to the need for highly accurate outcome supervision models. Additionally, for some domains such as creative writing or general question answering, such supervision can be ambiguous and difficult to collect effectively. Moreover, inference with an external PRM function is computationally expensive, requiring additional FLOPs and memory, particularly for mid-generation execution.

Nevertheless, for the rebuttal/final version of the paper, we will attempt to train a PRM for math contest problems using domain-specific datasets, such as PRM800K or Math Shepherd. Additionally, we are open to rephrasing the title to temper any overstated claims.

We thank the reviewer for their comments and for engaging with the paper. To address the concerns, we add a more complex reasoning domain, Math 500 and add additional clarity to the computational efficiency of our algorithms with the measure of the number of FLOPs. We are happy to clarify any questions you may have. Please let us know if your concerns are addressed and if so we would be grateful if you would be willing to raise your score. We would be happy to discuss if you have any concerns.

Q1: Why did you choose the LMSYS dataset for fine-tuning, and how does the selection of the fine-tuning dataset affect the generalization of the model? Would using a different or more complex dataset improve performance?

A1: Thank you for your question. We selected the LMSYS dataset for fine-tuning because it consists of real-world prompts. Our goal was to enable the model to adaptively allocate samples across a wide range of domains, and the broad distribution of prompts in LMSYS supports this objective.

We evaluated the fine-tuned model on diverse domains, including AlpacaEval, GSM8K, and now Math 500, which demonstrates its generalization capabilities. Since our performance already closely approaches the theoretical upper bound set by the underlying reward model used in constructing the preference dataset, we believe that training on a different dataset would likely yield minimal improvement. However, fine-tuning on more specific, complex datasets could enhance performance on similar in-distribution tasks, and this is something we may explore in a future revision.

Q2: Can you evaluate your method on more diverse and complex datasets, especially those requiring longer reasoning chains, to demonstrate its generalizability on reasoning tasks?

A2: To address this concern, we added the Math 500 (Hendrycks et al, 2020) as an additional domain as seen in our new figure which can be found here. This domain is more complex, requiring longer reasoning chains to perform harder math contest problems. We find that the results are relatively consistent with what we see on GSM8K.

Q3: Can you show the distribution of generated token lengths?

A3: Certainly. Below are the token length statistics for the evaluated datasets:

For AlpacaEval 2.0, the mean response length is 453 tokens with a standard deviation of 305 tokens. For GSM8K, the mean response length is 251 with a standard deviation of 170 tokens. For MATH 500, the mean response length is 537 tokens with standard deviation of 337 tokens. We will add these metrics and a histogram of the generated tokens to our appendix by the end of the rebuttal period or for the final version of the paper.

Q4: The paper maps the number of samples or tokens directly to computational cost, but in practice, batch inference can make multiple samples cost similar to one sample. Can you address how your method impacts computational cost and latency in real-world LLM serving scenarios?

A4: Thank you for this important question. Please refer to the sections “What are we improving?” and “How much are we improving it by?” in the overall response. In those sections, we provide concrete metrics and explanations on the cost-effectiveness of our approach, including its impact on computational cost in practical settings.

Q5: Considering that inserting a prefill (self-evaluation prompt) is not free and may introduce more latency than naive batch decoding, can you clarify how your method affects end-to-end latency?

A5: We appreciate your question. Please see the section “How much are we improving it by?” in the overall response. There, we provide concrete metrics and a detailed explanation of the tradeoff between the cost of a reward model and the self-evaluation process, as well as their respective impacts on end-to-end latency.

We sincerely thank the reviewer for their comments and for engaging with the paper. To address the concerns, we add a more complex reasoning domain, Math 500 and add additional clarity to the computational efficiency of our algorithms with the measure of the number of FLOPs. We are happy to clarify any questions you may have. Please let us know if your concerns are addressed and if so we would be grateful if you would be willing to raise your score. We would be happy to discuss if you have any concerns.

Q1: How does the proposed method compare with the baselines in terms of compute time and FLOPs, considering the cost of probing the model mid-generation?

A1: Thank you for this question. Please see the sections “What are we improving?” and “How much are we improving it by?” in the overall response for detailed insights.

In brief, self-evaluations and early pruning reduce costs (computation, memory, and energy consumption) by a factor of 4, without introducing any additional latency. Adaptive sampling further reduces costs by approximately a factor of 6 to 8, though it comes with the tradeoff of additional latency. Additional details on how these improvements are measured are included in the overall response.

Q2: Since pruning longer sentences improves performance, wouldn't it be better to fully generate sentences before selecting them? Why is pruning mid-generation necessary?

A2: This is a good question. Please refer to the sections “What are we improving?” and “How much are we improving it by?” in the overall response.

To clarify, our approach does not aim to improve the maximum theoretical performance of Best-of-N but rather to make Best-of-N far more cost-effective. Specifically, early pruning with self-evaluations reduces costs by a factor of 4. Generating a larger number of responses and pruning some early—rather than generating fewer full responses—is shown to be computationally more efficient. This method can be viewed as improving downstream task performance under fixed compute, memory, or energy budgets. The details of these measurements can be found in the overall response.

Q3: Why does your method hit a plateau while the Best-of-N method still shows improvement over more samples? Would Best-of-N surpass your method when evaluated at 32 and 64 fully generated samples?

A3: Thank you for raising this point. Please refer to the section “What are we improving?” in the overall response for more details.

As shown in Figure 1, we plot accuracy against the number of samples, with a fixed maximum sample budget. Best-of-16 serves as the oracle or upper-bound performance for a given reward model. Our methods, such as Adaptive Sampling and Pruning, do not increase the theoretical maximum performance of Best-of-N or other search algorithms but instead significantly reduce the costs associated with achieving that performance (in terms of computation, memory, and energy consumption). We will evaluate the performance at 32 and 64 samples and will include these results either by the end of the rebuttal period or in the final version of the paper.

Q4: Can you move important ablation studies from the Appendix to the main paper for better understanding of your method's parameters? Can you restructure your figures to present one observation per figure instead of combining many into Figure 1 to improve readability? Could you revise the writing to make it more succinct and avoid long, complex sentences?

A4: Thank you for these suggestions. We will take the following steps to improve clarity and readability in the final revision of the paper:

1. We will add clarity as you have suggested by reducing text and complex sentences in the methods section of the paper.
2. Additionally, we will move Figure 1 to the experiments section of the paper and break it into two separate figures to improve readability.
3. Finally, we will move key Tables and Figures from the appendix to give better understanding of the method's parameters.

We would be happy to make further changes that lead to a better understanding of this paper.

There are competitive methods out there, as highlighted by Pqum, that have not been considered nor compared against in this paper. Such comparison is important for the type of paper presented here.

Thank you for this comment. We now include the baseline the reviewer suggested to the paper. We copy the response below for your convenience:

This is a very helpful suggestion! This allows us to compare our method to a strong value model instead of a domain-specific PRM. Following what you described, we finetuned ArmoRM (the same reward model we used to create our preference data) to predict the reward given truncated generations. This is a very strong baseline, as ArmoRM was already trained on 1 million preferences and we further finetuned it to predict the reward directly instead of having to predict $P(Win \cup Tie)$ as our model is trained to do. Compared to Best of N, we found that this baseline does allows for some saving in FLOPs, with the pruning of less desirable responses. However, this gain is completely overshadowed by the additional cost of having to process all tokens a second time due to the lack of KV cache. This results in our method still saving almost 2x more FLOPs. You can find an updated performance vs FLOPs graph here. We will add these results to our paper in the final revision.

We would be most grateful if you would consider upgrading your score, given that your concerns have now been addressed.

We thank the reviewer for their comments and for engaging with the paper. To address the concerns, we add a more complex reasoning domain, Math 500 and add additional clarity to the computational efficiency of our algorithms with the measure of the number of FLOPs. We are happy to clarify any questions you may have. Please let us know if your concerns are addressed and if so we would be grateful if you would be willing to raise your score. We would be happy to discuss if you have any concerns.

Q1: Are the performance gains marginal for Alpaca Eval?

A1: Thank you for your question. Please refer to the sections “What are we improving?” and “How much are we improving it by?” in the overall response for a detailed explanation.

Specifically, for AlpacaEval, we employed self-evaluations and early pruning to significantly reduce the cost (in terms of computation, memory, and energy consumption) of Best-of-16 by approximately a factor of 4. Additional details on how these improvements are measured can be found in the overall response.

Q2: Can you quantify the actual computational efficiency and latency of your methods compared to Best-of-N, considering that your method requires serial processing while Best-of-N can be executed in parallel? How does your method perform on datasets that require more complex reasoning?

A2: Thank you for this suggestion. The section “How much are we improving it by?” in the overall response contains the necessary details that we will add to our final revision.

In summary, self-evaluations and early pruning reduce costs by roughly a factor of 4 without incurring additional latency. Meanwhile, adaptive sampling achieves cost reductions of approximately 6 to 8 times, though it introduces some additional latency as a tradeoff. Further details about these measurements can be found in the overall response.

Regarding complex reasoning tasks, we evaluated the method on MATH 500, a much more challenging dataset, and observed results consistent with GSM8K. This evaluation is also discussed in the “How much are we improving it by?” section.

Q3: Can you study the effect of self-evaluation training in isolation by conducting Best-of-N after self-evaluation training to separate the effect of this training from the adaptive inference method?

A3: To observe the effect of self-evaluation training, please refer to Table 1. The table highlights the difference in performance between zero-shot (LLM-as-a-Judge), reward modeling (Bradley-Terry), and our capability-aware self-evaluations. To clarify, here we utilize the same samples from Llama 3.1 8B-Instruct per question and evaluate solely the reward modeling/self-evaluation capabilities of the model. Here, we show that Best-of-16 performance improves from 24.4% to 33.8% on AlpacaEval and improves from 86.7% to 91% on GSM8K when comparing LLM-as-a-Judge to capability-aware self-evaluation. Additionally, we see 33.2% to 33.8% and 87.7% to 91% in the same domains when comparing a Bradley Terry reward model to capability-aware self-evaluation. In short, the self-evaluation training allows us to provide a significant performance benefit over traditional Best-of-N with a reward model or using the zero-shot ability of the model in the LLM as a Judge formulation.

Q4: What is the performance of Llama 3.1 8B Instruct when using Best-of-16 on AlpacaEval and GSM8K?

A4: As shown in Table 1, with Llama 3.1 8b Instruct, capability-aware self-evaluation gets a Best-of-16 performance of 33.8% on Best-of-16 on AlpacaEval and 91% on Best-of-16 with GSM8K. For the learned reward model, the Best-of-16 on AlpacaEval is 33.2% and the Best-of-16 with GSM8K is 87.7%.

Q5: How does the proportion of ties (the epsilon for deciding ties) in the preference dataset construction for self-evaluation training affect the performance of the model?

A5: This is an excellent question. Many responses in an on-policy preference dataset will be semantically similar as they come from the same model as seen in prior work (Tajwar et al, 2024). Therefore, relaxing the Bradley-Terry Model to account for ties can make the prediction problem easier. Intuitively, with a large proportion of ties, we increase the semantic distance between the winning and losing response, making self-evaluation easier. This is the reason our model can routinely make confident predictions such as P(Win or Tie) = 0.99, indicating that the probability of a significant improvement is very low.