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

Authors: Rohin Manvi, Anikait Singh, Stefano Ermon
Primary area: foundation or frontier models, including LLMs
Keywords: LLMs, inference-time, inference-time efficiency, Best-of-N, self-evaluation

Abstract:
Inference-time computation is a powerful paradigm to enhance the performance of large language models (LLMs), with Best-of-N sampling being a widely used technique. However, this method is computationally expensive, requiring both (1) an external reward model and (2) the generation of multiple samples. In this work, we introduce a new generative self-evaluation scheme designed to adaptively reduce the number of generated samples while maintaining or even improving performance. We use a generative reward model formulation, allowing the LLM to predict mid-generation the probability that restarting the generation will yield a better response. These predictions are obtained without an external reward model and can be used to decide whether or not to generate more samples, prune unpromising samples early on, or to pick the best sample. This capability is very inexpensive as it involves generating a single predefined token. Trained using a dataset constructed with real unfiltered LMSYS user prompts, Llama 3.1 8B's win rate against GPT-4 on AlpacaEval increases from 21\% to 34\% with 16 samples and math performance on GSM8K improves from 84\% to 91\%. By sampling only when the LLM determines that it is beneficial to do so and adaptively adjusting temperature annealing, we demonstrate that 74\% of the improvement from using 16 samples can be achieved with only 1.2 samples on average. We further demonstrate that 50–75\% of samples can be pruned early in generation with minimal degradation in performance. Overall, our methods enable more efficient and scalable compute utilization during inference for LLMs.

Decision: reject
Reviewer scores: 6, 3, 6, 5

Reviews and discussion:
<thread>
creation date:
1730396765624

thread_type:
official_review

<thread_object>
creator: reviewer_pqum

<content>
summary:
The paper proposes a new strategy to improve the inference of LLMs, involving two key components: adaptive sampling (selecting the N in best-of-N as a function of the prompt) and early pruning (pruning unpromising mid-generations). This relies on a new ability from LLMs; given a prompt and a potentially unfinished generation, they can learn to detect whether they will be better with an additional sampling round. This ability is trained explicitly in supervised way, with feedbacks obtained from an external RM; the key point being that this RM is then not required during deployment. This strategy improves performances on math and gsm8k benchmarks.

strengths:
* The paper is clearly written, properly highlighting the different ideas and results.
* Improving inference strategies for LLms is an important topic for the ICLR community, with potential significant impact.
* The two contributions, the discovered "trainable self-evaluations" and its applications (adaptive sampling and early pruning) make sens.
* The experiments, althought succinct (see weakness), show convincing performances on standard benchmarks.

weaknesses:
* The main limitation of the work is that the proposed ideas are arguably small improvements over existing baselines. For example, the paper does not mention "process-based reward models" or the "reward-guided decoding" literature (such as MCTS), that seeks the same objective. The only benefit of the paper would be to include those abilities inside the LLM itself, thus removing the need for an external value network. 

* In terms of practicality, the need to generate sequentially is a key limitation in real-time deployment, thus hindering some of the efficiency benefits. Though, I acknowledge that the idea of "exponentially increased batch size" is an interesting but limited answer.

* With your training, and given a prompt and generation, you actually train the model to estimate the quantile of this generation according to the external reward model; a discussion on this topic would be needed. Then similarly, it would be interesting to consider extension where, be given N>2 generations, your model learns to detect whether the generation is actually the best-of-N.

* Some ablations are missing. For example, could you provide the performances of applying temperatre annealing to basic Best-of-N? What if you train an external RM to evaluate the quantile on mid-generations? Or about a simple strategy that would detect, be given the prompt, its complexity, and then allocate automatically a prompt-aware compute.

* I fully agree that "We need to keep the underlying policy or model unchanged when learning to self-evaluate." (l.232) Though I disagree with the next sentence: "To do so, we need the model’s prior distribution over tokens to be similar to what we expect after training" (l.233). Could you clarify this? 
From my understanding, finetuning for self evalaution should decrease other abilities because of catastrophic forgetting.

