Mind the Ambiguity: Aleatoric Uncertainty Quantification in LLMs for Safe Medical Question Answering

AI summary

Overview

Research area: Natural language processing — specifically uncertainty quantification and representation engineering applied to large language models for medical question answering (Medical QA).

Technical level: Intermediate. The paper assumes familiarity with transformer internals (residual-stream activations, layers), uncertainty decomposition (aleatoric vs. epistemic), and standard metrics such as AUROC, ECE, and Brier score, but the core idea — that ambiguity is linearly encoded in hidden states and can be read off by a small linear probe — is explained in an accessible way.

Scope: The paper formalizes ambiguous medical queries as a source of aleatoric uncertainty (AU), builds a benchmark of clear/ambiguous question pairs (CV-MedBench), shows AU is linearly separable in LLM hidden states, and uses a lightweight probe to drive a "Clarify-Before-Answer" pipeline that improves answer accuracy across four open LLMs.

What This Paper Is About

Users frequently ask medical questions that are underspecified — for example, "Can I take ibuprofen for pain?" with no mention of demographics or symptom details. These ambiguous queries reduce answer accuracy and create safety risk in high-stakes healthcare settings. The paper's goal is to quantify this input-driven uncertainty (aleatoric uncertainty) accurately and cheaply enough to run at inference time, and to use that signal to proactively request clarification before the model answers.

Key Contributions

1. CV-MedBench (Clear-to-Vague Medical QA Benchmark). The first benchmark designed for studying input ambiguity in Medical QA. Clear questions from MedQA, MedMCQA, and MedExQA are rewritten by an LLM (e.g., GPT-4o) into underspecified variants using three transformation types — context omission, semantic vagueness, and logical inconsistency — producing contrastive pairs (clear, ambiguous) that differ only in ambiguity level. The dataset undergoes human verification, with agreement rates exceeding 95% on topic fidelity, ambiguity validity, and linguistic fluency.

2. AU-Probe (Aleatoric Uncertainty Probe). A lightweight plug-in module that estimates AU directly from an LLM's internal hidden states using a per-layer linear probe trained with binary cross-entropy on binary labels (clear = 0, ambiguous = 1). It requires no LLM fine-tuning and no multiple forward passes. Sigmoid outputs give continuous scores in (0,1) despite binary supervision.

3. AU-Guided "Clarify-Before-Answer" Pipeline. A two-stage framework: Stage 1 computes the AU score before answer generation and, if it exceeds a threshold τ, flags the input as ambiguous and requests clarification; Stage 2 generates the answer once AU is at or below τ. Built on the finding that AU is linearly encoded in activation space and therefore separable by a simple hyperplane.

4. Empirical validation across four open-source LLMs (Qwen2.5-7B-Instruct, Llama-3.1-8B-Instruct, Bio-Medical-Llama-3-8B, BioMistral-7B), covering AU quantification quality, in-distribution (ID) and out-of-distribution (OOD) robustness, low-data behavior, latency, and a case study.

Main Findings

• Near-perfect ambiguity discrimination: AU-Probe achieves near-perfect AUROC across all models and datasets (e.g., 0.9998 on CV-MedQA and 0.9999 on CV-MedMCQA for Qwen2.5-7B-Instruct; 0.9997 on CV-MedQA for BioMistral-7B), compared with baselines such as MSP (0.5573 on CV-MedQA for Qwen2.5-7B-Instruct) and Semantic Entropy (0.5479). The paper reports an average AUROC improvement of 49.24% over baselines.

• Better calibration: AU-Probe generally shows lower Brier scores than competing methods; the paper notes SAR sometimes yields lower ECE in OOD settings because of more conservative AU estimates, but AU-Probe's Brier scores are consistently superior.

• Input ambiguity clearly hurts accuracy: The paper reports a significant drop in response accuracy across LLMs and datasets when queries are ambiguous (Figure 2).

