---
title: "WISER: multimodal variational inference for full-waveform inversion without dimensionality reduction"
author:
  - name: Ziyi Yin
    orcid: "0000-0002-5024-8771"
    email: ziyi.yin@gatech.edu
    url: https://ziyiyin97.github.io
    affiliations:
      - name: Georgia Institute of Technology
  - name: Rafael Orozco
    orcid: "0000-0003-0917-2442"
    email: rorozco@gatech.edu
    affiliations:
      - name: Georgia Institute of Technology
  - name: Felix J. Herrmann
    orcid: "0000-0003-1180-2167"
    email: felix.herrmann@gatech.edu
    affiliations:
      - name: Georgia Institute of Technology
bibliography: WISER.bib
lightbox: True
format:
  html:
    link-external-newwindow: True
    number-sections: true
    toc: true
  pdf:
    number-sections: true
filters:
  - pseudocode
  - SLIM
execute:
  freeze: auto
abstract: |
  We present a semi-amortized variational inference framework designed for computationally feasible uncertainty quantification in 2D full-waveform inversion to explore the multimodal posterior distribution without dimensionality reduction. The framework is called **WISER**, short for full-**W**aveform variational **I**nference via **S**ubsurface **E**xtensions with **R**efinements. WISER leverages the power of generative artificial intelligence to perform approximate amortized inference that is low-cost albeit showing an amortization gap. This gap is closed through non-amortized refinements that make frugal use of acoustic wave physics. Case studies illustrate that WISER is capable of full-resolution, computationally feasible, and reliable uncertainty estimates of velocity models and imaged reflectivities.
---

## Introduction

Full-waveform inversion (FWI) aims to estimate unknown multi-dimensional ($\mathrm{D}\geq 2$) velocity models, denoted as $\mathbf{x}$, from noisy seismic data, $\mathbf{y}$, by inverting the nonlinear forward operator, $\mathcal{F}$, which relates $\mathbf{x}$ and $\mathbf{y}$ via $\mathbf{y} = \mathcal{F}(\mathbf{x})+\boldsymbol{\epsilon}$ with $\boldsymbol{\epsilon}$ measurement noise [@virieux2009]. FWI poses significant challenges, as it requires solving a high-dimensional, non-convex, and ill-posed inverse problem, with a computationally demanding forward operator in multiple dimensions. In addition, the inherent nonuniqueness of FWI results leads to multiple possible Earth models compatible with the observed data, underscoring the need for uncertainty quantification (UQ) to handle this multimodality.

The trade-off between accuracy and computational cost is a critical consideration in any high-dimensional inference routine with expensive forward operators [@cranmer2020frontier], including FWI. To circumvent the costs associated with global optimization, several approaches have attempted localized UQ [@fang2018uncertainty; @keating2021null; @izzatullah2023physics; @hoffmann2024local] based on the Laplace approximation. However, these approaches may not capture the full complexities of multimodal parameter spaces. In contrast, a Bayesian inference approach offers a costly but comprehensive resolution of the posterior distribution, $p(\mathbf{x}\mid\mathbf{y})$.

Bayesian inference algorithms are broadly categorized into two groups. The first, sampling-based methods, like Markov-chain Monte Carlo [McMC, @cowles1996markov], struggle with high-dimensional parameter spaces. To meet this challenge, they often rely on too restrictive low-dimensional parameterizations to reduce the number of sampling iterations [@fang2020deep; @liang2023uncertainty; @wei2023quantifying; @wei20243d; @wei2024reconstruction; @dhabaria2024hamiltonian], which could bias the inference results, rendering them impractical for multi-D UQ studies especially when solutions are nonunique.

