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

## Abstract

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.

## 1 Introduction

Figure: Figure 1: Adaptive Inference-Time Compute: Strategies such as adaptive sampling (left) and early pruning (right) make compute utilization during inference far more efficient and scalable.
Refer to caption: extracted/5899596/assets/adaptive_and_pruning_slim.jpg

Figure: Figure 2: Capability-Aware and Mid-Generation Self-Evaluations enable adaptive inference-time compute strategies. They are obtained without an external reward model and can determine whether or not to generate more samples, prune unpromising samples early on, and pick the best sample.
Refer to caption: x1.png

As large language models (LLMs) continue to advance, delivering high-quality responses across diverse applications becomes increasingly important. One promising direction to enhance response quality is the strategic use of inference-time computation, particularly through methods like Best-of-N sampling Snell et al. (); Charniak and Johnson (); Cobbe et al. (), which selects the best response from multiple candidates. However, this method incurs substantial inference cost from querying an external reward model and producing a large, fixed number of samples.

In this work, we introduce a new reward modeling paradigm, which we denote as capability-aware self-evaluations. This paradigm allows for adaptive allocation of inference-time compute, aiming to reduce the computational overhead while maintaining or improving LLM performance across various domains. We demonstrate that LLMs can directly model the probability that restarting generation yields in a better response, enabling informed decisions about whether to continue generating a response, initiate new ones, as well as rank responses. These predictions are obtained by simply appending a predefined self-evaluation prompt to the partially or fully generated response and generating a single predefined token whose likelihood is used as the prediction. This is in contrast to preference-based reward models which can primarily only be used to rank responses.

Our self-evaluation method is highly cost-effective, requiring no external reward models and incurring only the minimal cost of generating a single token. In contrast, an external reward model inherently requires more memory and storage. Additionally, it is unable to reuse the KV cache obtained when generating the response and would have to process the input and response from scratch.

To demonstrate adaptive inference-time compute allocation, we introduce two techniques: (1) adaptive sampling and (2) early pruning of unpromising samples. Adaptive sampling involves resampling a response for a given prompt until it is predicted that further samples will not yield additional improvements, thus conserving computation for complex tasks that will benefit from it. Furthermore, early pruning discards samples midway through generation if they are likely to result in suboptimal completions. These are not possible with standard reward models.

To give an LLM the ability to self-evaluate, one must construct an on-policy pairwise preference dataset with ties. In our experimental evaluation, we construct a dataset of approximately 30,000 preferences constructed using real unfiltered LMSYS (Chiang et al., ) user prompts and an existing reward model, ArmoRM (Wang et al., ) trained on roughly 1 million preferences or ratings. With this dataset, we fine-tune a Llama 3.1 8B Instruct (Dubey et al., ) model to self-evaluate and demonstrate significant performance improvements across both in-distribution and out-of-distribution tasks. Notably, as shown in Figure 1, the win rate against GPT-4 on AlpacaEval increases from 21% to 34% with 16 samples, and performance on held-out GSM8K math problems improves from 84% to 91%. We find that with adaptive sampling, using just 1.2 samples on average captures 74% of the improvement observed with 16 samples and 1.9 samples captures 84%. Additionally, early pruning can prevent 75% of unpromising samples from being fully generated, saving 56% of tokens generated with very minimal degradation in performance.

