Technical Summary: Improving Consistency-Based Uncertainties in LLMs via Beam Search

Problem Statement

Consistency-based uncertainty quantification (UQ) has established itself as a robust method for assessing the reliability of outputs from large language models (LLMs), particularly in black-box scenarios. These methods typically estimate uncertainty by measuring the level of agreement among multiple generations produced via multinomial sampling.

However, the authors identify critical limitations in applying multinomial sampling to short question-answering (QA) tasks:

1. High Redundancy: Due to sharp probability distributions for short answers, multinomial sampling frequently produces duplicate or near-identical generations. The work notes that for short outputs (2–4 tokens), redundancy rates can reach 30–50%, leading to a waste of computational resources.
2. High Variance: Since each sampling run generates a different set of candidates, the resulting uncertainty estimates exhibit significant variance across runs, undermining their robustness.
3. Inefficiency: Achieving stable estimates often requires a large number of samples, which is computationally expensive.

The core problem is that multinomial sampling fails to effectively explore the model's output space in scenarios with short outputs, leading to noisy confidence estimates and suboptimal uncertainty quantification.

Methodology

The work proposes a new family of UQ methods that replace multinomial sampling with Beam Search for generating candidate outputs. The methodology relies on three key insights:

1. Diverse Candidates: Beam Search inherently maintains the $k$ most probable partial sequences, thereby guaranteeing distinct candidate outputs and reducing redundancy.
2. Probability-Weighted Estimation: Treating beam candidates as uniform samples would overemphasize low-probability outputs. Instead, the authors introduce a weighted estimator. They normalize the sequence probabilities of the beam candidates ($b^{(i)}$) to form weights $w_i$:
   $$w_i = \frac{p(b^{(i)} | x)}{\sum_{j=1}^M p(b^{(j)} | x)}$$
   These weights are then used to calculate consistency-based metrics (e.g., Dissimilarity, Eccentricity) instead of using a simple average.
3. Theoretical Justification: The authors derive a theoretical lower bound for the probability mass ($m_B$) of the beam set required for the beam-weighted estimator to achieve a lower mean squared error (MSE) than the standard Monte Carlo estimator (multinomial). They prove that if the beam set covers a sufficient portion of the probability mass (specifically $m_B > 1 - \frac{1}{2\sqrt{M}}$ for a simplified condition), the beam search approach yields lower error.

Adapted Methods

The authors adapt several existing consistency-based UQ methods to utilize Beam Search:

• Dissimilarity: Measures the expected semantic difference between the primary answer and the beam candidates, weighted by $w_i$.
• Eccentricity: Uses a graph Laplacian on the similarity matrix of all samples (including the primary answer). The authors modify the embedding calculation to use weighted centroids.
• Eigenvector Dissimilarity: Measures the distance of the primary answer's embedding to the weighted mean of the sample embeddings in the Laplacian space.
• CoCoA: A hybrid method combining token-level probabilities with sample consistency. The authors replace the consistency signal with the beam-weighted estimator.

Main Contributions

1. Identification of Limitations: The work demonstrates empirically that multinomial sampling leads to high redundancy and variance in short QA tasks, making it a suboptimal choice for consistency-based UQ in these contexts.
2. Beam-Weighted Estimators: The proposal of a new family of UQ methods that utilize beam search candidates with importance weighting to better approximate the model's distribution without generating duplicates.
3. Theoretical Analysis: The derivation of a distribution-free sufficient condition (Theorem 1) guaranteeing that the beam-weighted estimator outperforms the multinomial Monte Carlo estimator when the beam set captures sufficient probability mass.
4. Empirical Validation: Comprehensive evaluation across six QA datasets (TriviaQA, WebQuestions, CoQA, HotpotQA, CommonsenseQA, ARC-Challenge) and six models (Gemma 3, Llama 3.1, Qwen 3).

Results

Experiments show that replacing multinomial sampling with Beam Search consistently improves the performance of consistency-based UQ methods:

• State-of-the-Art Performance: Beam-driven variants of Dissimilarity, Eccentricity, Eigenvector Dissimilarity, and CoCoA achieve the best or second-best Prediction-Rejection Ratio (PRR) scores across almost all model and dataset combinations.
• Robustness: Beam search methods exhibit significantly lower variance in uncertainty estimates compared to multinomial sampling.
• Sample Efficiency: Improvements are most pronounced with small sample budgets ($M \ge 2$). Beam search achieves high PRR scores with fewer candidates (saturating around $M=5$ for open-ended tasks and $M=2$ for multiple-choice tasks), whereas multinomial sampling improves more slowly.
• Advantage for Short Outputs: The performance gap between Beam Search and multinomial sampling is largest for short output lengths (where redundancy is highest) and narrows as output length increases.

Significance and Claims

The work claims that Beam Search should not be discarded as a decoding strategy for uncertainty quantification. Although often associated with deterministic, high-quality generation, the authors argue that Beam Search, combined with probability weighting, offers a more efficient and stable mechanism for exploring the output space in short QA tasks than stochastic multinomial sampling.

The significance of this work lies in:

• Practical Efficiency: It offers a way to improve UQ performance in white-box scenarios "essentially for free," as Beam Search is often already performed or can be executed with negligible additional cost compared to generating $M$ independent multinomial samples.
• Theoretical Foundation: It provides a formal condition under which Beam Search is provably superior for UQ, moving beyond heuristic improvements.
• Broad Applicability: The method improves both pure consistency-based metrics and hybrid methods (like CoCoA) that combine consistency with token probabilities, suggesting a generalizable improvement in uncertainty estimation for short text generation tasks.

The authors remain modest regarding limitations, noting that their evaluation is currently restricted to white-box scenarios (requiring access to model probabilities) and short QA task datasets. They acknowledge that generalizability to long-form generation and black-box scenarios (requiring empirical probability estimation) remains an open question.