• Clarification improves answer accuracy: Under threshold 0.5, AU-guided clarification improves accuracy over the no-clarification baseline for every model/dataset combination, from +4.79% (BioMistral-7B on CV-MedQA) to +17.50% (Bio-Medical-Llama-3-8B on CV-MedMCQA). The framework's overall reported average accuracy improvement is 9.48%.

• AU-Probe outperforms the strongest baseline for guiding clarification: ASK4CONF was the strongest baseline in the AUROC evaluation, but AU-Probe consistently beat it in driving clarification decisions, with higher accuracy in every reported model/dataset cell of Table 2.

• Layer placement of the AU signal varies by model: On Llama-3.1-8B-Instruct and BioMistral-7B, the per-layer probe peaks in deeper layers (layers 28–32); on Qwen2.5-7B-Instruct and Bio-Medical-Llama-3-8B, it peaks in intermediate layers (layers 8–12).

• Robust under distribution shift: Under OOD evaluation (trained on CV-MedQA or CV-MedMCQA, tested on CV-MedExQA), AUROC stays consistently above 0.85, and ID AUROC is close to 1.0. Qwen2.5-7B-Instruct is especially robust, maintaining AUROC near 1.0 in both ID and OOD settings.

• High data efficiency: Starting from the full CV-MedBench training splits of 1,600 instances, the probe reaches near-optimal ID AUROC with as little as 5–10% of the data for most models, with saturation by 30% at the latest; OOD performance stabilizes between 20% and 40%. The main results use 30% (960 samples), and even 2% (64 samples) remains substantially stronger than the baselines.

• Low latency: AU-Probe takes 1.05–1.22 seconds per sample across the four models on an NVIDIA A40 GPU, comparable to the fastest baselines (MSP, MTE, ASK4CONF at roughly 0.9–1.2 seconds), while consistency-based methods (SE, SAR) and RAUQ range from about 10 to over 35 seconds.

• Case study: On Llama-3.1-8B-Instruct, the query "I have a red itchy rash on my arm. What should I use to treat it?" received an AU score of 0.82 (above τ = 0.5), triggering clarification; after the user supplied details, the score dropped to 0.39 and the model answered safely.

• Threshold choice: Thresholds were searched in {0.1, 0.3, 0.5, 0.7, 0.9}; τ = 0.5 performed best across datasets and models. The threshold is described as tunable — raising it reduces false alarms, lowering it increases sensitivity.

Methodology in Plain English

The researchers started by assembling clear medical questions from three public datasets (MedQA, MedMCQA, MedExQA) and using an LLM to rewrite each one into a vague version that preserves the medical topic but removes or blurs key details. Human reviewers then checked the rewrites on three dimensions, and agreement exceeded 95%.

Next, they fed each clear/ambiguous pair through an LLM and recorded the internal activations (the numerical state of the model at the final prompt token, for each transformer layer). Projecting these activations into two dimensions with PCA showed that clear and ambiguous questions form visibly separate clusters in certain layers — evidence that ambiguity is encoded along an approximately linear direction.

Based on that, they trained a separate simple linear classifier for each layer, using binary labels for clear (0) and ambiguous (1), with a sigmoid to produce scores between 0 and 1. The layer with the best AUROC was selected as the AU-Probe; its score is the estimated aleatoric uncertainty.

Finally, they wrapped this probe into a two-stage pipeline. If the AU score exceeds a threshold, the system asks the user for clarification instead of answering immediately; only when the score falls to or below the threshold does the model generate the final medical answer. They evaluated this against baselines including MSP, Mean Token Entropy, Semantic Entropy, SAR, RAUQ, and ASK4CONF, measuring discrimination (AUROC), calibration (ECE, Brier), accuracy on clarified questions, robustness across datasets, data efficiency, and per-sample latency.

Why This Matters

