Don't Throw Away Your Beams: Improving Consistency-based Uncertainties in LLMs via Beam Search

Motivation

• Multinomial sampling can generate similar or identical candidates in short-form QA, making uncertainty estimates noisy and less representative of the output space.The problem is amplified when only a small number of samples is used for computational efficiency.Source: §1 · p. 2

Approach

• Beam search explores alternative decoding paths and generates distinct candidates that better capture the model’s output space.The proposed estimators weight beam outputs by their probabilities rather than treating them uniformly.Source: §1 · p. 2

Contributions

• The paper introduces an importance-weighted estimator of consistency-based uncertainty using beam-search output candidates.This replaces multinomially sampled outputs while retaining compatibility with existing consistency-based and hybrid UQ methods.Source: §1 · p. 2
• The authors provide a distribution-free sufficient condition under which the beam-weighted estimator has lower error than the expected multinomial-sampling error.The condition is described as a theoretical lower bound on the relevant beam-set probability mass.Source: §1 · p. 2

Results

• Beam-search estimators improve existing consistency-based UQ approaches on short-form QA tasks and achieve state-of-the-art results.The paper evaluates this approach across six QA datasets and six models.Source: §1 · p. 2

2.1 LANGUAGE MODEL DECODING

• Autoregressive LLMs generate sequences token by token, with each token conditioned on the prompt and the preceding sequence.Decoding converts the model’s conditional probability distributions into concrete outputs.Source: §2.1 · p. 2
• Greedy decoding selects the highest-probability token, multinomial sampling draws tokens from the distribution, and beam search retains the top-k partial sequences.These strategies trade off output quality and diversity differently.Source: §2.1 · p. 2

2.2 UNCERTAINTY QUANTIFICATION FOR LLMS

• UQ methods include information-based, reflexive, density-based, and sampling-based approaches, with consistency-based methods measuring agreement among generated outputs.Consistency-based UQ is especially suited to confidence in a specific prediction rather than overall predictive-distribution uncertainty.Source: §2.2 · p. 3
• Dissimilarity estimates predictive uncertainty by measuring semantic differences between a chosen answer and alternative generations.Higher dissimilarity indicates lower consistency and therefore greater predictive uncertainty.Source: §2.2 · p. 3
• Because multinomial sampling does not reliably produce distinct, high-probability, and stable candidates, it does not best serve consistency-based uncertainty estimation.The resulting samples provide a limited approximation to the model’s full output space.Source: §2.2 · pp. 3–4
• Multinomial sampling produces 30–50% duplicate outputs for shorter generations, wasting computation and increasing variance in uncertainty estimates.The issue is especially relevant when methods use small sample sets or draw many full generations solely for uncertainty estimation.Source: §2.2 · pp. 3–4

3 UNCERTAINTY QUANTIFICATION BASED ON CONSISTENCY OF BEAM SEARCH CANDIDATES

• Beam search provides distinct candidates, reduces variance, and supplies a candidate distribution for consistency-based uncertainty estimates.Probability weighting avoids overemphasizing lower-probability beam outputs while retaining the beam search distribution.Source: §3 · p. 4, §3.1 · p. 4

3.1 REPLACING MULTINOMIAL SAMPLING

• Top-M truncation introduces a small bias but typically reduces variance and duplication on peaked distributions, producing more stable estimates per unit budget.Source: §3.1 · p. 4

3.2 THEORETICAL ANALYSIS

• Theorem 1 states that the beam-weighted estimator has lower mean-squared error than the Monte Carlo estimator when the beam probability mass satisfies a sufficient condition.The condition compares the beam set mass and inside–outside dissimilarity structure against Monte Carlo estimation error.Source: §3.2 · pp. 4–5
• For M = 10, beam search is guaranteed to outperform multinomial sampling when the top-10 hypotheses capture more than 0.842 of the model probability mass.On TriviaQA, 22.7% of examples meet the sufficient condition overall, rising to 30–40% for generations of at most three tokens.Source: §3.2 · p. 5, Figure 3 · p. 5

3.3 ADAPTING OTHER UQ METHODS TO BEAM SEARCH

• The beam-weighted estimator can be incorporated into other consistency-based UQ methods, including CoCoA, Eccentricity, and a new Eigenvectors Dissimilarity measure.Eigenvectors Dissimilarity increases when the answer moves away from the embedding bulk and when candidate samples become more dispersed.Source: §3.3 · pp. 5–6

4.1 EXPERIMENTAL SETUP

• The evaluation covers six QA datasets, three model families with base and instruct variants, PRR as the primary metric, and multiple consistency and CoCoA baselines.The datasets include closed-book, open-book, and multiple-choice QA; PRR uses AlignScore as its quality metric, with ROC-AUC and PR-AUC reported additionally.Source: §4.1 · pp. 6–7

4.2 RESULTS AND DISCUSSION

