PaperlayerView exact revision on arXiv

Source-linked AI summary

Fast Best-of-N Decoding via Speculative Rejection

Hanshi Sun, Momin Haider, Ruiqi Zhang, Huitao Yang, Jiahao Qiu, Ming Yin, Mengdi Wang, Peter Bartlett, Andrea Zanette

arXiv:2410.20290v2cs.CL

TL;DR

Post-training alignment complicates deployment, while Best-of-N provides strong inference-time alignment but becomes computationally impractical at large sample counts. Speculative Rejection uses partial reward signals to stop unpromising generations early. It achieves comparable reward with 16–32 times greater computational efficiency than Best-of-N in the reported evaluation.

  • Problem

    Post-training adds deployment complexity, while Best-of-N alignment requires too many inference resources at large N to remain computationally viable.

    Source: §1 · pp. 1–2

  • Method

    Speculative Rejection repeatedly scores partial generations and halts responses unlikely to achieve high final rewards, dynamically reducing the batch size.

    Source: §1 · p. 2, §4.2 · p. 6

  • Results

    16–32 GPUs are required by Best-of-N to achieve a reward comparable to Speculative Rejection on a single GPU, with similar latency.

    Source: §1 · p. 3, §5.1 · p. 8

  • Takeaways & Limitations

    Speculative Rejection provides a general-purpose, simpler-to-implement framework for accelerating reward-oriented LLM decoding.

    Source: §6 · p. 10

    The correlation between partial and final rewards varies across prompts, so fixed rejection rates may not provide the best speedup and normalized score everywhere.

    Source: §6 · p. 10

Abstract

from arXiv · show

The safe and effective deployment of Large Language Models (LLMs) involves a critical step called alignment, which ensures that the model's responses are in accordance with human preferences. Prevalent alignment techniques, such as DPO, PPO and their variants, align LLMs by changing the pre-trained model weights during a phase called post-training. While predominant, these post-training methods add substantial complexity before LLMs can be deployed. Inference-time alignment methods avoid the complex post-training step and instead bias the generation towards responses that are aligned with human preferences. The best-known inference-time alignment method, called Best-of-N, is as effective as the state-of-the-art post-training procedures. Unfortunately, Best-of-N requires vastly more resources at inference time than standard decoding strategies, which makes it computationally not viable. In this work, we introduce Speculative Rejection, a computationally-viable inference-time alignment algorithm. It generates high-scoring responses according to a given reward model, like Best-of-N does, while being between 16 to 32 times more computationally efficient.

1 Introduction

Inference-time alignment avoids the deployment complexity of post-training, but Best-of-N becomes computationally impractical at large N. Speculative Rejection addresses this by using partial-response rewards to stop unlikely high-scoring generations early, achieving comparable performance with substantially fewer resources.

Best-of-N

  • Best-of-N can match strong post-training alignment while remaining simple, hyperparameter-light, and useful for generating high-quality data.Its only inference-time hyperparameter is N, and it is also used in Expert Iteration and rejection-sampling-based alignment.Source: §1 · p. 2
  • Large-N Best-of-N is computationally impractical because generating thousands of sequences requires dozens or hundreds of accelerators.Although sequences can be generated in parallel with largely unchanged latency, memory limits can require multiple GPUs; useful N values may reach 1000–60000.Source: §1 · p. 2

Speculative Rejection

  • Speculative Rejection uses the positive correlation between partial and final rewards to halt generations unlikely to score highly when completed.This enables performance comparable to Best-of-N for N > 1000 using one accelerator with similar latency.Source: §1 · p. 2
  • Speculative Rejection begins with a large batch and repeatedly rejects unpromising sequences, dynamically reducing batch size before GPU memory is exhausted.Reward-model queries during generation identify which responses should continue to completion.Source: §1 · pp. 2–3, Figure 1 · p. 2

Empirical evaluation

  • Speculative Rejection achieves reward comparable to Best-of-N using one GPU versus 16–32 GPUs for Best-of-N, with similar latency.The evaluation uses AlpacaFarm with varied generative and reward models and also examines win-rate and probability maximization.Source: §1 · p. 3

2 Related Literature

Prior work aligns language models through weight updates or decoding-time selection, while efficiency methods and search-based approaches motivate Speculative Rejection’s model-agnostic pruning perspective.

