Transformers Can Learn Posterior Predictive Distributions In-Context
Provides theoretical foundations for understanding how transformers approximate Bayesian posterior distributions in-context, addressing a gap in the theory of prior-data fitted networks.
This paper theoretically shows that transformers can learn posterior predictive distributions (PPDs) in-context for Gaussian process regression by implementing gradient descent and nonlinear mappings, with error bounds depending on attention depth and bin resolution. Simulations confirm the theoretical findings, highlighting the role of normalization and attention depth in extrapolation.
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.