Provable Test-Time Scaling for Beam Search in LLM Reasoning

Conference: NeurIPS 2026
arXiv: 2609.38672
Area: LLM Reasoning
Keywords: test-time compute, beam search, confidence filtering, coverage coefficient, prefix competitiveness

TL;DR

With sampling access to next tokens but no full logits, this paper filters low-frequency tokens before beam pruning and, under alignment between prefix likelihood and correctness, reduces the worst-step coverage dependence of search samples from worst-case quadratic to nearly linear; accuracy on an LLM matrix-multiplication task increases from 37.2% to 38.8%.

Background & Motivation

Increasing test-time compute does not necessarily require training a larger model: a fixed model can generate more candidates, followed by voting or reward-model selection. Best-of-N and Best-of-Majority face a difficulty because they usually generate complete reasoning chains before deciding which ones to retain. When the probability of a correct step is below 1, complete-path success probabilities multiply across the horizon, and substantial compute can be spent on paths that already went wrong. Beam search instead expands and prunes incrementally, concentrating computation on a few promising prefixes.

Early pruning, however, introduces irreversible errors. Rather than assuming access to exact next-token probabilities, this paper draws multiple samples at each prefix, estimates probabilities from frequencies, and ranks prefixes by accumulated log-probability. In a large vocabulary, a rare token observed once receives an estimated probability of one divided by the sample count. If that incorrect prefix subsequently has high-probability continuations, it can displace the correct path whose individual steps have modest probabilities. Even an accurate final reward model cannot recover an answer that has already been pruned.

The analysis therefore separates survival of the correct path from correct selection after survival. Local sampling, prefix ranking, and filtering govern the former, while reward-model error still affects the latter. Core Idea: filter low-frequency tokens lacking repeated sampling support before likelihood-based pruning, suppressing tail-estimation noise while analyzing local search failures separately from final reward-selection error.

Method

Overall Architecture

Inputs are a prompt, a fixed base model, beam width \(b\), maximum horizon \(L\), per-prefix sample count \(N\), and filtering thresholds. Starting from an empty prefix, each round independently samples \(N\) next tokens at every active prefix and merges repeated tokens into frequency counts. Only distinct tokens passing the threshold become extension candidates. Each candidate adds its empirical log-probability to its parent's score, and the highest-scoring at most \(b\) prefixes are retained globally.

After \(L\) rounds, standard CF-Beam uses an external reward model to choose one complete response from the final beam. The self-consistent variant instead chooses the complete path with the largest cumulative empirical log-likelihood, without a reward model. Neither uses intermediate rewards or updates base-model parameters; they are not new training frameworks with process-reward supervision.

The main contribution is a theoretical analysis of this simple filtering step, not a new network architecture. Prefix states and theorem conditions therefore explain the algorithm here without turning proof sections into an architecture diagram. If the beam becomes empty, the algorithm returns a fixed fallback response; ties follow a fixed rule. These edge cases are included in the search-failure event.

Key Designs

1. Sampling-based prefix scoring: localize path difficulty without assuming correctness alignment away

The empirical probability of a next token is its observed count divided by \(N\). Only distinct observed tokens enter the candidate pool; repeated occurrences do not occupy multiple beam slots. Prefix scores sum empirical log-probabilities along the path, comparing estimated likelihoods of entire prefixes rather than just their last steps. Sampling randomness is independent across prefixes, but paths sharing ancestors also share earlier scores, so path scores cannot all be treated as independent.

To characterize difficulty, the authors identify a reward-optimal path and define its token-level coverage coefficient as the inverse probability of its correct token at depth \(t\). The maximum summarizes the hardest step, while the product summarizes complete-path difficulty:

\[ C_t^{\star}(x)=\frac{1}{\pi_{\mathrm{ref}}(a_t^{\star}\mid s_{t-1}^{\star},x)},\qquad C_{\max}^{\star}(x)=\max_t C_t^{\star}(x),\qquad \overline C(x)=\prod_{t=1}^{L}C_t^{\star}(x). \]

For example, if each correct token has probability 0.8, the hardest-step coefficient is 1.25, whereas the inverse probability of the complete path is \(1.25^L\). Beam search aims to retain correct prefixes repeatedly rather than wait for an independent complete rollout to follow every correct step. This is a sampling advantage, not an automatic justification that high likelihood implies correctness.

Theorem 1 constructs a hard instance for vanilla beam search: each of the first two correct steps has probability \(1/C^{\star}\); many incorrect first tokens have tiny true probabilities but appear once when sampled; and each incorrect prefix subsequently has two continuations with probability 1/2 each. The true correct prefix remains more competitive, yet incorrect paths can have larger empirical probability products. This explains why widening the beam alone need not remove tail-candidate crowding. The quadratic sample lower bound is an existential worst-case result, not a required cost on every task.

