PaperlayerView exact revision on arXiv

Source-linked AI summary

Provable Test-Time Scaling for Beam Search in LLM Reasoning

Qijia He, Yu Huang, Yuan Cheng, Yuxin Chen, Yingbin Liang

arXiv:2609.38672v1cs.LGcs.AI

TL;DR

Although beam search is empirically effective for long-horizon LLM reasoning, its test-time scaling guarantees are not well understood. The paper analyzes sampling-access beam search and proposes confidence filtering, finding improved coverage dependence and greater robustness on hard and longer-horizon instances. Its guarantees assume sampling access to token probabilities and rely on reward-model estimation in the standard setting.

  • Problem

    Theoretical understanding of beam search’s test-time scaling is limited, including its guarantees under constrained compute and imperfect verifiers and its advantages over sequence-level methods.

    Source: §1 · p. 2

  • Method

    The paper analyzes beam search using sampled token-probability estimates for intermediate scoring and an external reward model only after complete responses, then introduces CF-Beam to filter low-probability samples.

    Source: §3 · p. 5

  • Results

    CF-Beam reduces sufficient coverage dependence from quadratic to nearly linear under fixed horizon, gap, and target accuracy, and experiments report greater robustness on hard instances and increasing horizons.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    The guarantees control the dominant search-induced term through token-level coverage that scales polynomially with horizon, unlike the potentially exponential coverage dependence of sequence-level methods.

    Source: §1 · p. 2

    The analysis assumes black-box sampling access to the reference policy, without full logits or exact next-token probabilities.

    Source: §2 · p. 3

Abstract

from arXiv · show

Beam-search-based test-time methods provide an effective way to improve large language model (LLM) performance on long-horizon generation by pruning invalid reasoning paths early, leading to significantly improved reasoning efficiency and more favorable test-time cost scaling. Despite strong empirical success, the theoretical understanding of beam search remains limited. In this paper, we study the test-time compute guarantee of the commonly used beam search framework that uses the model's internal log-likelihood for intermediate scoring, while relying on an external reward model only after a complete response is generated. We first establish a lower bound for vanilla beam search, showing that at least $\Omega(C^\star(x)^2)$ samples are required for the optimal response to survive, where $C^\star(x)$ is the token-level coverage coefficient for prompt $x$. This motivates our modified confidence-filtered beam search (CF-Beam), which reduces the sufficient coverage dependence from quadratic to nearly linear under prefix competitiveness, for fixed horizon, gap, and target accuracy. We then show that the regret of CF-Beam is upper-bounded by the probability of rare failure events and the reward estimation error scaled by a path-level coverage coefficient, where the rare-failure term vanishes as per-step sampling increases. Our results highlight a fundamental advantage of beam search over sequence-level inference methods such as Best-of-N and Best-of-Majority. While the guarantees of these approaches typically involve coverage coefficients that grow exponentially with the horizon $L$, CF-Beam controls the dominant search-induced term through a token-level coverage coefficient that scales polynomially with $L$. Our numerical experiments further confirm that beam search is more robust on hard instances and under increasing reasoning horizons.

1 Introduction

Test-time methods improve fixed-model LLM outputs through additional inference-time computation, but guarantees for beam search remain limited. This paper analyzes beam search under sampling access and imperfect verification, establishing sample-complexity bounds and comparing its behavior with sequence-level methods.

1 Introduction

  • Existing test-time methods such as Best-of-N and voting can require rapidly increasing samples as context length grows, while imperfect verifiers can misrank candidates.The paper identifies limited theoretical understanding of beam search’s test-time scaling as a central gap.Source: §1 · p. 2
  • Beam search expands reasoning token by token and prunes low-quality continuations, which can focus computation on promising prefixes in sparse-success, long-horizon settings.Intermediate verification may use an external reward model or internal self-evaluation.Source: §1 · p. 2
  • Stepwise pruning creates token-level dependencies, while heuristic intermediate scores may not align with one globally consistent objective.These properties make beam-search dynamics harder to analyze than sequence-level methods.Source: §1 · p. 2
  • The analysis studies sampling-access beam search that uses estimated token probabilities as an internal verifier and queries an external reward model only after a complete response.The motivating questions concern scaling under constrained compute and imperfect verifiers, and possible advantages over Best-of-Majority and Best-of-N.Source: §1 · p. 2
  • Beam search is compared with sequence-level methods through theoretical guarantees and numerical experiments, with CF-Beam reported as more robust on hard instances and at longer horizons.The paper reports that beam search has a polynomial-in-horizon dominant search term, whereas guarantees for Best-of-N and Best-of-Majority may involve exponentially growing coverage coefficients.Source: §1 · pp. 2–3

