Can Transformers Learn Full Bayesian Inference in Context?

TL;DR

This paper shows that transformers can perform full Bayesian inference in context, enabling them to generate posterior distributions for statistical models without additional training, matching traditional inference methods in quality.

Contribution

It introduces a novel framework combining ideas from prior fitted networks and normalizing flows, allowing transformers to perform complex Bayesian inference in context.

Findings

Transformers can generate high-quality posterior samples in context.

The approach matches the performance of MCMC and variational inference.

Experiments on real datasets validate the method's effectiveness.

Abstract

Transformers have emerged as the dominant architecture in the field of deep learning, with a broad range of applications and remarkable in-context learning (ICL) capabilities. While not yet fully understood, ICL has already proved to be an intriguing phenomenon, allowing transformers to learn in context -- without requiring further training. In this paper, we further advance the understanding of ICL by demonstrating that transformers can perform full Bayesian inference for commonly used statistical models in context. More specifically, we introduce a general framework that builds on ideas from prior fitted networks and continuous normalizing flows and enables us to infer complex posterior distributions for models such as generalized linear models and latent factor models. Extensive experiments on real-world datasets demonstrate that our ICL approach yields posterior samples that are similar in quality to state-of-the-art MCMC or variational inference methods that do not operate in context. The source code for this paper is available at https://github.com/ArikReuter/ICL_for_Full_Bayesian_Inference.