Research map / Probabilistic and scientific modeling / Uncertainty and Gaussian processes
Can Transformers Learn Full Bayesian Inference in Context?
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Trains a TabPFN-style transformer encoder with a flow-matching diffusion decoder on synthetic joint samples so it outputs full posterior samples for GLMs, factor analysis and GMMs in context, without parameter updates at inference.
Problem
Full Bayesian inference (complex, high-dimensional posteriors over latent variables) is slow with MCMC and biased by restrictive assumptions in VI; the paper asks whether in-context learning can perform this mapping from dataset to posterior samples.
Approach
Amortize posterior inference at the dataset level: pretrain on synthetic (dataset, latent) pairs from the generative model, using a TabPFN-style encoder for the dataset and a conditional flow-matching decoder (adaLN time conditioning) to generate posterior samples via an ODE solve, so a single forward pass conditions on the observed data.
Key results
- On GLM synthetic data, ICL reaches C2ST 0.657, MMD 0.183, W2 0.556, the best of all methods versus HMC/analytical reference (VI full: 0.714, 1.016, 1.601).
- On factor analysis synthetic data, ICL achieves C2ST 0.568 and MMD 0.057, versus 0.987 or higher C2ST for all VI methods.
- On GMM synthetic data, ICL achieves C2ST 0.825 and MMD 0.706, while VI methods saturate near C2ST 0.99.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use ICL instead of VI: Structured | GLM posterior inference · GLMs synthetic datasets | MMD: 0.929 -> 0.183 | Helped |
| Use ICL instead of VI: Structured | GLM posterior inference · GLMs real-world datasets | W2: 0.370 -> 0.387 | Hurt |
| Use ICL instead of VI: Full | Factor analysis posterior inference · Factor analysis synthetic datasets | MMD: 3.222 -> 0.057 | Helped |
| Use ICL instead of VI: IAF | Factor analysis posterior inference · Factor analysis real-world datasets | W2: 0.251 -> 0.583 | Hurt |
| Use ICL instead of VI: IAF | GMM posterior inference · GMM synthetic datasets | MMD: 2.308 -> 0.706 | Helped |
| Use ICL instead of VI: IAF | GMM posterior inference · GMM real-world datasets | MMD: 1.845 -> 1.051 | Helped |
| Use flow matching rather than diffusion objectives (ablation) | ICL posterior inference · Evaluated scenarios | posterior sample similarity | Helped |
| Use flow matching rather than a multivariate Gaussian posterior parameterization (ablation) | ICL posterior inference · Evaluated scenarios | posterior sample similarity | Helped |
| Increase latent dimensionality to K=20 or K=50 (ablation) | ICL posterior inference · GLM scenarios | sample similarity to HMC | No clear effect |
| Apply larger train-test distribution shifts (ablation) | ICL posterior inference · Out-of-distribution evaluation | posterior sample quality | Hurt |
| Replace the transformer encoder with an equally sized MLP encoder (ablation) | ICL posterior inference · Architecture evaluation | performance | Hurt |
| Use a nonlinear neural network classifier alongside the random forest for C2ST (ablation) | C2ST evaluation · Evaluated methods | relative method performance | No clear effect |
Surprising findings
- On real-world GLM data, ICL's Wasserstein-2 (0.387) is slightly worse than VI structured (0.370), and on FA real-world data VI IAF has the best W2 (0.251) despite ICL leading on C2ST and MMD.
Limitations
- Requires expensive up-front pretraining on modern GPUs; large contexts are computationally costly.
- Performance degrades under larger distribution shift between synthetic training and test data, and metrics lose sensitivity in high dimensions.
- Experiments use relatively simple posteriors and low-dimensional latent variables compared to HMC/VI references.
Experimental setup
Datasets: synthetic datasets from GLM, factor analysis, and GMM generative processes 17 real-world tabular regression datasets from a benchmark suite
Benchmarks: 50 synthetic datasets and 17 real-world tabular regression datasets
Metrics: C2ST (random forest ROC-AUC) MMD with exponential kernel Wasserstein-2 distance
Baselines: Hamiltonian Monte Carlo (NUTS) Laplace approximation ADVI diagonal normal ADVI full-covariance normal ADVI structured normal AutoIAF
Builds on: TabPFN flow matching diffusion transformer prior-data fitted networks automatic differentiation variational inference
Closest papers by meaning
- In-Context Learning through the Bayesian Prism (ICLR 2024 · similarity 0.91)
- Technical Debt in In-Context Learning: Diminishing Efficiency in Long Context (NeurIPS 2025 · similarity 0.89)
- Transformers are Minimax Optimal Nonparametric In-Context Learners (NeurIPS 2024 · similarity 0.89)
- Nonlinear transformers can perform inference-time feature learning (ICML 2025 · similarity 0.89)
- On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures (ICML 2025 · similarity 0.89)
- In-Context Deep Learning via Transformer Models (ICML 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.