Impact on research: The paper shows that aleatoric uncertainty — not just total predictive uncertainty — has a linear representation in LLM hidden states, extending the Linear Representation Hypothesis to a new attribute. It also supplies a benchmark (CV-MedBench) for studying input ambiguity in Medical QA, an under-explored source of risk compared with the field's dominant focus on knowledge coverage and diagnostic reasoning. The probe's design shows that high-quality uncertainty estimation can be achieved with a small linear module rather than fine-tuning or repeated sampling.

Real-world applications:

• Clinical decision-support tools that screen incoming patient questions and ask for missing details before offering guidance.
• Consumer-facing symptom checkers, where vague queries are the norm and unsafe advice carries direct health consequences.
• Triage and intake systems that route underspecified queries to a clarification step rather than a potentially misleading automated answer.
• Any latency-sensitive medical QA deployment, since the probe adds negligible overhead compared with sampling- or ensemble-based uncertainty methods.

Industry relevance: The method requires no LLM fine-tuning and no model replicas, is model-agnostic, and runs in about the same time as the cheapest baselines (1.05–1.22 seconds per sample on an NVIDIA A40). That combination of accuracy and speed is what makes it plausible for interactive clinical workflows, where multi-round sampling methods taking 10 to over 35 seconds are impractical.

Future Directions

• Factual and reality mismatch detection. The authors state as a limitation that AU-Probe only assesses ambiguity in the linguistic formulation of the query; it does not detect cases where a clearly phrased question rests on incorrect underlying medical information. They suggest external validation mechanisms (the text is truncated at "real-time physiologic") would be needed.
• A second limitation is announced but not detailed in the available content, so it is not reported here.
• Extending CV-MedBench generation to other datasets and LLMs. The paper notes that the prompt template enables the ambiguous-question generation to be extended to other datasets and models; building broader coverage of real-world ambiguity types is a natural next step.
• Tuning the safety/efficiency trade-off. Since τ is explicitly described as tunable — higher to reduce false alarms, lower to increase sensitivity — calibrating this threshold per clinical deployment setting is an open practical question.
• Generalizing the probe beyond the four evaluated models, including to larger and closed models, is untested in the reported results.

Target Audience

Researchers and practitioners working on LLM safety, uncertainty quantification, and representation engineering; medical AI developers building question-answering or clinical decision-support systems; and NLP researchers interested in benchmarks for input ambiguity. Readers focused on calibration, inference-time efficiency, or the geometry of LLM internal states will also find the layer-wise separability analysis and latency comparisons directly useful. A general audience would need some background in transformer representations and uncertainty metrics to follow the methodology.

Authors’ abstract

The deployment of Large Language Models in Medical Question Answering is severely hampered by ambiguous user queries, a significant safety risk that demonstrably reduces answer accuracy in high-stakes healthcare settings. In this paper, we formalize this challenge by linking input ambiguity to aleatoric uncertainty (AU), which is the irreducible uncertainty arising from underspecified input. To facilitate research in this direction, we construct CV-MedBench, the first benchmark designed for studying input ambiguity in Medical QA. Using this benchmark, we analyze AU from a representation engineering perspective, revealing that AU is linearly encoded in LLM's internal activation patterns. Leveraging this insight, we introduce a novel AU-guided "Clarify-Before-Answer" framework, which incorporates AU-Probe - a lightweight module that detects input ambiguity directly from hidden states. Unlike existing uncertainty estimation methods, AU-Probe requires neither LLM fine-tuning nor multiple forward passes, enabling an efficient mechanism to proactively request user clarification and significantly enhance safety. Extensive experiments across four open LLMs demonstrate the effectiveness of our QA framework, with an average accuracy improvement of 9.48% over baselines. Our framework provides an efficient and robust solution for safe Medical QA, strengthening the reliability of health-related applications. The code is available at https://github.com/yaokunliu/AU-Med.git, and the CV-MedBench dataset is released on Hugging Face at https://huggingface.co/datasets/yaokunl/CV-MedBench.