# Distribution Transformers: Fast Approximate Bayesian Inference With On-The-Fly Prior Adaptation

arXiv: https://arxiv.org/abs/2502.02463 · HTML: https://mlsanity.com/paper/2502.02463/

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 76 references that are in the map.
- First posted on arXiv 4 Feb 2025.

## Builds on (most cited first, 30 of 36)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Variational Inference: A Review for Statisticians](https://mlsanity.com/paper/1601.00670/) (2016)
- [Aleatoric and Epistemic Uncertainty in Machine Learning: An Introduction to Concepts and Methods](https://mlsanity.com/paper/1910.09457/) (2019)
- [Masked Autoregressive Flow for Density Estimation](https://mlsanity.com/paper/1705.07057/) (2017)
- [A Survey of Uncertainty in Deep Neural Networks](https://mlsanity.com/paper/2107.03342/) (2021)
- [TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second](https://mlsanity.com/paper/2207.01848/) (2022)
- [The frontier of simulation-based inference](https://mlsanity.com/paper/1911.01429/) (2019)
- [Conditional Neural Processes](https://mlsanity.com/paper/1807.01613/) (2018)
- [Are Transformers universal approximators of sequence-to-sequence functions?](https://mlsanity.com/paper/1912.10077/) (2019)
- [Transformers Can Do Bayesian Inference](https://mlsanity.com/paper/2112.10510/) (2021)
- [Black Box Variational Inference](https://mlsanity.com/paper/1401.0118/) (2013)
- [Neural Processes](https://mlsanity.com/paper/1807.01622/) (2018)
- [Attentive Neural Processes](https://mlsanity.com/paper/1901.05761/) (2019)
- [Sequential Neural Likelihood: Fast Likelihood-free Inference with Autoregressive Flows](https://mlsanity.com/paper/1805.07226/) (2018)
- [Automatic Posterior Transformation for Likelihood-Free Inference](https://mlsanity.com/paper/1905.07488/) (2019)
- [Expectation Propagation for approximate Bayesian inference](https://mlsanity.com/paper/1301.2294/) (2013)
- [Flexible statistical inference for mechanistic models of neural dynamics](https://mlsanity.com/paper/1711.01861/) (2017)
- [BayesFlow: Learning complex stochastic models with invertible neural networks](https://mlsanity.com/paper/2003.06281/) (2020)
- [Transformer Neural Processes: Uncertainty-Aware Meta Learning Via Sequence Modeling](https://mlsanity.com/paper/2207.04179/) (2022)
- [Flow Matching for Scalable Simulation-Based Inference](https://mlsanity.com/paper/2305.17161/) (2023)
- [A survey on high-dimensional Gaussian process modeling with application to Bayesian optimization](https://mlsanity.com/paper/2111.05040/) (2021)
- [All-in-one simulation-based inference](https://mlsanity.com/paper/2404.09636/) (2024)
- [Sequential Neural Score Estimation: Likelihood-Free Inference with Conditional Score Based Diffusion Models](https://mlsanity.com/paper/2210.04872/) (2022)
- [Can Transformers Learn Full Bayesian Inference in Context?](https://mlsanity.com/paper/2501.16825/) (2025)
- [Amortized Probabilistic Conditioning for Optimization, Simulation and Inference](https://mlsanity.com/paper/2410.15320/) (2024)

## Built on by (most cited first, 13 of 13)

- [Effortless, Simulation-Efficient Bayesian Inference using Tabular Foundation Models](https://mlsanity.com/paper/2504.17660/) (2025)
- [ALINE: Joint Amortization for Bayesian Inference and Active Data Acquisition](https://mlsanity.com/paper/2506.07259/) (2025)
- [Efficient Autoregressive Inference for Transformer Probabilistic Models](https://mlsanity.com/paper/2510.09477/) (2025)
- [One Operator for Many Densities: Amortized Approximation of Conditioning by Neural Operators](https://mlsanity.com/paper/2605.06873/) (2026)
- [End-to-End Analysis of Charge Stability Diagrams with Transformers](https://mlsanity.com/paper/2508.15710/) (2025)
- [PriorGuide: Test-Time Prior Adaptation for Simulation-Based Inference](https://mlsanity.com/paper/2510.13763/) (2025)
- [metabeta -- A fast neural model for Bayesian mixed-effects regression](https://mlsanity.com/paper/2510.07473/) (2025)
- [α-PFN: Fast Entropy Search via In-Context Learning](https://mlsanity.com/paper/2606.07134/) (2026)
- [Transformers can do Bayesian Clustering](https://mlsanity.com/paper/2510.24318/) (2025)
- [Multi-Task Bayesian In-Context Learning](https://mlsanity.com/paper/2606.20538/) (2026)
- [Guidance for Prior Change via Density Ratio Estimation](https://mlsanity.com/paper/2608.21729/) (2026)
- [Amortized Factor Inference Networks for Posterior Inference](https://mlsanity.com/paper/2605.26419/) (2026)
- [Spectra: Exact Component Transport for Test-Time Prior Adaptation in Simulation-Based Inference](https://mlsanity.com/paper/2610.08021/) (2026)

## Abstract (arXiv)

While Bayesian inference provides a principled framework for reasoning under uncertainty, its widespread adoption is limited by the intractability of exact posterior computation, necessitating the use of approximate inference. However, existing methods are often computationally expensive, or demand costly retraining when priors change, limiting their utility, particularly in sequential inference problems such as real-time sensor fusion. To address these challenges, we introduce the Distribution Transformer -- a novel architecture that can learn arbitrary distribution-to-distribution mappings. Our method can be trained to map a prior to the corresponding posterior, conditioned on some dataset -- thus performing approximate Bayesian inference. Our novel architecture represents a prior distribution as a (universally-approximating) Gaussian Mixture Model (GMM), and transforms it into a GMM representation of the posterior. The components of the GMM attend to each other via self-attention, and to the datapoints via cross-attention. We demonstrate that Distribution Transformers both maintain flexibility to vary the prior, and significantly reduces computation times-from minutes to milliseconds-while achieving log-likelihood performance on par with or superior to existing approximate inference methods across tasks such as sequential inference, quantum system parameter inference, and Gaussian Process predictive posterior inference with hyperpriors.