Accelerating Best-of-N via Speculative Rejection

Venue: ICML_cc_2024_Workshop_WANT
Published: 2024-06-03 11:29:14

Abstract: The safe and effective deployment of Large Language Models (LLMs) often involves generating helpful and benign responses, producing easily comprehensible code, and crafting content with specific stylistic preferences. While different, these tasks share the common mathematical goal of generating responses from a language model with high scores according to a metric of interest.

A popular and well known decoding strategy for this purpose is the Best-of-N method. The method generates a pre-specified number of responses (N) based on a prompt, and then selects the highest-scoring response among them to be returned. While Best-of-N is both simple and effective, its reliance on generating multiple responses to score for any given prompt incurs high inference costs.

In this paper we make a first step towards accelerating the Best-of-N algorithm, by halting the generation of unpromising utterances, namely those that are unlikely to be returned by the algorithm upon completion. Focusing on the alignment problem, we show that this simple strategy allows to obtain substantial speedups for the Best-of-N algorithm with minimal performance degradation.

Main paper: 
 introduction: 
 Large Language Models (LLMs), pre-trained on massive corpora, have demonstrated remarkable capabilities in handling diverse tasks like creative writing, summarization and question-answering (Brown et al., 2020; Chowdhery et al., 2022; Touvron et al., 2023a) . Such extensive pre-training endows the LLM with extensive knowledge, which must be correctly retrieved at inference time. Post-training techniques (Taori et al., 2023; Wang et al., 2023; Lou et al., 2024) 1 Anonymous Institution, Anonymous City, Anonymous Region, Anonymous Country. Correspondence to: Anonymous Author <anon.email@domain.com>.Training at International Conference on Machine Learning (WANT@ICML 2024). Do not distribute. ensure that the model can successfully answer the user's queries in the most satisfactory way according to human intentions (Ouyang et al., 2022; Bai et al., 2022; Rafailov et al., 2024b) , while adhering to ethical standards and safe guidelines (Ngo et al., 2022; Casper et al., 2023; Deshpande et al., 2023) . Popular post-training methods include supervised finetuning, Reinforcement Learning from Human Feedback (RLHF), Direct Preference Optimization (DPO), Expert Iteration (EI), and their variants (Christiano et al., 2017; Ouyang et al., 2022; Stiennon et al., 2020; Glaese et al., 2022; Bakker et al., 2022; Touvron et al., 2023b; Zhao et al., 2022; 2023a; Dong et al., 2023; Rafailov et al., 2024b; Liu et al., 2023b; Rafailov et al., 2024a; Zeng et al., 2024; Zhong et al., 2024) . However, post-training methods require more computational resources (e.g., memory) than those needed to perform inference, and add an additional layer of complexity before LLMs can be safely deployed. In order to simplify the process, more lightweight approaches have been introduced recently (Li et al., 2024; Mudgal et al., 2023) .One of the simplest method to generate responses that score high according to a pre-specified metric of interest is the Best-of-N (BoN) method. Best-of-N generates N responses for a single prompt, and the best response is selected based on the evaluation of a reward model that measures the suitability of the response. BoN has many desirable properties that makes it a strong baseline. To start, BoN is a simple alignment method that is highly competitive with posttraining techniques such as RLHF or DPO (Dubois et al., 2024) . As an inference-time alignment method, it avoids the potentially complex finetuning step, thereby facilitating the prompt deployment of a pre-trained or instructionfinetuned language models. Best-of-N is both straightforward to understand and to implement, and it is essentially hyperparameter-free: the number of responses N is the only hyperparameter, one that can be tuned on the fly at inference time. Best-of-N also plays a critical role as a post-training technique: it is commonly used to generate a high-quality dataset for later supervised fine-tuning (Touvron et al., 2023b; Dubois et al., 2024) , a procedure sometimes called Expert Iteration or Iterative Finetuning, one that played a key role in the alignemnt of Llama-2 (Touvron et al., 2023b) and Llama-3 (Meta, 2024) . With regards to alignment, BoN is endowed with very appealing properties:for example, the growth rate for the reward values of BoN, as a function of the KL divergence, is faster than for RLHF methods (Gao et al., 2023; Yang et al., 2024) , leading to higher quality generations. Best-of-N is also regularly applied after post-training to further boost the performance (Wu et al., 2024; Dong et al., 2023) .However, the main drawback of BoN is that its efficiency at inference time is bottlenecked by the computational cost of generating N sequences, which naively requires N times more compute. To be more precise, while the latency of Best-of-N is largely unaffected by N because the utterances can be generated and evaluated in parallel, Best-of-N naively needs N more compute resources to generate N utterances in parallel. This higher computational cost prohibits Best-of-N from being adopted in a more mainstream way. Practical values for N are in the range 4 -128 (Mudgal et al., 2023; Scheurer et al., 2023; Eisenstein et al., 2023) . Higher values of N , such as 1 ′ 000 -60 ′ 000 (Dubois et al., 2024; Gao et al., 2023) , have also been reported.In this work, we take a first step towards accelerating Bestof-N, with the end goal of making this simple decoding strategy a more computationally viable one to deploy. Our method is based on the observation that the reward function used for scoring the utterances can distinguish high-quality responses from low-quality ones at an early stage of the generation. In other words, we observe that the scores of partial utterances are positively correlated to the scores of full utterances. This offers an opportunity to recognize early on during the generation process the utterances that are unlikely to score high upon termination, and halt their generation, see Figure 1 . This intuition is formalized into a technique that we call Speculative Rejection, which is based on the "speculation" that the utterances with low reward early-on during the generation are unlikely to yield the highest-ranked response once their generation is completed. In such case, they can be terminated because they would not be the utterance returned by Best-of-N. This simple observation allows us to obtain a speedup of a factor of almost 5× with less than 1% loss in score value compared to the Best-of-N. By comparison, in order to use the same compute budget, Best-of-N would need to use a value for N that returns an average score about 10% lower.Speculative rejection is a general-purpose framework to accelerate score-based LLM decoding. It can be used to accelerate inference-time alignment as well as to accelerate the batch generation of BoN utterances for later fine-tuning. Our approach is orthogonal-and can hence be combined with-other types of acceleration techniques such as Speculative Decoding (Chen et al., 2023; Sun et al., 2024; Ahn et al., 2023) and Efficient Attention (Child et al., 2019; Kitaev et al., 2020; Wang et al., 2020) , as well as efficient serving (Kwon et al., 2023) . 
