Paper notes — Uncertainty Quantification for In-Context Learning of LLMs

Ling, C., Zhao, X., Zhang, X., Cheng, W., Liu, Y., Sun, Y., Oishi, M., Osaki, T., Matsuda, K., Ji, J., Bai, G., Zhao, L., & Chen, H. (2024). Uncertainty Quantification for In-Context Learning of Large Language Models. NAACL 2024 (main conference).

Motivation stated in the paper. Standard uncertainty-quantification tools for LLMs (predictive entropy, ensembling, MC dropout–style methods) were designed around two sources of randomness: the model’s parameters and randomness in generation. In-context learning introduces a third source that none of those tools were built to isolate: the prompt itself is a choice — which demonstrations are picked, how many, and in what order — and that choice can shift predictions on a fixed underlying task, in a way that’s a property of the prompt construction rather than of the model or the label noise.

Decomposition proposed. The paper defines the total predictive uncertainty of an ICL prediction and splits it into two named components:

  • Aleatoric uncertainty, attributed to randomness in which demonstrations were selected for the prompt — estimated by resampling many different demonstration sets (and/or orderings) for the same task and measuring how much the predictive distribution varies across them.
  • Epistemic (or “configuration”) uncertainty, the residual uncertainty that remains even after accounting for demonstration selection — attributed to the model’s own limitations rather than to which examples happened to be shown.

Formally, this follows the same law-of-total-variance style decomposition Kendall and Gal use for vision uncertainty: total variance = variance explained by demonstration-set choice, plus the expected residual variance conditioned on a fixed demonstration set.

How it’s estimated in practice. Because there is no gradient-based retraining step in ICL to get posterior samples from, the paper’s estimator works by repeated resampling and re-inference: draw many demonstration sets from the available pool of labeled examples for a task, run the model on each resulting prompt, and use the spread of the resulting output distributions as the aleatoric term, with the within-prompt predictive entropy (or similar) contributing the epistemic term.

Evaluation. The method is tested across multiple LLMs and classification/QA-style ICL benchmarks, comparing against uncertainty baselines that don’t separate the two sources (e.g., plain predictive entropy on a single fixed prompt). The reported result is that the decomposed estimate correlates better with actual prediction correctness/calibration than the undecomposed baselines, and that the aleatoric component specifically tracks how much accuracy varies when the demonstration set is swapped for a fixed model and task — confirming that a meaningful share of ICL’s apparent uncertainty is attributable to prompt construction rather than the model.

Where it sits relative to function-vector-based approaches. The estimator’s inputs are entirely the model’s outputs under many resampled prompts — it never inspects internal activations. This makes it model-agnostic (works on any black-box LLM) but requires many forward passes per query to get a reliable estimate, and it can only speak to uncertainty indirectly, through how much the final prediction moves under input perturbation, rather than through any direct read of what the model internally represents “the task” to be.