Amortized Energy-Based Bayesian Inference

TL;DR

This paper introduces a likelihood-free, transport-based amortized Bayesian inference method that efficiently approximates posterior distributions in inverse problems using neural operators and energy-distance minimization.

Contribution

It proposes a novel transport-based approach for amortized Bayesian inference that is likelihood-free, scalable to infinite dimensions, and capable of capturing complex posterior structures.

Findings

The method accurately captures multimodal and complex posterior distributions.

It enables fast sampling for new observations after training.

Demonstrated effectiveness on PDE-constrained inverse problems.

Abstract

We consider amortized Bayesian inference for nonlinear inverse problems in settings where only samples from the joint distribution of parameters and observations are available. Classical methods such as Markov chain Monte Carlo require solving a new inference problem for each observation, which can be computationally prohibitive when inference must be repeated many times. We propose a transport-based approach that learns an observation-dependent map pushing forward a reference measure to approximate the posterior distribution. The map is trained by minimizing an averaged energy-distance objective between the true posterior and the learned pushforward. This formulation is likelihood-free, requiring only joint samples, and avoids density evaluation, invertibility constraints, and Jacobian determinant computations. For function-space inverse problems with Gaussian priors, we parameterize the transport map as the identity plus a perturbation in the Cameron-Martin space of the prior, preserving absolute continuity with respect to the prior. In infinite-dimensional settings, the map is represented using neural operators. We illustrate the method on a finite-dimensional nonlinear inverse problem and two PDE-constrained inverse problems arising in porous medium flow and seismic inversion. The results show that the learned transport captures posterior structure, including multimodality and dominant modes, while enabling fast posterior sampling for new observations.