ASPIRE: Iterative Amortized Posterior Inference for Bayesian Inverse Problems

TL;DR

This work tackles Bayesian inverse problems in which uncertainty quantification is essential yet computationally challenging. It introduces ASPIRE, an approach that iteratively refines amortized posterior inferences by coupling physics-based score summaries with conditional normalizing flows, enabling offline training and fast online evaluation. Through a stylized linear-Gaussian example and a high-dimensional transcranial ultrasound tomography application, ASPIRE demonstrates improved posterior fidelity and calibrated uncertainty with substantially lower online computational costs than non-amortized methods. The results indicate meaningful practicality for real-time or near-real-time imaging scenarios where forward-model evaluations are expensive. Overall, ASPIRE provides a principled, scalable middle ground between amortized and non-amortized VI, with broad implications for Bayesian inference in high-dimensional, physics-driven inverse problems.

Abstract

Due to their uncertainty quantification, Bayesian solutions to inverse problems are the framework of choice in applications that are risk averse. These benefits come at the cost of computations that are in general, intractable. New advances in machine learning and variational inference (VI) have lowered the computational barrier by learning from examples. Two VI paradigms have emerged that represent different tradeoffs: amortized and non-amortized. Amortized VI can produce fast results but due to generalizing to many observed datasets it produces suboptimal inference results. Non-amortized VI is slower at inference but finds better posterior approximations since it is specialized towards a single observed dataset. Current amortized VI techniques run into a sub-optimality wall that can not be improved without more expressive neural networks or extra training data. We present a solution that enables iterative improvement of amortized posteriors that uses the same networks architectures and training data. The benefits of our method requires extra computations but these remain frugal since they are based on physics-hybrid methods and summary statistics. Importantly, these computations remain mostly offline thus our method maintains cheap and reusable online evaluation while bridging the approximation gap these two paradigms. We denote our proposed method ASPIRE - Amortized posteriors with Summaries that are Physics-based and Iteratively REfined. We first validate our method on a stylized problem with a known posterior then demonstrate its practical use on a high-dimensional and nonlinear transcranial medical imaging problem with ultrasound. Compared with the baseline and previous methods from the literature our method stands out as an computationally efficient and high-fidelity method for posterior inference.

Figures (51)

• Figure 1: Our algorithm ASPIRE is a middle ground between amortized and non-amortized variational inference.
• Figure 2: Offline training phase: During training, our proposed algorithm refines the fiducial point used to calculate the summary statistics. Importantly, this refinement is done in an amortized manner over a training dataset.
• Figure 3: Online inference phase: our algorithm refines the fiducial used to calculate the summary statistics. During inference, it maintains the ability to amortize to many observations and has low cost since $J$ is typically a low number (3-4).
• Figure 4: The quality of the proposed amortized posterior approximation improves at each iteration as measured by the estimated posterior mean with respect to the analytically known ground truth posterior mean.
• Figure 5: Comparison of full covariance matrix from our method as compared to the analytical ground truth posterior covariance. After three iterations of our method, the estimated posterior covariance is close to the ground truth covariance.

Theorems & Definitions (1)

• Lemma 3.1