PaperlayerView exact revision on arXiv

Source-linked AI summary

Sampling-Efficient Test-Time Scaling: Self-Estimating the Best-of-N Sampling in Early Decoding

Yiming Wang, Pei Zhang, Siyuan Huang, Baosong Yang, Zhuosheng Zhang, Fei Huang, Rui Wang

arXiv:2503.01422v3cs.CLcs.AI

TL;DR

Best-of-N scaling is costly because fully generating all samples consumes GPU memory, while reward models add latency, memory, and training-data burdens. ST-BoN uses early internal consistency in model states to truncate weaker samples without reward models, reporting major memory, latency, and cost savings alongside accuracy gains at similar cost.

  • Problem

    Best-of-N’s cost-performance trade-off is underexplored because full generation consumes substantial GPU memory and reward models add memory, latency, training-data, and generalization burdens.

  • Method

    ST-BoN uses early consistency in latent model states to identify and truncate suboptimal samples without fully generating all N samples or using reward models.

  • Results

    ST-BoN reduces dynamic GPU memory by over 80% and latency by up to 50%, saves 70% to 80% cost at Full-BoN performance, and improves accuracy by 3 to 4 points at similar cost.

  • Takeaways & Limitations

    ST-BoN provides a flexible cost-performance trade-off across domains, including objective and subjective tasks, under limited resources.

  • Takeaways & Limitations

    ST-BoN requires access to hidden states and still has room to adapt its window length to task complexity.

Abstract

from arXiv · show

Test-time scaling enhances large language model performance by allocating additional compute resources during inference. Best-of-N (BoN) sampling serves as a common sampling-based scaling technique, broadening the search space in parallel to find better solutions from the model distribution. However, its cost-performance trade-off is still underexplored. Two main challenges limit the efficiency of BoN sampling: (1) Generating N full samples consumes substantial GPU memory, reducing inference capacity under limited resources. (2) Reward models add extra memory and latency overhead, and training strong reward models introduces potential training data costs. Although some studies have explored efficiency improvements, none have addressed both challenges at once. To address this gap, we propose Self-Truncation Best-of-N (ST-BoN), a decoding method that avoids fully generating all N samples and eliminates the need for reward models. It leverages early sampling consistency in the model's internal states to identify the most promising path and truncate suboptimal ones. In terms of cost, ST-BoN reduces dynamic GPU memory usage by over 80% and inference latency by 50%. In terms of cost-performance trade-off, ST-BoN achieves the same performance as Full-BoN while saving computational cost by 70%-80%, and under the same cost, it can improve accuracy by 3-4 points.

1 Introduction

Best-of-N expands inference-time search by sampling and reranking multiple responses, but Full-BoN memory costs and reward-model overhead limit efficient deployment. ST-BoN addresses both issues by using early internal consistency to truncate weaker samples without reward models.

  • Best-of-N Sampling: BoN samples N responses and reranks them to select the best candidate, commonly using self-consistency or a trained reward model.Self-consistency selects the most frequent answer, while reward models score candidates for richer tasks.
  • Efficiency Challenges: Full-BoN must fully generate all N samples, creating substantial memory overhead for long reasoning sequences.GPU parallelization can reduce latency but cannot remove the additional memory requirement.
  • Efficiency Challenges: Reward models add memory and latency costs, require costly high-quality feedback data, and may generalize poorly across domains.The passage specifically contrasts domains such as mathematics and open-ended question answering.
  • Related Work: Existing efficiency methods reduce some generation overhead but still depend on reward-model inference or assumptions that limit suitability for reasoning tasks.The cited approaches score or discard prefixes, yet do not overcome both identified challenges simultaneously.
  • ST-BoN: ST-BoN uses early sampling consistency in latent states to preserve a promising sample and truncate suboptimal samples without reward-model intervention.The method is designed to reduce memory use and accelerate generation.
  • Cost: Over 80% lower dynamic GPU memory and up to 50% lower inference latency are reported for ST-BoN versus Full-BoN.These gains exclude the overhead of additional reward models.
  • Cost-Performance Trade-offs: At matched Full-BoN performance, ST-BoN saves 70% to 80% computational cost; at similar cost, it improves accuracy by 3 to 4 points.The reported results also state applicability across a wide range of domains.