Early Stopping Algorithms

  • Unlike model-specific early-exit methods, the proposed paradigm is not confined to particular architectures and applies across several scenarios.Early-exit methods add internal classifiers and confidence-based stopping rules to specific models such as Shallow-Deep Network and FastBERT.Source: §2 · p. 3

Alignment and Use of Best-of-N

  • The method resembles beam-search expansion but completes only a fraction of N utterances, reducing costs associated with linear KV-cache memory and quadratic reward-model evaluation.Source: §2 · p. 3

Inference Efficiency in LLMs

  • Speculative decoding uses small models for proposal generation and large models as verifiers, and can be combined with Speculative Rejection for reward maximization.Source: §2 · p. 3

Alignment and Use of Best-of-N

  • Best-of-N is an effective decoding-time alignment strategy with strong reward-KL tradeoffs, but it is computationally impractical even for moderate N.It samples multiple responses, scores them, and returns the highest-scoring response.Source: §2 · p. 3

Pruning in Games

  • Its early stopping also resembles game-tree pruning by rejecting suboptimal trajectories, but it lacks an adversary and preserves Best-of-N generation quality.Source: §2 · p. 4
  • Compared with MCTS-based LLM approaches, the approach is potentially simpler to implement, focuses on preserving Best-of-N quality, and does not require the training used by some alignment methods.Source: §2 · p. 4

3 Preliminaries

The preliminaries define autoregressive response generation and reward-based inference-time alignment, culminating in Best-of-N selection and its computational limitation.

Language Model Generation

  • An autoregressive language model generates response tokens conditionally from a prompt and previously generated tokens until producing an EOS token.Source: §3 · p. 4

Inference-time Alignment

  • A reward model assigns a real-valued score to a prompt-response pair based on qualities such as helpfulness, harmlessness, coherence, relevance, and fluidity.Source: §3 · p. 4
  • Inference-time alignment uses a frozen autoregressive model to generate responses with as high a reward-model score as possible.Source: §3 · p. 4
  • Best-of-N samples N independent responses, scores each, and returns the response with the highest reward.Source: §3 · p. 4
  • Best-of-N is extremely effective but computationally impractical even for moderate values of N.Source: §3 · p. 4

4 SPECULATIVE REJECTION

SPECULATIVE REJECTION uses early partial-reward signals to stop unlikely-best responses, dynamically shrinking the batch while preserving promising candidates. It thereby simulates large-N Best-of-N decoding without exhausting GPU memory.

4.1 Observation

  • Partial rewards are positively correlated with final rewards, enabling early stopping of responses whose decision-token scores fall below the best-response threshold.The example compares partial and completed utterance scores and uses the threshold defined by the best response’s partial reward.Source: §4.1 · p. 5, Figure 2 · p. 5
  • 85.5% of tokens can be saved while preserving Best-of-N output quality when early stopping does not terminate the best response.The observation relies on partial and final rewards being sufficiently correlated for the example.Source: §4.1 · p. 5
  • The method addresses unknown prompt-specific thresholds and decision tokens by determining stopping points from GPU memory capacity during decoding.It continues only the most promising utterances until the next decision token or completion.Source: §4.1 · p. 5

4.2 Algorithm

  • SPECULATIVE REJECTION begins with an oversized batch, evaluates partial responses near the memory limit, and rejects a fraction of the lowest-scoring sequences to avoid out-of-memory failure.Each rejection round repeats when the GPU approaches its memory limit.Source: §4.2 · p. 6
  • The rejection threshold is the α-th lower quantile of partial rewards, and sequences above it continue while rejected sequences terminate.The rejection rate α controls the fraction of trajectories terminated, and the batch size is updated from the accepted set.Source: §4.2 · p. 6
  • The algorithm returns the highest-final-reward response among sequences not halted during decoding, dynamically reducing batch size after each rejection round.This simulates a larger-N Best-of-N initial phase while using available GPU memory more efficiently with minimal latency increase.Source: §4.1 · p. 6, §4.2 · p. 7

5 Experiments

The experiments evaluate SPECULATIVE REJECTION across efficiency, generation quality, and probability maximization, finding that it can match Best-of-N rewards with substantially fewer GPU resources while preserving quality.