## Nit

* The fig1 actually shows the opposite of what the legend "Rewards are not interpretable" states; indeed, actually the reward provides the lowest score to the most unconfident generation.

* Related work: some papers that may be worth discussing
"Don't throw away your value model! Generating more preferable text with Value-Guided Monte-Carlo Tree Search decoding"
"ARGS: Alignment as Reward-Guided Search"
"Efficient Controlled Language Generation with Low-Rank Autoregressive Reward Models"
"Critic-Guided Decoding for Controlled Text Generation"

* The discussion on ties ends up being unecessary, as it does not impact performance. Thus, you might not want to state "Accounting for ties in reward modeling is especially important" (l.243).

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730452919897

thread_type:
official_review

<thread_object>
creator: reviewer_6jzh

<content>
summary:
The paper shows that one of the challenging problem at inference-time computation scaling (Best-of-N) is to identify if a specific generation is "good", or if it is worth allocating more compute to continue its generation. For example, Best-of-N is considered computationally expensive because it requires external reward model to evaluate the quality of the generation, and it often requires multiple such generation samples to obtain a good quality result. 

To address this problem, the paper introduces capability-aware self-evaluation, a technique that allows LLM to self-predict if the current generation is promising, or in other words, if restarting the generation will yield better result or not. In the Best-of-N setting, the paper also introduces adaptive sampling and annealing to further improve the quality of the generation, and early prune unpromising samples using the self-evaluation methodology.

In the evaluation section, the paper presents a finetuned Llama3.1-8B-Instruct model trained on 30k preference pairs constructed from unfiltered LMSYS using a reward model, and shows an increasing win rate against GPT-4.

strengths:
**Elegant solution**. The evaluation section shows that the proposed fine-tuning method is able to generalize from LMSys to AlpacaEval and GSM8K. The solution is elegant and intuitive, and it works well.

The method shows a promising way to improve the quality of generation in Best-of-N setting by early pruning unpromising samples, and adaptively sampling and annealing the remaining samples. The proposed method is simple, intuitive, and shown to be effective.

**Writing**. The writing is clear and the paper is well-organized. The paper is well-written and easy to follow.

weaknesses:
**Finetuning dataset construction - what is a good dataset to finetune on?**. Why is LMSYS dataset chosen? Would another dataset other than LMSYS do better or worse? Would using a more complex dataset do better or worse? It is not as convincing if there is no ablation study on the dataset choice.

**Limited evaluation**. Only 2 datasets are used to perform the evaluation. GSM8K is not a hard dataset to see performance gain (e.g. using hyperparam tuning). To show the generalizability on reasoning tasks, it would be better to evaluate on more diverse and complex datasets, especially on slightly longer reasoning chains to show the increasing size of sampling and prunning is effective. Showing the distribution of generated token lengths would be very useful to understand "how far" are we talking about.

**Limited setting of inference time algorithm**. As mentioned in conclusion / limitation section, the paper only discuss Best-of-N setting. But many reasoning techniques such as self-refinement / chain-of-thought are not generating parallel requests. It would be better to evaluate on more diverse inference settings.

**Naively mapping number of samples / tokens to computational cost**. The paper did not include any statement about computational cost in terms of number of tokens generated. In actual LLM serving with batch inference (e.g. vLLM, DeepSpeed), the computational cost of 16 samples performing decoding is roughly the same as 1 sample performing decoding. Thus, if the metric is "end-to-end latency", then the bottleneck is actually the straggler - meaning what is the longest request in the batch that needs to be finished before the generation is complete. The proposed method does not seem to address this problem. In addition, as mentioned in discussion section, introducing a prefill (inserting prompt into the middle of generation) is not free, and may actually introduce more latency than naive batch decoding setting. Although this is not the focus of the paper, it is important to properly address this point in writing as it is not very clear at reading.

