Can Transformers Learn Full Bayesian Inference In Context?

Conference: ICML 2025
arXiv: 2501.16825
Code: https://github.com/ArikReuter/ICL_for_Full_Bayesian_Inference
Area: Optimization/Bayesian Inference
Keywords: In-Context Learning, Bayesian Inference, Posterior Sampling, Diffusion Transformer, Normalizing Flows

TL;DR

This paper demonstrates that Transformers can perform full Bayesian inference in context. By pre-training an encoder-decoder architecture (TabPFN encoder + diffusion Transformer decoder) on synthetic data, the model can generate posterior samples of comparable quality to HMC for statistical models like GLMs and Gaussian mixture models during deployment, without requiring any parameter updates.

Background & Motivation

Background: In-context learning (ICL) has become a fundamental capability in NLP, enabling LLMs to adapt to tasks using in-context information without fine-tuning. TabPFN has also demonstrated the powerful capabilities of ICL in tabular data classification. However, existing ICL methods only output point estimates or univariate distributions of the posterior predictive distribution.

Limitations of Prior Work: - Full Bayesian inference (obtaining high-dimensional continuous posteriors \(P^{z|x}\)) is crucial in many domains (medicine, physics, neuroscience). - Traditional methods (MCMC/HMC) suffer from slow inference speed, requiring Markov chains to be run from scratch for every new dataset. - Variational inference (VI) requires parametric assumptions and can be inaccurate. - Can a Transformer "understand" data and directly output posterior samples, just like an LLM understands text?

Key Challenge: Full Bayesian inference requires handling complex, high-dimensional posterior distributions, whereas current ICL mostly handles low-dimensional or discrete outputs.

Goal: To achieve full Bayesian inference via ICL, mapping \(x \mapsto P^{z|x}\) to take data as input and output samples of the posterior distribution.

Key Insight: Combine an encoder (TabPFN to process context data) with a generative decoder (diffusion Transformer to generate posterior samples). Train the model on synthetic samples of the joint distribution \((x, z)\), and perform posterior sampling on new data at deployment.

Core Idea: Just as LLMs learn conditional distributions of text from text datasets, this method learns the conditional posterior of parameters from synthetic statistical data. The key is to represent and sample the high-dimensional posterior using continuous normalizing flows (a diffusion process trained via flow matching).

Method

Overall Architecture

Two-stage design: 1. Encoder (TabPFN architecture): Processes the input dataset \(x = \{(x_1, y_1), ..., (x_n, y_n)\}\) to generate context representations. 2. Decoder (Diffusion Transformer + Flow Matching): Conditioned on the encoder output, generates posterior samples \(z \sim P^{z|x}\) by solving a neural ODE. Training: End-to-end training using synthetic samples from the joint distribution \((x, z)\).

Key Designs

1. Synthetic Data Pre-training Paradigm:

   • Function: Samples parameters \(z\) from a prior distribution, samples data \(x\) from the model \(P(x|z)\), and trains the model on \((x, z)\) pairs.
   • Mechanism: The training data covers a large number of different datasets (each \(z\) corresponds to one dataset), enabling the model to learn to infer posteriors from arbitrary datasets.
   • Prior Inclusion Strategy: To enhance robustness, TabPFN's "universal prior"—which covers a wider range of data patterns—is mixed into the pre-training data.
   • Design Motivation: Eliminates the need for real data annotation through fully synthetic training while generalizing effectively to real-world data.
   • Theoretical Foundation: Proves that a flow matching model trained on joint distribution samples can learn the conditional distribution \(P^{z|x}\).

2. Continuous Normalizing Flow Decoder:

   • Function: Maps standard Gaussian noise to the target posterior distribution.
   • Mechanism: Trains a diffusion Transformer using the flow matching objective to learn a velocity field \(v_t(z_t, x)\) that maps \(z_0 \sim N(0,I)\) to \(z_1 \sim P^{z|x}\).
   • Cross-Attention: The decoder accesses the encoder's context representation via cross-attention.
   • Design Motivation: Continuous normalizing flows can represent arbitrarily complex posterior distributions, unlike VI which is limited by parametric forms.

3. Model-Data Flexibility:

   • Function: A single trained model can handle datasets of varying sizes and dimensions.
   • Mechanism: The encoder's self-attention does not restrict the number of data points \(n\), and positional encodings adapt to different feature dimensions.
   • Supported Models: Generalized Linear Models (GLMs), Gaussian Mixture Models (GMMs), Factor Analysis (FA).
   • Design Motivation: Analogous to how LLMs handle texts of different lengths, leveraging the generalization capability of ICL.

Loss & Training