5.1 Efficiency Evaluation

  • With Llama3-8B and reward model RM-Mistral-7B, SPECULATIVE REJECTION achieves a reward score that would require Best-of-N to use between 16 and 32 GPUs; Best-of-120 is the single-80GB-GPU baseline.The evaluation uses AlpacaFarm on a DGX node with H100 GPUs.Source: §5 · p. 7, §5.1 · p. 8
  • SPECULATIVE REJECTION uses fewer GPU resources to achieve higher scores than Best-of-N across evaluated generative-model and reward-model pairs.Figure 3 reports relative GPU compute and improvement score, with Best-of-120 as the single-GPU baseline.Source: §5 · p. 7, §5.1 · p. 8, Figure 3 · p. 8
  • The method provides less improvement for Llama-3-8B-Instruct than for base models because its stronger alignment produces shorter responses and fewer rejection rounds.The comparison specifically names Mistral-7B and Llama-3-8B as base-model references.Source: §5.1 · p. 9
  • Lower rejection rates preserve more candidate responses and yield higher final scores, but slightly increase latency through reward-model computation and larger intermediate batches.Higher rejection rates are more aggressive and require fewer rejection rounds, whereas lower rates retain responses that may later score highly.Source: §5.1 · p. 9

5.2 Win-rate Evaluation

  • SPECULATIVE REJECTION maintains generation quality while achieving notable speedups on most model and reward-model combinations under GPT-4-Turbo evaluation.Quality is assessed with win-rate and length-controlled win-rate against the Bo120 baseline.Source: §5.2 · p. 9, Table 1 · p. 9

5.3 Maximization of the Probability of the Generated Utterances

  • SPECULATIVE REJECTION is a general reward-maximizing decoding strategy that can use probability as its reward function, beyond human-preference alignment.The method rejects the top α fraction of responses with the lowest average probability during each rejection round.Source: §5.3 · p. 9
  • SPECULATIVE REJECTION outperforms Best-of-N when maximizing model probability, producing higher-probability responses with remarkable speedup.The reward is average response log-probability, and Table 2 reports perplexity results across models.Source: §5.3 · p. 9, Table 2 · p. 10

6 Limitations and Conclusions

The paper presents Speculative Rejection as a general-purpose method for accelerating reward-oriented decoding, while identifying opportunities to improve prompt adaptivity and reward prediction.

  • Speculative Rejection accelerates reward-oriented decoding and is presented as simple to implement with substantial speedups over Best-of-N.Source: §6 · p. 10

Prompt-dependent Stopping

  • Prompt-dependent correlations between partial and final rewards may require adaptive rejection rates to balance speedup and normalized score across prompts.Source: §6 · p. 10

Reward Models as Value Functions

  • Reward models trained as value functions could predict completion scores more accurately and, according to the paper, enable optimal speedup.Source: §6 · p. 10

A Detailed Authors’ Contributions

The authors’ contributions span idea conception, theory, implementation, experiments, statistical analysis, writing, and project support.

B Correlation between partial and final rewards

This section examines whether partial rewards correlate with final rewards for responses generated for a single prompt, using Pearson and Kendall’s tau measures.

NeurIPS Paper Checklist

The checklist records the paper’s claimed contributions, reproducibility materials, experimental disclosures, limitations discussion, and ethics-related responses.

  • The authors state that the paper’s abstract and introduction accurately reflect its contributions and scope.Source: NeurIPS Paper Checklist · p. 17
  • The paper discusses limitations in its final section, while the checklist emphasizes documenting assumptions, robustness, scope, and computational efficiency.Source: NeurIPS Paper Checklist · p. 17

3. Theory Assumptions and Proofs

5. Open access to data and code

  • The authors provide code, evaluation datasets, public model checkpoints, and README instructions intended to support reproduction of the experiments.Source: NeurIPS Paper Checklist · pp. 18–19

4. Experimental Result Reproducibility

  • The experimental reproducibility responses address setup, statistical significance, compute resources, and the absence of reported error bars because experiments are costly.Source: NeurIPS Paper Checklist · p. 19
  • The authors report no direct societal impact, do not release data or models, cite existing assets, include a README, and state that the work involves neither human subjects nor crowdsourcing.Source: NeurIPS Paper Checklist · pp. 20–21
Loading 2410.20290v2…