Efficient semantic uncertainty quantification in language models via diversity-steered sampling

AI summary

Overview

Research area: Natural language processing; uncertainty quantification for large language model generation, with a focus on free-form question answering.

Technical level: Advanced. The paper builds on semantic entropy, importance sampling, control variates, and natural language inference, and assumes familiarity with autoregressive and masked diffusion decoding.

Scope: The paper introduces a diversity-steered decoding sampler that injects an entailment-based semantic-similarity penalty into an LLM's proposal distribution, then corrects and variance-reduces the resulting uncertainty estimates, for both autoregressive models (ARMs) and masked diffusion models (MDMs).

What This Paper Is About

Estimating uncertainty in free-form LLM answers usually requires drawing many samples, and those samples are often semantically redundant, which wastes compute and produces unstable estimates. The authors propose a sampler that actively pushes each new generation away from the meanings of previously generated samples, then uses importance reweighting to undo the bias this introduces, so that semantic uncertainty (both aleatoric and epistemic) can be estimated accurately with fewer samples.

Key Contributions

1. A diversity-steered sampler for both decoding paradigms. A continuous semantic-similarity penalty is subtracted from the model's logits during decoding (Equation 5 for ARMs, Equation 7 for MDMs), using an NLI model finetuned to score partial prefixes or intermediate diffusion states. The method requires no gradient access to the base LLM or the NLI model and is model-agnostic.

2. A lightweight NLI adaptation for partial sequences. Starting from DeBERTa-large-MNLI, the authors add a special [TRUNC] token for ARMs or reuse [MASK] for MDMs, freeze all existing parameters, and update only the special-token embedding, [CLS] embedding, and classification head, which corresponds to only 0.3% of model parameters (around 3M parameters) in the DeBERTa-large-MNLI case.

3. Importance-reweighted uncertainty estimators. Because sampling comes from a biased proposal, the authors derive self-normalized importance weights (Equation 8) and apply them to an importance-weighted semantic entropy estimator (Equation 10) and an importance-reweighted mutual information (MI) estimator (Equation 13) built on the MI proxy of Yadkori et al. (Algorithm 2, restricted to the pairwise case n=2).

4. Adaptive control variates and practical sample-efficiency tools. Control variates using base-model log probabilities reduce estimator variance (Equations 14 and 15) with an adaptively estimated coefficient and no extra inference cost; the paper also describes adaptive tuning of the diversity hyperparameter λ and online stopping based on estimator stability.

Main Findings

• Higher AUROC across models and datasets. On AUROC of semantic entropy with N=16 sequences and a ROUGE-L (F1) correctness threshold of 0.3, the method is best on average by mean. Examples: CoQA with OPT-6.7B, 0.75 ± 0.02 versus 0.59 ± 0.06 for vanilla sampling; TruthfulQA with LLaDA 8B-Instruct, 0.94 ± 0.02 versus 0.85 ± 0.04 for vanilla. On TriviaQA with OPT-13B, SDLG scored 0.86 ± 0.03 versus 0.85 ± 0.03 for the proposed method.

• More semantic clusters at the same sample budget. By design the sampler covers more semantic clusters than the baselines while using the same number of samples, illustrated on the CoQA dataset in Figure 4.

• Stronger on ambiguous, free-form tasks. The advantage over baselines is more evident in free-form and ambiguous datasets like CoQA and AmbigQA than in TriviaQA, which has a single, usually short, unambiguous answer.

• Threshold-invariant trends. Results hold at ROUGE-L thresholds of 0.1 and 0.5, and the threshold-free Spearman ρ between negated ROUGE-L scores and estimated uncertainties was 5% and 6% greater on average than those of DBS and SDLG, respectively.

• Acceptable importance-weight variance. The ratio ESS/N stays above 0.4, which the authors say suggests acceptable variance even without control variates.

• Partial sequences already carry semantic signal. The finetuned NLI converges to the ground-truth label long before the final token appears; its predicted entailment probability peaks sharply at the "random guess" value of 1/3 when only one token is revealed, which motivates the default λ schedule (λ is increased when the bidirectional entailment score with the most similar existing generation exceeds 1/3).

• Semantics-agnostic heuristics fall short. Consistent with Aichberger et al., simple temperature scaling or diverse beam search are insufficient to fully explore the semantic space.

Methodology in Plain English

The core idea is to make each new generation as different in meaning as possible from the ones already drawn. While decoding, the model's next-token log probabilities are penalized by a term proportional to the maximum bidirectional entailment score between the current partial sequence and any previously generated sequence. The penalty is controlled by a diversity parameter λ, which the authors tune adaptively; the aggregation over previous samples uses a max, though mean or median could be used for softer guidance. The same logic is applied to masked diffusion models, where the penalty is computed against an intermediate sequence constructed by substituting the current proposal into the masked positions.

To score partial text, a pretrained NLI model (DeBERTa-large-MNLI) is lightly finetuned. Truncation is applied to one side of each MNLI training pair for ARMs, and random masking is applied for MDMs; only the special-token embedding, the [CLS] embedding, and the classification head are updated.