2 Preliminaries

The preliminaries formulate test-time decoding for a fixed reference LLM policy and sparse outcome rewards. They define the sampling-access setting, query budget, and coverage quantities used to analyze whether an optimal response can be found.

2 Preliminaries

  • A response is a length-L token sequence generated from a fixed reference LLM policy, and test-time algorithms use extra inference compute to seek an improved output.The response space is represented as token sequences over vocabulary A.Source: §2 · p. 3
  • The complete response receives a ground-truth outcome reward only after generation, while an approximate reward model may provide a surrogate score.The text describes this as a sparse outcome-reward setting.Source: §2 · p. 3
  • The analysis assumes black-box sampling access to the reference policy, so next-token probabilities must be estimated from sampled continuations rather than read from full logits.The policy is treated at the token level because beam search processes information token by token.Source: §2 · p. 3
  • The per-prefix sample complexity counts token-level policy queries at each prefix, while the total query budget sums those queries across the instance.The paper also introduces coverage coefficients to quantify how difficult it is for the reference policy to sample the optimal trajectory.Source: §2 · pp. 3–4
  • Token-level coverage reflects the inverse probability of sampling the optimal token at each depth, whereas sequence-level coverage captures compounding difficulty along an autoregressive trajectory.The supplied definition identifies the unique optimal response and its optimal prefixes.Source: §2 · p. 4

3 Algorithm

Vanilla-Beam estimates token probabilities from repeated samples, scores trajectories by cumulative estimated log-likelihood, and keeps the top-b candidates at each depth. CF-Beam adds a confidence filter before beam selection to reduce the sample complexity needed for the optimal response to survive.

3 Algorithm

  • Vanilla-Beam expands each beam trajectory with sampled next tokens, estimates their probabilities by empirical frequencies, and ranks trajectories by cumulative estimated log-likelihood.At each depth it retains the top b trajectories, then selects among final candidates using the external reward model.Source: §3 · p. 4
  • Finite-sample tail candidates can crowd out the optimal prefix under hard top-b pruning, and an early elimination is irreversible.CF-Beam addresses this fragility by filtering samples whose estimated sampling probability falls below a threshold before selecting the top b.Source: §3 · p. 5
  • For 2 ≤ b ≤ N/2 and N ≤ (C⋆(x))^2/16, a constructed instance makes Vanilla-Beam eliminate the correct answer with probability at least 0.5.The theorem states Vanilla-Beam needs sample size greater than Ω((C⋆max)^2) for the correct response to survive to the end.Source: §3 · p. 5
  • With fixed horizon, prefix gap, and target accuracy, CF-Beam’s sufficient per-prefix sample complexity is eO(C⋆max), compared with Vanilla-Beam’s quadratic dependence on token-level coverage.The confidence filter removes low-probability samples before top-b selection.Source: §3 · p. 5

4 Regret Analysis

The regret analysis bounds CF-Beam’s losses through search failures and reward-model error, under assumptions on reward estimation and prefix competitiveness. It compares search costs with sequence-level methods, establishes a lower bound, and extends the method to reward-model-free selection.

4.1 Technical Assumptions

  • The guarantees assume controlled reward-model error and prefix competitiveness; without the latter, beam search can discard the optimum with probability at least 3/4 at any sample size when b = 1.The proposition establishes this failure on an instance without prefix competitiveness.Source: §4.1 · pp. 5–6