The second category, optimization-based methods, like variational inference [VI, @zhang2021introduction], seek to approximate the posterior distribution using classes of known parameterized distributions. VI can be subdivided into amortized and non-amortized methods. Amortized VI involves a computationally intensive offline training phase, leveraging advances in generative artificial intelligence (genAI), particularly with models like conditional diffusion [@baldassari2023conditional] and conditional normalizing flows [CNFs, @winkler2019learning]. After training, amortized VI provides rapid sampling during inference [@siahkoohi2023; @orozco2023amortized], exemplified by the WISE framework [@yin2024wise] for multi-D FWI problems. However, these methods may suffer from an *amortization gap* --- implying that the amortized networks may only deliver suboptimal inference for a single observation at inference time, particularly when trained with limited examples or when there exists a discrepancy between training and inference [@marino2018iterative]. Conversely, non-amortized VI dedicates entire computational resources to the online inference [@zhao2022bayesian; @zhang20233; @zhang2023bayesian]. They result in more accurate inference, but the costly optimization has to be carried out repeatedly for new observations, and integrating realistic priors can be challenging since the prior needs to be embedded involving density evaluations [@kruse2021hint].

This paper introduces WISER as a semi-amortized VI framework to facilitate computationally feasible and reliable UQ for multi-D FWI without dimensionality reduction. Building on WISE, we train CNFs for efficient, suboptimal amortized inference, but then follow up with a crucial refinement step that only needs frugal use of the forward operator and its gradient. The refinement step aligns the posterior samples with the observation during inference, effectively bridging the amortization gap and enhancing inference accuracy.

Our contributions are organized as follows. We begin by outlining WISER in @alg-wiser. We explore the algorithm by explaining amortized VI with WISE, followed by computationally feasible multi-D physics-based refinement. The performance of WISER is demonstrated through realistic synthetic 2D case studies using the Compass model [@e.jones2012], showcasing improvements over WISE for both in- and out-of-distribution scenarios.