questions:
1. Does the selection of finetuning datasets matters, or to what extent it affect the generalization of the model? 
2. What is the distribution of the generation length? 
3. Would the proposed method still be effective when problem is more difficult or unseen?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730569633335

thread_type:
official_review

<thread_object>
creator: reviewer_pbqs

<content>
summary:
A method for predicting whether the answer to a query being generated by a LLM is worth continuing or not.
The method fine-tunes a pre-trained LLM to predict whether the current (partially) generated answer is worth   continuing or not. The fine-tuning involves pre-pending the prompt 'Would you do better if you started over?   (“Yes.” or “No.”)' to the answer of a query and SFT the model to predict "Yes" if the answer is prefered over  the alternative answer.
By deciding to abort the answer generation, the authors claim that the LLM may save compute at inference-time  compared to a method that do not abort the answer generation.

strengths:
The paper tackles an important challenge and chose to do so using an open-weight model. Advances on open-      weight models may benefit everyone and reducing the compute and inference time may increase the adoption of    foundation models. The method itself was understandable (althought I believe that the writing of the paper     could generally be made clearer and more succinct).

weaknesses:
The paper proposes a method that improves inference-time compute, yet it never compares the compute-time and   the FLOPs with baseline methods. Having to probe the model mid generation definitely incurs a cost, which is   not being discussed in the paper. Moreover, the model finds that pruning longer sentence improves performance, in which case perhaps it would be better to fully generate the sentences before selecting them and thus the    idea of pruning mid-generation is unclear.

In Figure 1, the x-axis stops at 16 Fully Generated Samples, where the method of the authors has hit a plateau and the Best-of-N method is still showing improvement. The authors should increase the amount of fully         generated samples to observe and compare the effect of more generated sample of their method and the Best-of-N baseline.

In Table 1 the author compares baseline methods with the proposed method. However, it is not at all clear what are the parameters of the capability-aware method and whether it is comparable to the baselines. The authors   have to clarify how they compare with the baseline methods. E.g. of questions I have (and this is a subset of  the questions I have -- the authors should make sure that the revised version clearly explain the experimental methodology): How many fully generated samples are used? (Line 404 the authors mention Best-of-16 sampling, is this 16 fully generated samples?) If this is indeed 16 fully generated samples, then why should a pratitioner  use that method instead of the best-of-N baseline which performs comparably well acccording to Figure 1.

The ablations that are important for understanding the parameters of the method are in the Appendix while they should be in the main paper.

The authors continually refer to Figure 1 throughout the paper, forcing the reader to continually scroll up    and down. Instead of packing many important observations within one figure, the authors are encouraged to      break down the observations into indivudal figure (i.e. one observation per figure).

Finally, I encourage the authors to revise the writing making it more succinct and avoid long winded sentences.

questions:
* How does the proposed method compares with the baselines in term of compute-time?
* How does the proposed method compares iwth the baselines in term of FLOPs?
* How does Figure 1 looks like when evaluating 32 and 64 fully generated samples?
* What are the parameters of the methods and the baselines used to produce table 1?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730694712832

thread_type:
official_review

<thread_object>
creator: reviewer_pws6

<content>
summary:
- The paper proposes to simple way to train a LLM to do self-evaluation (with additional tokens) and uses the logits for the trained additional tokens to decide whether to continue sampling or not. They show that they can approximately match the performance of Best-of-16 at, on average, 1/4 of the number of samples on GSM8K.
- The authors claim on the efficiency and latency of the proposed method, but since it requires serial, instead of parallel, generations when doing Best-of-N, the latency may actually increase by a non-negligible amount. This needs to be properly addressed in the paper as this is a paper about making Best-of-N efficient.

strengths:
- The authors propose a very simple and concrete way to adaptively scale generations for Best-of-N and show an noticeable improvement on GSM8K where they can approximately match the performance of Best-of-16 at 1/4 of the number of samples.