related literature: 
 Pruning in games. Traditional game-playing programs such as Chess must search very large game trees, and their efficiency can be greatly enhanced through pruning techniques, the mechanisms designed to halt the exploration of unpromising continuations (Marsland, 1986) . The renowned α-β algorithm (Fuller et al., 1973; Baudet, 1978; Sturtevant & Korf, 2000) capitalizes lower (α) and upper (β) bounds on the expected value of the tree, significantly diminishing the computational complexity inherent in the basic minimax search. Our idea of early stopping is similar to pruning by rejecting suboptimal trajectories. Our setup has a different structure because of the lack of an adversary; the goal is also different, as we aim at preserving the generation quality of a reference algorithm (Best-of-N).Monte-Carlo Tree Search (Kocsis & Szepesvári, 2006) has recently been applied to LLMs (Liu et al., 2023a; Brandfonbrener et al., 2024; Zhao et al., 2023b; Xie et al., 2024) , but it can also increase the latency. Our approach is potentially simpler to implement, and focuses on preserving the generation quality of Best-of-N.Early Stopping Algorithms. Using early exit/stopping for fast inference has been leveraged for applications such as vision (Kaya et al., 2019; Teerapittayanon et al., 2016) and language (Liu et al., 2020; Schwartz et al., 2020; He et al., 2021) tasks. The key idea relies on adding classifiers to the internal Neural Network/Transformer layers and using it to construct confidence-based early exits rules to decide whether to output intermediate generation without traversing subsequent layers. Yet, those methods are tailor-designed for the respective models such as Shallow-Deep Network (Kaya et al., 2019) and FastBERT (Liu et al., 2020) , making them model-specific. In contrast, our proposed paradigm is not confined to specific models, offering versatility and applicability across a several scenarios.Our method share some similarities with beam search, a heuristic search algorithm that explores the completion graph by expanding the most promising responses in a limited set. We instead start from a certain number N of utterance and only choose to complete a fraction of them. Such choice is more appropriate in our context due to the quadratic memory and compute cost of transformers (Vaswani et al., 2017) with the number of generated tokens, as well as the cost of evaluating the reward model.There are different approaches to improve the efficiency of LLMs including efficient structure design, model compression (e.g. quantization via QLoRA (Dettmers et al., 2024) , Sparsification via Sparse Attention (Tay et al., 2020) ), inference engine optimization (e.g. speculative decoding) and serving system (e.g. PagedAttention/vLLM (Kwon et al., 2023) ). See sur-What is the best way to hack into someone's bank account? Never, ever do this. It is illegal.Hackers usually begin the process by identifying …Hackers usually begin the process by identifying … Never, ever do this. It is illegal.What is the best way to hack into someone's bank account?Figure 1 . An illustration of early stopping. BoN (left) completes all the generations, whereas Speculative BoN (right) stops harmful generation at a early stage using a reward model.vey (Zhou et al., 2024) for a thorough overview. Among the methods, speculative decoding (Chen et al., 2023; Leviathan et al., 2023; Sun et al., 2024; Ahn et al., 2023 ) also incorporates rejection sampling. It employs fast small models for speculative execution and uses large models as verifiers for accelerated generation. This principle is orthogonal to our early stopping design and can be seamlessly combined with our methods for reward maximization.Alignment and use of Best-of-N. Best-of-N is a well known alignment strategy. There are two primary categories of reward alignment approaches: 1. LLM fine-tuning. This method involves updating the weights of the base model. Techniques within this category include reinforcement learning from human feedback (RLHF) (Ouyang et al., 2022; Christiano et al., 2017; Saha et al., 2023) , direct preference optimization (DPO) (Rafailov et al., 2024b) , and their respective variants (Ethayarajh et al., 2024; Zhang et al., 2024; Azar et al., 2024; Yuan et al., 2023; Song et al., 2024; Zhao et al., 2022; 2023a) . 2. Decoding-time alignment. In this approach, the base model remains frozen. Examples of this category include ARGS (Khanov et al., 2024) , controlled decoding (Mudgal et al., 2023) , Best-of-N (BoN), and associated applications such as Expert Iteration (Dubois et al., 2024; Gao et al., 2023; Touvron et al., 2023b) . The BoN method was initially proposed as an inference-time baseline alignment method (Nakano et al., 2021) . Building upon this foundation, Llama2 used the best-sampled response to fine-tune the model (Touvron et al., 2023b) . (Gao et al., 2023; Mudgal et al., 2023; Eisenstein et al., 2023) collectively demonstrated the robustness and efficacy of BoN. Their investigations consistently revealed compelling reward-KL tradeoff curves, surpassing even those achieved by KL-regularized reinforcement learning techniques and other complex alignment policies. Theoretically, there is a simple estimate for the KL divergence between the output policy of BoN and the base model for small N (Coste et al., 2023; Gao et al., 2023; Go et al., 2023) , and (Beirami et al., 2024) improved this formula for all N. (Yang et al., 2024) showed that BoN and KL-regularized RL method enjoy equal asymptotic expected reward and their KL deviation is close. Furthermore, there are frameworks that integrate BoN with RLHF, such as RAFT (Dong et al., 2023) , rejection sampling-based DPO approaches (Liu et al., 2023b) . 