2 Preliminary

Autoregressive language-model decoding generates token sequences step by step, while sampling methods introduce diversity by drawing tokens from a truncated model distribution. Best-of-N evaluates multiple independently sampled sequences and selects the highest-scoring one.

  • Autoregressive Decoding: An autoregressive language model generates a token sequence Y_1:T by sampling each token sequentially until the end-of-sequence token.The sequence has generation length T and ends with y_T = [EOS].
  • Autoregressive Decoding: Greedy decoding chooses the highest-probability token, whereas top-k and top-p sampling draw diverse tokens after distribution truncation.These decoding choices determine how candidate sequences are produced.
  • Best-of-N Sampling: Best-of-N sampling independently generates N sequences, assigns each a scalar score, and selects the sequence with the maximum score.The procedure broadens the candidate set before reranking.
  • Scoring and Consistency: A reward model provides an external scoring approach, while self-consistency selects candidates using agreement among extracted answers.Self-consistency is described as an unsupervised alternative associated with answer accuracy.

3 ST-BoN: Self-Truncation Best-of-N Decoding

ST-BoN estimates which of N samples is most promising during early decoding, using internal consistency signals rather than full generation or reward models. Its theory links early similarity to final consistency, while its implementation uses latent-state distances and a buffer window to stabilize selection.

  • Motivation: ST-BoN relies on the premise that early consistency can foreshadow final consistency, allowing weaker samples to be truncated before completion.The method targets efficiency bottlenecks from full generation and reward-model use.
  • Theoretical Support: If early partial sequences are closer under a metric, the probability that their completed sequences satisfy a consistency threshold increases.The comparison is made between a primary sample and the other N−1 samples.
  • Theoretical Support: Under local Lipschitz and bounded-increment assumptions, expected final deviations satisfy E[S_T] ≤ Γ^(T−t)S_t, yielding a high-probability consistency bound.The proof uses maximal coupling and Markov’s inequality to connect early and final distances.
  • Theoretical Support: As early consistency increases, the lower bound on the probability of final consistency also increases, supporting early estimation of a promising path.This is the stated theoretical feasibility for using sampling consistency during early decoding.
  • Estimation Time: ST-BoN begins self-estimation when samples become pairwise inconsistent, checking sequence equality during parallel inference and terminating any sample that reaches [EOS].The earliest qualifying time is denoted c.
  • Internal Consistency Measure: Because early text differences are limited, ST-BoN represents samples with informative hidden states from the model’s internal computation.This latent representation supports consistency measurement before semantic information is fully developed.
  • Internal Consistency Measure: Using Chain-of-Embedding features, ST-BoN measures pairwise latent-space distance and selects the sample with the lowest average distance to the other samples.The lowest-score sample becomes the optimal estimate at time c.
  • Buffer Window: A buffer window extends decoding τ steps after c, with τ proportional to c, and chooses the most frequent optimal estimate across the window.The window is intended to reduce randomness from making a decision at one early time point.

4 Cost Analysis: Towards Efficient Best-of-N Sampling

ST-BoN’s cost analysis compares its GPU-memory and latency overhead with Full-BoN. Early estimation occurs well before full completion, enabling memory reductions that exceed 80% for N ≥5 while shifting later generation to single-sample inference.

  • GPU Memory: GPU memory is mainly constrained by KV-cache accumulation, whose peak usage grows with sampling size N and generation length T.The analysis excludes basic model-weight memory and focuses on peak KV-cache overhead.
  • GPU Memory: ST-BoN’s earliest estimation time c occurs much earlier than the full generation length T across MATH, TheoremQA, and MMLU.The observed Pc distributions shift left relative to PT.
  • GPU Memory: 80%: The memory reduction rate RD surpasses 80% for sampling sizes N ≥5.RD increases with N toward an upper limit determined by the relative expected estimation and full-generation lengths.
  • GPU Memory: RD increases with sampling size N and task complexity, with MATH attaining a higher optimization upper limit because its full generations are longer.Variation in expected full-generation length contributes more to domain differences than variation in expected estimation time.
  • Wall-Clock Time: ST-BoN resumes single-sample generation after c + τ, while equality checks and self-estimation within the buffer window may add serial-processing latency.Latency is evaluated using wall-clock time, including the trade-off between early truncation and buffer-window processing.

