Transformers Can Learn Posterior Predictive Distributions In-Context

Prior-data fitted networks (PFNs) have recently emerged as a powerful approach for Bayesian prediction tasks, approximating the posterior predictive distribution (PPD) through in-context learning. Despite their strong empirical performance and ability to go beyond point predictions, theoretical understandings of the algorithmic capability of transformers to learn distributions in context are still lacking. Focusing on Gaussian process regression problems, we show by construction that transformers can implement a gradient descent algorithm targeting the posterior predictive mean and variance, followed by nonlinear mappings that yield binned probabilities of PPD. We study the error bounds of the approximated PPD in terms of attention depth and bin resolution. Based on these results, we further demonstrate the key role of normalization and the choice of attention depth in enabling the extrapolation abilities of transformers beyond the pretraining sample size range. We conduct simulations that corroborate our findings, providing insight into the expressivity of PFNs targeting PPDs and how architectural choices may influence generalization capabilities.

Prior-data fitted networks are transformer-based models that aim to make Bayesian predictions after being trained on many simulated datasets. Unlike models that output only a single estimate, they can produce a full predictive distribution, which describes both the likely value and the uncertainty around it. However, it has been unclear how a transformer can carry out the Bayesian computation needed to produce such a distribution. In this paper, we study this question using Gaussian process regression, a classical setting where the correct Bayesian predictive distribution is known. We show that a transformer can be constructed so that its attention layers run an iterative algorithm for computing the predictive mean and variance, while its final neural network layer converts these quantities into probabilities over output ranges. We also show that normalizing attention is important for making the model work reliably on larger datasets than those seen during training. This gives a concrete explanation of how transformer-based probabilistic predictors can perform Bayesian prediction in context, and provides a template for understanding practical PFNs with more complex priors.