Uncertainty Quantification for In-Context Learning of Large Language Models

TL;DR

This paper tackles predictive uncertainty in large language models during in-context learning by decomposing it into aleatoric uncertainty (from demonstrations) and epistemic uncertainty (from model configurations). It formulates in-context learning within a Bayesian framework with latent concepts and model parameters, and introduces an entropy-based, unsupervised estimation pipeline to separate EU and AU. Extensive experiments on LLaMA-2 and OPT-13B across sentiment, COLA, SST-2, and AG_News demonstrate that EU and AU can outperform baselines in misclassification detection and are robust to domain shifts, with EU proving especially effective for detecting out-of-domain and semantic out-of-distribution scenarios. The work provides a practical tool for diagnosing and mitigating uncertainty in ICL, guiding data selection and model usage for more reliable AI systems.

Abstract

In-context learning has emerged as a groundbreaking ability of Large Language Models (LLMs) and revolutionized various fields by providing a few task-relevant demonstrations in the prompt. However, trustworthy issues with LLM's response, such as hallucination, have also been actively discussed. Existing works have been devoted to quantifying the uncertainty in LLM's response, but they often overlook the complex nature of LLMs and the uniqueness of in-context learning. In this work, we delve into the predictive uncertainty of LLMs associated with in-context learning, highlighting that such uncertainties may stem from both the provided demonstrations (aleatoric uncertainty) and ambiguities tied to the model's configurations (epistemic uncertainty). We propose a novel formulation and corresponding estimation method to quantify both types of uncertainties. The proposed method offers an unsupervised way to understand the prediction of in-context learning in a plug-and-play fashion. Extensive experiments are conducted to demonstrate the effectiveness of the decomposition. The code and data are available at: https://github.com/lingchen0331/UQ_ICL.

Figures (4)

• Figure 1: Uncertainty in LLM's prediction can stem from two aspects: a) Demonstration Quality: LLMs are likely to make wrong predictions if the demonstrations are inappropriate; b) Model Configuration: different decoding strategies (e.g., beam_search and top_k sampling) and their parameter settings may return different predictions.
• Figure 2: Uncertainty Quantification of In-context Learning Pipeline: we want to quantify the uncertainty that comes from 1) different in-context demonstrations $\mathbf{x}_{1:T}$; and 2) different model configurations $\Theta_l$.
• Figure 3: Framework of entropy-based uncertainty estimation, which consists of 1) generating $M$ sequences based on a set of $\mathbf{x}_{1:T-1}$; 2) selecting token(s) that is relevant to the answer and extract the probabilities; 3) aggregating the token probabilities of $M$ sequences into a distribution of predicted labels; 4) iterating the process $L$ times corresponding to $L$ different demonstration sets and form a probability matrix $\mathcal{M}$, where the column denotes different demonstration sets and the row denotes labels of the dataset.
• Figure 4: The performance of misclassification rate using two backbone LLMs: OPT-13B and LLaMA-2-13B on Emotion dataset. (a) and (b) demonstrate the precision-recall curves (x-axis is the recall and y-axis is the precision) for OPT-13B and LLaMA-2-13B; (c) and (d) demonstrate the ROC curve (x-axis is the false positive rate and y-axis is the true positive rate) for OPT-13B and LLaMA-2-13B, respectively.