5 Experiments

Experiments evaluate ST-BoN across objective and subjective tasks, models, and cost metrics. ST-BoN largely preserves Full-BoN performance while substantially improving cost-performance trade-offs and maintaining effectiveness across domains.

  • Experimental Setup: The evaluation covers objective datasets in mathematics, theorem application, science reasoning, and general knowledge, plus subjective summarization and instruction-following tasks.Models include Qwen2.5-7B-Instruct, Llama3-8B-Instruct, Mistral-7B-Instruct-v0.3, and Qwen2.5-72B-Instruct.
  • Experimental Setup: ST-BoN uses early truncation to reduce memory and time costs, with computational cost evaluated relative to greedy decoding.The evaluation measures memory and wall-clock time, while performance uses task-specific metrics including Accuracy, Rouge-L, Human Scoring, and Win Rate.
  • Objective Tasks: At N = 20, ST-BoN shows no significant performance drop compared with Full-BoN without a reward model when cost is ignored on the objective-task comparison.The comparison uses Qwen2.5-7B-Instruct on four objective tasks.
  • Objective Tasks: 70–75% lower costs let ST-BoN match Full-BoN without a reward model at N = 20 using N = 20 ∼30, while 75–80% lower costs achieve parity with Full-BoN with a reward model using N = 40 ∼80.The comparison is reported when Full-BoN uses N = 20.
  • Objective Tasks: Under the same computational cost, ST-BoN reaches nearly N = 80 and achieves 3-4 point higher accuracy than the Full-BoN baselines across all objective datasets.The cited examples compare against Full-BoN without a reward model at N = 10 and with a reward model at N = 3.
  • Subjective Tasks: On subjective tasks, ST-BoN reaches Full-BoN with a reward model’s instruction-following performance while reducing computational cost by approximately 50%.Full-BoN without a reward model performs poorly on subjective tasks, while its reward-model variant underperforms on summarization because of domain mismatch.

6 In-depth Analysis

The analysis examines ST-BoN’s components, window length, sampling robustness, and behavior in ambiguous cases. Results support CoE and buffering, with task- and temperature-dependent cost-performance effects.

  • Component ablation: ST-BoN with CoE and a buffer window achieves stronger early-final consistency than random, semantic, or string-based estimation.The buffer window consistently outperforms its removal, while semantic and string variants remain near the random baseline.
  • Component ablation: CoE consistency under full generation outperforms self-consistency with majority voting, supporting its effectiveness for measuring early consistency.
  • Window length τ: When τ < c, performance gain relative to cost increases faster; beyond c, the gain slows significantly, making c the experimentally preferred balance.Ignoring cost, performance can continue fluctuating upward with larger τ, especially on MATH, where longer outputs may require more information.
  • Sampling strategy: Temperature changes affect performance and cost more than top-k or top-p changes in the MMLU sampling-strategy ablation.Increasing temperature makes the distribution more uniform and random, causing earlier estimation and lower computational cost.
  • Case studies: ST-BoN better addresses ambiguous cases where samples are inconsistent or majority voting is highly random.

7 Related Work and Comparisons

Related work improves BoN efficiency through depth pruning, breadth pruning, or training-time distillation. ST-BoN differs by targeting inference-time pruning without reward models while preserving parallel sampling advantages.

  • Scaling paradigms: BoN expands reasoning through parallel sampling, complementing sequential scaling methods that lengthen a single reasoning path.
  • BoN foundations: BoN commonly selects among sampled responses using self-consistency or reward-model scoring.
  • Depth pruning: Depth-pruning methods discard suboptimal prefixes or segments, but retain reward-model latency and memory overhead.Their validation has focused on preference alignment rather than reasoning tasks.
  • Breadth pruning: Breadth-pruning methods reduce the effective sample count under a fixed budget, but serial adaptive stopping compromises BoN’s parallelism benefits.
  • Training-time distillation: Training-time distillation methods replace multiple inference samples with one, shifting computational cost from inference to training.