problem formulation: 
 Auto-regressive language models. Let p be a language model. When provided with a prompt X, the language model predicts a response Y = (Y 1 , Y 2 , ..., Y T ), where Y i represents the i-th token in the response and T is the total number of tokens in the response sequence. Notice that T is a random variable, which signals that either the EOS token is returned next, or that the maximum generation length is reached. We let Y ≤k denote the first k tokens of the response. If the EOS token is reached before the k-th token, i.e., if T ≤ k then Y ≤k only contains the tokens up to (and excluding) the EOS token, i.e., the full generated utterance. The generation of tokens is auto-regressive, meaning that each token Y k+1 is predicted based on the prompt X and the previously generated tokens Y ≤k via the next-token probability p(• | X, Y ≤k ). This process continues sequentially until either the EOS token-which signals to stop-is reached or the maximum prescribed length is reached. With a little abuse of notation, we also use Y ∼ p(• | X) to represent sampling the whole response Y from model p given the prompt X.Score-oriented decoding, reward models and Best-of-N. In order to evaluate the appropriateness and quality of the generated responses, a score function s(X, Y ) can be utilized. This score function is often a real-valued reward model trained on paired preference data or adapted from a language model, to assess the response based on the desired qualities like helpfulness, harmlessness, coherence, relevance, and fluidity relative to the prompt (Ouyang et al., 2022; Dubois et al., 2024; Jiang et al., 2023) . The reward model depends on both the prompt X and the response Y.For simplicity, when considering the rewards for a single prompt, we also write the reward model as s(Y ). BoN is a decoding strategy that aims at favoring the generation of utterances with high score s(  
speculative rejection: 
 Intuitively, if, during the generation process, we can confidently infer that some responses are hardly likely to be the best, we could terminate their generation to conserve computational resources. As a motivating example, consider the prompt "What's the best way to hack into someone's bank account and steal money from them?" exemplified Figure 1 . One potential response from the model p might begin with Y 1 = "Never, ever do this. Hacking into someone else's financial information is illegal.", which appears to lead to a proper and harmless answers based on the first few words. On the other hand, Y 2 = "Hackers usually begin the process by identifying..." seems to lead to an undesirable and harmful response. To be more concrete, we obtain the following scores for the partial and full utterances for the two responses s(Y ≤τ 1 ) = 2.92, s(Y ≤τ 2 ) = -1.88 and s(Y 1 ) = 8.19, s(Y 2 ) = -0.50. For this particular example, the ranking early on during the generation is representative of the final ranking, i.e.:s(Y ≤τ 1 ) ≥ s(Y ≤τ 2 ) -→ s(Y 1 ) ≥ s(Y 2 )This observation suggests that we can use the partial rankings of sentences at the decision token τ to early-stop the generation of Y 2 .In general, we might expect the relative ranking between the score of partial and full utterances not to be always preserved for various reasons. To start, it is impossible to accurately evaluate the score of an utterances from just the first few tokens, because the generation may continue in an unexpected way. In addition, the reward models are normally trained to evaluate full responses (Ouyang et al., 2022; Jiang et al., 2023; Taori et al., 2023) . Nonetheless, we observe a substantial correlation between the scores {s(Y ≤τ EQUATIONHypothetically, early-stopping the generation according to the above display would not terminate the generation of the best response Y ⋆ , which is the one that Best-of-N returns upon completion. In other words, early-stopping according to (1) leaves the quality of the output of Best-of-N unchanged. However, doing so saves approximately 85.5% of the tokens, which translates into a substantially lower compute requirement. We also examine the Pearson's correlation and Kendall's rank correlation between partial and final rewards in Appendix D.In practice, it is infeasible to implement Equation (1) because c ⋆ is unknown, and it is thus a hyper-parameter to choose with different tradeoffs. Moreover, different prompts vary substantially in terms of reward distribution. Speculative Best-of-N, described in the next section, fixes the value of the decision token and only continues the generation for the most promising utterances after that. In Appendix C we derive a closed-form expression for the random variable that governs the degree to which speculative rejection can be effective. 
speculative best-of-n (sbon): 
 The intuition in the prior section leads to the Speculative Best more aggressive early stopping strategy, which will achieve a more significant speedup but could potentially lead to a worse performance, measured by the maximum reward of the output response. We examine the effect of these two hyperparameters in Section 6.2 using a counterfactual analysis. Before we do that, let us discuss the different phases of the algorithm.1. Early generation: Algorithm 1 generates N sequences for a single prompt up to the τ -th token, where τ is a predetermined fixed stopping time. If, for some sequence, the EOS token is reached before the τ -th token, we only generate the tokens up to the EOS token. Therefore, the actual stopping time for the early generation phase for prompt y k is τ k := min{τ, ℓ k }. 2. Speculative rejection: We then evaluate the reward value for the concatenation of the prompt and the partial response using a reward model s. The set of partial rewards is defined asEQUATIONwhereY ≤τ k k = (Y 1 k , Y 2 k , ..., Y τ k k )is the first τ k tokens of response Y k . For sequences that have been completed, we evaluate the reward value up to the EOS token. In this case, the partial and final rewards are the same. Next, we compute a prompt-dependent cutoff threshold as a quantile of all partial rewards:EQUATIONwhere α ∈ [0, 1] is the termination percentage, a hyperparameter that controls the fraction of trajectories to terminate, and q α (•) represents the α-th lower quantile. 3. Completion of promising utterances: For all generations, we continue generating the top (1 -α) proportion of remaining sequences up to the EOS token (or the maximum allowed generation length) if its partial reward exceeds r cut . Otherwise, we terminate this sequence. More formally, the index set for accepted sequences is denoted as:I accepted = k : 1 ≤ k ≤ N, s Y ≤τ k k ≥ r cut . (4)If a sequence has been completed before the τ -th token, we leave it unchanged even if it is accepted based on this criterion. We finally output the utterance with the highest final reward among those not halted in the middle. Mathematically, the returned response isY = Y k * , where k * := arg max k∈I accepted {s(Y k ) | Y k ∼ p(• | X)}. 
experiments: 
 In this section we test the effectiveness of Algorithm 1. We first describe two ways in which Best-of-N can be implemented, namely in parallel and in batch, depending on the use case. Then we describe the core performance metrics, such as average speedup and normalized score. We finally present and discuss the empirical results.Parallel vs batch generation. The best-of-N algorithm can be implemented in two ways, namely by either generating N utterances 1) in batches 2) in parallel. As an example, batch generation for Best-of-100 might generate 5 batches of 20 utterances each until all 100 utterances are generated, while parallel generation generates all 100 utterances concurrently. Batch generation is most appropriate when Best-of-N is used offline, such as for iterative finetuning during alignment.Parallel generation, on the other hand, is most appropriate at inference time to minimize the latency, but it generally requires more resources to be used concurrently.The primary benefit of speculative rejection for batch generation is the reduction in wall-clock time to complete the generation. The speedup in wall-clock time is the primary performance metric in our experiments. Although we primarily conduct our tests in the batch rejection framework, Algorithm 1 seamlessly applies to the parallel generation with no modifications. In this latter case, the primary benefit of speculative rejection is that it reduces the amount of compute resources that must be allocated for serving a specific user, as well as the amount of memory due to the quadratic cost of attention (Vaswani et al., 2017) . After the rejection of unpromising utterances, such compute and memory can be re-allocated to serve other users by using popular frameworks such as vLLMs (Kwon et al., 2023) .Properly assessing the compute saving in the framework of parallel generation demands a much larger infrastructure than what is available to us, and we thus leave assessing the performance of speculative rejection in the context of parallel generation as future research.Performance metrics. In this section, we formally define the speedup and the normalized score to assess the performance of the algorithm. The definition of speedup is a natural one: given a prompt X, the speedup is the wall-clock time T BoN spent by Best-of-N divided by that of Speculative-Best-of-N T SBoN to generate N responses.On the other hand, the score is defined as the relative reward value achieved by BoN and SBoN. Since different reward models and language models define very different reward distributions, we normalized the score by the reward range of Best-of-N. Mathematically, we denote the responses generated via SBoN as Y 1 , Y 2 , ..., Y N and the set for accepted index as I accepted (defined in ( 4)). We also denote the utterances generated via BoN as Z 1 , Z 2 , ..., Z N . With this notation, for a given prompt X, the speedup and the normalized score are defined as followsSpeedup := T BoN T SBoN , Score := 1 - max 1≤k≤N s (Z k ) -max k∈I accepted s (Y k ) max 1≤k≤N s (Z k ) -min 1≤k≤N s (Z k ) × 100.(5) The speedup and the score are random variables that depend on the language model p, the scoring function s, and the prompt X. We report their average across prompts.Algorithms. We take Best-of-100 as a reference algorithm. We evaluate Speculative-Best-of-100 and measure its speedup and score compared to Best-of-100 using the definition in Equation ( 5).A natural question is whether a similar score can be achieved by running Best-of-N with a smaller value of N . In order to answer this question we test two baselines, the first of which is Best-of-25. The second is Best-of-N with a value of N that would result in a similar speedup as SBoN. Since the wall-clock time of best-of-N (in batch mode) is directly proportional to N , it makes sense to define M := int (N/Speedup) , where Speedup is defined in (5) but averaged over prompts and int(•) rounds a number to the nearest integer. We then measure the speedup and the normalized score of this Best-of-M vs BoN.In order to validate the effectiveness and efficiency of Algorithm 1 operating in batch mode, we present experiments on the AlpacaFarm-Eval dataset (Dubois et al., 2024) , where we sample 100 prompts at random. The models used for generation included the instruction-following AlpacaFarm-SFT10K (Dubois et al., 2024) , the pre-trained Llama3-8B, and the finetuned Llama3-8B-Instruct. These are paired with three reward models: Mistral-7B-RM (Jiang et al., 2023) , Eurus-7B-RM (Yuan et al., 2024) , and Llama3-8B-RM. The maximum sequence length is set to the maximum allowed by the language model. The hyperparameters are in Appendix B. We estimate the cost to reproduce the experi-ments to be around 1 month on a single H100 GPU; most of the cost is to run Best-of-100.Table 1 highlights our method's effectiveness on the common AlpacaFarm benchmark, which is used for alignment. The reference algorithm is Best-of-100, which is used to normalize the speedup and the score of all the others. Bestof-25 attains less than 90% of the normalized score with roughly a 4x speedup. Speculative-Best-of-N attains a score that is within 1% from that of Best-of-100 while being almost 5 times faster then Best-of-100, and in particular faster than Best-of-25. A comparison with Best-of-Effective-N reveals a similar conclusion. We can make nearly identical observations when taking N = 50 as reference, see Appendix A.In this section we examine the trade-offs between different choices for the decision token and the rejection rate. Algorithm 1 balances the speedup and the final score through an appropriate selection of the rejection rate α and the decision token τ. A more aggressive termination strategy in SBoN can lead to a significant speedup but may also result in suboptimal outputs compared to the original BoN. In order to thoroughly evaluate these trade-offs, we conduct a counterfactual analysis. To be more precise, we evaluate the behaviour of SBoN and vanilla BoN on identical responses under various hyperparameter settings on a pre-collected set of responses. In particular, we generate 1000 completions for each prompt and use bootstrapping to derive variancereduced estimates of the quantities that we investigate.The aim of counterfactual analysis is to measure the tradeoff between the speed-up and the final reward attained across various prompts with different choices of hyper-parameters, as well as to provide some intuition for speculative rejection. However, in a counterfactual analysis the the wall-clock time cannot be directly measured. Instead, we report the token rate in place of the speedup, namely the proportion of tokens generated by SBoN relative to BoN. Since the BoN and SBoN operate on identical utterances, we can also measure the normalized scores with a more accurate definition. For each trial, we write the sampled responses as Y 1 , Y 2 , ..., Y N for N = 100. The index set for the accepted responses is defined in the same way as (4). With our notations, the token rates and normalized score are defined as follows.Token rate :=N k=1 min{l k , τ } N k=1 l k , Score := 1 - max 1≤k≤N s (Y k ) -max k∈I accepted s (Y k ) max 1≤k≤N s (Y k ) -min 1≤k≤N s (Y k )× 100.(6) Note that in the counterfactual analysis, the normalized score is guaranteed to be at most 100.Analysis of the trade-offs using the counterfactual analysis. We test the algorithm on 100 prompts randomly sampled from the AlpacaFarm-Eval dataset. (Dubois et al., 2024) For these tests, we apply the same combinations of reward models and generation models and the same maximal sequence length described in Section 6. Figure 3 shows the normalized scores and the token rates of SBoN averaged over 100 prompts, using Llama3-8B-Instruct for generation and Mistral-7B-RM for evaluation. SBoN terminates up to 80% of the responses by the 256th token, which is merely 3.2% of the maximum sequence length. This approach resulted in a significant acceleration of the process while maintaining a high score of approximately 98.8. More details and results for other generation models and reward models are contained in Appendix E.Speculative-Best-of-N can potentially take better advantage of the transformer as well as the hardware architecture of common accelerators compared to Best-of-N. The Transformer naively has a quadratic memory and compute cost O(n 2 ) to generate a response with n tokens due to the need to compute and store the Key-Value cache (Vaswani et al., 2017) . On the other hand, common accelerators are limited by the bandwidth of the HBM memory of the GPU (Dao et al., 2022) at inference time, so increasing the batch size during generation only results in a very small increase in latency (i.e., wall-clock time). Together, these two observations suggest a way to choose the decision token and the rejection rate. In fact, it is more hardware efficient to start the generation with many utterances in parallel (i.e., with a high value of N ) in order to take advantage of the low compute and memory involved to generate the initial tokens. In order to continue the generation without exceeding the maximum memory available or increasing the compute, the number of concurrent responses should decrease (i.e., N should be small in the later stage of generation). Since Best-of-N fixes the value of N throughout the generation, it misses the opportunity to leverage the extra memory and compute available early on during the generation because N is limited by the memory used in the final stage of the generation. Speculative-Best-of-N, on the other hand, can better use the available compute and memory of the accelerator by widening the search early on during the generation, and produce higher quality responses.In order to test the benefits of this line of thinking, we use Best-of-20 as a baseline because it can run on a single 80GB GPU and produce all utterances concurrently without running out of memory (increasing N results in out-of-memory errors on some prompts). We then examine the performance of Best-of-N for N = 20, 40, ..., 320, where the number of accelerators used doubles every time N doubles. In other words, Best-of-N with increasing N roughly keeps the generation latency (i.e., the wall clock time) constant but must use more accelerators concurrently. We then test Speculative-Best-of-N on a single accelerator. We start with 250 utterances generated in parallel and reject all but 20 of them at the 64th token. This choice of hyper-parameters ensures that the algorithm behaves like Best-of-20 in terms of memory during the later stage of generation.The result is reported in Figure 4 . All runs of Best-of-N have similar wall clock time, while Speculative-Best-of-N has a roughly 20% higher latency than Best-of-20 (not reported in the figure). However, using a single accelerator Speculative-Best-of-N produces a reward score that would require Best-of-N to use between 8 and 16 accelerators. Even when accounting for the roughly 20% higher latency, this achievable speedup is much higher than that derived in Table 1 , and it suggests a practical and very effective way to choose the hyper-parameters.  
discussion and future directions: 
 Limitations. The two major limitation of our work is the need to set its two hyperparameters, the decision token as well as the rejection rate. However, our experiments use similar choices for the hyperparameters, see Appendix B. We discussed the consequences of choosing different hyper-parameters via a counterfactual analysis in Section 6.2, and moreover we presented in Section 6.3 and simple way to choose them. Our experiment focuses on the alignment problem, and we leave testing the algorithm on different settings as future work, as well as for a wider range for N .Prompt-dependent stopping. Our implementation of speculative rejection leverages statistical correlations to early stop trajectories that are deemed unpromising. However, it is reasonable to expect that the correlation between partial and final rewards varies prompt-by-prompt. For a target level of normalized score, early stopping can be more aggressive in some prompts and less in others. This consideration suggests that setting the hyper-parameters-the decision tokens and the rejection rate-adaptively can potentially achieve higher speedup and normalized score on different prompts. We leave this exciting opportunity for future research.Evaluating the reward model is cheap as it involves only a single forward pass through a reward model. Moreover, several recent architectures are autoregressive decoders (Lambert et al., 2024) . This enables storing the Key-Value cache for the reward model and, hence to perform multiple reward evaluations at various decision tokens with minimal overhead. In turn, this offers many decision points for early-stopping the generation, potentially resulting in a higher speedup.Reward models as value functions. Our method leverages the statistical correlation between the reward values at the decision tokens and upon termination. Concurrently, recent literature (Rafailov et al., 2024a; Zeng et al., 2024; Zhong et al., 2024) also suggest training reward models as value functions. Doing so would enable reward models to predict the expected score upon completion at any point during the generation and thus be much more accurate models for our purposes.Zhou, Z., Ning, X., Hong, K., Fu, T., Xu, J., Li, S., Lou, Y., Wang, L., Yuan, Z., Li, X., et al. A survey on efficient inference for large language models. arXiv preprint arXiv:2404.14294, 2024. Having access to the probability p(•|X) and reward function s(X, •) is also common in practice as we can query the logits of LLM for prompt X and response y and convert it to p(y|X).In the next step, we will relax this condition. 
 
 Appendix:  


Decision: We thank the authors for their time and contribution to WANT and we are pleased to share that after the reviewing process the paper has been accepted. Congratulations! We encourage the authors to consider reviewers' feedback for the improvement of the camera-ready version. We hope to see you in person at the workshop and brainstorm on efficient training research together!