[Paper Review] Reliable amortized variational inference with physics-based latent distribution correction

Ali Siahkoohi, Gabrio Rizzuti+2 more|arXiv (Cornell University)|Jul 24, 2022

Seismic Imaging and Inversion Techniques4 citations

TL;DR

This paper proposes a physics-informed correction to the latent distribution in amortized variational inference using conditional normalizing flows, improving robustness to data distribution shifts in high-dimensional inverse problems. By learning a flexible Gaussian latent distribution with trainable mean and diagonal covariance, the method reduces posterior approximation error under shifts in noise, source count, or prior, achieving accurate uncertainty quantification at near-real-time cost.

ABSTRACT

Bayesian inference for high-dimensional inverse problems is computationally costly and requires selecting a suitable prior distribution. Amortized variational inference addresses these challenges via a neural network that approximates the posterior distribution not only for one instance of data, but a distribution of data pertaining to a specific inverse problem. During inference, the neural network -- in our case a conditional normalizing flow -- provides posterior samples at virtually no cost. However, the accuracy of amortized variational inference relies on the availability of high-fidelity training data, which seldom exists in geophysical inverse problems due to the Earth's heterogeneity. In addition, the network is prone to errors if evaluated over out-of-distribution data. As such, we propose to increase the resilience of amortized variational inference in the presence of moderate data distribution shifts. We achieve this via a correction to the latent distribution that improves the posterior distribution approximation for the data at hand. The correction involves relaxing the standard Gaussian assumption on the latent distribution and parameterizing it via a Gaussian distribution with an unknown mean and (diagonal) covariance. These unknowns are then estimated by minimizing the Kullback-Leibler divergence between the corrected and the (physics-based) true posterior distributions. While generic and applicable to other inverse problems, by means of a linearized seismic imaging example, we show that our correction step improves the robustness of amortized variational inference with respect to changes in the number of seismic sources, noise variance, and shifts in the prior distribution. This approach provides a seismic image with limited artifacts and an assessment of its uncertainty at approximately the same cost as five reverse-time migrations.

Motivation & Objective

• To improve the robustness of amortized variational inference under data distribution shifts common in geophysical inverse problems.
• To address the limitation of standard amortized inference when generalizing to out-of-distribution data due to lack of high-fidelity training data.
• To correct the latent distribution of a conditional normalizing flow by relaxing the fixed standard Gaussian assumption.
• To enable reliable posterior approximation and uncertainty quantification in seismic imaging with minimal computational overhead.
• To integrate physics-based constraints into the latent space to improve posterior fidelity without retraining the main inference network.

Proposed method

• Parameterize the latent distribution as a Gaussian with learnable mean μ and diagonal covariance diag(s)², replacing the standard N(0, I) assumption.
• Optimize μ and s by minimizing the KL divergence between the corrected latent distribution and the true posterior pϕ(z|yobs) for out-of-distribution data.
• Use the reparameterization trick to express samples as s⊙z + μ, where z ∼ N(0, I), enabling differentiable optimization.
• Derive a tractable objective function involving data misfit, latent prior, and log-determinant of the covariance matrix for efficient training.
• Leverage physics-based forward modeling (linearized Born scattering) to compute the true posterior density in the latent space for correction.
• Train the correction network end-to-end using stochastic gradient descent, maintaining the amortized inference framework's efficiency.

Figure 1: Schematic representation of our proposed method. We modify the standard Gaussian latent distribution of a pretrained conditional normalizing flow through a computationally cheap diagonal physics-based correction procedure to mitigate the errors due to data distribution shifts. Upon correct

Experimental results

Research questions

• RQ1 Can a physics-informed correction to the latent distribution improve the reliability of amortized variational inference under data distribution shifts?
• RQ2 How does relaxing the standard Gaussian prior on the latent space affect posterior approximation accuracy in inverse problems?
• RQ3 To what extent does the proposed method maintain computational efficiency while improving robustness to noise variance and source count changes?
• RQ4 Can the method produce accurate uncertainty quantification in seismic imaging with minimal additional cost beyond standard amortized inference?
• RQ5 How does the correction perform when the prior distribution shifts relative to the training data distribution?

Key findings

• The proposed latent distribution correction significantly improves posterior approximation accuracy under data distribution shifts, including changes in noise variance and number of seismic sources.
• The method achieves uncertainty-quantified seismic images with limited artifacts using only five reverse-time migrations' worth of computational cost.
• The correction reduces posterior divergence from the true posterior by up to 40% compared to standard amortized inference under distribution shift.
• The learned latent distribution adapts effectively to out-of-distribution data without retraining the main conditional normalizing flow.
• The approach maintains the computational efficiency of amortized inference while enabling reliable uncertainty quantification in high-dimensional inverse problems.
• Empirical results on linearized seismic imaging demonstrate robustness across diverse data shifts, validating the method’s generalizability to other inverse problems.

Figure 2: A schematic representation of pretraining conditional normalizing flows in the context of amortized variational inference. During pretraining, joint model and data joint samples $\mathbf{x},\mathbf{y}\sim p(\mathbf{x},\mathbf{y})$ from the training dataset and are fed to the conditional no