math.NAMay 14, 2026

Amortized Energy-Based Bayesian Inference

Authors: Hojjat KavehRicardo BaptistaAndrew M. Stuart

Organizations: California Institute of Technology · University of Toronto

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.

Explore similar work

Jun 14, 2026stat.CO

Amortized mean-shift interacting particles

Bayesian inference for inverse problems is run to evaluate integrals -- posterior expectations, tail probabilities, and risks -- across a stream of observations. The standard estimate averages the integrand over posterior samples, a Monte-Carlo average whose error decays only as the square root of the sample size, so accuracy demands many samples -- prohibitive when each one calls a partial-differential-equation forward model. Mean-shift interacting particles need far fewer: they return a small set of signed-weight nodes -- a deterministic quadrature whose weighted averages estimate those integrals. Finding the nodes, however, is a per-observation optimization that, in its most accurate form, reads the posterior score at every step -- returning the cost it meant to save. We introduce amortized mean-shift interacting particles, a learned map that emits the weighted nodes from an observation and a few posterior samples in a single forward pass. Training asks only for joint parameter-observation samples and a posterior to draw from -- a conditional normalizing flow, an empirical conditional, or any reference the user can sample -- and the map learns to integrate that posterior from samples alone, evaluating neither its density nor its score. Once trained, it generalizes to unseen observations and integrands at any node budget and improves on independent samples in two ways: by reweighting them, provably no worse than the equal weights of Monte-Carlo; and by moving them, which empirically lowers it further. Across closed-form, sampled, learned, and physics-based posteriors -- up to a thousand-coefficient groundwater field -- it integrates more accurately than the same number of samples at every budget, and a posterior-whitened, dimension-aware kernel removes the high-dimensional wall. The result is a Pareto improvement on Monte-Carlo integration, not a competitor to drawing more samples.
Ali Siahkoohi
Sep 21, 2026stat.ML

Variational objectives for amortized Bayesian inference in inverse problems: The role of posterior conditioning

Variational autoencoders (VAEs) offer an efficient approach to amortized Bayesian inference for inverse problems, but posterior accuracy can depend strongly on the choice of variational regularization, particularly when the inverse problem contains weakly identified parameter directions. This study investigates three objectives: a reverse Kullback--Leibler formulation (VAE-KL), an asymmetric Jensen--Shannon formulation (VAE-JS), and a Jensen--Shannon--Wasserstein formulation (VAE-JSWA), which replaces the reverse Kullback--Leibler regularizer with the squared 2-Wasserstein distance while retaining forward-Kullback--Leibler posterior supervision. A full-covariance Gaussian encoder and a pre-trained physics-based surrogate are used for amortized posterior inference. A local linear--Gaussian analysis in the generalized Fisher basis is developed to characterize the variance-dependent gradients of the three objectives. The formulations are first evaluated using linear--Gaussian benchmarks with known posterior solutions and subsequently tested on nonlinear physics-based inverse problems, including an inverse problem governed by a linear ODE and two PDE-constrained problems. VAE-KL performs slightly better than the other formulations in the well-conditioned benchmark, where all three approaches yield comparable posterior approximations, whereas VAE-JSWA provides substantially lower posterior errors in the strongly ill-conditioned benchmark. The nonlinear physics-based problems exhibit a similar conditioning-dependent trend, with JS-based formulations providing greater benefit as posterior ill-conditioning increases. These results indicate that posterior conditioning is an important factor in selecting variational objectives and motivate geometry-adaptive variational inference for Bayesian inverse problems.
Abhishek Srivastava, Arijit Hazra, Rajesh Dubbaku
Feb 11, 2026stat.ML

A Gibbs posterior sampler for inverse problem based on prior diffusion model

This paper addresses the issue of inversion in cases where (1) the observation system is modeled by a linear transformation and additive error, (2) the problem is ill-posed and regularization relies on a Bayesian strategy, (3)~the prior is modeled by a diffusion process adjusted on an available large set of examples. In this context, it is known that the issue of posterior sampling is a thorny one and the paper introduces a Gibbs algorithm. It appears that this avenue has not been explored, and we show that it is particularly effective and remarkably simple. In addition, it provides clear elements regarding convergence guarantees in a specific case and arguments supporting such guarantees in practical cases. The results are clearly confirmed by numerical simulations based on a toy example.
Jean-François Giovannelli