weaknesses:
- The performance gain is marginal for Alpaca Eval. Thus, more comprehensive evaluations on a broader range of benchmarks that perhaps require more complex reasoning will be beneficial. However, I fully acknowledge the limits of what can be done during the rebuttal.
- The paper/authors argue that their proposed adaptive inference-time compute is scalable and efficient based on the performance vs. number of samples comparison. They make an argument on latency. However, an important aspect that is underplayed is the actual computational efficiency when implemented. Although mentioned, because of the need to self-evaluate, the adaptive inference-time compute is allocated in serial. On the other hand, Best-of-N can be executed in parallel, which is its advantage over this method. Although inference speed is very dependent on the inference stack engine, it would be important to quantify this cost of serial processing.
- Some further analysis is warranted to analyze the effect of learning to self-evaluate alone. Interestingly, the authors note that a pretrained (and instruction tuned) model has a nontrivial ability to conduct self-evaluations and that a Bradley Terry reward model further improves it. It would be constructive to study the effect of self-evaluation training in isolation and essentially conduct Best-of-N after this self-evaluation training to separate the effect of this training on newly annotated data and the adaptive inference method.

questions:
- For Section 4.2 experiments on AlpacaEval and GSM8K, what is the performance of Llama 3.1 8B Instruct when using Best-of-16?
- The authors focus on showing the efficiency in terms of the reduction of the number of samples vs performance. However, an interesting experiment would be to, fix the amount of compute budget (e.g., FLOPS), and to see whether one can adaptively allocate higher # of samples for more difficult questions and fewer for easier questions, while using everything. That would essentially show the effectiveness of this adaptive compute allocation method compared to a naive Best-of-N that allocates the same amount of compute to any prompt.
- How does the proportion of ties (the epsilon for deciding ties) in the preference dataset construction for self-evaluation training change the performance of the model?
- I am not sure how realistic this is given the time constraint of the rebuttal, but a natural question that arises is how this changes for datasets that require more complex reasoning such as ArenaHard, Math Hard, etc.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732389878798

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Overall Response

comment:
We thank all the reviewers for their helpful feedback and suggestions. We appreciate that the reviewers recognize that our solution is simple and elegant (pws6, 6jZH), provides substantial cost reduction (pws6, 6jZH, Pqum), and that our paper is well-written (6jZH, Pqum).

In this rebuttal we make the improvements our methods offer more precise, concrete, and easy to understand. This was not clear to many reviewers as we made improvements to multiple aspects of Best-of-N which we measured in different ways. 

Reviewers asked for evaluation with a more difficult benchmark that requires more reasoning, so we added evaluations with the MATH 500 benchmark. We also implemented a process reward model (PRM) baseline for requested comparison to our mid-generation self-evaluations. Finally, we added and restructured our figures with concrete metrics (FLOPs, wall-time) and additional ablations for improved clarity.

1. **What are we improving?**

Inference-time methods such as Best-of-N are expensive as they require multiple samples and an external reward model. The primary objective of our paper is to make this far more **cost-effective through the substantial reduction of total computation, memory usage, and energy consumption.** These metrics translate directly to cost in most modern LLM serving scenarios as queries are processed in continuous batches (Daniel et al, 2023) and minimal compute or memory is sitting idle. Our methods improve performance with fixed computation, memory, or energy.

We proposed methods that reduce cost without sacrificing latency (self-evaluations and early pruning) and methods that reduce more cost with some additional latency (adaptive sampling).

2. **How are we improving it?**

We introduce a generative reward model formulation, allowing LLMs to predict mid-generation the probability that restarting the generation will yield a better response. These capability aware and mid-generation self-evaluations enable adaptive inference-time compute strategies. Specifically, these predictions are obtained without an external reward model and can be used to decide whether or not to generate more samples, prune unpromising samples early on, or to pick the best sample. 

We propose a two strategies to take advantage of this capability, namely adaptive sampling and early pruning. Adaptive sampling takes advantage of the capability-aware aspect of the self-evaluations to determine if generating more samples is beneficial. Early pruning takes advantage of mid-generation self-evaluations to stop unpromising samples early in generation. 

