# Energy-based Transport for Amortized Bayesian Inference

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 12 of its 32 references that are in the map.
- First posted on arXiv 14 May 2026.

## Builds on (most cited first, 12 of 12)

- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Sequential Neural Likelihood: Fast Likelihood-free Inference with Autoregressive Flows](https://mlsanity.com/paper/1805.07226/) (2018)
- [BayesFlow: Learning complex stochastic models with invertible neural networks](https://mlsanity.com/paper/2003.06281/) (2020)
- [Flow Matching for Scalable Simulation-Based Inference](https://mlsanity.com/paper/2305.17161/) (2023)
- [Optimal experimental design: Formulations and computations](https://mlsanity.com/paper/2407.16212/) (2024)
- [Conditional Sampling with Monotone GANs: from Generative Models to Likelihood-Free Inference](https://mlsanity.com/paper/2006.06755/) (2020)
- [On the representation and learning of monotone triangular transport maps](https://mlsanity.com/paper/2009.10303/) (2020)
- [An Optimal Transport Formulation of Bayes' Law for Nonlinear Filtering Algorithms](https://mlsanity.com/paper/2203.11869/) (2022)
- [Conditional simulation via entropic optimal transport: Toward non-parametric estimation of conditional Brenier maps](https://mlsanity.com/paper/2411.07154/) (2024)
- [Machine Learning for Inverse Problems and Data Assimilation](https://mlsanity.com/paper/2410.10523/) (2024)
- [LazyDINO: Fast, scalable, and efficiently amortized Bayesian inversion via structure-exploiting and surrogate-driven measure transport](https://mlsanity.com/paper/2411.12726/) (2024)
- [Learning to solve Bayesian inverse problems: An amortized variational inference approach using Gaussian and Flow guides](https://mlsanity.com/paper/2305.20004/) (2023)

## Abstract (arXiv)

We consider amortized Bayesian inference for nonlinear inverse problems using only samples from the joint distribution of parameters and observations, including problems with unknown functions in a Banach space. Classical methods such as Markov chain Monte Carlo solve a new inference problem for each observation, making repeated posterior inference computationally prohibitive, particularly in infinite dimensions. Amortized Bayesian inversion instead learns a reusable map that rapidly generates posterior samples for new observations. We learn an observation-dependent transport map that pushes a reference measure to an approximate posterior. Training minimizes the average energy distance between the posterior and the learned pushforward. Averaging over observations allows generalization across observation instances and efficient amortized inference. Furthermore, the formulation is likelihood-free, requiring only samples from the joint distribution and avoiding likelihood evaluation. In addition, the use of an energy-distance objective removes the need for invertibility of the transport map and for computation of Jacobian determinants, enabling flexible parameterizations in high- and infinite-dimensional settings. Moreover, when the posterior has a density with respect to a Gaussian prior measure, we construct transport maps as the identity plus a learnable map valued in the prior's Cameron--Martin space. This guarantees that the learned posterior remains absolutely continuous with respect to the prior. In infinite dimensions, the transport map is parameterized using neural operators, enabling use at different grid resolutions. We demonstrate the approach on a finite-dimensional problem and PDE-based porous-medium flow and seismic inverse problems. The learned transport captures multimodality and dominant posterior modes while enabling fast sampling.