```pseudocode
#| label: alg-wiser
#| html-indent-size: "1.2em"
#| html-comment-delimiter: "//"
#| html-line-number: true
#| html-line-number-punc: ":"
#| html-no-end: false
#| pdf-placement: "htb!"
#| pdf-line-number: true

\begin{algorithm}
\caption{WISER: full-\textbf{W}aveform variational \textbf{I}nference via \textbf{S}ubsurface \textbf{E}xtensions with \textbf{R}efinements}
\begin{algorithmic}
\State {\Large\textbf{Offline training phase}}
\State
\State \textbf{Dataset generation}
\For{$i = 1:N$}
  \State $\mathbf{x}^{(i)} \sim p(\mathbf{x})$
  \State $\boldsymbol{\epsilon}^{(i)} \sim p(\boldsymbol{\epsilon})$
  \State $\mathbf{y}^{(i)} = \mathcal{F}(\mathbf{x}^{(i)})+\boldsymbol{\epsilon}^{(i)}$
  \State $\overline{\mathbf{y}}^{(i)} = \overline{\nabla\mathcal{F}}(\mathbf{x}_0)^{\top}\mathbf{y}^{(i)}$
\EndFor
\State
\State \textbf{Network training}
\State $\displaystyle \boldsymbol{\theta}^{\ast}=\underset{\boldsymbol{\theta}}{\operatorname{argmin}} \quad \frac{1}{N}\sum_{i=1}^{N} \left(\frac{1}{2}\|f_{\boldsymbol{\theta}}\left(\mathbf{x}^{(i)};\overline{\mathbf{y}}^{(i)}\right)\|_2^2-\log\left|\det\mathbf{J}_{f_{\boldsymbol{\theta}}}\right|\right)$
\State
\State {\Large\textbf{Online inference phase}}
\State
\State $\overline{\mathbf{y}}_{\mathrm{obs}} = \overline{\nabla\mathcal{F}}(\mathbf{x}_0)^{\top}\mathbf{y}_{\mathrm{obs}}$
\For{$i = 1:M$}
\State $\mathbf{z}_i \sim \mathrm{N}(\mathbf{0}, \mathbf{I})$
\State $\mathbf{x}_i = f_{\boldsymbol{\theta}^{\ast}}^{-1}\left(h_{\boldsymbol{\phi}}\left(\mathbf{z}_i\right);\overline{\mathbf{y}}_{\mathrm{obs}}\right)$
\EndFor
\State
\State \textbf{Physics-based refinements}
\For{$ii = 1:\mathrm{maxiter}_1$}
  \For{$i=1:M$}
        \State $\displaystyle\mathbf{g}_i=\nabla_{\mathbf{x}_i}\left[\frac{1}{2\sigma^2}\|\mathcal{F}(\mathbf{x}_i)-\mathbf{y}_{\mathrm{obs}}\|_2^2 + \frac{1}{2\gamma^2}\|\mathbf{x}_i-f_{\boldsymbol{\theta}^{\ast}}^{-1}\left(h_{\boldsymbol{\phi}}\left(\mathbf{z}_i\right);\overline{\mathbf{y}}_{\mathrm{obs}}\right)\|_2^2\right]$
        \State $\mathbf{x}_i = \mathbf{x}_i - \tau \mathbf{g}_i$
  \EndFor
  \For{$iii= 1:\mathrm{maxiter}_2$}
  \State $\displaystyle\mathcal{L}(\boldsymbol{\phi})=\sum_{i=1}^M \left[\frac{1}{2\gamma^2}\|\mathbf{x}_i-f_{\boldsymbol{\theta}^{\ast}}^{-1}\left(h_{\boldsymbol{\phi}}\left(\mathbf{z}_i\right);\overline{\mathbf{y}}_{\mathrm{obs}}\right)\|_2^2 + \frac{1}{2}\|h_{\boldsymbol{\phi}}\left(\mathbf{z}_i\right)\|_2^2 - \log\left|\det\mathbf{J}_{h_{\boldsymbol{\phi}}}\right|\right]$
  \State $\boldsymbol{\phi} \leftarrow \mathrm{ADAM}\left(\mathcal{L}(\boldsymbol{\phi})\right)$
  \EndFor
\EndFor
\State
\State {\Large\textbf{Output}}: $\{f_{\boldsymbol{\theta}^{\ast}}^{-1}\left(h_{\boldsymbol{\phi}}\left(\mathbf{z}_i\right);\overline{\mathbf{y}}_{\mathrm{obs}}\right)\}_{i=1}^M$ as samples of $p(\mathbf{x}|\mathbf{y}_{\mathrm{obs}})$
\end{algorithmic}
\end{algorithm}
```

## Amortized VI with WISE (lines 1---20)

WISER starts with an offline training phase that leverages conditional generative models to approximate the posterior distribution. This is achieved by WISE [@yin2024wise], which involves generating a training dataset (lines 3---9 of @alg-wiser) and training the CNFs (line 11---12).

### Dataset generation (lines 3---9)

We begin by drawing $N$ velocity models from the prior distribution, denoted by $p(\mathbf{x})$ (line 5). For each sample, $\mathbf{x}^{(i)}$, we simulate the observed data, $\mathbf{y}^{(i)}$, by performing the wave modeling and adding a random noise term (lines 6---7). Next, we compute common-image gathers [CIGs, @hou2016] for each observed data with an initial smooth 1D migration-velocity model, $\mathbf{x}_0$, which can be rather inaccurate. These CIGs, represented by $\overline{\mathbf{y}}^{(i)}$, are produced by applying the adjoint of the extended migration operator, $\overline{\nabla\mathcal{F}}(\mathbf{x}_0)^{\top}$, to the observed data. Using CIGs as the set of physics-informed summary statistics not only preserves information from the observed seismic data [@ten2023omnidirectional] but also enhances the training of CNFs in the next stage [@radev2020bayesflow; @orozco2023], as they help to decode the wave physics, translating prestack data to the image (subsurface-offset) domain.

### Network training (lines 11---12)