• Beam-guided consistency methods generally improve PRR across six models averaged over six QA datasets, often ranking first or second against baselines and original methods.Dissimilarity with beam search is best for all base models and second-best for two instruct models; CoCoA beam variants also achieve leading rankings.Source: §4.2 · p. 8, Table 3 · p. 8

4.3.1 EFFECT OF SAMPLE COUNT

• Beam search generally achieves higher PRR than multinomial sampling for every candidate budget M ≥2 and reaches high PRR with small budgets before saturating.Multinomial sampling improves more gradually and remains below beam search throughout the tested budgets; at M = 1, beam search reduces to greedy decoding and Dissimilarity is nearly zero.Source: §4.3.1 · p. 8, Figure 5 · p. 9

4.3.2 EFFECT OF OUTPUT LENGTH

• Beam search most clearly outperforms sampling for short generations, where multinomial duplicate rates are ∼30–50% for 2–4 tokens; the gap becomes negligible around 7 tokens and above.The comparison uses weighted beam search versus unweighted multinomial sampling on TriviaQA and CoQA with Gemma 3 4B base.Source: §4.3.2 · p. 9, Figure 4 · p. 9

4.3.3 PREDICTION-REJECTION CURVES

• Across three estimators, beam-search rejection curves dominate multinomial curves across nearly the full rejection range and become increasingly better as rejection rises.Beam-guided estimates remain stable at high rejection rates, while multinomial estimates flatten or degrade.Source: Figure 6 · p. 9, §4.3.3 · p. 10

4.3.4 ADDITIONAL ABLATIONS

• Additional ablations examine diverse and hybrid candidate generation, restricted-mass normalization, other sampling objectives, and top-1 beam decoding.These analyses include a probability floor for stabilizing low-mass beam weights and beam-weighted formulations of Semantic Entropy and Degree Matrix.Source: §4.3.4 · p. 10

Consistency-based uncertainty estimation

• Consistency-based UQ methods estimate uncertainty from agreement among sampled outputs, including similarity-matrix, lexical-similarity, Semantic Entropy, CoCoA, and SAR approaches.These methods are relevant in black-box settings, while hybrid white-box methods combine consistency signals with token probabilities.Source: §5 · p. 10

Uncertainty and decoding

• Decoding-focused UQ studies examine token-probability scores or incorporate uncertainty into generation, leaving stochastic decoding and non-likelihood consistency methods less studied.Uncertainty-aware decoding methods generally optimize generation quality rather than improving the uncertainty estimate itself.Source: §5 · p. 10

• The proposed beam-weighted estimators reduce variance in dissimilarity and increase candidate-answer diversity compared with multinomial sampling.Source: §6 · p. 10
• A theoretical lower bound identifies beam-set probability mass conditions under which the multinomial Monte Carlo estimator has larger error.Source: §6 · p. 10
• Evaluation on six QA datasets and six models demonstrates state-of-the-art uncertainty-quantification performance.Source: §6 · p. 10

LIMITATIONS

• The evaluation is limited to white-box short-form QA and relies on existing neural metrics, leaving longer-form generalization and specialized similarity measures open.The methods require model probability distributions, and AlignScore plus a pretrained NLI model are used for generation quality and consistency.Source: LIMITATIONS · p. 11

A.1 DIFFERENT SAMPLING STRATEGIES

• Plain beam search with probability weighting is a reasonable default because temperature and diverse beam strategies underperform it, while hybrid gains are sensitive to tuning.No single strategy dominates across datasets, models, or estimators; hybrid multinomial–beam can reach top-1 for specific B but gains are not systematic.Source: §A.1 · p. 16

A.2 RESTRICTED-MASS NORMALIZATION

• No clear best restricted-mass floor ϵ emerges, so selecting ϵ remains case-dependent across methods and dataset–model pairs.ϵ = 0 recovers probability normalization, whereas ϵ = 1 yields uniform weights and intermediate values trade distributional fidelity against robustness to noisy, length-biased tails.Source: §A.2 · p. 16, Table 5 · p. 16

A.3 OTHER SAMPLING-BASED METHODS UNDER BEAM SEARCH

• Beam search substantially helps Semantic Entropy but provides little or no improvement for Degree Matrix, whose absolute PR-AUC remains below the main answer-focused methods.The paper attributes this gap to target mismatch: these objectives score uncertainty about input x rather than correctness of the produced answer y∗.Source: §A.3 · pp. 16–17, Table 6 · p. 17

A.4 GRAPH LAPLACIAN EMBEDDING PARAMETERS

• Thresholding graph-Laplacian eigenvalues is more robust than fixing embedding size: α ∈[0.7, 0.9] performs consistently strongly, supporting α = 0.9.A fixed K can underfit or include noisy directions because the number of informative directions varies between candidate sets.Source: §A.4 · p. 18, Table 7 · p. 18

A.5 CROSS-ENCODER SIMILARITY

