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

TL;DR

The paper introduces the Distribution Transformer, a neural architecture that performs fast, approximate Bayesian inference by transforming prior distributions into posterior distributions using attention mechanisms, suitable for real-time applications.

Contribution

It presents a novel architecture that learns to map priors to posteriors efficiently, reducing computation time significantly while maintaining high inference quality.

Findings

Reduces inference time from minutes to milliseconds.

Achieves comparable or superior log-likelihood performance to existing methods.

Demonstrates effectiveness across diverse tasks like sequential inference and quantum parameter estimation.

Abstract

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.