2. Confidence filtering: remove accidental tail tokens before hard pruning

CF-Beam adds one step to Vanilla-Beam: a next token whose empirical probability is below the threshold is not expanded. Remaining candidates still use accumulated log-likelihood and top-\(b\) selection. “Confidence” here refers to empirical-frequency thresholding, not a calibrated confidence interval computed for every token. Filtering occurs during expansion rather than as deduplication or voting after complete responses have been generated.

The main theorem sets the threshold to \(1-1/L\) times the true correct-token probability, namely \(\beta_t=(1-1/L)/C_t^{\star}(x)\). This suppresses tail tokens while allowing some downward sampling fluctuation of the correct token. It nevertheless uses an unknown correct-token probability, making it an oracle threshold rather than a directly computable deployment rule. The reward-free variant uses \(\beta_t=1/(2C_t^{\star}(x))\), which is also an oracle choice.

The practical variant sets the threshold at each visited prefix to a fixed fraction of the largest empirical next-token frequency:

\[ \beta_t(s)=\gamma\max_{a\in\mathcal A}\hat\pi(a\mid s,x),\qquad \gamma\in(0,1). \]

Experiments use \(\gamma=0.3\) by default. This rule does not require the correct answer or coverage coefficients, but its motivation assumes that the correct continuation is locally competitive, with a probability comparable to the local maximum. Global prefix competitiveness does not directly replace this local proxy condition. The paper does not establish the same oracle theorem for the empirical threshold. Excessively high thresholds discard useful paths, while very low thresholds approach vanilla beam search.

3. Prefix competitiveness and regret decomposition: identify when pruning is reliable and when more sampling cannot help

The authors require the optimal path to have an average cumulative log-likelihood advantage at every depth against every other prefix at that depth. This is neither a statement that the final answer is more frequent nor a requirement that each correct token is greedily optimal. It is a likelihood ordering condition for entire prefixes:

\[ \frac{1}{t}\log\frac{\pi_{\mathrm{ref}}(s_t^{\star}\mid x)}{\pi_{\mathrm{ref}}(s\mid x)}\geq\kappa,\qquad s\neq s_t^{\star},\quad t\in[L],\quad \kappa\in(0,1]. \]

The gap \(\kappa\) lets empirical rankings approach the correct prefix ranking as sampling increases. Without this condition, Proposition 1 gives a one-step counterexample: the correct token has probability 1/8, the incorrect token has probability 7/8, and beam width is 1. For every sample count, the correct answer is pruned with probability at least 3/4. More accurate estimation of an incorrect likelihood ordering does not fix that ordering.

The proof separates failure into filtering out the correct token and top-\(b\) pruning of its prefix despite passing the filter. Binomial lower tails and a union bound over depths control the first event. For the second, the analysis counts filtered competitors whose scores match or exceed the correct prefix. Truncated binomial moments and inverse moments control score fluctuations, while summation over true prefix probability mass avoids paying directly for exponentially many candidates.

To handle sampling dependencies induced by pruning, the appendix analytically pre-generates an independent frequency histogram for every possible prefix, revealing it only when the algorithm visits that prefix. This coupling defines empirical frequencies even along an optimal path pruned earlier. It does not require the actual algorithm to pre-sample the full tree or increase its query budget. Cauchy–Schwarz handles correlations from shared prefixes instead of assuming competing paths are independent.

For rewards in \([0,1]\), optimal reward 1, reward mean-squared error at most \(\epsilon_{\mathrm{RM}}^2\) under the reference distribution, and pointwise optimal-response error at most \(\epsilon_{\mathrm{opt}}\), Theorem 2 gives:

\[ \mathrm{Reg}(x)\leq L\exp\!\left(-\frac{N}{2C_{\max}^{\star}L^2}\right) +\frac{2C_{\max}^{\star}}{b}\exp\!\left(-\frac{\kappa^2N}{48C_{\max}^{\star}}\right) +\epsilon_{\mathrm{opt}}(x) +2\sqrt{\overline C(x)\epsilon_{\mathrm{RM}}^2(x)}. \]

Regret is the optimal true reward minus the expected true reward of the algorithm's output. The theorem also requires \(L\geq2\) and \(N\geq(48C_{\max}^{\star}/\kappa^2)\max\{2,\log(2C_{\max}^{\star})\}\). The first two terms decrease with sampling, whereas the last two do not decrease with \(N\) or \(b\) in this bound. Removing complete-path coverage from the search term does not remove exponential difficulty from the entire regret guarantee.