• Flow matching loss: \(\mathcal{L} = \mathbb{E}[\|v_t(z_t, x) - u_t(z_t|z_1)\|^2]\)
• End-to-end training of the encoder and decoder.
• On-the-fly generation of synthetic data (each batch is a newly synthesized dataset).

Key Experimental Results

Main Results

GLM posterior inference (real datasets, compared with MCMC/VI):

Method	\(W_2\) Distance vs HMC ↓	Inference Time	Model-Specific?
HMC (NUTS)	0 (Baseline)	~Minutes	✓ Requires sampler setup
ADVI	0.15	~Seconds	✓ Requires defining the model
Mean-Field VI	0.23	~Seconds	✓ Requires defining the model
ICL (Ours)	0.08	<1s	✗ Out-of-the-box

Gaussian Mixture Model Posterior

Method	Clustering Accuracy ↑	Posterior Coverage ↑	Time
EM Algorithm	90.2%	N/A (Point Estimate)	Fast
Gibbs Sampling	92.1%	94.8%	Slow
ICL (Ours)	91.5%	93.2%	<1s

Ablation Study

Configuration	\(W_2\) vs HMC	Description
Without TabPFN Prior	0.18	Trained only on synthetic data from the target model
Without Cross-Attention	0.35	Decoder cannot see the data
Full Model	0.08	TabPFN + Diffusion + Cross-Attention
Standard VAE Decoder	0.22	VAE posterior is not flexible enough
Flow Matching Decoder	0.08	Continuous flows are more flexible

Key Findings

• The quality of ICL posterior samples is close to HMC (\(W_2 = 0.08\) vs \(0.15\) for ADVI) but 1-2 orders of magnitude faster.
• More accurate than variational inference methods (ADVI, mean-field VI) because it is not limited by parametric assumptions.
• Incorporating the TabPFN prior significantly improves generalization to real-world data (reducing \(W_2\) from 0.18 to 0.08).
• The flow matching decoder vastly outperforms the VAE decoder; the flexibility of flows is required to capture the multi-modality and irregular shapes of posteriors.
• Good scale adaptability: trained on datasets with \(n < 1000\) but generalizes successfully to larger datasets.

Highlights & Insights

• "Understanding data like LLMs understand text"—A grand perspective generalizes ICL from NLP to statistical inference.
• The paradigm of "training on synthetic data + deploying on real data" is highly elegant, requiring no real annotated data.
• The combination of flow matching and cross-attention enables the model to represent arbitrarily complex posteriors, bypassing the Gaussian assumptions in VI.
• Clever reuse of the TabPFN prior, enhancing the robustness of specialized models with a general "data-understanding prior."
• Transformative potential for practical data analysis—statisticians no longer need to configure MCMC samplers; the posterior is produced instantly upon feeding dataset inputs.

Limitations & Future Work

• Only validated on GLM, GMM, and FA; more complex models (e.g., Bayesian neural networks) remain to be explored.
• Posterior dimensionality is currently constrained (<100D), and scalability to high-dimensional posteriors has not been verified.
• Prior distributions for synthetic data must be manually designed; prior mismatch can lead to performance degradation.
• Robustness to model misspecification still has room for improvement.
• Lack of theoretical guarantees (such as convergence rates) for inference quality.

Related Work & Insights

• vs TabPFN: TabPFN performs posterior predictive inference (\(P(y|x)\)), whereas this work obtains the full posterior (\(P(z|x)\)), making it more general.
• vs Simulation-Based Inference (SBI): SBI is trained on simulated data but typically targets a single model, whereas ICL can generalize across models.
• vs MCMC/HMC: Accurate but slow, whereas ICL is fast but approximate, offering a new trade-off between accuracy and speed.
• Insight: The capabilities of ICL may extend far beyond current understanding—from text comprehension to statistical inference, ICL serves as a general "conditional distribution learner."

Rating

• Novelty: ⭐⭐⭐⭐⭐ The intersection of ICL and full Bayesian inference represents a pioneering combination.
• Experimental Thoroughness: ⭐⭐⭐⭐ In-depth evaluation across multiple models and datasets, with comprehensive comparisons against MCMC and VI.
• Writing Quality: ⭐⭐⭐⭐⭐ Clear and elegant, with highly inspiring figures illustrating the analogy with LLMs.
• Value: ⭐⭐⭐⭐⭐ Has the potential to change the way statistical inference is practiced.

Related Papers

• [ICML 2025] How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Bias
• [NeurIPS 2025] Multi-head Transformers Provably Learn Symbolic Multi-step Reasoning via Gradient Descent
• [ICML 2025] In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention
• [ICML 2025] Training Dynamics of In-Context Learning in Linear Attention
• [ICML 2025] On Understanding Attention-Based In-Context Learning for Categorical Data