Reliable amortized variational inference with physics-based latent distribution correction

Bayesian inference for high-dimensional inverse problems is computationally costly and requires selecting a suitable prior distribution. Amortized variational inference addresses these challenges by pretraining a neural network that acts as a surro-gate conditional distribution that approximates the posterior dis-tribution not only for one instance of the observed data but also for the distribution of the data pertaining to a specific inverse problem. When fed previously unseen data, the neural network - in our case, a conditional normalizing flow - provides the posterior samples at virtually no cost. However, the accuracy of amortized variational inference relies on the availability of high-fidelity training data, which seldom exist in geophysical inverse problems because of the earth's heterogeneous subsurface. In addition, the network is prone to errors if evaluated over data that are not drawn from the training data distribution. As such, we have aimed to increase the resilience of amortized variational inference in the presence of moderate data distribution shifts. We achieve this via a correction to the conditional normalizing flow's latent distribution that improves the approximation to the posterior distribution for the data at hand. The correction in-volves relaxing the standard Gaussian assumption on the latent distribution and parameterizing it via a Gaussian distribution with an unknown mean and (diagonal) covariance. These un-knowns are then estimated by minimizing the Kullback-Leibler divergence between the corrected and the (physics-based) true posterior distributions. Although generic and applicable to other inverse problems by means of a linearized seismic imaging ex-ample, we find 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, given noisy seismic data simulated via the linearized Born modeling, provides a seismic image with limited artifacts and an assessment of its uncertainty at approx-imately the same cost as five reverse time migrations.