# Reliable amortized variational inference with physics-based latent distribution correction

arXiv: https://arxiv.org/abs/2207.11640 · HTML: https://mlsanity.com/paper/2207.11640/

## Key facts

- Cited by 25 papers in the map (data through 14 Oct 2026).
- Builds on 39 of its 95 references that are in the map.
- First posted on arXiv 24 Jul 2022.

## Builds on (most cited first, 30 of 39)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [Variational Inference: A Review for Statisticians](https://mlsanity.com/paper/1601.00670/) (2016)
- [How transferable are features in deep neural networks?](https://mlsanity.com/paper/1411.1792/) (2014)
- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [The frontier of simulation-based inference](https://mlsanity.com/paper/1911.01429/) (2019)
- [Compressed Sensing using Generative Models](https://mlsanity.com/paper/1703.03208/) (2017)
- [BayesFlow: Learning complex stochastic models with invertible neural networks](https://mlsanity.com/paper/2003.06281/) (2020)
- [Preconditioned Stochastic Gradient Langevin Dynamics for Deep Neural Networks](https://mlsanity.com/paper/1512.07666/) (2015)
- [Invertible generative models for inverse problems: mitigating representation error and dataset bias](https://mlsanity.com/paper/1905.11672/) (2019)
- [Coupling-based Invertible Neural Networks Are Universal Diffeomorphism Approximators](https://mlsanity.com/paper/2006.11469/) (2020)
- [Semi-Amortized Variational Autoencoders](https://mlsanity.com/paper/1802.02550/) (2018)
- [Simulation Intelligence: Towards a New Generation of Scientific Methods](https://mlsanity.com/paper/2112.03235/) (2021)
- [Task Agnostic Continual Learning Using Online Variational Bayes](https://mlsanity.com/paper/1803.10123/) (2018)
- [Deep Bayesian Inversion](https://mlsanity.com/paper/1811.05910/) (2018)
- [Detecting Model Misspecification in Amortized Bayesian Inference with Neural Networks](https://mlsanity.com/paper/2112.08866/) (2021)
- [Composing Normalizing Flows for Inverse Problems](https://mlsanity.com/paper/2002.11743/) (2020)
- [Trumpets: Injective Flows for Inference and Inverse Problems](https://mlsanity.com/paper/2102.10461/) (2021)
- [HINT: Hierarchical Invertible Neural Transport for Density Estimation and Bayesian Inference](https://mlsanity.com/paper/1905.10687/) (2019)
- [Universal approximation property of invertible neural networks](https://mlsanity.com/paper/2204.07415/) (2022)
- [Unsupervised Learning of Full-Waveform Inversion: Connecting CNN and Partial Differential Equation in a Loop](https://mlsanity.com/paper/2110.07584/) (2021)
- [Preconditioned training of normalizing flows for variational inference in inverse problems](https://mlsanity.com/paper/2101.03709/) (2021)
- [Task adapted reconstruction for inverse problems](https://mlsanity.com/paper/1809.00948/) (2018)
- [Conditional Injective Flows for Bayesian Imaging](https://mlsanity.com/paper/2204.07664/) (2022)
- [Invertible Neural Networks versus MCMC for Posterior Reconstruction in Grazing Incidence X-Ray Fluorescence](https://mlsanity.com/paper/2102.03189/) (2021)
- [Deep Bayesian inference for seismic imaging with tasks](https://mlsanity.com/paper/2110.04825/) (2021)
- [Conditional Prior Networks for Optical Flow](https://mlsanity.com/paper/1807.10378/) (2018)

## Built on by (most cited first, 25 of 25)

- [Neural Methods for Amortized Inference](https://mlsanity.com/paper/2404.12484/) (2024)
- [Conditional score-based diffusion models for Bayesian inference in infinite dimensions](https://mlsanity.com/paper/2305.19147/) (2023)
- [BayesFlow: Amortized Bayesian Workflows With Neural Networks](https://mlsanity.com/paper/2306.16015/) (2023)
- [Conditional Optimal Transport on Function Spaces](https://mlsanity.com/paper/2311.05672/) (2023)
- [ASPIRE: Iterative Amortized Posterior Inference for Bayesian Inverse Problems](https://mlsanity.com/paper/2405.05398/) (2024)
- [Robust Amortized Bayesian Inference with Self-Consistency Losses on Unlabeled Data](https://mlsanity.com/paper/2501.13483/) (2025)
- [WISE: full-Waveform variational Inference via Subsurface Extensions](https://mlsanity.com/paper/2401.06230/) (2023)
- [Does Unsupervised Domain Adaptation Improve the Robustness of Amortized Bayesian Inference? A Systematic Evaluation](https://mlsanity.com/paper/2502.04949/) (2025)
- [Solving multiphysics-based inverse problems with learned surrogates and constraints](https://mlsanity.com/paper/2307.11099/) (2023)
- [BayesFlow 2: Multi-Backend Amortized Bayesian Inference in Python](https://mlsanity.com/paper/2602.07098/) (2026)
- [Misspecification-robust Sequential Neural Likelihood for Simulation-based Inference](https://mlsanity.com/paper/2301.13368/) (2023)
- [Amortized Normalizing Flows for Transcranial Ultrasound with Uncertainty Quantification](https://mlsanity.com/paper/2303.03478/) (2023)
- [Learned multiphysics inversion with differentiable programming and machine learning](https://mlsanity.com/paper/2304.05592/) (2023)
- [InvertibleNetworks.jl: A Julia package for scalable normalizing flows](https://mlsanity.com/paper/2312.13480/) (2023)
- [Inference of CO2 flow patterns -- a feasibility study](https://mlsanity.com/paper/2311.00290/) (2023)
- [Refining Amortized Posterior Approximations using Gradient-Based Summary Statistics](https://mlsanity.com/paper/2305.08733/) (2023)
- [Amortized Bayesian Workflow](https://mlsanity.com/paper/2409.04332/) (2024)
- [A Primer on Variational Inference for Physics-Informed Deep Generative Modelling](https://mlsanity.com/paper/2409.06560/) (2024)
- [Conditional neural control variates for variance reduction in Bayesian inverse problems](https://mlsanity.com/paper/2602.21357/) (2026)
- [Dual-space posterior sampling for Bayesian inference in constrained inverse problems](https://mlsanity.com/paper/2603.00393/) (2026)
- [On the role of memorization in learned priors for geophysical inverse problems](https://mlsanity.com/paper/2603.19629/) (2026)
- [Posterior sampling with CNN-based, Plug-and-Play regularization with applications to Post-Stack Seismic Inversion](https://mlsanity.com/paper/2212.14595/) (2022)
- [Amortized quadrature for posterior expectations in inverse problems](https://mlsanity.com/paper/2606.15871/) (2026)
- [Priors learned from legacy reconstructions inherit undetectable overconfidence](https://mlsanity.com/paper/2607.21721/) (2026)
- [MoRF-AST: Calibrated Probabilistic Virtual Sensing for Structural Monitoring under Changing Operating Conditions](https://mlsanity.com/paper/2608.24531/) (2026)

## Abstract (arXiv)

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.