The reward-error term reflects distribution shift from the reference model to search-selected outputs. Even if a reward model has low mean-squared error on common responses, search may favor responses the reference model rarely generates. Oracle thresholds bound joint output probabilities by at most \(4\overline C\) times reference probabilities, enabling transfer of the reward-error guarantee. Theorem 3 uses small-gap hard instances to show that this path-coverage dependence cannot simply be deleted in the worst case. It does not establish that search can never isolate the correct answer under a fixed positive gap.

4. Self-consistent final selection: remove the reward model while retaining likelihood–correctness alignment

Self-consistent CF-Beam chooses the complete path with the largest empirical likelihood instead of the highest reward score. Without reward-error transfer, the filtering threshold can stay at half the correct-token probability rather than approach its mean as \(L\) increases. Corollary 2 consequently permits \(N\gtrsim(C_{\max}^{\star}/\kappa^2)\log(eLC_{\max}^{\star}/\delta)\) samples per prefix to achieve regret at most \(\delta\). For fixed beam width, coverage, gap, and target accuracy, the total query budget grows nearly linearly with the horizon.

However, retaining the correct path in the beam does not suffice for this final selection rule. The proof must additionally rule out any complete competitor whose score is at least that of the correct answer. This is an extra event beyond survival required by reward-based selection. The variant is also not conventional voting across different reasoning chains grouped by final answer: it directly maximizes cumulative empirical likelihood over complete paths and requires a top-1 prefix gap.

Appendix B provides two extensions. Reward-optimal responses need not be unique if a witness path satisfies the gap and optimal-response error conditions, but that witness must still beat distinct prefixes of other correct paths. The theory does not exploit their aggregate probability mass. For reward-based selection, at most \(k-1\) gap-violating prefixes per depth can be allowed when \(1\leq k\leq b\), replacing \(b\) by \(b-k+1\) in the pruning bound. Reward-free maximum-likelihood selection does not directly inherit this relaxation.

Loss & Training

The paper adds no training loss or fine-tuning stage; the base policy and final reward model are treated as given. Reward-model error is an assumption, not a guarantee learned by CF-Beam itself. AceMath-7B-RM in the experiments likewise scores only complete responses.

A query means drawing one next token at a prefix. The actual count is at most \(N+(L-1)bN\leq LbN\). Before pruning, each round processes at most \(bN\) distinct candidates; filtering only reduces this pool and introduces no extra branching. These counts are not neural-network forward passes, latency, or KV-cache size. A local model can reuse one forward-pass distribution for multiple draws, while API and batching costs differ.

Key Experimental Results

Main Results

The LLM experiment evaluates Qwen3-1.7B-Base at sampling temperature 1.3 on 500 synthetic \(4\times4\) integer matrix-multiplication problems. Methods requiring reward-based selection use AceMath-7B-RM. Both beam methods set \(N=12\) and \(b=2\), and CF-Beam uses the empirical threshold with \(\gamma=0.3\). Results below come from Appendix H.2, Table 2; accuracy is a proportion and budgets are thousands of token-level policy queries.

Method	Accuracy	Query budget (thousands)
Best-of-N	0.332	16.8
Majority Voting	0.304	16.8
Best-of-Majority	0.326	16.8
Vanilla-Beam	0.372	18.6
Empirical CF-Beam	0.388	16.2

CF-Beam improves on Vanilla-Beam by 1.6 percentage points while the reported budget decreases from 18.6 thousand to 16.2 thousand queries. Its gain over Best-of-N is 5.6 percentage points. These are comparable, not identical, measured budgets and do not establish equal-FLOPs or equal-latency gains.

Ablation Study

Appendix H.2, Table 3 fixes \(N=12\) and varies the filtering coefficient and beam width; \(\gamma=0\) disables filtering. The subset below illustrates the benefit of moderate filtering and harm of excessive filtering rather than reproducing every threshold.

Beam width	\(\gamma=0\)	\(\gamma=0.15\)	\(\gamma=0.3\)	\(\gamma=0.6\)	\(\gamma=0.9\)
\(b=2\)	0.372	0.378	0.388	0.318	0.290
\(b=4\)	0.482	0.512	0.472	0.374	0.302
\(b=8\)	0.536	0.550	0.480	0.364	0.300

The best threshold is not fixed: it is 0.3 for \(b=2\) but 0.15 for \(b=4,8\). Relative to no filtering, the best accuracy gains at the three widths are 1.6, 3.0, and 1.4 percentage points. Wider beams improve the best accuracy, but the table has no corresponding budget column and is not an equal-total-budget comparison of beam widths.

The theoretical results also provide an analysis table. Their nearly linear claims retain conditions; LLM accuracy alone does not verify the theorems.

