Variational Uncertainty Decomposition for In-Context Learning

I. Shavindra Jayasekera*, Jacob Si*, Filippo Valdettaro, Wenlong Chen, A. Aldo Faisal, Yingzhen Li
Imperial College London
*Denotes Equal Contribution
Paper Code

Abstract

As large language models (LLMs) gain popularity in conducting prediction tasks in-context, understanding the sources of uncertainty in in-context learning becomes essential to ensuring reliability. The recent hypothesis of in-context learning performing predictive Bayesian inference opens the avenue for Bayesian uncertainty estimation, particularly for decomposing uncertainty into epistemic uncertainty due to lack of in-context data and aleatoric uncertainty inherent in the in-context prediction task. However, the decomposition idea remains under-explored due to the intractability of the latent parameter posterior from the underlying Bayesian model. In this work, we introduce a variational uncertainty decomposition framework for in-context learning without explicitly sampling from the latent parameter posterior, by optimising auxiliary queries as probes to obtain an upper bound to the aleatoric uncertainty of an LLM's in-context learning procedure, which also induces a lower bound to the epistemic uncertainty. Through experiments on synthetic and real-world tasks, we show quantitatively and qualitatively that the decomposed uncertainties obtained from our method exhibit desirable properties of epistemic and aleatoric uncertainty.

VUD Fig 1

Figure 1. Uncertainty Decomposition with Auxiliary Data (Above). Decomposition Example for Two-Moons Dataset (Below).

VUD Fig 1

Figure 2. Variational Uncertainty Decomposition (VUD) Framework.

VUD Visualizations

The following figures depict the decomposition of total uncertainty into aleatoric and epistemic components via VUD. These diagrams illustrate the uncertainty decomposition in classification tasks with an in-context learning (ICL) dataset size of 15. Blue scatter points represent the total uncertainty while pink scatter points represent aleatoric uncertainty. The difference (gap) between the blue and pink curve indicates the epistemic uncertainty. Green dashed lines in the background marks ICL examples with label y = 0, while gray dashed lines indicate those with label y = 1. Notice how uncertainty changes in dense and sparse regions.

Aleatoric & Epistemic Decomposition

Figure 3. Uncertainty Decomposition Algorithm on Logistic Regression Dataset.

Total Uncertainty Visualization

Figure 4. Uncertainty Decomposition Algorithm on Uniform Interval Dataset.

Decomposition while Adding ICL Points

Figure 5. Uncertainty Decomposition Algorithm on Logistic Regression Dataset with Increasing ICL Examples.

BibTeX

Please consider citing our paper if you find it helpful. Thank you!


      @misc{jayasekera2025variationaluncertaintydecompositionincontext,
            title={Variational Uncertainty Decomposition for In-Context Learning}, 
            author={I. Shavindra Jayasekera and Jacob Si and Filippo Valdettaro and Wenlong Chen and A. Aldo Faisal and Yingzhen Li},
            year={2025},
            eprint={2509.02327},
            archivePrefix={arXiv},
            primaryClass={stat.ML},
            url={https://arxiv.org/abs/2509.02327}, 
      }