[論文レビュー] 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.

研究の動機と目的

• 地球物理学的逆問題に一般的に見られるデータ分布シフトに対して、アモアタイズド変分推論のロバスト性を向上させること。
• 高精度なトレーニングデータが不足する状況において、標準的なアモアタイズド推論が分布外データに一般化できないという限界を克服すること。
• 固定された標準ガウス分布仮定を緩和することで、条件付きノーマライジングフローの潜在分布を補正すること。
• 最小限の計算コスト増加で、地震画像処理における信頼性の高い後方分布近似と不確実性評価を可能にすること。
• 物理的制約を潜在空間に統合することで、主な推論ネットワークの再トレーニングなしに後方分布の忠実度を向上させること。

提案手法

• 標準的なN(0, I)仮定に代えて、学習可能な平均μと対角共分散diag(s)²を用いて潜在分布をガウス分布としてパラメータライズする。
• 分布外データにおける真の後方分布pϕ(z|yobs)とのKLダイバージェンスを最小化することで、μとsを最適化する。
• 再パラメータライゼーショントリックを用い、サンプルをs⊙z + μ（z ∼ N(0, I)）として表現することで、微分可能な最適化を可能にする。
• データ不一致、潜在事前分布、および共分散行列の行列式の対数を含む、効率的なトレーニングに適した解析的最適化関数を導出する。
• 線形化されたボーン散乱に基づく物理的前向きモデリングを活用し、補正のための潜在空間内での真の後方分布密度を計算する。
• 確率的勾配降下法を用いて、アモアタイズド推論フレームワークの効率を保ちつつ、補正ネットワークをエンドツーエンドでトレーニングする。

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 トレーニングデータの分布とは異なる事前分布にシフトした場合、補正手法の性能はどの程度維持されるか？

主な発見

• 提案手法による潜在分布補正は、ノイズ分散の変化や地震源数の変化を含む、さまざまなデータ分布シフト下でも、後方分布近似精度を顕著に向上させた。
• わずか5回分のリバースタイムミグレーションに相当する計算コストで、アーチファクトが少ない不確実性を伴う地震画像を生成できた。
• 分布シフト下では、標準的なアモアタイズド推論に比べて、後方分布の真の分布からの乖離を最大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