4.2 Regret Upper Bound

  • CF-Beam’s regret separates search error from reward-model error; filtering and pruning failures decay exponentially with per-prefix sampling, but reward-selection error does not decrease with sample size or beam width.The reward-model term uses sequence-level coverage C(x), while the search terms depend on token-level coverage, horizon, and prefix gap.Source: §4.2 · p. 7
  • For fixed prefix gap and target accuracy, sufficient per-prefix sampling is nearly linear in token-level coverage, reducing Vanilla-Beam’s quadratic coverage requirement to nearly linear.The total CF-Beam query budget is eO(bL^3C⋆max) for fixed gap and confidence; the filtering improvement is stated for fixed horizon, gap, and accuracy.Source: §4.2 · pp. 7–8

4.3 Comparison with Other Test-time Methods

  • CF-Beam’s search-induced guarantee depends on L C⋆max rather than the sequence-level product C(x), yielding polynomial rather than potentially exponential horizon dependence.Sequence-level methods may be more sample-efficient when token-level coverage is near one and the horizon is moderate.Source: §4.3 · p. 8

4.4 Regret Lower Bound

  • The lower bound matches the upper bound’s reward-estimation-error term, establishing minimax optimality with respect to that error.The construction uses arbitrarily small prefix gaps, making it information-theoretically difficult to prune all suboptimal prefixes before the final step.Source: §4.4 · pp. 8–9

4.5 Extension to Reward Model Free Setting

  • The reward-model-free variant selects the most self-consistent candidate in the final beam and needs no reward-model accuracy assumption.Its total query budget is eO(LbC⋆max/κ²), with only logarithmic dependence on L in the per-prefix sample requirement.Source: §4.5 · p. 9

5 Experiments

CF-Beam performs best in the hardest simulated setting and maintains the strongest accuracy as reasoning horizons grow. Experiments compare two CF-Beam implementations with four baselines using a noisy reward model.

5 Experiments

  • The comparison tests Oracle and Empirical CF-Beam against Vanilla-Beam, Best-of-N, Majority Voting, and Best-of-Majority for optimal-response identification.The synthetic LLM simulator uses a noisy reward model with error level 0.01; Empirical CF-Beam sets its threshold using γ = 0.3 without requiring the optimal-token probability.Source: §5 · pp. 9–10
  • At p = 0.01, CF-Beam has its largest advantage; as p increases, the gap narrows and Vanilla-Beam catches up with sufficiently large budgets.The three difficulty levels use p ∈ {0.01, 0.05, 0.3}, with larger p indicating easier instances.Source: §5 · p. 10
  • As horizon L grows, CF-Beam consistently attains the highest accuracy while beam-based methods remain robust and sequence-level baselines degrade rapidly.The comparison uses a per-step success probability of 0.8; the observed pattern aligns with polynomial-in-L dominant sample-complexity for CF-Beam versus potentially exponential dependence for sequence-level methods.Source: §5 · p. 10

6 Related Works

The section reviews theoretical work on test-time methods and beam search, then identifies beam-search test-time scaling guarantees as an unaddressed research gap. The paper presents its work as the first such guarantee.

6 Related Works

  • Prior analyses cover CoT estimation, Best-of-N, Majority Voting, and Best-of-Majority, including convergence or sample-complexity results.The section also notes Yu et al. (2025c), which combines Best-of-N with rewind-and-repeat.Source: §6 · p. 10
  • Prior beam-search theory studied implicit objectives and self-consistency-based uncertainty quantification, while this work provides the first test-time scaling guarantee for beam-search decoding.The cited studies are Meister et al. (2020) and Fadeeva et al. (2025).Source: §6 · p. 10

7 Conclusion and Limitations

The conclusion presents CF-Beam’s scaling advantage as conditional on prefix alignment between likelihood and final reward. The surrounding discussion reviews related methods, clarifies computational-cost measures, and extends the analysis to multiple optimal responses.

7 Conclusion and Limitations

  • CF-Beam improves search-induced scaling over sequence-level methods on hard and long-horizon instances, but its advantage relies on prefix alignment between likelihood and final reward.If alignment fails, reward-optimal prefixes may be pruned before final reward evaluation.Source: §7 · p. 10