Result	Conditions and guarantee	Scope
Theorem 1: vanilla beam lower bound	There exists an instance with $50\leq C^{\star}\ll	\mathcal A
Theorem 2 / Corollary 1: reward-based CF-Beam	Oracle thresholds, prefix gap, and reward-error assumptions; for fixed gap and accuracy, total budget is \(\widetilde O(bL^3C_{\max}^{\star})\)	Controls search failure; reward error still involves \(\overline C\)
Corollary 2: self-consistent CF-Beam	Oracle half-probability thresholds and a top-1 prefix gap; total budget is \(\widetilde O(LbC_{\max}^{\star}/\kappa^2)\)	No reward-error assumption, but maximum final likelihood must align with optimal reward
Proposition 1: likelihood-misalignment example	\(b=1\), correct/incorrect probabilities 1/8 and 7/8; failure probability at least 3/4 for every \(N\)	More sampling alone cannot repair likelihood–correctness misalignment

Key Findings

• Simulator experiments run 300 independent trials per instance and use symmetric label flips with probability 0.01 to model reward error. Filtering has its largest advantage at the hardest complete-path probabilities among 0.01, 0.05, and 0.3.
• The horizon experiment considers \(L\in[2,40]\) and sets complete-path success probability to \(0.8^L\). The authors report faster degradation of sequence-level methods, but readable pointwise curve values are absent from the cache, so no specific curve accuracies are reconstructed.
• The empirical variant can exceed the oracle variant at a fixed finite budget. The oracle threshold leaves a relative margin of only \(1/L\) below the correct-token probability, becoming more sensitive to downward sampling fluctuations as the horizon grows. “Oracle” does not mean optimal at every budget.
• Figure captions label the difficulty experiment as Figure 2 and the horizon experiment as Figure 3, while the text and H.1 repeatedly cite both as Figure 3. This note distinguishes them by experiment content without silently repairing figure numbers. Algorithm cross-references also use Algorithm 3 despite the displayed Algorithm 1.

Highlights & Insights

• The analysis connects accidental observation of tail tokens with highly concentrated continuations of incorrect prefixes. Vanilla beam search is therefore limited by more than whether the correct token was sampled; filtering breaks this chain before incorrect paths occupy the beam.
• The most useful distinction is between local search coverage and distribution shift in final reward selection. The former can avoid the product of complete-path inverse probabilities, whereas rare-path reward error can still damage the latter.
• The proof does not treat adaptively pruned candidates as independent samples. Pre-sampling coupling, truncated moments, and probability-mass summation provide reusable tools for decoders that estimate scores from samples before hard selection.

Limitations & Future Work

• Prefix competitiveness is strong: correct responses must maintain cumulative likelihood superiority at every depth. Rare but correct reasoning steps may violate it, so these guarantees do not extend to arbitrary mathematical reasoning.
• The deployed empirical threshold differs from the main theorem's oracle threshold. Experiments support the former, but do not close the theoretical gap; observable gap diagnostics and adaptive thresholds are worthwhile directions.
• The LLM evaluation includes one base model, one matrix-multiplication task, and 500 problems, without confidence intervals accompanying the reported results. The gains do not yet establish reliability across tasks, temperatures, or models.
• Reward-based guarantees still amplify error by complete-path coverage. Reward calibration, intermediate verification, or switching to sequence-level methods on likelihood-misaligned prompts may be more targeted than unconditional increases in sampling.
• The cost measure counts sampling queries rather than actual runtime performance. Future evaluations should also report forward passes, generated lengths, batching, latency, and memory, while theory could exploit shared prefixes among multiple correct paths.

Related Work & Insights

• vs Best-of-N / Best-of-Majority: these methods generate complete responses before selection, with coverage difficulty typically involving complete-path inverse probability. CF-Beam maintains prefixes early and improves the search term under likelihood alignment, but introduces irreversible pruning and gap assumptions.
• vs self-consistency voting: conventional methods aggregate different complete reasoning chains by answer frequency. This paper's reward-free variant selects by empirical path likelihood, so its nearly linear budget guarantee should not be transferred directly to answer voting.
• vs process-reward-guided search: work by Lightman et al. and on process verifiers can supply intermediate correctness signals. This paper deliberately permits rewards only at completion, exposing both the capabilities and boundaries of likelihood-only pruning.
• Transferable insight: assess whether candidate-ranking signals align with the final objective before choosing early pruning over delayed selection. Filtering coefficients should also adapt to beam width and local sampling reliability rather than be treated as universally optimal constants.

Rating

• Novelty: 4/5 — Simple filtering, but a clear theoretical problem with a vanilla-beam counterexample and improved coverage dependence.
• Experimental Thoroughness: 3/5 — Simulator evidence, LLM comparisons, and threshold ablations are present; model and task coverage remains limited.
• Writing Quality: 3/5 — Search and reward errors are clearly separated, but figure and algorithm cross-references are inconsistent.
• Value: 4/5 — Clarifies when beam search genuinely saves computation and why larger budgets do not automatically ensure reliable reasoning.