Our approach allows for more efficient and scalable use of compute resources during inference. By enabling models to dynamically allocate compute during inference based on task complexity and the model’s capability, we optimize resource usage, ensuring efficiency in processing all types of prompts and tasks. This adaptability makes using inference-time compute far more practical and ready for the real world where LLMs are used for a wide variety of applications. (^1^11Code for reproducibility can be found at in our [public github implementation](https://github.com/rohinmanvi/Capability-Aware_and_Mid-Generation_Self-Evaluations).)

## 2 Preliminaries and Notation

An autoregressive language model generates a sequence $\mathbf{y}=(y_{1},y_{2},\dots,y_{T})$ given an input context $\mathbf{x}$ by predicting tokens sequentially. Assuming the model is parameterized by $\theta$, the conditional probability distribution of generating a sequence $\mathbf{y}$ given context $\mathbf{x}$ is

$$ $p_{\theta}(\mathbf{y}|\mathbf{x})=\prod_{t=1}^{T}p_{\theta}(y_{t}|\mathbf{x},y _{<t}),$ (2.1) $$

with the convention $y_{<t}=(y_{1},y_{2},\dots,y_{t-1})$. For ease of notation, we define $p_{\theta}(y_{t}|\mathbf{x}):=p_{\theta}(y_{t}|y_{<t},\mathbf{x})$. For a vocabulary size $M$, the probability of predicting the $t$-th token $y_{t}$ is determined using a softmax with temperature $\gamma$ on logit scores $z$ of all the tokens:

$$ $p_{\theta}(y_{t}|\mathbf{x})=\frac{\exp(z_{t}/\gamma)}{\sum_{i=1}^{M}\exp(z_{i }/\gamma)},$ (2.2) $$

where $z_{t}=\text{logit}_{\theta}(y_{t}|\mathbf{x},y_{<t})$. Higher values of $\gamma$ introduce more randomness; as the temperature $\gamma$ approaches zero, the distribution becomes concentrated on the token with the highest logit.

Next-token prediction is a typical approach used for pre-training and fine-tuning of LLMs. In particular, supervised fine-tuning (SFT) minimizes the cross-entropy loss between the model’s predicted next token and the target token in a given sequence. Given a dataset $\mathcal{D}=\{(\mathbf{x},\mathbf{y})\}$ of input context $\mathbf{x}$ and target response $\mathbf{y}$, the SFT loss is given by:

$$ $\mathcal{L}_{\text{SFT}}(\theta,\mathcal{D})=-\mathbb{E}_{(\mathbf{x},\mathbf{ y})\sim\mathcal{D}}\left[\sum_{t=1}^{|\mathbf{y}|}\log p_{\theta}(y_{t}| \mathbf{x},y_{<t})\right].$ (2.3) $$

On-policy pairwise preference dataset is a preference dataset that consists of responses generated by a single model:

$$ $\mathcal{D}_{\text{preference}}=\left\{\left(\mathbf{x},\mathbf{y}_{1},\mathbf {y}_{2},l\right)_{i}\right\}_{i=1}^{N},$ (2.4) $$

where $\mathbf{x}$ is the input, $\mathbf{y}_{1}$ and $\mathbf{y}_{2}$ are two responses generated by the model, and $l$ is the preference label that indicates the outcome of the comparison between $\mathbf{y}_{1}$ and $\mathbf{y}_{2}$:

$$ $l=\begin{cases}1&\text{if }\mathbf{y}_{1}\text{ resulted in a }\textit{Win} \text{ and }\mathbf{y}_{2}\text{ resulted in a }\textit{Loss}\leavevmode \nobreak\ (\mathbf{y}_{1}\succ\mathbf{y}_{2}),\\ 0&\text{if }\mathbf{y}_{1}\text{ and }\mathbf{y}_{2}\text{ resulted in a } \textit{Tie}\leavevmode\nobreak\ (\mathbf{y}_{1}\approx\mathbf{y}_{2}),\\ -1&\text{if }\mathbf{y}_{1}\text{ resulted in a }\textit{Loss}\text{ and } \mathbf{y}_{2}\text{ resulted in a }\textit{Win}\leavevmode\nobreak\ (\mathbf{ y}_{1}\prec\mathbf{y}_{2}).\end{cases}$ (2.5) $$

To simplify our analysis, we assume that the preference labels $l$ are based on an underlying reward function $r(\mathbf{x},\mathbf{y})$ and a threshold $\epsilon$ that defines if one completion is more preferred:

$$ $l=\begin{cases}1&\text{if }r(\mathbf{x},\mathbf{y}_{1})-r(\mathbf{x},\mathbf{y }_{2})>\epsilon,\\ 0&\text{if }\left|r(\mathbf{x},\mathbf{y}_{1})-r(\mathbf{x},\mathbf{y}_{2}) \right|\leq\epsilon,\\ -1&\text{if }r(\mathbf{x},\mathbf{y}_{1})-r(\mathbf{x},\mathbf{y}_{2})<- \epsilon.\end{cases}$ (2.6) $$

Bradley-Terry (BT) (Bradley and Terry, ) reward models consist of a base model $\theta_{B}$ and a shallow MLP reward head $\theta_{R}$. They predict a scalar reward score $\hat{R}=r_{\theta_{B},\theta_{R}}(\mathbf{x},\mathbf{y})$ given an input and response. They are initialized from a pretrained base model and a randomly initialized reward head, then trained on a dataset of $N$ examples, $\mathcal{D}=\{(\mathbf{x},\mathbf{y}_{\text{win}},\mathbf{y}_{\text{loss}})_{i
}\}_{i=1}^{N}$, where $\mathbf{y}_{\text{win}}$ is the preferred response and $\mathbf{y}_{\text{loss}}$ is the dispreferred response. They are trained to predict a higher reward for $\mathbf{y}_{\text{win}}$ than for $\mathbf{y}_{\text{loss}}$ under the Bradley-Terry model. This is achieved by minimizing:

$$ $\mathcal{L}_{RM}(\theta_{B},\theta_{R},D)=-\mathbb{E}_{(\mathbf{x},\mathbf{y}_ {\text{loss}},\mathbf{y}_{\text{win}})\sim D}\left[\log\left(\sigma(r_{\theta_ {B},\theta_{R}}(\mathbf{x},\mathbf{y}_{\text{win}})-r_{\theta_{B},\theta_{R}}( \mathbf{x},\mathbf{y}_{\text{loss}}))\right)\right],$ (2.7) $$

Best-of-N (Snell et al., ) is a widely used test-time compute approach to improve the performance of LLMs. Specifically, given a prompt, $N$ candidate responses from an LLM are scored using a reward model, and the highest-scoring response is filtered as the final response.

## 3 Adaptive Inference-Time Compute via Capability-Aware and Mid-Generation Self-Evaluations

A simple, yet effective method for allocation of test-time compute is Best-of-N, allowing the model to improve upon its responses over greedy sampling. However, the generation of a large, fixed number of samples is computationally expensive and the efficacy of this approach relies heavily on the quality of an external reward model, which can additionally incur substantial computational cost. In the following section, we address the cost and robustness of the reward model by establishing a general framework for self-evaluation via token prediction. We show how to train capability-aware and mid-generation self-evaluators to reduce the cost of generating a large, fixed number of samples.

### 3.1 Capability-Aware and Mid-Generation Self-Evaluations

In the classic reward modeling paradigm, a pre-trained LM is used as an initialization for reward modeling. However, this approach also relies on a newly added reward head to output the reward, which diverges from the token prediction task for which the LLM was trained, which can hurt performance for reward modeling. Additionally, generally preference-based reward models learn only to rank arbitrary responses. This may be limiting for comparing on-policy samples (such as those from Best-of-N) during inference to determine how good the response is with respect to the model’s own capabilities. Furthermore, reward models are generally only able to evaluate full responses. Responses that are clearly unpromising very early on in generation still have to be generated until the end for evaluation, wasting computational resources. Can we use an alternate paradigm to model rewards to address these limitations?

We propose capability-aware and mid-generation self-evaluations, a paradigm in which we query the reward model to predict the probability that restarting generation will not yield a better response. This probability can be used to elicit adaptive test-time computation, where the model can dynamically allocate compute dependent on the difficulty of a query and its own capability. This is because to decide whether or not to allocate more compute to a given input, we need to know if it is fruitful to do so. For example, for easy queries, the model may have already outputted the best response it can, so resampling is unnecessary. Additionally, for difficult queries, sampling more may lead to a minor improvement in performance, leading to unfruitful resampling. We illustrate this paradigm in  Figure 2. Modeling this probability effectively has three components: (1) reward modeling using token prediction to better leverage the pretrained model’s existing knowledge and capabilities (2) relying on on-policy pairwise preferences and ties to model the probability that the model cannot generate a more preferred response (3) using truncated responses to model the probability that restarting generation will not yield a more preferred response. Traditional reward models can suffer at this task as they are model agnostic. Though these models can rank how different responses perform for a given task, they are unable to quantify how effective resampling would be (not capability aware). Furthermore, a desirable property for a reward model is to early stop the generation of samples that are not promising in order to give the model another chance without wasting compute.

Reward Modeling with Token Prediction. To obtain a self-evaluation of a response from an LLM, we simply append a predefined self-evaluation prompt $I$ to the generated response $\mathbf{y}$ and obtain a score in the form of the likelihood of a predefined token $t_{\text{good}}$ corresponding to the likelihood of the response being good. This approach allows us to acquire rewards without any external reward model, making it highly cost-effective as we can reuse the KV cache obtained during the generation of the response. Also, this leverages the existing zero-shot capability of the models to judge its own responses as a prior, which has been shown to be effective in works such as Madaan et al. ().

Formally, given a preference dataset $\mathcal{D}_{\text{pref}}=\{x,\mathbf{y}_{\text{good}},\mathbf{y}_{\text{bad}}\}$, which contains input context $x$ and response pairs $y$, we can train a self-evaluation model by maximizing the likelihood of the good token $\log p_{\theta}(t_{\text{good}}\mid(\mathbf{x},\mathbf{y}_{\text{good}}))$ for good responses $\mathbf{y}_{\text{good}}$, and the likelihood of the bad token $\log p_{\theta}(t_{\text{bad}}\mid(\mathbf{x},\mathbf{y}_{\text{bad}}))$ for bad responses $\mathbf{y}_{\text{bad}}$. To achieve this, we minimize the SFT loss in Equation 2.3 on a modified variant of the preference dataset $\mathcal{D}_{\text{self-evaluation}}$

$$ $\displaystyle\mathcal{D}_{\text{self-evaluation}}=\{((\mathbf{x},\mathbf{y}_{ \text{good}},I),t_{\text{good}})\}\cup\{((\mathbf{x},\mathbf{y}_{\text{bad}},I ),t_{\text{bad}})\}.$ (3.1) $$

During inference, we normalize the likelihood of $t_{\text{good}}$ as the score to rank responses:

$$ $\displaystyle\text{Score}=\frac{p_{\theta}(t_{\text{good}}\mid\mathbf{x}, \mathbf{y},I)}{\sum_{t\in\{t_{\text{good}},t_{\text{bad}}\}}p_{\theta}(t\mid \mathbf{x},\mathbf{y},I)}.$ (3.2) $$

Given that the language model has a fixed vocab size $|V|$, probability mass can be spuriously assigned to tokens $t\notin\{t_{\text{good}},t_{\text{bad}}\}$, resulting in this normalization to be desirable.

Note: it is beneficial to establish within the self-evaluation prompt that the LLM has to perform a classification task and that the target tokens are natural responses to effectively leverage the zero-shot capability of the model. We find that with a poorly designed self-evaluation prompt $I$, the model will overgeneralize during training and respond with $t_{\text{good}}$ or $t_{\text{bad}}$ for all queries, regardless of relevance. We need to keep the underlying policy or model unchanged when learning to self-evaluate. To do so, we need the model’s prior distribution over tokens to be similar to what we expect after training where $p_{\theta}(t_{\text{good}}\cup t_{\text{bad}}\mid\mathbf{x},\mathbf{y},I)\approx
1$. Fortunately, this can easily be done by explicitly stating that the LLM should respond with only $t_{\text{good}}$ or $t_{\text{bad}}$ in the self-evaluation prompt $I$.

The probability that the model cannot generate a more preferred response can be easily derived from on-policy pairwise preferences and ties. It is the probability that the current sample $\mathbf{y}$ results in a Win or a Tie against another sample $\mathbf{y}^{\prime}$. More formally:

$$ $\displaystyle P_{\mathbf{y}^{\prime}\sim p_{\theta}(\mathbf{y}\mid\mathbf{x})} \left(r(\mathbf{x},\mathbf{y})-r(\mathbf{x},\mathbf{y}^{\prime})\geq-\epsilon \right)=P(\textit{Win}\cup\textit{Tie}\mid\mathbf{x},\mathbf{y}).$ (3.3) $$

Accounting for ties in reward modeling is especially important for on-policy pairwise data where responses coming from the same model are very likely to be similar. Ties commonly occur when generating responses from the same model (e.g. 40% of the time) and are even more common with simple tasks. Since ties indicate that the model cannot do better, they are crucial when learning to do capability-aware self-evaluations. If ties are not specified, it is also harder to model the reward as the model has to distinguish between extremely similar responses. Furthermore, since $(100-40)/2=\leavevmode\nobreak\$30% of samples result in a Loss, we see that the model can do significantly better on a query only roughly 30% of the time.

Notice that $P(\textit{Win}\cup\textit{Tie}\mid\mathbf{x},\mathbf{y})$ monotonically increases with the sample’s underlying reward $r(\mathbf{x},\mathbf{y})$. This means that it can also be used to determine if one sample has a higher reward than another. Since inference-time compute strategies only care about being able to rank responses, it turns out that modeling $P(\textit{Win}\cup\textit{Tie}\mid\mathbf{x},\mathbf{y})$ is more useful than $r(\mathbf{x},\mathbf{y})$ since it can also be used inform decisions on whether or not to allocate more compute.

The probability that restarting generation will not yield a more preferred response can be used to evaluate the quality of a partial response $\mathbf{y}_{1:t}$. In the context of adaptive inference-time compute, if we find that this probability is low, we can stop spending additional inference resources on generating the remainder of the partial response, pruning this poor partial response. Again, in the context of on-policy pairwise preferences or ties, this probability is simply $P(\textit{Win}\cup\textit{Tie})$ conditioned on a partial response $\mathbf{y}_{1:t}$:

$$ $P_{\begin{subarray}{c}\mathbf{y}^{\prime}\sim p_{\theta}(\mathbf{y}\mid\mathbf {x})\\ \mathbf{y}_{t+1:T}\sim p_{\theta}(\mathbf{y}_{t+1:T}\mid\mathbf{x},\mathbf{y}_ {1:t})\end{subarray}}\left(r(\mathbf{x},\mathbf{y}_{1:T})-r(\mathbf{x},\mathbf {y}^{\prime})\geq-\epsilon\right)=P(\textit{Win}\cup\textit{Tie}\mid\mathbf{x} ,\mathbf{y}_{1:t}).$ (3.4) $$

Modeling these probabilities with an on-policy pairwise preference dataset with ties. To train LLMs to make capability-aware self-evaluations, we construct a dataset derived from an on-policy pairwise preference dataset with ties, where good responses are those that resulted in a Win or Tie ($\mathbf{y}_{\text{win}}$, $\mathbf{y}_{\text{tie}}$), and bad responses are those that resulted in a Loss ($\mathbf{y}_{\text{loss}}$). To train LLMs to make capability-aware self-evaluations mid-generation, we simply include the same examples but with responses randomly truncated ($\mathbf{y}_{\text{win,trunc}},\mathbf{y}_{\text{tie,trunc}},\mathbf{y}_{\text{
loss,trunc}}$).

We use the following self-evaluation prompt and target tokens:

$$ $I=\text{`Would you do better if you started over? (``Yes." or ``No.")'}$ $$

$$ $t_{\text{good}}=\text{`No'},\quad t_{\text{bad}}=\text{`Yes'}$ $$

Our final dataset $\mathcal{D}_{\text{capability-aware}}$, used to minimize the SFT loss (Equation 2.3), is constructed as:

$$ $\mathcal{D}_{\text{capability-aware}}=\{((\mathbf{x},\mathbf{y}_{\text{win}},I ),t_{\text{good}})\}\cup\{((\mathbf{x},\mathbf{y}_{\text{tie}},I),t_{\text{ good}})\}\cup\{((\mathbf{x},\mathbf{y}_{\text{loss}},I),t_{\text{bad}})\}$ (3.5) $$

$$ $\cup\leavevmode\nobreak\ \{((\mathbf{x},\mathbf{y}_{\text{win,trunc}},I),t_{ \text{good}})\}\cup\{((\mathbf{x},\mathbf{y}_{\text{tie,trunc}},I),t_{\text{ good}})\}\cup\{((\mathbf{x},\mathbf{y}_{\text{loss,trunc}},I),t_{\text{bad}})\}.$ $$

During inference, we compute the normalized likelihood of the $t_{\text{good}}$ (‘No’) token to score full or partial responses:

$$ $p_{\theta}(\textit{Win}\cup\textit{Tie}\mid\mathbf{x},\mathbf{y}_{1:t})=\frac{ p_{\theta}(t_{\text{good}}\mid\mathbf{x},\mathbf{y}_{1:t},I)}{\sum_{t\in\{t_{ \text{good}},t_{\text{bad}}\}}p_{\theta}(t\mid\mathbf{x},\mathbf{y}_{1:t},I)}$ (3.6) $$

### 3.2 Making Best-of-N Efficient and Scalable

Best-of-N is very computationally expensive, as it requires the generation of a large, fixed number of samples. This leads to wasted computation on queries that do not require more compute, and underutilization of compute on queries that could benefit from it. To remedy this and make Best-of-N more practical, we introduce two new primitives that leverage the capability-aware and mid-generation self-evaluations that we introduced in the last section. The first is adaptive sampling, where additional samples are allocated to a query only if they are predicted to be beneficial. The second is early pruning, where unpromising samples are stopped from being generated further to save inference computation.

#### 3.2.1 Adaptive Sampling and Annealing

We introduce adaptive sampling as a technique to allocate inference-time compute only when it is beneficial to do so. We further introduce exponentially increasing parallel sampling to mitigate the main disadvantage of adaptive sampling, latency. Finally, we introduce a temperature annealing strategy to balance exploration and exploitation while adaptively sampling and boost efficiency.

Resampling until meeting a threshold. We adaptively sample by resampling only when the model predicts it can produce a better response. We do so by computing the likelihood that the model cannot generate a better response $p_{\theta}(\textit{Win}\cup\textit{Tie}\mid\mathbf{x},\mathbf{y})$ for every sample. If this probability exceeds a predefined threshold $\tau$ for any sample, we stop resampling and select the sample with the highest $p_{\theta}(\textit{Win}\cup\textit{Tie}\mid\mathbf{x},\mathbf{y})$ as the final response. Otherwise, we resample, repeating this process until the threshold is met or a maximum number of samples $N_{\text{max}}$ is reached. This approach concentrates computational resources on prompts where improvement is likely, avoiding unnecessary sampling when further gains are improbable.

Increasing number of parallel samples exponentially. To reduce latency, we sample in exponentially increasing batch sizes. Specifically, for the $k$-th sampling iteration, the batch size $N_{k}$ is defined as:

$$ $N_{1}=1,\quad N_{k}=2^{k-2}\quad\text{for}\ k>1.$ (3.7) $$

This ensures that the cumulative number of samples by the $k$-th iteration is $2^{k-1}$. This exponential increase minimizes latency, reducing the number of iterations needed to meet $N_{\text{max}}$, while allowing for enough self-evaluations to determine if larger batches of samples are necessary.

Temperature annealing schedule based on number of samples generated so far. To balance exploitation and exploration, we vary the temperature $\gamma$. The temperature for the $k$-th iteration is given by:

$$ $\gamma_{k}=1-2^{-(k-1)}.$ (3.8) $$

Initially, the temperature starts low (e.g., $\gamma_{1}=0$) to prioritize high-probability responses. As $k$ increases, $\gamma_{k}$ quickly approaches 1, encouraging more diverse and creative sampling. This annealing schedule allows the model to first focus on exploiting the most likely responses, then explore alternative options as more samples are generated.

Figure: Algorithm 1 Adaptive Sampling and Annealing

#### 3.2.2 Early Pruning of Unpromising Samples

One downside to adaptive sampling requires the generation of samples in series, which introduces latency. However, can we still leverage parallel sampling to avoid increasing the latency of response generation but still enable adaptive test-time compute allocation? One approach to reducing computational costs in parallel generation is to early prune unpromising samples based on mid-generation self-evaluations.

When to prune. Pruning too early risks discarding samples that could improve, while pruning too late offers minimal savings. Thus, we balance this trade-off when selecting the fixed number of initial tokens (e.g., 64 or 128) before making any pruning decisions.

Which samples to prune. After generating an initial number of tokens, we compute the intermediate self-evaluations for each partially generated sample. After ranking samples by the resulting scores, we stop the generation of the bottom $x\%$ (e.g., 50% or 75%) to conserve computation. This ensures that only the most promising partial samples continue to be generated.

## 4 Experiments

We evaluate our methods by first examining the sample efficiency of self-evaluation. We then evaluate its impact on optimizing inference-time compute with adaptive sampling and early pruning of unpromising samples.

### 4.1 Training Data and Evaluations

Construction of On-Policy Pairwise Preference Dataset with Ties. Training reward models typically requires human-labeled preferences, which are costly to obtain, especially for on-policy data generated by the model being trained. To mitigate this, we utilize an existing reward model, ArmoRM (Wang et al., ) trained on approximately 1,000,000 preferences. This reward model serves as our underlying reward $r(\mathbf{x},\mathbf{y})$.

We conduct our experiments using the Llama 3.1 8B Instruct model, which is fine-tuned on a preference dataset of 50,000 real user prompts from LMSYS (Chiang et al., ) and pairs of on-policy responses (sampled from the same Llama 3.1 8B Instruct model). These responses are scored using ArmoRM, the underlying reward. Preferences or ties between responses are determined using a threshold $\epsilon$ of 0.01 on the reward difference, resulting in an on-policy pairwise preference dataset with approximately 40% ties.

Evaluation Protocol
We evaluate the performance of our self-evaluation model and adaptive test-time compute approaches in two domains: (1) the AlpacaEval (Dubois et al., ) benchmark and (2) the GSM8K dataset (Cobbe et al., ). AlpacaEval 2.0 is an automatic benchmark that compares model responses to those generated by GPT-4 across approximately 800 representative prompts. The final metric is the win rate against GPT-4, adjusted to increase correlation with human preferences (Spearman correlation of 0.98). While highly correlated with human judgments, this metric is relative—it measures success based on outperforming another model rather than achieving absolute human satisfaction. GSM8K is a collection of 8.5K grade school math word problems involving multi-step reasoning and basic arithmetic operations. We use GSM8K as it is a popular benchmark for reasoning as well as absolute measure of performance that is not relative to another LLM. This is particularly useful in evaluating adaptive sampling where it allocates compute resources to queries that benefit from it.

### 4.2 Performance and Sample Efficiency

**Table 1: Best-of-16 Performance on AlpacaEval (Win Rate vs GPT-4, %) and GSM8K (%) using capability-aware self-evaluations using token prediction, LLM-as-a-Judge which also uses token prediction but without finetuning, the commonly used Bradley-Terry external reward model, and the underlying reward model used to create our experimental preference data.**
| Benchmark | Random | Zero-Shot | Bradley-Terry | Capability-Aware | Underlying |
| --- | --- | --- | --- | --- | --- |
|  | (1 sample) | (LLM-as-a-Judge) | Reward Model | $p_{\theta}(\textit{Win}\cup\textit{Tie})$ | $r(\mathbf{x},\mathbf{y})$ |
| AlpacaEval | 21.2 | 24.4 | 33.2 | 33.8 | 36.3 |
| GSM8K | 84.2 | 86.7 | 87.7 | 91.0 | 92.6 |

Baselines. To evaluate performance with Best-of-N, we use two baselines. The Zero-Shot (LLM-as-a-Judge) baseline uses Llama 3.1 8B Instruct without additional training. We prompt and get scores for zero-shot predictions in the exact same manner that we do for capability-aware self-evaluations. The Bradley-Terry Reward Model also uses Llama 3.1 8B Instruct as the base model and is trained using the same on-policy pairwise preference dataset with ties.

As shown in Table 1, the zero-shot approach performs modestly on both benchmarks, confirming that LLMs have a non-trivial ability to self-evaluate. The Bradley-Terry reward model unsurprisingly does much better. Our token prediction method outperforms both. On AlpacaEval, it achieves a 33.8% win rate against GPT-4 using Best-of-16 sampling, slightly surpassing the Bradley-Terry model’s 33.2% and nearing the underlying reward model’s 36.3% (trained on 1 million preferences). On GSM8K, token prediction attains 91.0% accuracy, significantly outperforming the Bradley-Terry model’s 87.7%. This suggests that our method is more sample-efficient and generalizes better, especially on queries that might be uncommon among real user prompts like those in GSM8K.

These results demonstrate that token prediction effectively leverages the pre-trained model’s priors for self-evaluation, approaching the performance of the larger Llama 3.1 70B, which achieves 38.1% on AlpacaEval and 95.1% on GSM8K. With additional data and hyperparameter tuning, we anticipate that we can easily further close the gap with the underlying reward.

We also examined other probabilities to model with token prediction, specifically $p_{\theta}(\textit{Win})$ or $p_{\theta}(\textit{Win}\mid\neg\textit{Tie})$. Our findings, as shown in Table 3, indicate that including or removing ties does not significantly impact performance, but modeling $p_{\theta}(\textit{Win})$ is significantly more difficult as it forces the model to distinguish between samples that result in wins or ties, which are both cases in which the model performs relatively well for its capability.

Notably, using Best-of-N sampling with our methods allows the Llama 3.1 8B model to approach the performance of Llama 3.1 70B, a model roughly $10\times$ larger that gets 38.1% and 95.1% respectively. This highlights the efficacy of Best-of-N sampling in enhancing model performance. In the following section, we address the expensive nature of this method with adaptive compute strategies.

**Table 2: Adaptive Sampling and Annealing performance gains and efficiency on GSM8K.**
| Win-or-Tie Probability Threshold | 0.92 | 0.96 | 0.98 | 0.99 | 1.00 |
| --- | --- | --- | --- | --- | --- |
| Average Samples Used | 1.2 | 1.9 | 3.7 | 6.8 | 16.0 |
| Average Batches Used (Latency) | 1.1 | 1.4 | 2.0 | 2.9 | 5.0 |
| GSM8K Pass@1 (%) | 89.2 | 89.9 | 90.8 | 90.8 | 91.0 |
| Percent of Maximum Improvement | 73.5 | 83.8 | 97.1 | 97.1 | 100.0 |

### 4.3 Efficiency and Scaling of Adaptive Sampling and Annealing

We evaluate the efficiency gains and performance trade-offs of our adaptive sampling and annealing strategy on the GSM8K dataset. Our objective is to retain the performance benefits of Best-of-16 sampling while significantly reducing the average number of samples used and minimizing latency. We want to evaluate the methods’ ability to allocate samples when necessary. For our experiments, the capability-aware self-evaluations are always used to select the best sample. This is so that the final selection method is not a confounding factor.

As a baseline, we first assess the performance Best-of-N. The best sample is selected using self-evaluation with token prediction. Figure 1 shows that increasing the number of samples from 1 to 16 incrementally improves the GSM8K Pass@1 accuracy from 84.2% to 91.0%. This represents the maximum achievable performance with our method, which we define as 100% of the maximum improvement. However, generating a fixed large number of samples per query is expensive.

To mitigate this, we introduce adaptive sampling controlled by a threshold $\tau$ and a maximum number of samples $N_{\text{max}}$. Initially, we test adaptive sampling using the underlying reward model to decide whether to generate additional samples. As shown in Figure 1, this approach does not significantly outperform random selection. While the underlying reward model effectively ranks responses, it does not consider the model’s capability to improve upon its own outputs, substantially limiting its effectiveness in informing resampling decisions.

In contrast, our adaptive sampling method utilizing capability-aware self-evaluations via token prediction yields substantial efficiency gains. Table 4 demonstrates that by adjusting the threshold $\tau$ for the win-or-tie probability $p_{\theta}(\textit{Win}\cup\textit{Tie}\mid\mathbf{x},\mathbf{y})$, we can control the average number of samples. For instance, setting $\tau=0.98$ results in 97.1% of the maximum performance while using only an average of 4.1 samples, compared to the 16 samples required for maximum performance.

Moreover, incorporating the annealing schedule further enhances performance, particularly at lower thresholds. As illustrated in Table 2, with annealing, we achieve 73.5% of the maximum improvement using an average of just 1.2 samples when $\tau=0.92$. This indicates that our annealing strategy effectively balances exploitation and exploration during sampling.

Our adaptive sampling approach offers significant compute savings compared to Best-of-N. By adjusting $\tau$ and $N_{\text{max}}$, we can fine-tune the trade-off between performance and efficiency. The latency introduced by adaptive sampling is minimal due to the exponentially increasing batch sizes.

### 4.4 Efficiency Gains of Pruning Unpromising Samples

Figure: Figure 3: Early Pruning performance gains and efficiency on AlpacaEval.
Refer to caption: extracted/5899596/assets/prune_at_different_tokens_slim.jpg

We evaluate the efficiency gains from early pruning of unpromising samples on the AlpacaEval benchmark. We begin with 16 samples and prune from them. Our primary metrics are the percent of maximum improvement achieved and the number of tokens generated per prompt. The latter is a more accurate measure of computational cost than the number of full samples generated, as partial responses also consume computational resources.

Our experiments control two variables: the number of tokens generated before making pruning decisions (64 or 128 tokens), and the percentage of samples pruned (e.g., pruning 75%). Again, the capability-aware self-evaluations are used the final sample selection method.

Impact of Evaluation Timing on Performance. As shown in  Figure 1, pruning at 64 tokens improves performance over random pruning at 0 tokens. For instance, pruning 50% of samples at 64 tokens achieves a win rate of 33.4%, closely approaching the maximum of 33.8% without pruning. The timing of when we perform self-evaluations significantly affects performance. Pruning at 128 tokens yields even better results. Pruning 75% of samples at 128 tokens attains a win rate of 33.2%, nearly matching the maximum performance without pruning. This improvement is expected, as longer partial responses provide more context for accurate self-evaluation.

Overall, there are substantial savings if one prunes at the right time and the right number of samples. In terms of computational cost measured by number of tokens, pruning 75% of 16 samples at 128 tokens reduces the average number of tokens generated per prompt from 7,259 to 3,220 as shown in Figure 3, resulting in significant efficiency gains.

## 5 Related Work

Reward Modeling in RLHF. Traditionally, reward models (RMs) within Reinforcement Learning with Human Feedback (RLHF) have been trained using discriminative techniques, often drawing on ranking models like the Bradley-Terry (BT) framework (Bradley and Terry, ). These RMs are generally treated as binary classifiers, distinguishing between preferred and dispreferred completions for a given prompt (Stiennon et al., ; Ouyang et al., ). The preference signal is predicted by generating a continuous score, typically derived from a linear head appended to an autoregressive language model. An alternative line of work departs from this discriminative paradigm by directly modeling preferences through joint probability estimations, as seen in methods that predict the likelihood of one response being preferred over another (An et al., ).

Recent advancements in RLHF have introduced implicit reward models that circumvent the necessity for a distinct reward classifier by learning the distribution of preferred and dispreferred completions via objectives like Direct Preference Optimization (DPO) (Rafailov et al., ) and Implicit Preference Optimization (IPO) (Gheshlaghi Azar et al., ). These methods embed preference learning directly into the policy optimization process, unlike explicit preference modeling.

Another approach to reward modeling is the application of prompting techniques, such as the “LLM as a Judge” framework (Zheng et al., ), where large language models (LLMs) are prompted to act as evaluators without additional finetuning. Despite the potential of these methods, even advanced models like GPT-4o underperform when compared to dedicated RMs in more complex evaluations, such as those found in the RewardBench benchmark (Lambert et al., ).

Techniques for Inference-time Compute. During inference, the integration of a reward model with a proposal distribution (LLM) can be employed to refine the output responses to a given prompt. One notable paradigm in this context is Self-Consistency (Wang et al., ), which is designed for factual queries with extractable answers. In this approach, the language model selects the response it generates with the highest frequency across multiple samples. Optimizations such as Early Stopping Self-Consistency have been proposed, which terminate the sampling process early if a subset of responses shows strong consistency. However, these approaches face limitations due to their dependence on identifying the most “consistent" response from a large pool of discrete answers, restricting their applicability to tasks like multiple-choice or mathematical problem-solving.

In addition, search algorithms such as Best-of-N and Beam Search have been explored in works such as Snell et al. (), which leverage reward models to select the most promising candidate samples in reasoning tasks. In this work, we examine how to enhance the efficiency of token generation relative to the accuracy of the outputs, with the complexity of the query dictating the amount of inference compute allocation. Furthermore, for multi-step reasoning tasks, this work presupposes that the problem can be decomposed into discrete semantic steps—a potentially strong assumption in domains outside of well-structured fields like mathematics.

Generative Verifiers. Recent concurrent work, such as GenRM (Zhang et al., ) and Cloud (Ankner et al., ), demonstrate that token-prediction-based reward modeling outperforms traditional reward modeling techniques. In this work, we further introduce capability-aware self-evaluation, which allows a model to understand its own generation capability to enable dynamic allocation of computational resources during test-time inference. We also introduce mid-generation self-evaluations that can allow for the pruning of unpromising samples.

## 6 Discussion, Conclusion, and Limitations

As large language models (LLMs) evolve, enhancing response quality through inference-time computation becomes critical. Best-of-N sampling, a traditional approach, generates multiple response candidates and selects the best, but it incurs high computational costs due to its reliance on external reward models and fixed sample sizes. This work introduces a cost-effective alternative: capability-aware self-evaluations performed mid-generation. By appending a self-evaluation prompt to responses, LLMs predict the likelihood of generating a better response without requiring external reward models. Two adaptive inference-time techniques—adaptive sampling and early pruning—are also proposed, allowing LLMs to resample or discard suboptimal responses during generation dynamically. These methods significantly improve efficiency, with fewer samples yielding performance gains similar to a larger set, and early pruning reducing unnecessary computation. Experimental results demonstrate notable improvements in performance across tasks, achieving a 34% win rate against GPT-4 on AlpacaEval with 16 samples and increasing accuracy on GSM8K math problems from 84% to 91%. Overall, this approach optimizes inference-time compute allocation, making it more scalable and practical for diverse real-world applications.

The main limitation of our method is that adaptive sampling introduces latency, which otherwise would not be a problem with only parallel sampling. We minimized this with exponentially increasing batch sizes. We also introduced early pruning to save computation even in parallel sampling. For future work, it may be possible to predetermine the number of samples one needs to allocate to a given query based on the probabiliy that samples result in a Tie. While this would be significantly less efficient than adaptive sampling, it would eliminate latency for adaptive sampling. It may be possible to do different types of search with active capability-aware self-evaluations. Beam search or other similar types of search may be possible on general prompts with this new capability.

## Appendix A Appendix

### A.1 Reproducability

For reproducability, we provide the following details so that readers can replicate our results. Firstly, we provide details on how the dataset is constructed in Sections 3 and 4 to enable Capability-Aware and Mid-Generation Self-Evaluations for any model. Additionally, we provide algorithm pseudocode as seen in Algorithm 1, giving the reader transparency in how to replicate the adaptive sampling and annealing algorithm. Additionally, we provide details on how the dataset is constructed in Sections 3 and 4 to enable Capability-Aware and Mid-Generation Self-Evaluations for any model. Finally, we provide evaluation details in Section 4. We release a [public github implementation](https://github.com/rohinmanvi/Capability-Aware_and_Mid-Generation_Self-Evaluations) for reproducability.

### A.2 Comparing likelihoods

We briefly compare alternative probabilities one might consider to model using an on-policy pairwise preference dataset with ties.

**Table 3: Performance of different likelihoods with token prediction trained on a much smaller dataset of roughly 600 preferences.**
| Benchmark (Best-of-16) | $p_{\theta}(\textit{Win})$ | $p_{\theta}(\textit{Win}\mid\neg\textit{Tie})$ | $p_{\theta}(\textit{Win}\cup\textit{Tie})$ |
| --- | --- | --- | --- |
| Win Rate vs GPT-4 on AlpacaEval (%) | 27.7 | 30.1 | 29.6 |
| GSM8K Pass @ 1 (%) | 86.9 | 88.5 | 88.6 |

### A.3 Adaptive Sampling Results

We show adaptive sampling performance with various aspects of the algorithm removed.

**Table 4: Performance on GSM8K with adaptive sampling using capability-aware self-evaluations. This is without the annealing schedule.**
| Win-or-Tie Probability Threshold | 0.92 | 0.96 | 0.98 | 0.99 | 1.00 |
| --- | --- | --- | --- | --- | --- |
| Average Samples Used | 1.4 | 2.2 | 4.1 | 6.9 | 16.0 |
| Average Batches Used (Latency) | 1.3 | 1.6 | 2.0 | 2.9 | 5.0 |
| GSM8K Pass@1 (%) | 88.0 | 89.9 | 90.8 | 91.0 | 91.0 |
| Percent of Maximum Improvement | 55.9 | 83.8 | 97.1 | 100.0 | 100.0 |

**Table 5: Performance on GSM8K with adaptive sampling using the underlying reward model.**
| Reward Threshold | 0.119 | 0.133 | 0.150 | 0.163 | Inf |
| --- | --- | --- | --- | --- | --- |
| Average Samples Used | 1.4 | 2.2 | 4.1 | 6.9 | 16.0 |
| Average Batches Used (Latency) | 1.2 | 1.4 | 2.0 | 2.8 | 5.0 |
| GSM8K Pass@1 (%) | 86.9 | 87.9 | 89.0 | 89.9 | 91.0 |
| Percent of Maximum Improvement | 39.7 | 54.4 | 70.6 | 83.8 | 100.0 |

**Table 6: Performance on GSM8K with Best-of-N at varying number of samples.**
| Samples Used | 1 | 2 | 4 | 8 | 16 |
| --- | --- | --- | --- | --- | --- |
| GSM8K Pass@1 (%) | 84.2 | 87.7 | 88.9 | 89.9 | 91.0 |
| Percent of Maximum Improvement | 0.0 | 51.5 | 69.1 | 83.8 | 100.0 |

### A.4 Early Pruning Results

We show the results associated with early pruning shown in Figure 1 as well as tokens generated.

**Table 7: Win Rate vs GPT-4 on AlpacaEval with varying pruning strategies.**
| Win Rate vs GPT-4 | Samples Used (Pruned Samples) |  |  |  |  |
| --- | --- | --- | --- | --- | --- |
| on AlpacaEval (%) | 1 (Prune 15) | 2 (Prune 14) | 4 (Prune 12) | 8 (Prune 8) | 16 (Prune 0) |
| Prune @ 0 tokens (Random) | 21.2 | 25.1 | 29.1 | 31.4 | 33.8 |
| Prune @ 64 tokens | 25.3 | 28.3 | 31.2 | 33.4 | 33.8 |
| Prune @ 128 tokens | 28.3 | 30.4 | 33.2 | 33.8 | 33.8 |
| No Pruning | 33.8 | 33.8 | 33.8 | 33.8 | 33.8 |

**Table 8: Percent of Maximum Improvement achieved on AlpacaEval with different pruning strategies.**
| Percent of Maximum | Samples Used (Pruned Samples) |  |  |  |  |
| --- | --- | --- | --- | --- | --- |
| Improvement (%) | 1 (Prune 15) | 2 (Prune 14) | 4 (Prune 12) | 8 (Prune 8) | 16 (Prune 0) |
| Prune @ 0 tokens (Random) | 0.0 | 31.0 | 62.7 | 81.0 | 100.0 |
| Prune @ 64 tokens | 32.5 | 56.3 | 79.4 | 96.8 | 100.0 |
| Prune @ 128 tokens | 56.3 | 73.0 | 95.2 | 100.0 | 100.0 |
| No Pruning | 100.0 | 100.0 | 100.0 | 100.0 | 100.0 |

**Table 9: Average tokens generated per prompt on AlpacaEval with different pruning strategies.**
| Tokens Generated | Samples Used (Pruned Samples) |  |  |  |  |
| --- | --- | --- | --- | --- | --- |
| per Prompt | 1 (Prune 15) | 2 (Prune 14) | 4 (Prune 12) | 8 (Prune 8) | 16 (Prune 0) |
| Prune @ 0 tokens (Random) | 454 | 907 | 1,815 | 3,629 | 7,259 |
| Prune @ 64 tokens | 1,370 | 1,763 | 2,544 | 4,113 | 7,259 |
| Prune @ 128 tokens | 2,214 | 2,551 | 3,220 | 4,566 | 7,259 |
| No Pruning | 7,259 | 7,259 | 7,259 | 7,259 | 7,259 |

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
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.
soundness: 2
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.
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.
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.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
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?
soundness: 2
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).
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.
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.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
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?
soundness: 3
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.
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.
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.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: See weaknesses.
soundness: 3
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.
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.
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).

Meta-review: 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).

Decision: Reject