CNFs are trained with pairs of velocity models and CIGs via minimization of the objective in line 12. The symbol $f_{\boldsymbol{\theta}}$ denotes the CNFs, characterized by their network weights, $\boldsymbol{\theta}$, and the Jacobian, $\mathbf{J}_{f_{\boldsymbol{\theta}}}$. The term "normalizing" within CNFs implies their functionality to transform realizations of velocity models, $\mathbf{x}^{(i)}$, into Gaussian noise from a standard multivariate normal distribution (as defined by the $\ell_2$ norm), conditioned on the summary statistics (CIGs).

### Online inference (lines 14---20)

The aforementioned data generation and CNF training procedures conclude the offline training phase. During online inference, amortized VI is enabled by leveraging the inherent invertibility of CNFs. For a given observation, $\mathbf{y}_{\mathrm{obs}}$, the online cost is merely generation of a single set of CIGs (line 16). Subsequently, the posterior samples are generated by applying the inverse of the CNFs to Gaussian noise realizations, conditioned on these CIGs (lines 18---19)[^1].

[^1]: We slightly abuse the notation to assume $h_{\boldsymbol{\phi}}$ as an identity operator here.

<!--WISE provides low-cost amortized inference despite an *amortization gap* that hampers its accuracy. To bridge this amortization gap, we subsequently propose a refinement procedure that reintegrates the wave physics.-->

## Physics-based refinements (lines 22---32)

Consider a single observation, $\mathbf{y}_{\mathrm{obs}}$, and its corresponding posterior samples, $\mathbf{x}_i \sim p(\mathbf{x} \mid \overline{\mathbf{y}}_{\mathrm{obs}})$. The latent representations generated by the trained CNFs, ${\hat{\mathbf{z}}_i} = f_{\boldsymbol{\theta}^{\ast}}(\mathbf{x}_i; \overline{\mathbf{y}}_{\mathrm{obs}})$, may not conform exactly to the standard Gaussian distribution during inference. To address this issue, we follow @siahkoohi2023 to apply latent space corrections to fine-tune the trained CNFs. This involves integrating a shallower, yet invertible, network[^2], specifically trained to map realizations of true Gaussian noise to the corresponding latent codes, ${\hat{\mathbf{z}}_i}$. Adhering to a transfer learning approach, we maintain the weights of the trained CNFs while solely updating the weights of the shallower network by minimizing the following objective:

$$
\underset{\boldsymbol{\phi}}{\operatorname{minimize}} \quad \mathbb{E}_{\mathbf{z}\sim\mathrm{N}(\mathbf{0}, \mathbf{I})}\left[\frac{1}{2\sigma^2}\|\mathcal{F}\circ f_{\boldsymbol{\theta}^{\ast}}^{-1}\left(h_{\boldsymbol{\phi}}\left(\mathbf{z}\right);\overline{\mathbf{y}}_{\mathrm{obs}}\right)-\mathbf{y}_{\mathrm{obs}}\|_2^2  + \frac{1}{2}\|h_{\boldsymbol{\phi}}\left(\mathbf{z}\right)\|_2^2- \log\left|\det\mathbf{J}_{h_{\boldsymbol{\phi}}}\right|\right].
$$ {#eq-strong}

Here, the refinement network, $h_{\boldsymbol{\phi}}$, mitigates the amortization gap by adjusting the latent variable $\mathbf{z}$ before feeding it to the inverse of the trained CNFs, $f_{\boldsymbol{\theta}^\ast}^{-1}$. Intuitively, minimizing the first terms ties the posterior samples closer to the observed data. The second and third terms prevent the corrected latent space from being far from the Gaussian distribution, which implicitly takes advantage of the prior information existing in the amortized training phase.

[^2]: For linear inverse problems in seismic imaging, @siahkoohi2023 show that an elementwise scaling and shift mechanism is adequate to bridge the gap. However, given the complex, non-convex nature of FWI, we employ $h_{\boldsymbol{\phi}}$ as a generic invertible network.

@eq-strong offers a fine-tuning approach that leverages the full multi-D wave physics to refine the amortized VI framework for a single observation at inference phase. However, it introduces notable computational demands because it necessitates the coupling of the modeling operator and the networks. Specifically, every update to the network weights, $\boldsymbol{\phi}$, requires costly wave modeling operations. Given that network training typically involves numerous iterations, these computational demands can render it impractical for realistic FWI applications.

To relieve this computational burden, we adopt a strategy from @siahkoohi2020weak to reformulate @eq-strong into a weak form by allowing the network output to be only weakly enforced (in an $\ell_2$ sense) to be the corrected velocity models. The objective function for this weak formulation reads:

$$
 \underset{\mathbf{x}_{1:M}, \boldsymbol{\phi}}{\operatorname{minimize}} \quad \frac{1}{M}\sum_{i=1}^{M} \left[\frac{1}{2\sigma^2}\|\mathcal{F}(\mathbf{x}_i)-\mathbf{y}_{\mathrm{obs}}\|_2^2
+ \frac{1}{2\gamma^2}\|\mathbf{x}_i-f_{\boldsymbol{\theta}^{\ast}}^{-1}\left(h_{\boldsymbol{\phi}}\left(\mathbf{z}_i\right);\overline{\mathbf{y}}_{\mathrm{obs}}\right)\|_2^2 + \frac{1}{2}\|h_{\boldsymbol{\phi}}\left(\mathbf{z}_i\right)\|_2^2 - \log\left|\det\mathbf{J}_{h_{\boldsymbol{\phi}}}\right| \right].
$$ {#eq-weak}

We strategically decouple the computationally expensive forward operator, $\mathcal{F}$, from the more cheap-to-evaluate networks, $f_{\boldsymbol{\theta}^\ast}$ and $h_{\boldsymbol{\phi}}$. This is achieved in a penalty form with the assumption that the misfit between the network outputs and the posterior samples adheres to a Gaussian distribution, $\mathrm{N}(\mathbf{0}, \gamma^2\mathbf{I})$, where $\gamma$ is a hyperparameter dictating the trade-off between data misfit and regularization. Setting $\gamma$ to $0$ equates this weak formulation to the constrained formulation in @eq-strong. This weak formulation also supports optimization strategies for updating the velocity models with physical constraints [@esser2018total; @peters2019projection] and multiscale optimization techniques [@bunks1995multiscale].

WISER takes full computational advantage of this weak formulation by employing a nested loop structure. The outer loop is dedicated to updating $M$ velocity models, $\mathbf{x}_i$, through costly gradient descent steps (lines 24---27 of @alg-wiser), while the inner loop (lines 28---31) focuses on more updates (with the ADAM optimizer) to network weights, $\boldsymbol{\phi}$, without computationally expensive physics modeling. To achieve a balance, we conduct $\mathrm{maxiter}_2=128$ iterations in the inner loop. After refinements, WISER first evaluates the refined network on the latent variables to obtain refined latent codes. Subsequently, the amortized network uses the refined codes conditioned on the CIGs to compute the corrected posterior samples (line 34).

## Case studies

Evaluation of WISER is conducted through synthetic case studies utilizing 2D slices of the Compass model and 2D acoustic wave physics. The parameter of interest is discretized into $512 \times 256$ gridpoints with a spatial resolution of $12.5\,\mathrm{m}$, resulting in over $10^5$ degrees of freedom. The forward operator, $\mathcal{F}$, simulates acoustic data with absorbing boundaries. A Ricker wavelet with a central frequency of $15\,\mathrm{Hz}$ and an energy cut below $3\,\mathrm{Hz}$ is employed. We use $512$ sources towed at $12.5\,\mathrm{m}$ depth and $64$ ocean-bottom nodes (OBNs) located at jittered sampled horizontal positions [@hennenfent2008simply]. We employ source-receiver reciprocity during the modeling and sensitivity calculations. The observed data, $\mathbf{y}^{\mathrm{obs}}$, is perturbed with band-limited Gaussian noise to achieve a signal-to-noise ratio (S/N) of $12\,\mathrm{dB}$. The training of the CNFs uses $N=800$ pairs of velocity models and CIGs. To demonstrate WISER's efficacy in mitigating the amortization gap, we compare results from WISE and WISER under two scenarios during inference:

(i) observed shot data is generated using an in-distribution velocity model with the same forward operator;
(ii) observed shot data is produced by an out-of-distribution (OOD) velocity model and also a slightly altered forward operator.

### Case 1: in distribution

The ground-truth velocity model is an unseen 2D slice from the Compass model, shown in @fig-true. Following @alg-wiser, we initiate WISER by drawing $M=16$ Gaussian noise realizations to create the initial set of $16$ velocity models, depicted in @fig-wise-v-3d. To minimize computational demands, stochastic gradients [@herrmann2013frugal] are calculated in line 25 of @alg-wiser. Each particle's gradient is estimated using only $1$ randomly selected OBN gather from the observed data. We also add box constraints to the velocity models to restrict their range to $1.48$ to $4.44\,\mathrm{km/s}$. Following $\mathrm{maxiter}_1=80$ outer iterations—equivalent to $20$ data passes or $2560$ PDE solves[^3]—--we obtain the posterior samples from WISER in @fig-wiser-v-3d.

[^3]: 1 PDE solve means solving the wave equation for a single source. A gradient requires 2 PDE solves (forward and adjoint).

::: {#fig-wise-vs-wiser layout="[[50], [25, 25], [25, 25], [25, 25]]"}
![](./figs/true-v-colorbar.png){#fig-true width=100%}

![](./figs/WISE-v-3D.png){#fig-wise-v-3d width=100%}

![](./figs/WISER-v-3D.png){#fig-wiser-v-3d width=100%}

![](./figs/WISE-rtm-3D.png){#fig-wise-rtm width=100%}

![](./figs/WISER-rtm-3D.png){#fig-wiser-rtm width=100%}

![](./figs/WISE-rtm-zoom2.png){#fig-wise-rtm-zoom width=100%}

![](./figs/WISER-rtm-zoom2.png){#fig-wiser-rtm-zoom width=100%}

Comparison between WISE and WISER for an in-distribution case. *(a)* Unseen ground-truth velocity model. *(b)* Estimated velocity models from WISE. The conditional mean estimate (CM) is shown in the center. For posterior samples, horizontal traces at $Z = 2.7\,\mathrm{km}$ and vertical traces at $X = 3.6\,\mathrm{km}$ are displayed on the top and on the right, respectively. *(d)* Imaged reflectivity samples from WISE. *(f)* Zoom-in views of *(d)* overlaying on the CM of WISE. *(c)(e)(g)* are the counterparts from WISER, showcasing significant improvements.
:::

#### Observations

The conditional mean estimate (CM) from WISE lacks finer details, particularly beneath the unconformity at depths below $2.4\,\mathrm{km}$ (in red). This is attributed to the excessive variability in structural details of the posterior samples, visible on the right panel of @fig-wise-v-3d.

In contrast, WISER generates more consistent and accurate posterior samples. In @fig-wiser-v-3d, the right panel shows that the uncertainty from WISER is reduced below the unconformity. The upper panel illustrates that the uncertainty is more focused on the dipping events at the unconformity, highlighting areas of poor illumination.

#### Impact on imaging

To assess the impact of uncertainty in velocity models on downstream tasks, we conduct high-frequency imaging using a Ricker wavelet with central frequency of $30\,\mathrm{Hz}$ and compare the imaged reflectivities derived from the posterior samples of both WISE and WISER, shown in @fig-wise-rtm and @fig-wiser-rtm, respectively.

The imaged reflectivities produced by CM from WISER exhibit superior continuity and a better correlation with the CM migration-velocity model, particularly noticeable in @fig-wiser-rtm-zoom under the unconformity. Also, reflectivity samples produced by WISER demonstrate improved alignment among themselves compared to those produced by WISE. In addition, notable vertical shifts observed in the imaged reflectivities from WISE to WISER indicate significant adjustments in the positioning of subsurface reflectors, underlining the necessity of the refinement procedure for precisely estimating migration-velocity models that locate subsurface reflectors more accurately.

### Case 2: out of distribution

To test the robustness and adaptability of WISER when faced with unexpected variations at inference, we also evaluate WISER's performance under OOD scenarios. We introduce alterations to the velocity model depicted in @fig-true through an elementwise perturbation shown in @fig-perturb-curve. This manipulation modifies the velocity values across different depth levels, resulting in a significant shift in their statistical distribution, illustrated in @fig-hist. We use the perturbed velocity as the unseen ground-truth velocity model in this case study, shown in @fig-v-ood. To further expand the amortization gap, we modify the encoding of the forward operator by introducing a higher amplitude of band-limited Gaussian noise (S/N $0\,\mathrm{dB}$).

These complexities present substantial challenges for WISE, leading to biased inference results as depicted in @fig-wise-v-3d-ood. The yellow histograms in @fig-hist show that the velocity values of the posterior samples from WISE closely resemble those of the original velocity model, despite the different distribution of the ground-truth velocity model. This indicates that WISE tends to incorporate an inductive bias from the training samples. In WISER, we conduct $\mathrm{maxiter}_1=160$ outer iterations, using $M=16$ particles and $1$ OBN per gradient. We also employ the frequency continuation method [@bunks1995multiscale] to compute the gradient in line 25 of @alg-wiser, transitioning gradually from low-frequency to high-frequency data. This results in $40$ datapasses or $5120$ PDE solves in total.

::: {#fig-wise-vs-wiser-ood layout="[[20, 30], [49], [25, 25], [25, 25], [25, 25], [25, 25]]"}

![](./figs/perturbv.png){#fig-perturb-curve width=100%}

![](./figs/histogram.png){#fig-hist width=100%}

![](./figs/v-ood.png){#fig-v-ood width=100%}

![](./figs/WISE-v-3D-ood.png){#fig-wise-v-3d-ood width=100%}

![](./figs/WISER-v-3D-ood.png){#fig-wiser-v-3d-ood width=100%}

![](./figs/WISE-ood-rtm-3D.png){#fig-wise-ood-rtm width=100%}

![](./figs/WISER-ood-rtm-3D.png){#fig-wiser-ood-rtm width=100%}

![](./figs/WISE-ood-rtm-zoom2.png){#fig-wise-rtm-zoom-ood width=100%}

![](./figs/WISER-ood-rtm-zoom2.png){#fig-wiser-rtm-zoom-ood width=100%}

![](./figs/fwi-post-1.png){#fig-wise-fwi width=100%}

![](./figs/WISER-post-1.png){#fig-wiser-post width=100%}

OOD case study. *(a)* Curves for velocity-value perturbations; *(b)* histograms of values at the depth of $0.5\,\mathrm{km}$ and $2.8\,\mathrm{km}$ in the original velocity model (@fig-true), perturbed velocity model (@fig-v-ood), posterior samples of WISE, and WISER, shown in red, blue, yellow and green color, respectively. *(c)---(i)* Comparison between WISE and WISER. The ordering remains the same as in @fig-wise-vs-wiser. *(j)* FWI result starting with a posterior sample from WISE. *(k)* A posterior sample from WISER.
:::

#### Observations

WISER produces more accurate posterior samples shown in @fig-wiser-v-3d-ood. Furthermore, the statistical distribution of the velocity values in the WISER posterior samples (green histogram in @fig-hist) aligns better with the distribution of the unseen ground-truth velocity values (blue histogram in @fig-hist), demonstrating WISER's robustness against potential distribution shifts during inference.

To further showcase WISER's robustness under severe measurement noise, we compare a posterior sample from WISER (@fig-wiser-post) to a velocity model estimated by FWI (@fig-wise-fwi), derived by minimizing only the data likelihood (the first term in line 25 of @alg-wiser), while starting from the same initial model as WISER. The FWI result is significantly impacted by noise, while the posterior samples from WISER remain relatively noise-free and capture all pertinent geological structures.

#### Impact on imaging

The imaging results from WISE (@fig-wise-ood-rtm) and WISER (@fig-wiser-ood-rtm) reveal noticeable discrepancies in quality. The CM migration-velocity model from WISE leads to discontinuities in the imaged reflectivities, particularly at the horizontal layer around $1.8\,\mathrm{km}$ depth and more so below the unconformity. In contrast, the CM from WISER significantly improves the continuity of the imaged reflectivities across the entire seismic section. The imaged reflectivity samples from WISER not only show better consistency among themselves but also align more accurately with the estimated CM migration-velocity model, particularly visible in @fig-wiser-rtm-zoom-ood.

## Discussion and conclusions

The primary contribution of WISER is to leverage both genAI and physics to achieve a semi-amortized VI framework for scalable ($\mathrm{D}\geq 2$) and reliable UQ for FWI even in situations where local approximations are unsuitable. At its core, WISER harnesses the strengths of both amortized and non-amortized VI: the amortized posterior obtained through offline training provides a low-fidelity but fast mapping, and the physics-based refinements offer reliable and accurate inference. Both approaches benefit from information preservation exhibited by CIGs, rendering our inference successful where conventional FWI fails due to cycle skipping.

Compared to McMC methods that rely on low-dimensional parameterizations, WISER does not impose intrinsic dimensionality reductions or simplifications of the forward model. Therefore, WISER is capable of delivering full-resolution UQ for realistic multi-D FWI problems. CNFs are primed for large-scale 3D inversion thanks to their invertibility, which allows for memory-efficient training and inference [@orozco2023invertiblenetworks].

Compared to non-amortized VI methods, WISER also requires significantly less computational resources during inference. This is because WISE, as a precursor of WISER, already provides near-accurate posterior samples, making the refinement procedure computationally feasible. @zhang20233 show that non-amortized VI, without access to realistic prior information, requires $\mathrm{O}(10^6)$ to $\mathrm{O}(10^8)$ PDE solves, while WISER only needs $\mathrm{O}(10^3)$ PDE solves. Apart from the computational cost reduction, WISER ensures that the posterior samples realistically resemble Earth models, thanks to the integration of the conditional prior information from WISE. Contrary to non-amortized VI, which needs density evaluations to embed the prior (i.e., $p(\mathbf{x})$ in line 5 of @alg-wiser) to produce realistic Earth models, WISER only needs access to samples of the prior distribution (i.e., $\mathbf{x}^{(i)}$ in line 5).

Opportunities for future research remain. Although case 2 demonstrates the robustness of WISER concerning OOD issues, these issues could be fundamentally addressed by diversifying the training set of WISE through a foundation model [@sheng2023seismic]. Also, our OOD case study has not yet explored scenarios where the likelihood term is more pathologically misspecified, such as the presence of unremoved shear wave energy outside the range of the forward operator, which calls for further investigations. Our approach will also benefit from calibration of the estimated posterior on which the authors report elsewhere, including application of WISE(R) in 3D.

In conclusion, this paper sets the stage for deploying genAI models to facilitate high-dimensional Bayesian inference with computationally intensive multi-D forward operators. Deep learning and AI have been criticized for their reliance on realistic training samples, but WISER alleviates this reliance and still offers computationally feasible and reliable inference through a blend of offline training and online frugal physics-based refinements, preparing our approach for large 3D deployment.

<!--
## Acknowledgement

This research was carried out with the support of Georgia Research Alliance and partners of the ML4Seismic Center.

## Declaration of generative AI and AI-assisted technologies in the writing process

During the preparation of this work, the authors used ChatGPT to refine sentence structures and improve the readability of the manuscript. After using this service, the authors reviewed and edited the content as needed and take full responsibility for the content of the publication.
-->

## References

::: {#refs}
:::
