Research map / Probabilistic and scientific modeling / Uncertainty and Gaussian processes
Variational Uncertainty Decomposition for In-Context Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes a variational uncertainty decomposition for LLM in-context learning that optimises auxiliary probe queries to upper-bound aleatoric uncertainty (and lower-bound epistemic), without sampling the latent parameter posterior.
Problem
Decomposing LLM in-context learning predictive uncertainty into aleatoric and epistemic parts is hard because the implicit Bayesian parameter posterior cannot be sampled explicitly and exchangeability often fails.
Approach
Append optimisable auxiliary queries Z (with LLM-generated answers U) to the context; the expected predictive entropy conditioned on (U, Z) upper-bounds the Bayesian aleatoric uncertainty, so minimising over Z gives a variational estimate, with epistemic uncertainty lower-bounded as the remainder. Permutation ensembling and KL filtering promote approximate exchangeability.
Key results
- On the Buttons bandit with p=0.7 and alpha=2, epistemic variance (EV) achieves worst-case regret 0.092 vs 0.199 for total variance (TV).
- In-context abstention on MMLU-M rejecting the top 10% uncertain questions raises accuracy from 68.03% to 72.22% using aleatoric uncertainty, versus 70.00% with total uncertainty.
- Visualisations on synthetic tasks show epistemic uncertainty decreasing near demonstrations and aleatoric uncertainty concentrated at the decision boundary.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use aleatoric uncertainty rather than total uncertainty to filter the highest 10% uncertain questions | VUD abstention · BoolQA | Accuracy (%): 85.56{ $\pm$.002} -> 85.56{ $\pm$.003} | No clear effect |
| Use aleatoric uncertainty rather than total uncertainty to filter the highest 10% uncertain questions | VUD abstention · HotpotQA | Accuracy (%): 87.78{ $\pm$.002} -> 90.00{ $\pm$.003} | Helped |
| Use aleatoric uncertainty rather than total uncertainty to filter the highest 10% uncertain questions | VUD abstention · PubMedQA | Accuracy (%): 88.89{ $\pm$.001} -> 90.00{ $\pm$.000} | Helped |
| Use aleatoric uncertainty rather than total uncertainty to filter the highest 10% uncertain questions | VUD abstention · MMLU-CS | Accuracy (%): 85.56{ $\pm$.002} -> 86.67{ $\pm$.001} | Helped |
| Use aleatoric uncertainty rather than total uncertainty to filter the highest 10% uncertain questions | VUD abstention · MMLU-M | Accuracy (%): 70.00{ $\pm$.003} -> 72.22{ $\pm$.002} | Helped |
| Use epistemic variance rather than total variance in UCB | LLM-UCB on Buttons · Buttons, p>0.5 | Mean Worst-Case Regret: p=0.6, α=2: 0.198{ $\pm$.001} -> 0.149{ $\pm$.039}; p=0.6, α=5: 0.199{ $\pm$.000} -> 0.140{ $\pm$.013}; p=0.7, α=2: 0.19 | Helped |
| Use epistemic variance rather than total variance in UCB | LLM-UCB on Buttons · Buttons, p>0.5 | Mean Regret: p=0.6, α=2: 0.035{ $\pm$.054} -> 0.068{ $\pm$.042}; p=0.6, α=5: 0.158{ $\pm$.065} -> 0.105{ $\pm$.027}; p=0.7, α=2: 0.07 | Mixed |
| Increase the in-context dataset size (ablation) | VUD logistic regression · Logistic regression, in-distribution and out-of-distribution test inputs | Epistemic uncertainty | Hurt |
| Increase the in-context dataset size (ablation) | VUD logistic regression · Logistic regression | Aleatoric uncertainty | No clear effect |
| Choose Z by perturbing the test input (ablation) | VUD logistic regression · Logistic regression | Aleatoric uncertainty estimate | Helped |
| Choose Z by random sampling rather than perturbing the test input (ablation) | VUD logistic regression · Logistic regression, low ICL-density regions | Aleatoric uncertainty estimate stability | Hurt |
| Choose Z by repeating the test input (ablation) | VUD logistic regression · Logistic regression | Aleatoric uncertainty estimate | Mixed |
Surprising findings
- Repeating the test input as the probe (Z = x*) gives the lowest variational aleatoric estimate, and random and BO sampling yield the lowest KL, though perturbation is preferred for scalability.
Limitations
- Approximate Bayesian behaviour: the exchangeability filtering is necessary but not sufficient, so ICL is only approximately Bayesian.
- Restricted to regression and classification with short outputs; natural-language responses not addressed.
- Bandit results have high variance across seeds.
Experimental setup
Datasets: synthetic 1-D logistic regression synthetic linear regression Two Moons Spirals BoolQA HotpotQA PubMedQA MMLU
Benchmarks: Buttons bandit BoolQA HotpotQA PubMedQA MMLU-CS MMLU-Moral in-context OOD detection
Metrics: mean regret worst-case mean regret median reward accuracy after abstention AUC
Baselines: UCB1 Greedy Instruct Baseline Martingale posterior decomposition total uncertainty thresholding
Builds on: Bayesian interpretation of in-context learning de Finetti's theorem Martingale posterior UCB bandits
Closest papers by meaning
- Decomposing Uncertainty for Large Language Models through Input Clarification Ensembling (ICML 2024 · similarity 0.91)
- Is In-Context Learning in Large Language Models Bayesian? A Martingale Perspective (ICML 2024 · similarity 0.89)
- Conformal Prediction with Paraphrase-Aware Scoring for LLM Uncertainty Quantification (NeurIPS 2026 · similarity 0.87)
- Inv-Entropy: A Fully Probabilistic Framework for Uncertainty Quantification in Language Models (NeurIPS 2025 · similarity 0.87)
- Distinguishing the Knowable from the Unknowable with Language Models (ICML 2024 · similarity 0.87)
- Training-Free Bayesianization for Low-Rank Adapters of Large Language Models (NeurIPS 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.