8 Conclusion

The paper concludes that ST-BoN self-estimates promising samples without fully generating all N candidates or using reward models, improving cost-performance trade-offs across evaluated tasks.

  • ST-BoN self-estimates the most promising sampling without fully generating N samples or using reward models.

Limitations

The limitations section identifies restricted model access and incomplete adaptation of the window length to task complexity.

  • ST-BoN requires access to hidden states, so it does not apply to closed-source models such as GPT-4 series models.The method also has room to adapt τ to task complexity, particularly for longer responses.

Societal Impact

The paper positions reward-model-based efficiency methods as constrained by computational overhead and domain-specific assumptions. It contrasts these limitations with ST-BoN’s internal-consistency approach.

  • Reward-model-based methods repeatedly score partial sequences, adding memory and inference-time overhead that can block parallel sampling.
  • Their effectiveness depends on prefix scores correlating with final-text scores, an assumption validated mainly in preference-alignment domains.
  • Process reward models may struggle with arbitrarily segmented reasoning text, limiting their suitability for reasoning tasks.
  • ST-BoN instead uses internal consistency signals and does not depend on reward-model overhead or domain-specific priors.

C Evaluation Setup

The evaluation measures reasoning accuracy and open-ended-task quality using automatic and human metrics. Instruction-following is evaluated through pairwise win rates against greedy decoding.

  • Evaluation Metrics: Reasoning scenarios use exact-match Accuracy against ground-truth answers.
  • Evaluation Metrics: Summarization uses Rouge-L plus human ratings of fluency, coherence, consistency, and relevance.
  • Evaluation Metrics: Human summary ratings average three annotators’ scores across four dimensions, each rated from 0 to 5.
  • Evaluation Metrics: Instruction-following uses Win Rate, computed from external-LLM judgments comparing each paradigm with greedy decoding.

D.1 Main Results of All LLMs

Across larger and smaller language models, the reported cost–performance conclusions remain valid, while ST-BoN’s early self-estimation correlates with final correctness across model families.

  • Main Results: ST-BoN remains effective as model size increases, preserving the reported cost–performance trade-off conclusions across Qwen2.5-72B, Llama3-8B, and Mistral-7B models.
  • Reward-Model Ablation: Additional experiments show that reward-model choice does not change ST-BoN’s cost–performance advantage over Full-BoN with reward models.
  • Consistency Analysis: Early self-estimation has a strong positive correlation with final answer correctness across Llama3-8B and Mistral-7B model series.
  • Consistency Ablation: ST-BoN’s absolute performance can be weaker than Full-BoN without reward models at the same sampling size, although the gap is not substantial.
  • Consistency Ablation: Under Full-BoN without reward models, chain-of-embedding outperforms majority voting in most evaluated cases at nearly identical computational cost.

E Case Study

The case studies illustrate ST-BoN’s selection behavior on reasoning examples and document the paper’s checklist responses about scope, reproducibility, limitations, and theory.

  • Case Studies: The case studies evaluate GPQA, TheoremQA, and MATH examples across Qwen2.5-7B, Llama3-8B, and Mistral-7B models with sampling size N = 5.
  • Case Findings: ST-BoN can select a correct answer that lacks the majority vote, whereas majority voting becomes random when all sampled answers are inconsistent.
  • Case Findings: Internal consistency assigns continuous scores that distinguish candidates more clearly in ambiguous cases than majority voting.
  • Paper Checklist: The checklist states that the paper discusses limitations and identifies Section Limitations as the justification.
  • Paper Checklist: The theory-checklist materials identify Section 3.1 and Appendix A as the locations for assumptions and proofs.
Loading 2503.01422v3…