Amortized Energy-Based Bayesian Inference

2026-05-14

AI Summary

This work addresses the computational burden of traditional Bayesian inference methods, such as MCMC, in nonlinear inverse problems where only joint samples of parameters and observations are available, particularly when repeated inference is required. The authors propose an amortized Bayesian inference framework based on transport maps, employing a neural operator to parameterize observation-dependent mappings that push a reference measure to an approximate posterior. The model is trained by minimizing the mean energy distance, eliminating the need for likelihood evaluations, density computations, invertibility constraints, or Jacobian calculations. The approach preserves absolute continuity with respect to Gaussian priors in function space, making it suitable for infinite-dimensional settings. Numerical experiments demonstrate its ability to efficiently capture multimodal structures and dominant modes, significantly accelerating posterior sampling for new observations, with validated performance on PDE-constrained inverse problems such as subsurface flow and seismic inversion.

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.

Problem

• amortized inference
• Bayesian inference
• nonlinear inverse problems
• likelihood-free
• posterior sampling

Innovation

• amortized inference
• transport map
• energy distance
• likelihood-free
• neural operators