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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Whittle, George, Ziomek, Juliusz, Rawling, Jacob, Osborne, Maike A.
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866911693040254976
author Whittle, George
Ziomek, Juliusz
Rawling, Jacob
Osborne, Maike A.
author_facet Whittle, George
Ziomek, Juliusz
Rawling, Jacob
Osborne, Maike A.
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2502_02463
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Distribution Transformers: Fast Approximate Bayesian Inference With On-The-Fly Prior Adaptation
Whittle, George
Ziomek, Juliusz
Rawling, Jacob
Osborne, Maike A.
Machine Learning
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.
title Distribution Transformers: Fast Approximate Bayesian Inference With On-The-Fly Prior Adaptation
topic Machine Learning
url https://arxiv.org/abs/2502.02463