[论文解读] 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 Techniques被引用 4

一句话总结

该论文提出了一种基于物理信息的潜在分布校正方法，应用于变分推断的摊销化框架中，利用条件归一化流提升高维反问题中数据分布偏移下的鲁棒性。通过学习一个具有可训练均值μ和对角协方差diag(s)²的灵活高斯潜在分布，该方法在噪声、源数量或先验分布发生变化时显著降低了后验近似误差，实现在接近实时成本下的精确不确定性量化。

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.

研究动机与目标

• 提升摊销变分推断在地球物理反问题中常见数据分布偏移下的鲁棒性。
• 解决标准摊销推断在缺乏高保真训练数据时泛化到分布外数据的局限性。
• 通过放松固定标准高斯假设，校正条件归一化流的潜在分布。
• 在计算开销极小的前提下，实现在地震成像中可靠的后验近似与不确定性量化。
• 将基于物理的约束整合到潜在空间，提升后验保真度，而无需重新训练主推理网络。

提出的方法

• 将潜在分布参数化为具有可学习均值μ和对角协方差diag(s)²的高斯分布，替代标准的N(0, I)假设。
• 通过最小化校正后潜在分布与真实后验pϕ(z|yobs)之间的KL散度，优化μ和s，以处理分布外数据。
• 利用重参数化技巧将样本表示为s⊙z + μ，其中z ∼ N(0, I)，从而实现可微优化。
• 推导出包含数据不拟合项、潜在先验项和协方差矩阵的对数行列式项的可计算目标函数，以实现高效训练。
• 利用基于物理的正向建模（线性化Born散射）计算潜在空间中真实后验密度，用于校正。
• 使用随机梯度下降端到端训练校正网络，同时保持摊销推理框架的效率。

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

实验结果

研究问题

• RQ1 基于物理信息的潜在分布校正是否能提升摊销变分推断在数据分布偏移下的可靠性？
• RQ2 放松潜在空间上的标准高斯先验对反问题中后验近似精度有何影响？
• RQ3 该方法在提升对噪声方差和源数量变化的鲁棒性的同时，能在多大程度上保持计算效率？
• RQ4 该方法是否能在仅增加极少额外计算成本的前提下，实现地震成像中精确的不确定性量化？
• RQ5 当先验分布相对于训练数据分布发生偏移时，该校正方法表现如何？

主要发现

• 所提出的潜在分布校正方法在数据分布偏移下显著提升了后验近似精度，包括噪声方差和地震源数量的变化。
• 仅需相当于五次逆时偏移的计算成本，即可生成具有有限伪影的不确定性量化地震成像结果。
• 在分布偏移条件下，与标准摊销推断相比，该方法将后验分布与真实后验之间的分歧降低了高达40%。
• 学习到的潜在分布能有效适应分布外数据，而无需重新训练主条件归一化流。
• 该方法在保持摊销推断计算效率的同时，实现了高维反问题中可靠的不确定性量化。
• 在线性化地震成像的实验结果表明，该方法在多种数据偏移下均表现出鲁棒性，验证了其在其他反问题中的泛化能力。

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