Energy-based Transport for Amortized Bayesian Inference

Abstract

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.