Can Generative AI Solve Your In-Context Learning Problem? A Martingale Perspective

TL;DR

This paper reframes in-context learning with conditional generative models as a Bayesian model-criticism problem, introducing the generative predictive $p$-value to evaluate when a CGM reliably solves an ICL task. By linking martingale predictive $p$-values to posterior predictive checks via Doob's theorem, the authors provide a practical procedure that uses dataset completions sampled from the CGM to assess model suitability without explicit latent explanations. Empirically, the method accurately predicts model capability across tabular, natural-language, and imaging tasks, and the interplay between NLML and NLL discrepancies reveals data sufficiency and informs computational trade-offs. The approach offers a principled, scalable framework for risk-aware deployment and potential model-selection extensions in generative AI systems.

Abstract

This work is about estimating when a conditional generative model (CGM) can solve an in-context learning (ICL) problem. An in-context learning (ICL) problem comprises a CGM, a dataset, and a prediction task. The CGM could be a multi-modal foundation model; the dataset, a collection of patient histories, test results, and recorded diagnoses; and the prediction task to communicate a diagnosis to a new patient. A Bayesian interpretation of ICL assumes that the CGM computes a posterior predictive distribution over an unknown Bayesian model defining a joint distribution over latent explanations and observable data. From this perspective, Bayesian model criticism is a reasonable approach to assess the suitability of a given CGM for an ICL problem. However, such approaches -- like posterior predictive checks (PPCs) -- often assume that we can sample from the likelihood and posterior defined by the Bayesian model, which are not explicitly given for contemporary CGMs. To address this, we show when ancestral sampling from the predictive distribution of a CGM is equivalent to sampling datasets from the posterior predictive of the assumed Bayesian model. Then we develop the generative predictive $p$-value, which enables PPCs and their cousins for contemporary CGMs. The generative predictive $p$-value can then be used in a statistical decision procedure to determine when the model is appropriate for an ICL problem. Our method only requires generating queries and responses from a CGM and evaluating its response log probability. We empirically evaluate our method on synthetic tabular, imaging, and natural language ICL tasks using large language models.

Figures (20)

• Figure 1: An example illustrating two ICL problems. One that the model $\bm{\theta}$ (Llama-2 7B touvron2023llama) can solve, and one that it cannot. On the left, we have (a) examples from the SST2 task socher-etal-2013-recursive comprising part of an ICL dataset $\mathrm{x}^{n}$, (b) a new query $\mathrm{z}$, and (c) some responses $\mathrm{y}$ sampled from the CGM $p_{\bm{\theta}}(\mathrm{y} \mid \mathrm{z}, \mathrm{x}^{n})$ when prompted with the dataset and query. The true label is "negative" and the CGM responds correctly. On the right, we have have the same format but the dataset (d) and query (e) are taken from the medical question pairs task mccreery2020effective. Here, the true label is "similar," but the model responds incorrectly with "different" half of the time.
• Figure 2: Examples of misaligned model and data combinations. Transformer models (pink) are fit to either linear or polynomial noisy data defined by reference Bayesian models (blue).
• Figure 3: Tabular data tasks.
• Figure 4: Natural language in-capability vs. out-of-capability tasks. Green solid line is the ICL error rate for Llama-2 7B. Gray dashed line is the random guessing error rate.
• Figure 5: Generative fill tasks using the test sets of SVHN, MNIST, and CIFAR-10.

Theorems & Definitions (9)

• Theorem 1
• proof
• Definition 1
• Theorem 2
• proof
• Lemma 3
• proof
• Theorem 2
• proof