• The study evaluates an alternative RoBERTa-large STS cross-encoder similarity function against the main NLI-based similarity setup on six datasets.The ablation keeps all other settings unchanged, but the supplied passages do not state its numerical outcome.Source: §A.5 · p. 19, Table 8 · p. 19

A.6 NUMBER OF SAMPLES ACROSS DIFFERENT TASKS

• For all budgets M > 1, beam search consistently outperforms sampling in PRR across three models on TriviaQA and on ARC-Challenge.Beam-search PRR plateaus around M = 5 on open-ended TriviaQA and M = 2 on multiple-choice ARC-Challenge.Source: Figure 8 · p. 19, §A.6 · p. 20
• Beam-guided Dissimilarity remains effective at small budgets: M ≈5 for open-ended short-form generation and M = 2 for multiple-choice settings.The constrained output space in multiple-choice tasks enables faster saturation.Source: §A.6 · p. 20

B.1 PROBABILITY MASS COVERAGE

• Beam search covers more probability mass than multinomial sampling, while its first few ranked beams capture most of that mass.This motivates probability weighting and helps explain effectiveness at small candidate budgets.Source: §B.1 · p. 21, Figure 9 · p. 21

B.2 EXAMPLES

• Qualitative examples across TriviaQA, WebQuestions, and CoQA illustrate reduced duplication and enhanced uncertainty with beam-search samples.The examples display greedy answers, ten multinomial samples, ten beam samples, autoregressive probabilities, and uncertainty scores.Source: §B.2 · p. 21, Figure 10 · p. 21, Figure 11 · p. 22

C DATASETS

• Accuracy is measured as the fraction of predictions whose AlignScore with the gold answer exceeds 0.5 across model–dataset pairs.Prompt templates are listed separately for base and instruct models.Source: C DATASETS · p. 23, Table 9 · p. 23, Table 10 · p. 23

D.1 SCORING TOP-BEAM OUTPUT

• When the top-1 beam is scored as the produced answer, beam-weighted methods achieve higher PRR than original methods and baselines in most cases.Table 11 reports PRR averaged over six datasets and marks improvements over multinomial counterparts.Source: §D.1 · p. 24, Table 11 · p. 24

D.2 ROC-AUC AND PR-AUC

• Across six datasets, beam-guided variants generally match or outperform multinomial sampling on ROC-AUC and PR-AUC for Gemma 3 4B base.Correctness is binarized at AlignScore > 0.5, with incorrect answers treated as the positive class for PR-AUC.Source: §D.2 · p. 24

D.3 DETAILED RESULTS FOR EACH DATASET

• Tables 14–19 provide separate PRR results for six datasets across Gemma, Llama, and Qwen base and instruct models.Each table marks the best and second-best methods and indicates improvements of beam-guided variants over multinomial baselines.Source: §D.3 · p. 24, Table 14 · p. 26, Table 15 · p. 26
• Tables 14–19 report PRR separately for six datasets across six base and instruct model variants, including Qwen 3 8B instruct, with Table 19 marking top-1 and second-best methods.Source: §D.3 · p. 24, Table 19 · p. 28

E DETAILED DESCRIPTION OF UNCERTAINTY QUANTIFICATION METHODS

• The study evaluates sequence probability, mean token entropy, perplexity, and claim conditioned probability as token-, sequence-, or claim-level uncertainty measures.Sequence Probability uses negative log-probability; Mean Token Entropy averages token entropy; Perplexity uses negative average token log-likelihood; CCP perturbs claim tokens with alternative generations.Source: E DETAILED DESCRIPTION OF UNCERTAINTY… · p. 29
• Shifting Attention to Relevance combines TokenSAR and SentenceSAR, replacing sequence-level generative probability with relevance-reweighted probability.TokenSAR uses token relevance functions to assign relevance weights to tokens in the sequence.Source: E DETAILED DESCRIPTION OF UNCERTAINTY… · p. 29
• The experiments also include P(True), lexical similarity, number of semantic sets, and Laplacian eigenvalue scores as alternative uncertainty measures.P(True) uses the negative log-probability of “True”; lexical similarity averages pairwise generation similarity; semantic sets cluster mutually entailing outputs; the Laplacian method uses eigenvalues of an output-similarity graph.Source: E DETAILED DESCRIPTION OF UNCERTAINTY… · pp. 29–30

F COMPUTATIONAL BUDGET

• The full study used approximately 34 GPU-days, including six-model evaluations across six datasets and additional ablations.The reported budget includes 24 GPU-days for the six models, 10 GPU-days for ablations, and two A100 GPUs with 80 GB each.Source: F COMPUTATIONAL BUDGET · p. 30

G THE USAGE OF LLMS

• LLMs are the primary object of analysis, while assistants such as ChatGPT are used only for limited grammar correction and code debugging under human monitoring.The stated assistant use is practical and restricted rather than part of the primary analysis.Source: G THE USAGE OF LLMS · p. 30