Note that these new strategies we propose are primitives that could be combined or enhanced in various ways for potentially greater efficiency gains. To improve adaptive sampling specifically, we also proposed a novel annealing schedule as well as exponentially increasing batch sizes.

3. **How much are we improving it by?**

We measure efficiency with respect to the number of FLOPs during inference (following Hoffman et al 2022). We also provide exact wall-clock times in a controlled environment using SGLang (Zheng et al.) on a consistent hardware setup (8×A100 GPUs) in Figure 3B. Below, we will discuss the efficiency gains we get from each component we propose.

Firstly, capability-aware self-evaluations alone cut FLOPs by 2x compared to using external reward models. This is because the KV cache can be shared during the evaluation process with infilling of our self-evaluation prompt (16 additional tokens).

Additionally, we see roughly a **4x reduction in total FLOPs** to match the performance of Best-of-16 by additionally leveraging early pruning. This is **without any additional latency**. Adaptive sampling achieves **roughly a 6-8x reduction in total FLOPs** with the tradeoff of roughly 2x latency (as batches are processed sequentially) as shown in Table 3 in our appendix. Using early pruning or adaptive sampling is a choice that can be made depending on constraints an application has on latency and computational efficiency. We provide the total FLOPs usage and downstream performance in the updated Figure 2 and Figure 3. **[Results](https://imgur.com/a/pdQ878w)** with the PRM baseline show that mid-generation self-evaluations are still roughly 2x more efficient simply due to KV cache.

Also, our proposed methods allow for savings in memory usage. Memory usage is determined by the number of active model parameters, as well as the size of the KV cache which is largely a function of the total number of tokens generated. Since we both remove the reward model (2x less total parameters to store) and reduce the total number of tokens generated (2x with early pruning and 4x with adaptive sampling), **the memory savings are roughly proportional to the savings in the total number of FLOPs**.

Since we save total FLOPs and memory usage, **we also proportionally save energy consumption** which accounts for about 5-10% of the total cost of operation for modern GPUs.
</content>
</thread_object>

</thread>


Metareview:
metareview:
In this paper, the authors propose a new strategy for efficient inference-time improvement of LLMs. Specifically, the authors exploit adaptive sampling (selecting the N in best-of-N as a function of the prompt) + early pruning (pruning unpromising mid-generations) upon the learned generative reward model. The authors conducted empirical comparison on GSM8k and MATH 500 to demonstrate the benefits in terms of FOLPS vs. Performance. 

The paper is well-motivated and easy-to-follow. The major issues raised by the reviewers mainly lies in about the novelty and significance. Specifically, using additional reward model to guide the inference-time generation and pruning is not novel, e.g., [1, 2, 3]. Moreover, the empirical study is limited.  Without discussion and comparison to these methods on more tasks, it is difficult to justify the claimed benefits. 

I suggest the authors to take the comments from the reviewers into account to emphasize the differences and benefits to the existing methods, especially the generative reward effect in accelerating inference in practice. 


[1] Mudgal, Sidharth, Jong Lee, Harish Ganapathy, YaGuang Li, Tao Wang, Yanping Huang, Zhifeng Chen et al. "Controlled decoding from language models." arXiv preprint arXiv:2310.17022 (2023).

[2] Sun, Haotian, Yuchen Zhuang, Wei Wei, Chao Zhang, and Bo Dai. "BBox-Adapter: Lightweight Adapting for Black-Box Large Language Models." arXiv preprint arXiv:2402.08219 (2024).

[3] Chakraborty, Souradip, Soumya Suvra Ghosal, Ming Yin, Dinesh Manocha, Mengdi Wang, Amrit Singh Bedi, and Furong Huang. "Transfer Q Star: Principled Decoding for LLM Alignment." arXiv preprint arXiv:2405.20495 (2024).