Because decoding now follows a biased proposal rather than the true model distribution, the authors correct downstream estimates with self-normalized importance weights, and then reduce remaining variance with control variates that exploit the correlation between base-model log probabilities and the estimator's target quantity. For MI, the epistemic proxy is built from pairs of responses drawn by iterative prompting and clustered with the same NLI-based entailment metric.

Evaluation uses the same estimation procedure on the final generations of every sampling scheme, so the comparison is fair. Baselines include standard IID sampling at temperatures 1 and 2, diverse beam search with a penalty hyperparameter of 0.5, and a re-implementation of SDLG; MDM baselines use Gumbel temperatures of 1 and 2. Clustering follows Kuhn et al.'s Algorithm 1 using the DeBERTa-large-MNLI model finetuned on MNLI.

Datasets: four QA benchmarks — 907 validation matched instances with shorter stories from CoQA, 1,000 instances from the validation no-context reading comprehension split of TriviaQA, 800 instances from the validation split of TruthfulQA, and the light validation split of AmbigQA.

Models: OPT-6.7B and OPT-13B (autoregressive), LLaMA 3 8B-Instruct (instruction-tuned autoregressive), and LLaDA 8B-Instruct (instruction-tuned masked diffusion). OPT-30B was omitted because the marginal AUROC improvement relative to OPT-13B has been insignificant in prior work.

Metrics: AUROC with correctness defined by ROUGE-L (F1) < 0.3 against the reference answer (maximum across multiple references), plus the average number of clusters and the effective sample size of the importance weights relative to N=16. Numbers are reported as mean ± std over 5 jackknife samples of size 200.

Why This Matters

Uncertainty estimates for LLM outputs are only useful if they can be produced cheaply and reliably. This work makes semantic uncertainty estimation more sample-efficient and extends it beyond autoregressive models to masked diffusion models, which the authors note have been largely overlooked for uncertainty quantification despite matching state-of-the-art autoregressive text quality.

Real-world applications:

• Risk-sensitive model deployment where a system must decide whether to answer or abstain, such as scientific or medical question answering.
• Hallucination detection and flagging of out-of-distribution or knowledge-gap responses, using the epistemic uncertainty proxy (MI).
• Selective prediction and routing, where high-uncertainty queries are escalated to a human expert or a stronger model.
• Benchmarking and monitoring of QA systems on ambiguous, multi-answer questions such as AmbigQA.

Industry relevance: The framework is modular and requires no gradient access to the base LLM, so it can be layered onto existing inference stacks. The authors describe it as a drop-in enhancement. The main practical cost is that sequential diversity steering scales linearly with the number of samples N, which trades against parallel throughput.

Future Directions

• Robustness of clustering. Real text often straddles several plausible semantic clusters and NLI scores are noisy; treating cluster assignments as random and marginalizing over them, potentially at the token level, the embedding level, or both, could improve downstream estimates.
• Prompt-level marginalization. The current pipeline conditions on a single prompt realization; sampling multiple paraphrased prompt templates (or generating question paraphrases with the base model) would yield uncertainty estimates robust to prompt wording.
• Hybrid sampling for throughput. The authors suggest investigating batch sampling and other hybrids to balance exploration against parallel generation, since sequential steering costs time that grows linearly with N.
• Integration with RL fine-tuning. The logit-repulsion plus self-normalized importance-weighting scheme could in principle be embedded inside on-policy RL methods such as PPO and its grouped-reward variant GRPO, driving roll-outs toward novel semantic clusters while keeping return estimates unbiased.

Target Audience

Researchers and practitioners working on LLM uncertainty quantification, hallucination detection, and decoding strategies, particularly those interested in semantic entropy, importance sampling, or masked diffusion models. It is also relevant to engineers building risk-sensitive QA or abstention systems who need compute-efficient uncertainty estimates, and to readers following work such as Aichberger et al.'s semantically diverse language generation. The paper assumes prior familiarity with importance sampling and entropy-based uncertainty measures, so it is best suited to readers with an intermediate-to-advanced background in probabilistic machine learning and NLP.

Authors’ abstract

Accurately estimating semantic aleatoric and epistemic uncertainties in large language models (LLMs) is particularly challenging in free-form question answering (QA), where obtaining stable estimates often requires many expensive generations. We introduce a diversity-steered sampler that discourages semantically redundant outputs during decoding, covers both autoregressive and masked diffusion paradigms, and yields substantial sample-efficiency gains. The key idea is to inject a continuous semantic-similarity penalty into the model's proposal distribution using a natural language inference (NLI) model lightly finetuned on partial prefixes or intermediate diffusion states. We debias downstream uncertainty estimates with importance reweighting and shrink their variance with control variates. Across four QA benchmarks, our method matches or surpasses baselines while covering more semantic clusters with the same number of samples. Being modular and requiring no gradient access to the base LLM, the framework promises to serve as a drop-in enhancement for uncertainty estimation in risk-sensitive model deployments.