A Related Works

  • Beam search is treated as parallel test-time computation; prior theory studied its implicit objective and uncertainty quantification, while this work supplies a test-time scaling guarantee.The related-work discussion also connects CF-Beam to self-consistency through its use of intermediate self-consistency scores to prune beams.Source: A Related Works · pp. 14–15

B Discussion of Assumptions and Computational Costs

B.1 Uniqueness of the Optimal Response

  • The analysis can track one qualifying optimal witness when multiple reward-optimal responses exist, but does not exploit their aggregate probability mass or shared-prefix structure.The prefix gap must hold against every distinct competing prefix, including prefixes of other correct trajectories.Source: §B.1 · p. 15

B.2 Prefix Competitiveness and Its Relaxation

CF-Beam’s likelihood-based pruning relies on a prefix advantage for the optimal trajectory, while the paper bounds token-level query and candidate-processing costs separately from real-world compute costs.

B.2 Prefix Competitiveness and Its Relaxation

  • Prefix competitiveness lets likelihood-based pruning exploit token-level coverage without independently sampling the entire trajectory, but higher-likelihood suboptimal prefixes can irreversibly eliminate the optimum.More sampling alone cannot compensate for arbitrary likelihood–reward misalignment; reward-free final selection still requires the top-1 gap.Source: §B.2 · pp. 15–16

B.3 Sampling Budget and Candidate Processing

  • For fixed N and beam width b, token-level queries and processed candidates grow linearly with horizon L, but these bounds are not direct measures of latency, forward passes, or KV-cache memory.Realized costs depend on local batching, API access, and cache management.Source: §B.3 · p. 16

C Proof of Theorem 1

The proof constructs a large-action instance where first-round sampling usually fills the beam, then shows suboptimal continuations can crowd out the optimal response. An auxiliary lemma establishes monotonicity of a binomial-tail bound across odd and even beam widths.

C Proof of Theorem 1

  • With a large action space, suboptimal first-step samples are almost surely distinct, so the optimal action is retained when sampled at least twice.The proof also argues that the first-round beam is almost surely filled with b distinct candidates under this regime.Source: C Proof of Theorem 1 · pp. 17–18
  • On the good event, at least b suboptimal continuations score above the threshold and crowd the optimal response out of the top-b candidates.The proof bounds the number of successful suboptimal prefixes using independent per-prefix events.Source: C Proof of Theorem 1 · pp. 19–20

C.1 Technical Lemmas

  • For N > 30, the binomial-tail bound F(b) is nondecreasing separately along the even and odd beam-width subsequences.This monotonicity is used to bound a probability in the proof of Theorem 1.Source: §C.1 · p. 20

D Proof of Proposition 1

The proposition proof constructs a one-step instance violating prefix competitiveness, where the optimal response survives only if its empirical frequency reaches half the samples. The technical lemmas analyze confidence-filtered output probabilities and survival under beam pruning.

D Proof of Proposition 1

  • The constructed instance violates prefix competitiveness, so no positive κ satisfies the condition there.Source: D Proof of Proposition 1 · p. 21
  • In the constructed one-step instance, the optimal response is retained only when its empirical frequency is at least one half.The score ordering is equivalent to comparing the empirical frequencies of the optimal and suboptimal actions.Source: D Proof of Proposition 1 · p. 21

E.1 Technical Lemmas: Filtering and Beam Survival

  • The filtering analysis bounds output probabilities using independence of empirical histograms along a fixed path and Markov’s inequality, without conditioning on output selection.Source: §E.1 · p. 22
  • When the optimal path is first pruned, at least b generated candidates have scores at least as high, so the count of filtered full-tree prefixes is at least b.Counting ties against the optimal prefix makes this implication valid for any tie-breaking rule.Source: §E.1 · pp. 23–24

E.2 Proof of Theorem 2

The proof bounds regret by separating the event that the optimal response survives from its complement, and derives sufficient sample conditions for the search-error terms. A related corollary controls maximum-score selection by bounding intermediate and final-depth competitors.

