TLDR: A new method called “diversity-steered sampling” enhances uncertainty quantification in large language models (LLMs) by discouraging semantically redundant outputs during text generation. It uses a fine-tuned natural language inference (NLI) model to apply a semantic similarity penalty, works for both autoregressive and masked diffusion models, and employs importance reweighting and control variates to produce more accurate and sample-efficient estimates of aleatoric and epistemic uncertainties across various QA tasks.
Large language models (LLMs) are incredibly powerful for generating text, but they often struggle with accurately expressing their uncertainty. This is particularly challenging in free-form question answering, where different responses can mean the same thing (semantic equivalence), and getting reliable uncertainty estimates usually requires generating many expensive outputs.
Researchers Ji Won Park and Kyunghyun Cho from Prescient Design, Genentech, along with Kyunghyun Cho also from New York University, have introduced a novel approach called “diversity-steered sampling” to address this problem. Their method significantly improves how efficiently LLMs can quantify both aleatoric (inherent ambiguity) and epistemic (lack of knowledge) uncertainties. You can read the full paper here.
The Challenge of Uncertainty in LLMs
LLMs can produce fluent text, but they are prone to two main types of uncertainty. Aleatoric uncertainty arises from the inherent ambiguity in a question, where multiple correct answers exist (e.g., “Name one way to measure the Hubble constant” has several valid methods). Epistemic uncertainty, on the other hand, reflects the model’s lack of knowledge, which can lead to “hallucinations” – confident but factually incorrect outputs.
Current methods for estimating uncertainty often rely on generating a large number of samples. However, these samples frequently contain semantically redundant outputs, wasting computational resources. While some diversity-oriented techniques exist, they often don’t consider the semantic meaning of the generated text.
Diversity-Steered Sampling: A Unified Approach
The core of Park and Cho’s innovation is a diversity-steered sampler that actively discourages semantically redundant outputs during the text generation process. This approach is versatile, working with both autoregressive models (ARMs), which generate text token by token, and masked diffusion models (MDMs), which iteratively refine masked spans.
The key idea is to inject a continuous semantic-similarity penalty into the model’s proposal distribution. This penalty is calculated using a natural language inference (NLI) model. This NLI model is lightly fine-tuned to understand semantic similarity even from partial prefixes (for ARMs) or intermediate masked states (for MDMs).
How It Works
During decoding, the sampler modifies the probability of generating the next token. It adds a penalty that pushes the current generation away from any semantically similar outputs that have already been generated. This ensures that the LLM explores a wider range of semantically distinct responses rather than repeating similar ideas.
To ensure accuracy, the method debiases downstream uncertainty estimates using importance reweighting. This corrects for the bias introduced by the diversity-steered sampling. Additionally, control variates are used to shrink the variance of these estimates, making them more stable and reliable.
A crucial aspect is the NLI model. It’s adapted to handle incomplete text by introducing special tokens like [TRUNC] for prefixes or [MASK] for diffusion masks. This fine-tuning is minimal, affecting only a small percentage of the NLI model’s parameters, and doesn’t require gradient access to the base LLM, making it a modular “drop-in” enhancement.
Also Read:
- String Seed of Thought: Empowering LLMs with Reliable Randomness and Creative Diversity
- Enhancing LLM Knowledge Distillation with Counterfactual Explanations in Low-Data Settings
Broader Applications and Benefits
The researchers evaluated their method across four question-answering benchmarks, including CoQA, TriviaQA, TruthfulQA, and AmbigQA. Their results show that diversity-steered sampling matches or surpasses existing baselines in quantifying uncertainty, while also covering a greater number of semantic clusters with the same number of samples. This means the model provides a more comprehensive view of possible answers and their associated uncertainties.
The framework is model-agnostic and can be applied to various LLMs, including modern instruction-finetuned models like LLaMA 3 8B-Instruct and LLaDA 8B-Instruct. This modularity and lack of gradient access requirements make it highly practical for integrating into risk-sensitive AI deployments.
By generating a more diverse set of semantically distinct outputs, this method offers a more accurate and sample-efficient way to understand what an LLM “knows” and “doesn’t know,” which is vital for building safer and more trustworthy AI systems.