E.2 Proof of Theorem 2

  • On the good event, the optimal path survives to the final beam and empirical reward selection is no worse than its estimate; off that event, regret is at most one.Source: §E.2 · p. 24

E.3 Proof of Corollaries 1

  • The corollary’s sufficient sample-size conditions make each search-error term at most δ/2, yielding the stated overall bound.Source: §E.3 · p. 25

E.4 Proof of Corollary 2

  • For maximum-score selection, bounding intermediate competitors and final-depth competitors, together with filtering failure, makes the total error at most δ.Source: §E.4 · p. 25

F Regret Lower Bound: Proof of Theorem 3

The proof constructs an instance where the optimal response is excluded from the final top-b set with probability at least 1/4. It combines this search failure with a misspecified reward estimator that selects a suboptimal response whenever one remains in that set.

F Regret Lower Bound: Proof of Theorem 3

  • The reward estimator ranks every suboptimal response above the optimum, so any final top-b set containing a suboptimal response yields an incorrect output.Source: F Regret Lower Bound: Proof of Theorem 3 · p. 26
  • The constructed instance excludes the optimal response from the final top-b set with probability at least 1/4.The proof establishes this bound by cases and uses it to bound the probability that the final set contains only the optimum.Source: F Regret Lower Bound: Proof of Theorem 3 · pp. 26–27
  • At each optimal-prefix step, the construction gives suboptimal children a constant probability of surviving confidence filtering, while the final set cannot contain only the optimum unless all suboptimal paths are filtered out.Source: F Regret Lower Bound: Proof of Theorem 3 · pp. 26–27

G Auxiliary Lemmas

  • For the auxiliary binomial bound, the proof obtains Pr(X < Np) ≥ 1/4 for the stated range of p and N.Source: G Auxiliary Lemmas · p. 30

H.1 Details of Numerical Experiments

The experiments compare confidence-filtered beam search with beam and sequence-level baselines in simulator settings and on a neural-model matrix-multiplication task. CF-Beam performs especially well on hard instances and as reasoning horizons grow, while filtering strength and beam width affect accuracy.

H.1 Details of Numerical Experiments

  • Empirical CF-Beam uses a prefix-adaptive cutoff based on the local maximum probability to filter low-frequency tails while retaining plausible continuations.The experiments set γ = 0.3; empirical CF-Beam performs comparably to, and sometimes slightly better than, Oracle CF-Beam.Source: §H.1 · p. 31
  • In the hardest setting (p = 0.01), both CF-Beam variants substantially outperform sequence-level baselines and Vanilla-Beam, especially under limited budgets.The performance gap narrows in easier settings, and Vanilla-Beam nearly matches CF-Beam at p = 0.3 with a sufficiently large budget.Source: §H.1 · p. 32
  • As horizon L grows, CF-Beam variants maintain the highest accuracy, while sequence-level baselines degrade rapidly and beam methods remain more robust.The experiments use a mean correct-action probability of 0.8 per step.Source: §H.1 · p. 32
  • The oracle threshold need not maximize accuracy at every finite budget; the horizon experiment cannot raise per-prefix sampling at the rate required by the sufficient filtering condition.The stated condition is N ≳ C⋆max L² log(2L/δ); with γ = 0.3, the empirical threshold can lower the cutoff, reducing premature elimination while potentially admitting more suboptimal candidates.Source: §H.1 · p. 32

H.2 Additional Experiments over LLMs

  • Empirical CF-Beam reaches 0.388 accuracy with a 16.2k query budget, beating Vanilla-Beam by 1.6 percentage points and Best-of-N by 5.6 percentage points.Vanilla-Beam uses 18.6k queries; CF-Beam uses a smaller budget.Source: Table 2 · p. 33, §H.2 · p. 33
  • On the LLM task, the best filtered configuration improves on Vanilla-Beam at each tested beam width, with best accuracy rising from 0.388 at b = 2 to 0.550 at b = 8.The best filtering coefficient is γ = 0.3 for b = 2 and γ = 0.15 for b = 4 and b = 8; wider beams generally require larger total budgets.Source: §H.2 · p. 33, Table 3 · p. 33
Loading 2609.38672v1…