A Unified Framework for Forward and Inverse Problems in Subsurface Imaging using Latent Space Translations

## Abstract

Abstract In subsurface imaging, learning the mapping from velocity maps to seismic waveforms (forward problem) and waveforms to velocity (inverse problem) is important for several applications. While traditional techniques for solving forward and inverse problems are computationally prohibitive, there is a growing interest in leveraging recent advances in deep learning to learn the mapping between velocity maps and seismic waveform images directly from data.
Despite the variety of architectures explored in previous works, several open questions still remain unanswered such as the effect of latent space sizes, the importance of manifold learning, the complexity of translation models, and the value of jointly solving forward and inverse problems.
We propose a unified framework to systematically characterize prior research in this area termed the Generalized Forward-Inverse (GFI) framework, building on the assumption of manifolds and latent space translations.
We show that GFI encompasses previous works in deep learning for subsurface imaging, which can be viewed as specific instantiations of GFI.
We also propose two new model architectures within the framework of GFI: Latent U-Net and Invertible X-Net, leveraging the power of U-Nets for domain translation and the ability of IU-Nets to simultaneously learn forward and inverse translations, respectively.
We show that our proposed models achieve state-of-the-art (SOTA) performance for forward and inverse problems on a wide range of synthetic datasets, and also investigate their zero-shot effectiveness on two real-world-like datasets.

## 1 Introduction

The goal of subsurface imaging is to identify the structure and geophysical properties of layers underneath the Earth’s surface such as velocity maps ($v$), which is useful for a number of applications including energy exploration, carbon capture and sequestration, and developing earthquake early warning systems (Zhang et al., ). A typical approach for subsurface imaging is to conduct seismic surveys, where the elastic wave energy generated from one or more controlled sources is made to propagate through the Earth’s layers that are refracted, reflected, and received as seismic waveforms ($p$) on the surface by multiple receivers. Mathematically, the relation between velocity maps and seismic waveforms is governed by the following acoustic wave equation expressed as a non-linear partial differential equation (PDE):

$$ $\nabla^{2}p(x,z,t)-\frac{1}{v^{2}(x,z)}\frac{\partial^{2}}{\partial t^{2}}p(x, z,t)=s(x,z,t),$ (1) $$

where $x$, $z$, and $t$ represent the horizontal direction, depth, and time, respectively, and $\nabla^{2}$ is the Laplacian operator. Using Equation 1, the forward problem in subsurface imaging is to compute seismic waveforms $p$ given velocity maps $v$ by learning the forward mapping $f_{v\rightarrow p}:\mathcal{V}\rightarrow\mathcal{P}$. The inverse problem, also referred to as Full Waveform Inversion (FWI), is then to learn the reverse mapping $f_{p\rightarrow v}:\mathcal{P}\rightarrow\mathcal{V}$ for inferring velocity maps $v$ given observations of $p$.

While solving the forward problem requires expensive numerical solutions (Tago et al., ) of Equation 1, learning the inverse is even more computationally demanding as it involves iterating over velocity maps $v$ to minimize the difference between simulated waveforms from the forward model and ground-truth observations of $p$ (Tarantola, ). As a result, traditional techniques for solving forward and inverse problems in subsurface imaging are difficult to scale in operational settings at the desired resolutions.
In response, there is a growing interest in the geophysics community to leverage deep learning for subsurface imaging (DL4SI), in particular, using image-to-image-translation methods to map seismic waveforms to velocity maps and back directly from data.

Figure: Figure 1: A unified framework for solving forward and inverse problems in subsurface imaging.
Refer to caption: extracted/5930152/figures/architectures/GFI_Framework.png

For example, several deep learning methods have recently been proposed for FWI using encoder-decoder architectures such as InversionNet (Wu & Lin, ) and VelocityGAN (Zhang et al., ), where the encoder maps transform seismic data to latent feature representations that are decoded back to the velocity domain (Lin et al., ). Techniques for solving the forward problem in PDEs include the seminal work on physics-informed neural networks (PINNs) (Raissi et al., ) and recent works in neural operators such as Fourier Neural Operators (FNOs) (Li et al., ). The basic idea of FNOs is to lift inputs to the frequency domain using Fourier transform and apply linear transforms before inverting them to the original space in every FNO layer. A recent work has also pursued the goal of solving both forward and inverse problems together using a common architecture termed Auto-Linear (Feng et al., ), where two auto-encoders are first separately trained to project velocity and waveform fields to higher-dimensional manifolds, which are then linked using linear transforms for translating between manifolds.

Despite the variety of architectures explored in previous works in DL4SI, several open questions still remain unanswered such as:
(1) How does the size of latent spaces in encoder-decoder frameworks influence the quality of domain translations in DL4SI? While previous works such as InversionNet used lower-dimensional embeddings in their bottleneck layers, recent methods such as Auto-Linear advocate for higher-dimensional mappings.
(2) Do we need complex architectures for latent space translations or are linear transforms sufficient?
(3) What is the role of manifold learning in DL4SI where latent spaces are used not just for translation but also for reconstruction?
(4) Can we simultaneously solve both forward and inverse problems using joint frameworks?

To answer these questions, we propose a unified framework to systematically characterize prior research in DL4SI termed the *Generalized Forward-Inverse* (GFI) framework (see Figure 1). The GFI framework builds on two key philosophies that are at the core of prior research in DL4SI. First, we assume that both velocity $v$ and seismic waveforms $p$ can be projected to latent space manifolds $\tilde{v}$ and $\tilde{p}$ using encoder-decoder pairs that allow reconstruction of the original domains, referred to as the manifold assumption. Second, we assume that it is possible to learn bidirectional translations between the two manifolds $\tilde{v}$ and $\tilde{p}$, referred to as the latent space translation assumption. We show that GFI encompasses previous works in DL4SI including InversionNet, VelocityGAN, and Auto-Linear, which can be viewed as specific instantiations of GFI. This unifying perspective helps us analyze the field of DL4SI going beyond specific architecture choices and also discover novel formulations in DL4SI such as those explored in this paper.

In particular, we propose two new architectures within the framework of GFI: *Latent U-Net* (see Figure 2(a)) and *Invertible X-Net* (see Figure 2(b)). Latent U-Net uses two separate U-Nets (Ronneberger et al., ) to model latent space translations in both directions. This takes inspiration from the success of U-Nets in mainstream computer vision applications to explore their benefits in DL4SI. Invertible-XNet uses Invertible U-Net (IU-Net) (Etmann et al., ) to simultaneously learn both forward and inverse translations in the latent space using a single architecture. Invertible U-Nets builds on the idea of Invertible Neural Networks (INNs) (Ardizzone et al., ) to learn bijective mappings between input and output domains in DL4SI.
We show that both Latent U-Net and Invertible X-Net achieve state-of-the-art (SOTA) performance for forward and inverse problems on a wide range of benchmark datasets commonly used in DL4SI. We show that jointly solving forward and inverse problems using Invertible X-Net especially helps in improving the accuracy on the forward problem.
We also investigate the zero-shot effectiveness of our proposed approaches on synthetic and real-world-like datasets, contributing to novel insights in the field of DL4SI.

Figure: (a) Latent U-Net architecture
Refer to caption: extracted/5930152/figures/architectures/LatentU-Net.png

## 2 Related Works

### 2.1 Deep Learning for Subsurface Imaging (DL4SI)

Solving inverse problems in DL4SI has received considerable attention in recent years, thanks to the seminal work by Deng et al. () in creating OpenFWI, an open-access synthetic dataset involving paired examples of velocity and waveforms with varying levels of complexity. This has enabled the formulation of FWI as an image-to-image translation problem, where the relationship between waveforms and velocity is learned through supervised, unsupervised, or self-supervised methods (Lin et al., ). One of the early works in supervised learning for FWI includes InversionNet (Wu & Lin, ), which uses an encoder-decoder architecture to project waveforms into flattened vector representations acting as the bottleneck layer, which are decoded back to velocity maps.
VelocityGAN (Zhang et al., ) builds on the idea of InversionNet by using an adversarial training scheme for learning the model parameters along with prediction loss.
These works have also been extended for estimating 3D velocity fields in Feng et al. ().

Since ground-truth observations of velocity maps are difficult to obtain in real-world settings, a number of unsupervised and self-supervised formulations have also been explored for FWI. This includes the framework of UPFWI (Jin et al., ) that uses the supervision of a physics-based forward model of the wave equation to enforce cycle-consistency with a data-driven inverse model. While the training of UPFWI is completely unsupervised and achieves comparable performance as InversionNet in some settings, it requires computationally expensive simulations of a forward model at every training epoch, restricting its scalability. A recent line of work in semi-supervised learning for FWI involves the framework of Auto-Linear (Feng et al., ). In this framework, both velocity and waveforms are first mapped to higher-dimensional manifolds using two auto-encoders separately trained to minimize reconstruction loss in their respective domains. This is followed by a supervised learning of translations between the manifolds using two independent linear transforms (one for forward and one for inverse), building on previous observations of linearity in high-dimensional spaces for FWI (Feng et al., ).

### 2.2 Domain Translation Methods

The goal in domain translation (Isola et al., ; Zhu et al., ; Choi et al., ) is to learn the mapping between input and output domains (e.g., the CT scan of a patient and the organ segmentation from the scan), either in a paired or unpaired setting.
A recent work on Latent Space Mapping (LSM) (Mayet et al., ) has introduced a unifying framework for domain translation based on the manifold assumption (i.e., both input and output domains reside in low-dimensional manifolds learned by minimizing reconstruction loss) and shared latent space assumption (i.e., both manifolds reside in a shared latent space making it possible to translate from one manifold to another). We take inspiration from LSMs in domain translation to develop our proposed framework of GFI.

## 3 Proposed Methodologies

### 3.1 Generalized Forward Inverse Framework

As shown schematically in Figure 1, our proposed framework of GFI unifies the learning of the forward problem $f_{v\rightarrow p}:\mathcal{V}\rightarrow\mathcal{P}$ and the inverse problem $f_{p\rightarrow v}:\mathcal{P}\rightarrow\mathcal{V}$ using a common architecture involving latent space translations.

The GFI framework is inspired by two key assumptions. First, according to the manifold assumption, we assume that the velocity maps $v\in\mathcal{V}$ and seismic waveforms $p\in\mathcal{P}$ can be projected to their corresponding latent space representations, $\tilde{v}$ and $\tilde{p}$, respectively, which can be mapped back to their reconstructions in the original space, $\hat{v}$ and $\hat{p}$. Note that the sizes of the latent spaces can be smaller or larger than the original spaces. Further, the size of $\tilde{v}$ may not match with the size of $\tilde{p}$ (e.g., in Auto-Linear). Second, according to the latent space translation assumption, we assume that the problem of learning forward and inverse mappings in the original spaces of velocity and waveforms can be reformulated as learning translations in their latent spaces. We provide theoretical justifications for the existence of these latent space translations in the Appendix A. Note that these assumptions can also be relaxed in specific instantiations of GFI. For example, the manifold assumption can be relaxed to learn latent spaces solely for translation and not reconstruction.

Formally, let $\mathtt{E}_{v}:\mathcal{V}\rightarrow\tilde{\mathcal{V}}$ and $\mathtt{E}_{p}:\mathcal{P}\rightarrow\tilde{\mathcal{P}}$ be the encoders that map velocity and waveforms to their respective latent representations, $\tilde{v}=\mathtt{E}_{v}(v)$ and $\tilde{p}=\mathtt{E}_{p}(p)$. Similarly, let $\mathtt{D}_{v}:\tilde{\mathcal{V}}\rightarrow\mathcal{V}$ and $\mathtt{D}_{p}:\tilde{\mathcal{P}}\rightarrow\mathcal{P}$ be the decoders that map the latent representations back to the original spaces.
Additionally, let $L_{\tilde{v}\rightarrow\tilde{p}}:\tilde{\mathcal{V}}\rightarrow\tilde{
\mathcal{P}}$ and $L_{\tilde{p}\rightarrow\tilde{v}}:\tilde{\mathcal{P}}\rightarrow\tilde{
\mathcal{V}}$ be the latent space translations between $\tilde{v}$ and $\tilde{p}$ and vice versa.
The forward and inverse mappings between velocity $v$ and waveform $p$ can then be formulated as:

$$ $\displaystyle\hat{p}=f_{v\rightarrow p}(v)=\mathtt{D}_{p}\circ L_{\tilde{v} \rightarrow\tilde{p}}\circ\mathtt{E}_{v}(v)$ (2) $\displaystyle\hat{v}=f_{p\rightarrow v}(p)=\mathtt{D}_{v}\circ L_{\tilde{p} \rightarrow\tilde{v}}\circ\mathtt{E}_{p}(p)$ (3) $$

Learning GFI with Paired Data: Let us assume that we have a paired i.i.d. dataset $(v,p)\sim q(v,p)$, where $q(v,p)$ is a joint probability distribution. We can then learn the parameters of the GFI framework in a supervised manner by optimizing the following loss functions:

$$ $\displaystyle\min_{\theta_{\mathtt{E}_{v}},\theta_{L_{\tilde{v}\rightarrow \tilde{p}}},\theta_{\mathtt{D}_{p}}}\mathbb{E}_{q(v,p)}\mathcal{L}(\mathtt{D}_ {p}\circ L_{\tilde{v}\rightarrow\tilde{p}}\circ\mathtt{E}_{v}(v),p)=\mathcal{L }(\hat{p},p)$ (4) $\displaystyle\min_{\theta_{\mathtt{E}_{p}},\theta_{L_{\tilde{p}\rightarrow \tilde{v}}},\theta_{\mathtt{D}_{v}}}\mathbb{E}_{q(v,p)}\mathcal{L}(\mathtt{D}_ {v}\circ L_{\tilde{p}\rightarrow\tilde{v}}\circ\mathtt{E}_{p}(p),v)=\mathcal{L }(\hat{v},v)$ (5) $$

where $\mathcal{L}$ denotes a loss function such as the Mean Squared Error (MSE) loss, Mean Absolute Error (MAE) loss, or even an adversarial loss if trained with a discriminator.

### 3.2 Prior Works in DL4SI as Special Cases of GFI

Table 1 summarizes how previous works in DL4SI can be viewed as specific realizations of the unifying GFI Framework, as described in the following. Schematic illustrations of all prior works in the framework of GFI are provided in the Appendix Figure 9.

**Table 1: Characterizing prior and proposed works using GFI Framework.**
| Model | Modeling Mode | Manifold<br>Learning | Latent Space<br>Translation | Sizes of<br>Latent Spaces |
| --- | --- | --- | --- | --- |
| InversionNet | Inverse | No | Identity | Low |
| VelocityGAN | Inverse | No | Identity | Low |
| FNO | Forward | No | Fourier Layers | N/A |
| Auto-Linear | Forward and<br>Inverse Disjointly | Yes | Linear | High and<br>Different |
| Latent U-Net<br>(ours) | Forward and<br>Inverse Disjointly | Possible | U-Net | Low and<br>Same |
| Invertible X-Net<br>(ours) | Forward and<br>Inverse Simultaneously | Implicit | IU-Net | Low and<br>Same |

InversionNet and VelocityGAN: Both InversionNet and VelocityGAN use the same encoder-decoder architecture for learning the inverse mapping from seismic waveforms to velocity maps. They can be viewed as using an encoder $\mathtt{E}_{p}$ to map waveforms to a lower-dimensional space (represented as a flat vector), which is decoded back to velocity using decoder $\mathtt{D}_{v}$ (both encoder and decoder are implemented as ConvNets).
The latent space translation is effectively the identity mapping, $L_{\tilde{p}\rightarrow\tilde{v}}=I$.
Also, since they operate in inverse-only mode, they do not involve $\mathtt{E}_{v}$, $\mathtt{D}_{p}$, and $L_{\tilde{v}\rightarrow\tilde{p}}$ components in their architecture.

As a result, they do not perform manifold learning since the latent spaces of $\tilde{p}$ and $\tilde{v}$ are solely used for translation and not for reconstruction.

Auto-Linear: Auto-Linear employs a two-stage approach for learning both forward and inverse mappings using two disjoint latent space translations. In the first stage, self-supervised learning using reconstruction loss is utilized to train encoder-decoder pairs independently on velocity maps (i.e., $\mathtt{E}_{v}$ and $\mathtt{D}_{v}$) and seismic waveforms (i.e., $\mathtt{E}_{p}$ and $\mathtt{D}_{p}$). Note that the sizes of $\tilde{v}$ and $\tilde{p}$ are kept different from each other and higher than that of $v$ and $p$. In the second stage, two linear transforms are separately trained for the forward and inverse translations in the latent space, $L_{\tilde{v}\rightarrow\tilde{p}}$ and $L_{\tilde{p}\rightarrow\tilde{v}}$, while keeping the previously trained encoders and decoders frozen. Auto-Linear leverages the manifold assumption as the latent manifolds $\tilde{v}$ and $\tilde{p}$ are useful for reconstruction.

### 3.3 Proposed Model: Latent U-Net

We propose a novel architecture for solving forward and inverse problems using two latent space translation models implemented using U-Nets, termed *Latent U-Net*. As shown in Figure 2(a), Latent U-Net uses ConvNet backbones for both encoder-decoder pairs: ($\mathtt{E}_{v}$, $\mathtt{D}_{v}$) and ($\mathtt{E}_{p}$, $\mathtt{D}_{p}$), to project $v$ and $p$ to lower-dimensional representations. We also constrain the sizes of the latent spaces of $\tilde{v}$ and $\tilde{p}$ to be identical,
i.e., $\text{dim}(\tilde{v})=\text{dim}(\tilde{p})$, so that we can train two separate U-Net models to implement the latent space mappings $L_{\tilde{v}\rightarrow\tilde{p}}$ and $L_{\tilde{p}\rightarrow\tilde{v}}$.
Note that while U-Nets cannot be applied directly in the original spaces of velocity and waveforms due to a mismatch in dimensionalities, Latent U-Net allows us to bring to bear the power of U-Net in the domain of subsurface imaging, which has been very successful in several applications of computer vision.

We train Latent U-Net in an end-to-end manner to directly learn latent spaces useful for translation. This is in contrast to the two-stage training process employed in Auto-Linear where latent spaces are first pre-trained for reconstruction, followed by translation.
Specifically, for the forward problem, the components $\mathtt{D}_{p}$, $L_{\tilde{v}\rightarrow\tilde{p}}$, and $\mathtt{E}_{v}$ are trained together by optimizing Equation 4, while for the inverse problem, the components $\mathtt{D}_{v}$, $L_{\tilde{p}\rightarrow\tilde{v}}$, and $\mathtt{E}_{p}$ are trained together by optimizing Equation 5.
Since the training of forward and inverse modeling components of Latent U-Net are disjoint from each other, Latent U-Net does not explicitly learns manifolds useful for reconstruction. However, it is possible to modify the training objective of Latent U-Net by adding reconstruction losses in both forward and inverse problems to perform manifold learning.

### 3.4 Proposed Model: Invertible X-Net

We propose another novel architecture within the GFI framework termed Invertible X-Net to answer the question: “can we learn a single latent space translation model that can simultaneously solve both forward and inverse problems?”
Recently, Invertible Neural Networks (INNs) (Ardizzone et al., ) have been introduced as a specialized class of neural networks capable of learning bijective mappings between inputs and outputs.
INNs are inherently invertible by design; once the INN is trained to learn a forward mapping $f:\mathcal{X}\rightarrow\mathcal{Y}$ between inputs $x\in\mathcal{X}$ and outputs $y\in\mathcal{Y}$, the inverse mapping $f^{-1}:\mathcal{Y}\rightarrow\mathcal{X}$ can be obtained for free. A variant of INN has also been developed for image-to-image translation problems using the U-Net backbone, termed IU-Net (Etmann et al., ).
However, one of the key constraints of INN and IU-Net is that they require $\text{dim}(\mathcal{X})=\text{dim}(\mathcal{Y})$, which is difficult to ensure in subsurface imaging problems (since velocity and waveforms reside in entirely different spaces).

This is where we can leverage the GFI framework to employ IU-Net in the latent spaces of velocity and waveforms, which can be constrained to be of the same size (just like Latent-UNets),
i.e., $\text{dim}(\tilde{v})=\text{dim}(\tilde{p})$. By using a single IU-Net, we can simultaneously learn both forward and inverse translations in the latent space, $L_{\tilde{p}\rightarrow\tilde{v}}$ and $L_{\tilde{v}\rightarrow\tilde{p}}$. We refer to this formulation as Invertible “X”-Net (owing to the cross-shape of translations between the latent spaces) defined as follows:

$$ $\displaystyle\tilde{p}=L_{\tilde{v}\rightarrow\tilde{p}}(\tilde{v})=f_{\text{ IU-Net}}(\tilde{v});\quad\tilde{v}=L_{\tilde{p}\rightarrow\tilde{v}}(\tilde{p} )=f^{-1}_{\text{IU-Net}}(\tilde{p})$ (6) $\displaystyle\hat{p}=\mathtt{D}_{p}\circ f_{\text{IU-Net}}(\tilde{v})\circ \mathtt{E}_{v}(v);\quad\hat{v}=\mathtt{D}_{v}\circ f^{-1}_{\text{IU-Net}}( \tilde{p})\circ\mathtt{E}_{p}(p)$ (7) $$

The objective function for training Invertible X-Net involves prediction losses on both velocity maps and waveforms defined as follows:

$$ $\displaystyle\mathcal{L}_{\text{X-Net}}$ $\displaystyle=\mathcal{L}_{\text{forward}}+\mathcal{L}_{\text{inverse}}$ (8) $\displaystyle=\mathbb{E}_{q(v,p)}\mathcal{L}(\mathtt{D}_{p}\circ f_{\text{IU- Net}}\circ\mathtt{E}_{v}(v),p)+\mathbb{E}_{q(v,p)}\mathcal{L}(\mathtt{D}_{v} \circ f^{-1}_{\text{IU-Net}}\circ\mathtt{E}_{p}(p),v)$ (9) $$

Note that Invertible X-Nets facilitate bidirectional training as both the forward and inverse modeling components are trained simultaneously. As a result, both encoder-decoder pairs of Invertible X-Nets implicitly learn manifolds useful for reconstruction. Also, Invertible X-Net offers several key advantages over baselines. First, it simultaneously addresses both the forward and inverse problems within a single model architecture, whereas other baselines typically require training separate models for each task (e.g., Latent U-Net and Auto-Linear), leading to greater parameter efficiency. Additionally, the use of IU-Net ensures that the mappings between the latent spaces of velocity maps and seismic waveforms are bijective, guaranteeing a one-to-one mapping between these representations – a property not necessarily true for other models such as Latent U-nets and Auto-Linear. Further, the bidirectional training of the forward and inverse problems introduces a strong regularization effect as the gradients of the forward loss $\mathcal{L}_{\text{forward}}$ affects the parameters of $f_{\text{IU-Net}}$, thereby affecting both forward and inverse performance. Similarly, the gradients of the inverse loss $\mathcal{L}_{\text{inverse}}$ also impact both forward and inverse performance by updating the parameters of $f_{\text{IU-Net}}$.

Cycle Loss for Training Invertible X-Net: Another advantage of Invertible X-net is that the architecture can be trained with unpaired examples using cycle-loss consistency (Zhu et al., ). We consider both variants of Invertible X-Net (with and without cycle loss) in our experiments to demonstrate its effect on generalization performance.

## 4 Experimental Setup

### 4.1 Datasets

We consider the OpenFWI collection of datasets (Deng et al., ), which encompasses multi-structural benchmark datasets for DL4SI organized into three groups: Vel Family, Fault Family, and Style Family. The Vel Family consists of four datasets representing simple geological patterns, while the Fault Family includes four datasets with fault-like deformations, presenting greater modeling challenges. The Style Family utilizes a style transfer method to create a range of highly complex geological settings.
The waveform data is structured as (# source $\times$ recording time $\times$ receiver length), resulting in a shape of ($5\times 1000\times 70$) for five uniformly spaced seismic sources and 70 receivers over 1000 milliseconds. In contrast, the velocity maps are represented as (1 $\times$ depth $\times$ receiver length), with a shape of ($1\times 70\times 70$) reflecting spatial dimensions of depth and horizontal coverage.
We also conduct out-of-distribution evaluations on the Marmousi (Martin et al., ) and Overthrust datasets (Aminzadeh et al., ). They are standard benchmarks that closely mimic real-world scenarios, significantly surpassing the complexity of the OpenFWI datasets.

### 4.2 Model Architectures

Latent U-Net and Invertible X-Net: We use two variants for Latent U-Net: (1) (Small) with 17M parameters, and (2) (Large) with 33M parameters. For Invertible X-Net, we used an architecture with 24M parameters trained with and without cycle loss. See Appendix for details.

Baselines: For the inverse problem, we considered InversionNet, VelocityGAN, and Auto-Linear as baselines. We obtained their trained model checkpoints to reproduce their results. For the forward problem, we implemented FNO (Li et al., ) and a U-Net model (termed WaveformNet) as baseline models (see Appendix for details), along with using Auto-Linear.

Training Details: The input seismic waveforms were standardized using a StandardScaler to center the data with a zero mean and unit standard deviation, while the velocity data were scaled through MinMax normalization. See Appendix for details on hyper-parameter settings. To make a fair comparison with baseline methods, we kept the original train-test splits and reported results in the unnormalized space of both velocity and waveforms. We used Mean Absolute Error (MAE), Mean Square Error (MSE), and Structured Similarity (SSIM) as evaluation metrics.

## 5 Results and Discussions

Quantitative Comparisons: Figure 3(a) compares MAE on velocity predictions while solving the inverse problem across different OpenFWI datasets. We can see that our proposed models, which capture complex non-linear relationships in the latent space using U-Net and IU-Net architectures, consistently outperform baseline methods InversionNet, VelocityGAN, & Auto-Linear across majority of datasets. Overall, the large Latent U-Net model shows the best performance in inverse modeling followed by the small Latent U-Net model. On complex datasets such as Style family and CFB, Invertible X-Net with and without cycle loss are also comparable to the small Latent U-Net model. In most cases, the Invertible X-Net model trained with the cycle loss performs at par with the Invertible X-Net trained with supervised loss except in a few cases like CFA and FFA. In addition to MAE, we also evaluate baselines using other metrics reported in Table 7 in the Appendix.

Figure 3(b) compares the performance of baseline methods on the forward problem (see Appendix Table 9 for other evaluation metrics). We can see that both Latent U-Net (Large) and Invertible X-Net show better performance than most baselines across all datasets. However, in contrast to the trends observed in the inverse problem, we observe that Invertible X-Net with and without cycle loss shows the best performance compared to Latent U-Net models. We further notice that for simple datasets such as FVA and FFA, FNO has significantly lower MAE compared to other models. We later qualitatively explain this pattern by showing that FNO tends to capture dominant modes of the waveform distribution while neglecting subtler reflection patterns.

Figure: (a) Inverse problem
Refer to caption: extracted/5930152/figures/supervised_results/Supervised_Inverse_MAE_flipped.png

Qualitative Comparisons: In Figure 4(a), we compare inverse model predictions on three datasets, CVB, CFB, and STA, choosing one from each family of OpenFWI. We observe that our proposed methods (Latent U-Net and Invertible X-Net) produce considerably better velocity maps than the two SOTA baselines, InversionNet and Auto-Linear. We chose these 3 datasets as they are complex and have high heterogeneity as the velocity rapidly changes in all directions. For heterogeneous velocity maps, the predictions of baseline methods are smooth and they fail to delineate sharp changes in velocities. On the other hand, Latent U-Net and Invertible X-Net are able to predict better velocity fields capturing rapid variations in both shallow and deeper parts (highlighted using black-boxes in the figure). Between Latent U-Net (Large) and Invertible X-Net, we can see that the performance of Latent U-Net is slightly better especially on CFB and STA datasets. Note that Invertible X-Net uses a smaller number of parameters (24M) compared to Latent U-Net (Large) (33M). We provide additional visualizations of more baseline methods and across other datasets in the Appendix.

Figure 4(b) shows visualizations of model predictions for the forward problem on three datasets. Since these datasets are heterogeneous, the recorded seismic waveforms exhibit convoluted interference patterns. In particular, we can see that in all ground truth visualizations, there is a direct arrival in the shallower layers represented as a ‘slanted line’ signature. While capturing the direct arrival is relatively simple, the primary challenge in solving forward problems is to capture the subtler reflections in the deeper layers, where even small magnitudes of errors can result in large differences in inferred velocity maps.
We can see that while all baseline methods are able to predict shallow waveforms easily (like direct arrivals), methods such as FNO and Auto-Linear are struggling to predict the seismic waveforms, especially in the deeper regions (highlighted using black-boxes in the Figure).
On the other hand, our proposed models consistently improve results over baselines across all datasets. Additional visualizations for the forward problem are provided in the Appendix.

Building on these quantitative and qualitative results, we provide key insights on four open questions in the field of DL4SI in the following.

Figure: (a) Inverse problem
Refer to caption: extracted/5930152/figures/visualizations/vel/main_paper/curvevel-b_main_paper_50.png

Figure: Figure 5: Effect of latent dimensions and skip connections on the performance of Latent U-Net (Large) (a and b) and Latent U-Net (Small) (c). ‘A’ and ‘B’ show ‘low’ and ‘high’ complexity scenarios for each data group.
Refer to caption: extracted/5930152/figures/latentdim_main/latent_dim_effect.png

### 5.1 Key Insights on Open Questions in DL4SI

#### 5.1.1 What is the Effect of Latent Space Sizes on Translation Performance?

In Figure 5 (a), we investigate the impact of varying latent space sizes on the quality of domain translations using the Latent U-Net (Large) model on the inverse problem. We compare MAE by reducing latent space dimensions from $70\times 70$ to $8\times 8$. We can see that the quality of translations is not affected much for ‘A’-family of OpenFWI datasets, while the ‘B’-family shows monotonically increasing trend in MAE as we reduce latent dimension. Since the ‘A’-family was intentionally created to represent simpler velocity distributions, even the smallest latent space ($8\times 8$) is sufficient to estimate velocity maps. However, the ‘B’-family of datasets represent more complex velocity distributions, requiring a sufficiently larger latent space for the translation. This indicates that the ideal size of the latent space should be decided based on the complexity of problem being solved.

#### 5.1.2 Do We Need Complex Architectures for Translations?

We further study the role of model size on tranlslation performance by comparing Latent U-Net (Small) and Latent U-Net (Large) with varying latent dimensions in Figure 5 (b). We find that for a given latent dimension, Latent U-Net (Large) outperforms Latent U-Net (Small), motivating architectures with sufficient complexity for domain translations.
We further study the impact of using skip connections in the latent space translation models on varying latent space dimensions in Figure 5 (c).. We can see that adding skip connections shows marginal improvements in performance for larger latent dimensions. However, when the number of dimensions is low (less than 20), we can see a sharp degradation in performance by removing skip-connections. This indicates that smaller latent representations are weaker and thus insufficient for domain translation in U-Net frameworks without being augmented with skip connections.

Figure: Figure 6: Comparing a two-stage “Reconstruct then Translate” variant of Latent U-Net (Large) with the default one-stage formulation, across different training fractions.
Refer to caption: extracted/5930152/figures/latentdim_main/translation_vs_reconstruction_MAE.png

#### 5.1.3 Should We Train Latent Spaces Solely for Reconstruction?

Auto-Linear and Latent U-Net employ two different schools of thought for learning latent spaces for domain translations. Auto-Linear uses a two-stage “Reconstruct then Translate” approach where both encoder-decoder pairs are first trained solely for reconstruction, followed by the training of latent translation models in the second stage while keeping encoders-decoders frozen. In contrast, Latent U-Net follows a single-stage “Directly Translate” appproach where encoder-decoder pairs are trained together with the translation models, to allow the learned latent spaces to be influenced by the translation objective.

To investigate the effect of using frozen latent spaces on translation performance, we consider a two-stage “Reconstruct then Translate” variant of Latent U-Net where encoder-decoder pairs are first trained for reconstruction followed by translation in Figure 6 for varying OpenFWI families as we reduce the training fraction from 100% to 5%. We can see that for complex families such as CVB and CFB, the MAE of “Reconstruct then Translate” variant is higher for lower training fractions, suggesting that it is useful to influence the training of the latent spaces based on the translation objectives.
We see a similar trend for Latent U-Net (Small) shown in Appendix Figure 21.

#### 5.1.4 Is it Useful to Jointly Solve Forward and Inverse Problems?

As discussed in the quantitative comparisons of Latent U-Net (Large) and Invertible X-Net on the forward problem in Figure 3(b), we see a surprising trend that Invertible X-Net generally outperforms Latent U-Net (Large) despite having smaller number of parameters.
To understand this trend, we analyzed the training behavior of Invertible X-Net to learn that X-Net first learns to solve the inverse problem in the initial epochs of training, followed by a gradual shift in learning the solution to the forward problem in later epochs (see Appendix Figures 11 and 12 for visualizations). Since the model optimizes the combined loss on both velocity and waveforms together, as soon as the model is able to achieve a good solution for velocity, the gradients from the combined loss start helping the model to achieve better forward solutions. This also suggests that Invertible X-Net is seemingly learning the connection between forward and inverse problems for effective bidirectional training, and therefore is able to outperform Latent U-Net (Large) despite lower number of model parameters. To verify this, we compared the performance of Latent U-Net (Large), Invertible X-Net and Invertible X-Net trained only in the forward direction (Appendix Figure 10). We observed that indeed Invertible X-Net with forward-only mode shows worst performance followed by Latent U-net (Large), while Invertible X-Net (joint training) shows best performance.

### 5.2 Zero-shot Generalization Results

Extrapolating across OpenFWI Datasets:
To assess the out-of-distribution generalizability of best-performing models on unseen datasets, in Figure 7, we evaluate zero-shot performance of models trained across all datasets (shown as rows) when tested across all datasets (shown as columns), for both forward and inverse problems. Here, the color intensity indicates difference in SSIM performance of two competing methods. We can see that Invertible X-Net outperforms other baselines for both forward and inverse problems across majority of train-test cases. The only exception is when the models are trained on FVB, a relatively simpler dataset than other data families such as CurveVel and FlatVel. From a geophysics perspective, we expect to have better generalizability for models trained on complex geological settings and tested on relatively simpler ones, as complex geological structures can be thought of as juxtapositions of simpler geological structures. Additionally, we compare the zero-shot generalization for all baselines and Latent U-Net models in detail in the Appendix (Figures 13 - 16).

Figure: (a) Invertible X-Net vs InversionNet (Inverse)
Refer to caption: extracted/5930152/figures/appendix/zero_shot_grids/Inverse_Invertible_X-Net-InversionNet_SSIM.png

Figure: (a) Inverse problem
Refer to caption: extracted/5930152/figures/visualizations/marmousi_vel/marmousi_crop.png

Extrapolating on Marmousi and Overthrust Datasets:
We show the zero-shot generalizability of our models on the more complex and real-world-like geological settings of Marmousi and Overthrust in Figure 8. Here, we chose models trained on the most complex OpenFWI datasets, i.e., STA (for Marmousi) and STB (for overthrust) datasets, and compared their zero-shot performance in predicting velocity maps (inverse problem). We observe that none of the baselines are able to invert velocity maps at high resolution. However, the predictions of Latent U-Net and Invertible X-Net are still able to delineate the sharp changes in velocity maps especially in the shallow part of the data.
On the forward problem, our models are able to accurately predict seismic waveforms for both Marmousi and Overthrust data relative to baselines.

## 6 Conclusions and Future Research Directions

In this work, we introduced a unified framework for solving both forward and inverse problems in subsurface imaging, termed the Generalized-Forward-Inverse (GFI) framework. Within this framework, we proposed two novel architectures, Latent U-Net and Invertible X-Net, that leverage the power of U-Nets and IU-Nets to perform latent space translations, respectively. Our study addresses several key questions left unanswered by prior research in DL4SI.
Future research directions include training a single model for all datasets together, given that models trained on complex geological environments can effectively extrapolate on unseen simpler geological structures.

### Ethics Statement

We would like to emphasize that the primary goal of this paper is to facilitate a scientific inquiry of DL4SI. While advancements in subsurface imaging have traditionally been used for energy exploration, there are several emerging applications of subsurface imaging with positive societal consequences such as Carbon Sequestration, Earthquake Hazard Warning, and Environmental Monitoring. Research in DL4SI thus has similar characteristics as research in mainstream AI applications such as computer vision and language modeling, with a diverse range of societal, environmental, and economic impacts.

### Acknowledgments

We are thankful for the support of computational resources provided by the Advanced Research Computing (ARC) Center at Virginia Tech. This work was supported in part by NSF awards #2239328 and #2107332 from the Division of Information and Intelligent Systems (CISE/IIS). This manuscript has been authored by UT-Battelle, LLC, under contract DE-AC05-00OR22725 with the US Department of Energy (DOE). The US government retains and the publisher, by accepting the article for publication, acknowledges that the US government retains a nonexclusive, paid-up, irrevocable, worldwide license to publish or reproduce the published form of this manuscript, or allow others to do so, for US government purposes. DOE will provide public access to these results of federally sponsored research in accordance with the DOE Public Access Plan (
https://www.energy.gov/doe-public-access-plan).

## Appendices

## Appendix A Theoretical Justifications for Latent Space Translation Assumption

In this Section, we provide two Lemmas to theoretically support our Latent Space Translation assumption for the Generalized-Forward Inverse (GFI) Framework.

###### Lemma A.1 (Forward Latent Space Translation Assumption) .

Let $f:\mathcal{V}\rightarrow\mathcal{P}$ be an arbitrary forward operator mapping velocity maps $\mathcal{V}$ to seismic waveforms $\mathcal{P}$. Let $E_{v}:\mathcal{V}\rightarrow\tilde{\mathcal{V}}$ and $D_{v}:\tilde{\mathcal{V}}\rightarrow\mathcal{V}$ denote the encoder and decoder for the velocity space $\mathcal{V}$, respectively. Similarly, let $E_{p}:\mathcal{P}\rightarrow\tilde{\mathcal{P}}$ and $D_{p}:\tilde{\mathcal{P}}\rightarrow\mathcal{P}$ denote the encoder and decoder for the seismic waveform space $\mathcal{P}$. Here, $\tilde{v}\in\tilde{\mathcal{V}}$ and $\tilde{p}\in\tilde{\mathcal{P}}$ represent the latent space encodings. If we assume that the auto-encoder for the velocity are optimal, i.e., $D_{v}\circ E_{v}(v)=\hat{v}\approx v$, then there exists a functional mapping in the latent space $L_{\tilde{v}\rightarrow\tilde{p}}:\tilde{\mathcal{V}}\rightarrow\tilde{
\mathcal{P}}$.

###### Proof.

Given the forward operator $f:\mathcal{V}\rightarrow\mathcal{P}$, by definition, for any $v\in\mathcal{V}$, there exists $p\in\mathcal{P}$ such that $p=f(v)$.

Let the latent space representations $\tilde{v}\in\tilde{\mathcal{V}}$ and $\tilde{p}\in\tilde{\mathcal{P}}$ be defined by the auto-encoders as follows:

$$ $\displaystyle\tilde{v}$ $\displaystyle=E_{v}(v),\quad\hat{v}=D_{v}(\tilde{v})$ (10) $\displaystyle\tilde{p}$ $\displaystyle=E_{p}(p),\quad\hat{p}=D_{p}(\tilde{p})$ (11) $$

To construct the latent space mapping $\tilde{p}=L_{\tilde{v}\rightarrow\tilde{p}}(\tilde{v})$, consider the sequence of compositions involving the encoders, decoders, and the forward operator $f$:

$$ $\displaystyle\tilde{p}$ $\displaystyle=E_{p}(p)$ $\displaystyle=E_{p}(f(v))\quad\text{(since }p=f(v)\text{)}$ $\displaystyle=E_{p}(f(\hat{v}))\quad\text{(assuming reconstruction: }\hat{v} \approx v\text{)}$ $\displaystyle=E_{p}(f(D_{v}(\tilde{v})))\quad\text{(since }\hat{v}=D_{v}( \tilde{v})\text{)}$ (12) $$

Thus, by definition of the composition of functions, the latent space mapping can be expressed as:

$$ $\displaystyle L_{\tilde{v}\rightarrow\tilde{p}}=E_{p}\circ f\circ D_{v}$ (13) $$

∎

###### Lemma A.2 (Inverse Latent Space Translation Assumption) .

Let $f^{-1}:\mathcal{P}\rightarrow\mathcal{V}$ be an arbitrary inverse operator mapping seismic waveforms $\mathcal{P}$ to velocity maps $\mathcal{V}$ that is unique. Let $E_{p}:\mathcal{P}\rightarrow\tilde{\mathcal{P}}$ and $D_{p}:\tilde{\mathcal{P}}\rightarrow\mathcal{P}$ denote the encoder and decoder for the seismic waveform space $\mathcal{P}$, respectively. Similarly, let $E_{v}:\mathcal{V}\rightarrow\tilde{\mathcal{V}}$ and $D_{v}:\tilde{\mathcal{V}}\rightarrow\mathcal{V}$ denote the encoder and decoder for the velocity space $\mathcal{V}$. Here, $\tilde{p}\in\tilde{\mathcal{P}}$ and $\tilde{v}\in\tilde{\mathcal{V}}$ represent the latent space encodings. If we assume that the auto-encoder for the seismic waveform space is optimal, i.e., $D_{p}\circ E_{p}(p)=\hat{p}\approx p$, then there exists a functional mapping in the latent space $L_{\tilde{p}\rightarrow\tilde{v}}:\tilde{\mathcal{P}}\rightarrow\tilde{
\mathcal{V}}$.

###### Proof.

Given the inverse operator $f^{-1}:\mathcal{P}\rightarrow\mathcal{V}$, by definition, for any $p\in\mathcal{P}$, there exists $v\in\mathcal{V}$ such that $v=f^{-1}(p)$.

Let the latent space representations $\tilde{p}\in\tilde{\mathcal{P}}$ and $\tilde{v}\in\tilde{\mathcal{V}}$ be defined by the auto-encoders as follows:

$$ $\displaystyle\tilde{p}$ $\displaystyle=E_{p}(p),\quad\hat{p}=D_{p}(\tilde{p})$ (14) $\displaystyle\tilde{v}$ $\displaystyle=E_{v}(v),\quad\hat{v}=D_{v}(\tilde{v})$ (15) $$

To construct the latent space mapping $\tilde{v}=L_{\tilde{p}\rightarrow\tilde{v}}(\tilde{p})$, consider the sequence of compositions involving the encoders, decoders, and the inverse operator $f^{-1}$:

$$ $\displaystyle\tilde{v}$ $\displaystyle=E_{v}(v)$ $\displaystyle=E_{v}(f^{-1}(p))\quad\text{(since }v=f^{-1}(p)\text{)}$ $\displaystyle=E_{v}(f^{-1}(\hat{p}))\quad\text{(assuming reconstruction: }\hat {p}\approx p\text{)}$ $\displaystyle=E_{v}(f^{-1}(D_{p}(\tilde{p})))\quad\text{(since }\hat{p}=D_{p}( \tilde{p})\text{)}$ (16) $$

Thus, by definition of the composition of functions, the latent space mapping can be expressed as:

$$ $\displaystyle L_{\tilde{p}\rightarrow\tilde{v}}=E_{v}\circ f^{-1}\circ D_{p}$ (17) $$

∎

**Table 2: Statistics on the number of samples, the size of the velocity and waveforms for each dataset in OpenFWI Deng et al. ().**
| Dataset | Examples | Velocity shape | Waveform shape |
| --- | --- | --- | --- |
| FlatVel-A | 30,000 | (1, 70, 70) | (5, 1000, 70) |
| FlatVel-B | 30,000 | (1, 70, 70) | (5, 1000, 70) |
| CurveVel-A | 30,000 | (1, 70, 70) | (5, 1000, 70) |
| CurveVel-B | 30,000 | (1, 70, 70) | (5, 1000, 70) |
| FlatFault-A | 60,000 | (1, 70, 70) | (5, 1000, 70) |
| FlatFault-B | 60,000 | (1, 70, 70) | (5, 1000, 70) |
| CurveFault-A | 60,000 | (1, 70, 70) | (5, 1000, 70) |
| CurveFault-B | 60,000 | (1, 70, 70) | (5, 1000, 70) |
| Style-A | 67,000 | (1, 70, 70) | (5, 1000, 70) |
| Style-B | 67,000 | (1, 70, 70) | (5, 1000, 70) |

###### Lemma A.1 (Forward Latent Space Translation Assumption) .

###### Proof.

###### Lemma A.2 (Inverse Latent Space Translation Assumption) .

###### Proof.

## Appendix B Dataset Description

The OpenFWI comprises multi-structural benchmark datasets of significant size that can be used for solving full waveform inversion (FWI) using machine learning techniques (Deng et al., ). In particular, the repository contains 3 major groups of data: (1) Vel Family, (2) Fault Family, and (3) Style Family. These groups represent simple to complex sub-surface geological settings with seismic velocity and waveforms information. The Vel family is the simplest geological patterns including four datasets - (1) FlatVel-A (FVA), (2) FlatVel-B (FVB), (3) CurveVel-A (CVA), and (4) CurveVel-B (CVB). The difference between FlatVel and CurveVel is that the former represents low-energy geological environments where the rock layers are deposited horizontally and the latter consists of curved layers which are formed due to structural deformation of flat layers. The Fault family also has four datasets - (1) FlatFault-A (FFA), (2) FlatFault-B (FFB), (3) CurveFault-A (CFA), and (4) CurveFault-B (CFB). Unlike Vel datasets, the Fault family contains fault-like deformations, which is fracturing of rocks under certain pressure conditions. Due to the presence of faults, the Fault family becomes more complicated and challenging to model. The Style family has two datasets - (1) Style-A (STA), and (2) Style-B (STB). This dataset is generated using the style transfer method where the COCO dataset (Lin et al., ) is set as the content images and the Marmousi dataset is set as the style image. This is the most complex OpenFWI dataset as it represents highly complex geological settings where the velocity is changing rapidly and abruptly. In summary, the details about the Vel and Fault datasets are described in table 2.

The waveform data is represented as (# source $\times$ recording time $\times$ receiver length) whereas the velocity follow (1 $\times$ depth $\times$ receiver length) shape. In seismic surveys, the wave arrival time is recorded at the surface and thus, the waveform data for a single source is represented as a function of time and receiver length. In total, there are 5 seismic sources uniformly spaced along the surface, and wavefields are recorded by 70 receivers (uniformly spaced) along the surface for 1000 milliseconds. Therefore, the seismic wavefields are of the shape ($5\times 1000\times 70$). On the other hand, velocity maps are represented as functions of spatial dimensions, depth and horizontal coverage, and thus have the shape ($1\times 70\times 70$).

## Appendix C Prior Works in DL4SI as special cases of GFI

Figure 9 provides additional schematic illustrations of prior works in DL4SI such as InversionNet, WaveformNet, and AutoLinear as special cases of our proposed GFI framework.

Figure: Figure 9: Prior Works exprerssed as GFI
Refer to caption: extracted/5930152/figures/architectures/InversionNet.png

## Appendix D Additional Experimental Details

### D.1 Cycle Loss for Invertible X-Net

Given the velocity maps $v$ and the waveforms $p$, the predictions for the waveform and velocity can be obtained using the Invertible X-Net as follows:

$$ $\displaystyle\hat{p}=f_{v\rightarrow p}(v)=\mathtt{D}_{p}\circ L_{\tilde{v} \rightarrow\tilde{p}}\circ\mathtt{E}_{v}(v)$ (18) $\displaystyle\hat{v}=f_{p\rightarrow v}(p)=\mathtt{D}_{v}\circ L_{\tilde{p} \rightarrow\tilde{v}}\circ\mathtt{E}_{p}(p)$ (19) $$

Now, the Invertible X-Net architecture can be further applied on $\hat{p}$ and $\hat{v}$ to create the following transformations:

$$ $\displaystyle\hat{\hat{p}}=f_{v\rightarrow p}(\hat{v})=\mathtt{D}_{p}\circ L_{ \tilde{v}\rightarrow\tilde{p}}\circ\mathtt{E}_{v}(\hat{v})$ (20) $\displaystyle\hat{\hat{v}}=f_{p\rightarrow v}(\hat{p})=\mathtt{D}_{v}\circ L_{ \tilde{p}\rightarrow\tilde{v}}\circ\mathtt{E}_{p}(\hat{p})$ (21) $$

The cycle-loss for Invertible X-Net can be mathematically defined as follows:

$$ $\displaystyle\mathcal{L}_{\text{cycle}}=\mathcal{L}(p,\hat{\hat{p}})+\mathcal{ L}(v,\hat{\hat{v}})$ (22) $$

where $\mathcal{L}$ is the loss function that can be MSE, MAE or Elastic Loss. Note that the formulation of Cycle-Loss does not rely on paired examples, and can be applied on un-paired data as well. The combined loss function including cycle-loss can be given as:

$$ $\displaystyle\mathcal{L}_{\text{X-Net (Cycle)}}=\mathcal{L}_{\text{forward}}+ \mathcal{L}_{\text{inverse}}+\mathcal{L}_{\text{cycle}}$ (23) $$

For training Invertible X-Net (Cycle) model, we use loss function shown in Equation 23.

### D.2 Additional Training Details

For training, we normalize the velocity using min-max normalization and seismic waveform using standard normalization to rescale the data to mean 0 and standard deviation as 1. Table 5 shows other hyperparameter details for training Latent U-Net and Invertible X-Net models on OpenFWI datasets.

Since the baseline models have different normalization schemes than our models, we compare model predictions during evaluation by unnormalizing the predictions to original domains. For example, AutoLinear uses min-max normalization for velocity and a combination of log-normalization $x_{\text{norm}}=(\log_{e}(1+|x|)*\mathtt{sign}(x)$ followed by min-max in the log-normalized domain for waveform. We unnormalized both the velocity and waveform predictions so that we can measure and visualize errors in the predictions in the original space, allowing easy comparison across models.

### D.3 Model Architecture

We provide more details related to model architecture, layers, and number of parameters related to seismic and velocity encoder-decoder architecture in Table 3, and regarding the latent models used for translation in Table 4. Further, we compare our model parameters with baseline models in Table 6.

**Table 3: Architecture Details of Seismic Waveform and Velocity Encoder-Decoder models.**
| Model | #Layers | #Embedding<br>Dim | Channels | #Params |
| --- | --- | --- | --- | --- |
| Velocity Encoder | 5 | 128x70x70 | [8, 16, 32, 64, 128] | 11632 |
| Velocity Decoder | 6 | 128x70x70 | [128, 64, 32, 16, 1, 1] | 27877 |
| Seismic Waveform<br>Encoder | 5 | 128x70x70 | [8, 16, 32, 64, 128] | 55680 |
| Seismic Waveform<br>Decoder | 6 | 128x70x70 | [128, 64, 32, 16, 5, 5] | 186497 |

**Table 4: Architecture details of Latent U-Net and IU-Net latent space translation models**
| Model | #Depths | #Conv blocks/<br>Coupling Blocks | #Params |
| --- | --- | --- | --- |
| Latent U-Net | 2 | 4 | 34.68M |
| IU-Net | 4 | 4 | 25.78M |

**Table 5: Hyperparameter details for training Latent U-Net and Invertible X-Net models.**
| Model | #Epochs | Optimizer | LR | LR Scheduler |
| --- | --- | --- | --- | --- |
| Latent U-Net | 450 | Adam | 2e-3 | StepLR |
| Invertible X-Net | 450 | Adam | 2e-3 | StepLR |

**Table 6: Comparison of encoder, decoder, and latent model parameters for our model (Latent U-Net and Invertible X-Net) with other baseline models.**
| Model | #Vel Encoder<br>Params | #Vel Decoder<br>Params | #Amp Encoder<br>Params | #Amp Decoder<br>Params | #Translation<br>Params | #Total<br>Params |
| --- | --- | --- | --- | --- | --- | --- |
| FNO | - | - | - | - | - | 7.38M |
| InversionNet | - | 9.34M | 35.76M | - | Identity | 24.41M |
| VelocityGAN | - | 9.34M | 35.76M | - | Identity | 24.41M |
| Autolinear | 12.98M | 9.98M | 2.29M | 10.18M | 16.5K | 35.45M |
| Latent U-Net(Small) | 11.6K | 27.88K | 55.68K | 186.5K | 17.86M | 18.13M |
| Latent U-Net(Large) | 11.6K | 27.88K | 55.68K | 186.5K | 34.68M | 34.96M |
| Invertible X-Net | 11.6K | 27.88K | 55.68K | 186.5K | 25.78M | 26.06M |

## Appendix E Additional Results

### E.1 Additional Inverse Modeling Results

We provide more detailed comparison of our models with other baseline models in Table 7. Our proposed models consistently outperform baseline models on multiple datasets indicating superior generalizability on in-distributions examples.

Additionally, we also show zero shot generalization of our models on the Marmousi and Overthrust dataset in Tables 8 and 10. Our model Invertible X-Net shows generalizability in SSIM indicating that overall prediction has better geological understanding than other baseline models.

**Table 7: Comparison of our models (Latent U-Net (Large) and Invertible X-Net) with other baseline models for the inverse problem across 10 OpenFWI datasets. The bold highlights the best performing model on that dataset.**
| Metric | Model | FVA | FVB | FFA | FFB | CVA | CVB | CFA | CFB | STA | STB |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| MAE $\downarrow$ | InversionNet | 19.67 | 52.77 | 102.77 | 224.57 | 25.80 | 158.31 | 38.90 | 246.94 | 93.96 | 103.37 |
| Auto-Linear | 12.16 | 70.05 | 110.72 | 273.02 | 24.55 | 181.39 | 41.38 | 268.47 | 107.85 | 95.63 |  |
| Latent U-Net | 9.01 | 28.11 | 48.05 | 123.56 | 30.91 | 92.91 | 27.05 | 185.22 | 67.00 | 71.41 |  |
| Invertible X-Net | 23.80 | 34.94 | 57.86 | 139.81 | 59.45 | 94.11 | 198.56 | 181.03 | 67.09 | 121.42 |  |
| MSE $\downarrow$ | InversionNet | 1002.74 | 17271.62 | 36438.33 | 188044.46 | 4081.10 | 68214.92 | 9490.73 | 136327.84 | 23626.76 | 58622.35 |
| Auto-Linear | 1053.03 | 33907.57 | 42391.39 | 236457.48 | 5952.25 | 81789.80 | 13715.42 | 154713.12 | 31274.60 | 21819.14 |  |
| Latent U-Net | 216.92 | 6464.06 | 12609.65 | 91784.50 | 2348.22 | 33935.12 | 3771.38 | 93249.92 | 14353.55 | 14564.94 |  |
| Invertible X-Net | 1245.52 | 6969.42 | 14659.46 | 96121.90 | 7719.55 | 32559.95 | 77105.41 | 86512.36 | 13106.96 | 31550.22 |  |
| SSIM $\uparrow$ | InversionNet | 0.9894 | 0.9461 | 0.8073 | 0.6726 | 0.9765 | 0.7208 | 0.9566 | 0.6136 | 0.8858 | 0.6314 |
| Auto-Linear | 0.9887 | 0.9044 | 0.8056 | 0.6169 | 0.97 | 0.6865 | 0.9424 | 0.5695 | 0.8422 | 0.7274 |  |
| Latent U-Net | 0.9967 | 0.9809 | 0.9273 | 0.8156 | 0.991 | 0.8515 | 0.98 | 0.6930 | 0.9298 | 0.8064 |  |
| Invertible X-Net | 0.9917 | 0.9769 | 0.9135 | 0.8076 | 0.9826 | 0.8532 | 0.9316 | 0.7116 | 0.9360 | 0.7913 |  |

**Table 8: Comparison of our models (Latent U-Net (Large) and Invertible X-Net) on real world like datasets - Marmousi, Overthrust, Marmousi Smooth, Overthrust Smooth for the Inverse Problem. The bold highlights the best performing model on that dataset.**
| Metric | Model | Marmousi | Overthrust | Marmousi Smooth | Overthrust Smooth |
| --- | --- | --- | --- | --- | --- |
| MAE $\downarrow$ | InversionNet | 282.39 | 273.80 | 179.34 | 148.34 |
| Auto-Linear | 285.38 | 298.89 | 206.38 | 207.41 |  |
| Latnet U-Net | 322.13 | 264.40 | 242.71 | 198.53 |  |
| Invertible X-Net | 298.39 | 308.50 | 245.86 | 258.26 |  |
| MSE $\downarrow$ | InversionNet | 160084.5 | 135988.73 | 67978.40 | 37804.58 |
| Auto-Linear | 159517.41 | 157434.56 | 78207.07 | 71128.92 |  |
| Latnet U-Net | 217508.17 | 122398.64 | 112685.39 | 69122.58 |  |
| Invertible X-Net | 180250.47 | 179093.10 | 107344.46 | 124700.75 |  |
| SSIM $\uparrow$ | InversionNet | 0.46 | 0.4519 | 0.6044 | 0.7217 |
| Auto-Linear | 0.438 | 0.4097 | 0.5423 | 0.6447 |  |
| Latnet U-Net | 0.438 | 0.4807 | 0.6371 | 0.7031 |  |
| Invertible X-Net | 0.504 | 0.4827 | 0.6633 | 0.6952 |  |

### E.2 Additional Forward Modeling Results

Similar to the inverse problem, we provide detailed comparison of our models with other baseline models in Table 9. Our proposed models consistently outperform baseline models on multiple datasets indicating superior in-distributions generalizability.

**Table 9: Comparison of our models (Latent U-Net (Large) and Invertible X-Net) with other baseline models for the forward problem across 10 OpenFWI datasets. The bold highlights the best performing model on that dataset.**
| Metric | Model | FVA | FVB | FFA | FFB | CVA | CVB | CFA | CFB | STA | STB |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| MAE $\downarrow$ | FNO | 0.0077 | 0.0385 | 0.0700 | 0.1405 | 0.0282 | 0.0829 | 0.0605 | 0.1075 | 0.0292 | 0.0379 |
| Auto-Linear | 0.0340 | 0.0906 | 0.0744 | 0.1537 | 0.0541 | 0.1048 | 0.0703 | 0.1356 | 0.0492 | 0.0584 |  |
| Latent U-Net | 0.0389 | 0.0473 | 0.0581 | 0.1098 | 0.0512 | 0.059 | 0.0576 | 0.0891 | 0.0269 | 0.0354 |  |
| Invertible X-Net | 0.0257 | 0.0527 | 0.0532 | 0.0887 | 0.0457 | 0.0477 | 0.0404 | 0.0671 | 0.0235 | 0.0341 |  |
| MSE $\downarrow$ | FNO | 0.0004 | 0.0066 | 0.0259 | 0.0825 | 0.0059 | 0.0358 | 0.0251 | 0.0471 | 0.0041 | 0.0053 |
| Auto-Linear | 0.0084 | 0.0427 | 0.0303 | 0.1045 | 0.0217 | 0.0625 | 0.033 | 0.0866 | 0.0174 | 0.0173 |  |
| Latent U-Net | 0.0173 | 0.0127 | 0.0160 | 0.0489 | 0.0344 | 0.0175 | 0.0442 | 0.0297 | 0.0053 | 0.0082 |  |
| Invertible X-Net | 0.0062 | 0.0176 | 0.0175 | 0.0305 | 0.0185 | 0.0117 | 0.0126 | 0.0163 | 0.004 | 0.0087 |  |
| SSIM $\uparrow$ | FNO | 0.9967 | 0.9781 | 0.8881 | 0.8354 | 0.9667 | 0.8702 | 0.9166 | 0.8160 | 0.9655 | 0.9417 |
| Auto-Linear | 0.8274 | 0.6081 | 0.5971 | 0.4262 | 0.7661 | 0.5808 | 0.6898 | 0.4667 | 0.6922 | 0.6631 |  |
| Latent U-Net | 0.9764 | 0.9779 | 0.9237 | 0.89 | 0.9659 | 0.9283 | 0.9571 | 0.8636 | 0.9819 | 0.9702 |  |
| Invertible X-Net | 0.9887 | 0.9782 | 0.9559 | 0.9221 | 0.9744 | 0.954 | 0.9757 | 0.9156 | 0.989 | 0.9805 |  |

**Table 10: Comparison of our models (Latent U-Net (Large) and Invertible X-Net) on real world like datasets - Marmousi, Overthrust, Marmousi Smooth, Overthrust Smooth for the Forward Problem. The bold highlights the best performing model on that dataset.**
| Metric | Model | Marmousi | Overthrust | Marmousi Smooth | Overthrust Smooth |
| --- | --- | --- | --- | --- | --- |
| MAE $\downarrow$ | FNO | 0.1484 | 0.2726 | 0.1077 | 0.1892 |
| Auto-Linear | 0.2818 | 0.3018 | 0.2821 | 0.2730 |  |
| Latnet U-Net | 0.1338 | 0.2502 | 0.1013 | 0.1875 |  |
| Invertible X-Net | 0.1425 | 0.2311 | 0.1056 | 0.2062 |  |
| MSE $\downarrow$ | FNO | 0.1110 | 0.580 | 0.0811 | 0.3495 |
| Auto-Linear | 0.4227 | 0.6203 | 0.5325 | 0.5593 |  |
| Latnet U-Net | 0.1116 | 0.5494 | 0.0927 | 0.4133 |  |
| Invertible X-Net | 0.1168 | 0.4026 | 0.08130 | 0.5049 |  |
| SSIM $\uparrow$ | FNO | 0.4728 | 0.4242 | 0.5690 | 0.5475 |
| Auto-Linear | 0.2235 | 0.2311 | 0.2569 | 0.2748 |  |
| Latnet U-Net | 0.4866 | 0.4567 | 0.5935 | 0.5787 |  |
| Invertible X-Net | 0.4408 | 0.4056 | 0.5719 | 0.5243 |  |

### E.3 Importance of Combined Loss Function for Invertible-XNet

In this section, we focus on the training of the Invertible X-Net model using a combined loss function (incorporating both the forward and inverse problems), as opposed to training it solely with a forward loss function. Figure 10 shows that when model is trained using only forward loss, then its performance falls short compared to the Latent U-Net (Large) model. This discrepancy can be attributed to the fact that the Latent U-Net has a higher model complexity than Invertible X-Net, despite their similar sizes. Nonetheless, when the Invertible X-Net model is trained with combined loss (forward and inverse), the model is able to outperform Latent U-Net model with a good margin. This highlights the value of joint training, demonstrating that simultaneously learning both the forward and inverse problems can lead to better results than learning the two models separately.

In Figures 11 and 12, we illustrate asymmetry in learning the translation for the forward and inverse problems using CVB and CFA datasets. From the figures, we observe that the model learns the inverse mapping in the initial epochs and gradually starts to learn the solution to the forward problem in later epochs. Since the network optimizes the combined loss function on both velocity and waveform together, the gradients from the combined loss help the model to achieve better forward solution. This corroborates with our hypothesis that the model trained on combined loss is able to learn the connection between forward and inverse problem.

Figure: Figure 10: Comparison of Latent U-Net (Large), Invertible X-Net, Invertible X-Net (Forward Only), and Invertible X-Net (Cycle) on the forward problem across various OpenFWI datasets.
Refer to caption: extracted/5930152/figures/appendix/IUnet_discuusion_forward_Supervised_Forward_MAE.png

Figure: (a) Epoch: 1
Refer to caption: extracted/5930152/figures/appendix/training_asymmetry/cvb/cvb_results_epoch0.jpg

Figure: (a) Epoch: 1
Refer to caption: extracted/5930152/figures/appendix/training_asymmetry/cfa/cfa_results_epoch0.jpg

### E.4 Zero-shot Performance

In this part, we provide detailed insights into zero shot performance of our models (Latent U-Net (Large) and Invertible X-Net) with other baselines across all datasets. This investigation helps us understand overall out-of-distribution generalization of a model and underscore the importance of learning the translation problem in the latent space.

### E.5 Inverse Problem:

Figure 13 and 14 show generalization performance of Invertible X-Net and Latent U-Net (Large) models respectively on the inverse problem using MAE, MSE, and SSIM metrics. As described in the main paper, we show models trained across all dataset as rows and evaluated across all test datasets as columns. For comparison, we evaluate metrics such as MAE, MSE, and SSIM and calculate its difference between our models and other baseline models. For MAE and MSE, when the color intensity is blue, our model show better generalizability and vice-versa whereas, for SSIM, when the color intensity is red indicates better generalization of our model and vice-versa.

From Figure 13, we observe that the Invertible X-Net shows superior generalization over the baseline models (AutoLinear, InversionNet, and VelocityGAN) across all metrics except on the FVB dataset. Figure 13 (d) compare Invertible X-Net with Latent U-Net (Large) model where we see that Invertible X-Net shows better generalization on complex datasets such as CFB, STA, and STB datasets, while Latent U-Net is better on relatively simpler datasets such as CVA, CFA, and more.
Similarly, Figure 14 shows the comparison of Latent U-Net (Large) model with other baseline models. As expected, we observe that our model is able to generalize much better than all the baseline models.

#### E.5.1 Forward Problem:

Similar to the inverse problem, we analyze out-of-distribution generalization of our models against baseline models across all evaluation metrics.

Figures 15 and 16 compares the performance of our models Invertible X-Net and Latent U-Net (Large) against baselines. In Figure 15, we observe strong generalization of Invertible X-Net over all baselines - AutoLinear, FNO, and WaveformNet. Figure 15 (d) shows the comparative performance of Invertible X-Net against Latent U-Net (Large) where we see that Invertible X-Net dominates overall across all metrics. In Figure 16, we compare the Latent U-Net (Large) model against same baselines as above. The figure indicates Latent U-Net (Large) model has much stronger generalizability than baseline models consistently.

Figure: (a) Invertible X-Net - AutoLinear
Refer to caption: extracted/5930152/figures/appendix/zero_shot_grids/Inverse_Invertible_X-Net-Auto-Linear.png

Figure: (a) Latent U-Net (Large) - AutoLinear
Refer to caption: extracted/5930152/figures/appendix/zero_shot_grids/Inverse_Latent_U-Net_Large-Auto-Linear.png

Figure: (a) Invertible X-Net - AutoLinear
Refer to caption: extracted/5930152/figures/appendix/zero_shot_grids/Forward_Invertible_X-Net-Auto-Linear.png

Figure: (a) Latnet U-Net (Large) - AutoLinear
Refer to caption: extracted/5930152/figures/appendix/zero_shot_grids/Forward_Latent_U-Net_Large-Auto-Linear.png

## Appendix F Ablation Studies

### F.1 Effect of varying Latent Space Sizes

Figure 17(a) (a) shows how the performance of Latent U-Net model is affected with change in the latent space size. As the latent space size is reduced, the MAE and MSE metrics for mapping seismic waveform to velocity is increasing, while the SSIM is decreasing. The impact is more pronounced on complex datasets such as B group of datasets. These datasets represent geologically complex datasets and therefore may require larger latent space to encode geological heterogeneity. In Figure 17(b), we compare the large Latent U-Net with the small Latent U-Net model, while latent space size is also reduced. We observe that the performance gap between the two models reduces as the latent space size is also reduced. This underscore the importance of latent space size for encoding the geological features for an effective translation.

Further, we provide visualizations of the first two primary PCA and t-SNE components of the velocity latent space in Figure 18 and 19 respectively. We take the encoder trained on a dataset and get the latent space encoding on several datasets. This visualization shows the in-distribution and out-of-distribution generality of our encoders and highlights the importance of the manifold for latent space translation. Overall, we observe that the larger latent space $70\times 70$ have better structure than $8\times 8$, indicating that the ideal size of latent space should be decided based on the complexity of dataset and problem being solved.

Figure: (a) Effect of Latent dimension on U-Net Large
Refer to caption: extracted/5930152/figures/appendix/latent_dim/effect_latent_dim_unet_Large_MAE_crop.png

Figure: Figure 18: Visualizing the PCA projection of Velocity latent space for $8\times 8$ and $70\times 70$.
Refer to caption: extracted/5930152/figures/appendix/latent_space_viz/FlatVel-A_pca_70.png

Figure: Figure 19: Comparing the t-SNE plot of latent space for velocity encoder.
Refer to caption: extracted/5930152/figures/appendix/latent_space_viz/FlatVel-A_tsne_8.png

### F.2 Effect of Skip Connections for Latent U-Net

In this section, we study the impact of skip connections of Latent U-Net model for the inverse problem. Figure 20 shows how the MAE and MSE are increasing and SSIM is decreasing when the latent space size is decreasing from $70\times 70$ to $8\times 8$. We observe that the impact of skip connections is more pronounced at smaller latent space size as opposed to larger space.

Figure: Figure 20: Effect of the size of latent space and skip vs no skip connections on the performance of large and small Latent U-Net models across OpenFWI datasets.
Refer to caption: extracted/5930152/figures/appendix/latentdim_skip/latent_dim_skip_vs_noskip_unet_Large_MAE_crop.png

### F.3 Effect of Manifold Learning

Figure 21 for Latent U-Net Small further elaborates on our finding that direct translation learning consistently outperforms the two-stage approach, where translation follows reconstruction.

Figure: Figure 21: Comparison of Latent U-Net Small’s performance when learning translation directly versus after reconstruction, across different training fractions.
Refer to caption: extracted/5930152/figures/appendix/latent_dim/translation_vs_reconstruction_MAE_unet_small.png

## Appendix G Additional Visualizations

Here, we provide additional visualization of the waveform and velocity predictions for baselines and our models, namely Latent-UNet (small), Latent-UNet (large), Invertible-XNet, and Invertible-XNet (cycle), for both forward and inverse problems. Please note that we show the prediction of seismic waveform and velocity in the original space by unnormalizing the predictions for every model.

Figure: Figure 22: Zero shot generalization results of model trained on Style-A dataset on Marmousi and Overthrust dataset samples and their smoothened versions.
Refer to caption: extracted/5930152/figures/appendix/marmousi_overthrust/style-a/marmousi_appendix.jpg

Figure: Figure 23: Zero shot generalization results of model trained on Style-B dataset on Marmousi and Overthrust dataset samples and their smoothened versions.
Refer to caption: extracted/5930152/figures/appendix/marmousi_overthrust/marmousi_appendix_styleb.jpg

Figure: Figure 24: Visualization of predictions for forward and inverse problems on FVA dataset.
Refer to caption: extracted/5930152/figures/appendix/visualizations/forward/flatvel-a_appendix_50_forward.jpg

Figure: Figure 25: Visualization of predictions for forward and inverse problems on FVB dataset.
Refer to caption: extracted/5930152/figures/appendix/visualizations/forward/flatvel-b_appendix_50_forward.jpg

Figure: Figure 26: Visualization of predictions for forward and inverse problems on CVA dataset.
Refer to caption: extracted/5930152/figures/appendix/visualizations/forward/curvevel-a_appendix_50_forward.jpg

Figure: Figure 27: Visualization of predictions for forward and inverse problems on CVB dataset.
Refer to caption: extracted/5930152/figures/appendix/visualizations/forward/curvevel-b_appendix_50_forward.jpg

Figure: Figure 28: Visualization of predictions for forward and inverse problems on FFA dataset.
Refer to caption: extracted/5930152/figures/appendix/visualizations/forward/flatfault-a_appendix_50_forward.jpg

Figure: Figure 29: Visualization of predictions for forward and inverse problems on FFB dataset.
Refer to caption: extracted/5930152/figures/appendix/visualizations/forward/flatfault-b_appendix_50_forward.jpg

Figure: Figure 30: Visualization of predictions for forward and inverse problems on CFA dataset.
Refer to caption: extracted/5930152/figures/appendix/visualizations/forward/curvefault-a_appendix_50_forward.jpg

Figure: Figure 31: Visualization of predictions for forward and inverse problems on CFB dataset.
Refer to caption: extracted/5930152/figures/appendix/visualizations/forward/curvefault-b_appendix_50_forward.jpg

Figure: Figure 32: Visualization of predictions for forward and inverse problems on STA dataset.
Refer to caption: extracted/5930152/figures/appendix/visualizations/forward/style-a_appendix_50_forward.jpg

Figure: Figure 33: Visualization of predictions for forward and inverse problems on STB dataset.
Refer to caption: extracted/5930152/figures/appendix/visualizations/forward/style-b_appendix_50_forward.jpg

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
questions: 1. Could you comment on the complexity of the problems and how to "calculate" it when choosing the latent space size? How to choose this size for real-world applications and are there any limitations in real-world scenarios?
2. Why have you decided to remove Invertible X-Net from the experiments in Section 5.1.2.? Have you tried different model sizes for this architecture?  What is the relation between the number of parameters in encoders and decoders and the U-Net and how does it influence performance? 
3. It is not clear from the paper what are the benefits of having such a unified framework. Could you comment on the limitations? What would be your recommendation on which architecture to use and when?
4. Could you comment on the examples in Figure 4. and Figure 8. and the performance of the models achieved in terms of MAE?
soundness: 3
strengths: - The paper is clearly presented and well-organized.
- Even though the GFI framework is not the first idea that unifies seismic imaging forward and inverse problems (e.g., Auto-Linear [1]), it tries to systematically characterize and unify prior research in this area. The paper answers a few questions that were open in the research field. The paper concludes the following: the size of the latent space should be decided based on the complexity of the problem, the training of the latent spaces should be influenced based on the translation objectives and jointly solving forward and inverse problems is helping the model to achieve better forward solutions. In my opinion, these are answers to important questions that were backed up by the experiments in the paper.
- The paper introduces U-Net and IU-Net architectures for learning translations between velocity and seismic waveforms in the latent space. This leads to two new architectures for unified seismic imaging, namely Latent U-Net and Invertible X-Net. In comparison, the Auto-Linear framework's architecture comprises two separate linear layers trained for the forward and inverse translations in the latent space.
- Extensive and systematic experiments. The proposed approach was compared to the existing state-of-the-art methods both for the forward and inverse seismic imaging problems. The proposed method employing Latent U-Net architecture outperforms the existing methods in most of the datasets in both forward and inverse directions.

---
[1] Y. Feng, Y. Chen, P. Jin, S. Feng, Y. Lin, "Auto-Linear Phenomenon in Subsurface Imaging
." ICML, PMLR:235, 2024, pp. 13153-13174.
summary: Subsurface imaging is a technique for identifying the geophysical properties of layers underneath the Earth's surface. There are two directions: learning the mapping from velocity maps to seismic waveforms, called the forward problem, and the inverse problem which is the mapping from seismic data to velocity maps. Most of the previous deep learning-based research in the field focuses only on the inverse problem also referred to as full waveform inversion (FWI). This paper proposes a generalized forward-inverse (GFI) framework that unifies both directions, i.e., FWI and the forward problem. The framework builds on the assumption of manifolds and latent space translations, proposing two model architectures for the latter, namely Latent U-Net and Invertible X-Net. Latent U-Net architecture employs two U-Nets for translation between the velocity and waveform latent representations and vice versa, while Invertible X-Net uses a single IU-Net that simultaneously learns forward and inverse translations. The GFI framework encompasses previous works in deep learning for subsurface imaging, at the same time trying to answer questions such as the effect of latent space size, the importance of manifold learning and the value of jointly solving forward and inverse problems. The models were evaluated on the synthetic OpenFWI dataset and their generalization ability was tested on two real-world-like datasets.
weaknesses: - Section 5.1.1. "What is the effect of latent space sizes on translation performance?" is kind of vague. It states that the ideal size of the latent space should be decided based on the complexity of the problem being solved. It is not clear how to determine the complexity in real-world applications. I think the authors should be more specific in this section.
- In Section 5.1.2. "Do we need complex architectures for translations?", there is a comparison between large and small Latent U-Net. The only difference was the complexity (size) of the U-Nets for the latent space translation and the influence of having skip connections. What is the relation between the number of parameters in encoders and decoders and the U-Net and how does it influence performance? The Invertible X-Net was not included in the experiments and only one size was used throughout the paper. It would be useful to see how its size influences the performance.
- The Invertible X-Net model achieves worse results than the Latent U-Net model for the inversion problem and often fails to outperform other state-of-the-art methods. In contrast, it outperforms most of the methods for the forward problem. It is not clear from the paper what are the benefits of having such a unified framework. I think the limitations should be stated more clearly in the paper.
- In the results section, in qualitative comparison, Figure 4. and Figure 8. are missing a measure of error (at least MAE). The picked images should have approximately the same MAE as the reported results in the tables presenting a realistic case and fair comparison. It is not clear whether the presented results were in a sense cherry-picked, e.g., the proposed method achieved better results than the average and other methods below the average for the illustrated example.
- Instead of having a qualitative comparison of the methods on the two real-world-like datasets, I think it would be much more representative to have a quantitative comparison. Here, the presented examples might also be cherry-picked.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: please see my concerns in the Weakness
soundness: 3
strengths: The paper presents a unified architecture for addressing both forward and inverse subsurface imaging problems, while previous studies typically focus only on inversions. Modeling these problems together within a single framework could offer valuable insights, though it may also introduce potential challenges.
summary: This paper explores deep learning methods to address the forward and inverse problems in subsurface imaging by mapping between velocity maps and seismic waveforms. The authors introduce the Generalized Forward-Inverse (GFI) framework, which unifies past approaches and addresses open questions around latent spaces, manifold learning, and model complexity. They propose two novel architectures, Latent U-Net and Invertible X-Net, which achieve state-of-the-art results on synthetic datasets and show promising zero-shot performance on real-world-like datasets.
weaknesses: 1. The paper’s narrow focus on subsurface imaging may limit its relevance for a broader audience at ICLR. If there are analogous problems in other fields where their framework could be applied? 
2. Comparisons to prior work are outdated, with main competing methods from 2019, like InversionNet and VelocityGAN. The authors are encouraged to consider more recent advances, such as “Physics-Informed Robust and Implicit Full Waveform Inversion Without Prior and Low-Frequency Information (2024)” and other state-of-the-art methods from 2023, for a fairer and more relevant assessment.
3. While the authors aim to model the forward problem using neural networks, it is uncertain if the learned mapping can capture complex physical dynamics, as seen in the governing equations of seismic data. Extensive out-of-domain experiments are recommended to assess the model’s robustness and generalizability to real-world data.

4. The paper's structure could be streamlined to highlight the proposed methodology. Currently, a substantial background section precedes a brief description of the methods. Specific sections of the introduction and related work and preliminary knowledge could be condensed.

5. The model specifications lack detail; adaptations made to standard UNet for subsurface imaging and the mechanisms enabling X-Net’s invertibility are unclear. Clarifying these aspects and specifying when each model might be preferred would improve understanding.
6. Experimental comparisons with recent baselines (published within the last 2-3 years) would add rigor to the evaluation, as several recent methods could serve as relevant benchmarks. To name a few, “Physics-Informed Robust and Implicit Full Waveform Inversion Without Prior and Low-Frequency Information 2024”; “Full-waveform inversion using a learned regularization 2023”; “Wasserstein distance-based full-waveform inversion with a regularizer powered by learned gradient 2023”; “Full-waveform inversion using a learned regularization 2023”, and more.
7. Key numerical results in Section 5 are deferred to the appendix, which reduces clarity. Summarizing key findings directly in the main text would enhance readability.
8. The Invertible X-Net significantly underperforms other models, raising questions about its ideal use cases. Further context on scenarios where the X-Net's design offers advantages would be valuable.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: See weaknesses.
soundness: 3
strengths: 1.	(Quality) High quality of experimental evaluation: Experiments are extensive over a range of datasets including an extensive evaluation of out-of-distribution performance and qualitative results are provided in addition to quantitative results. Hence, the conclusion of generally improved performance seems sound. 
2.	(Significance) The improvements in performance seem to be significant, however interpretability of how significant could be improved (see weaknesses point 5)
3.	(Significance) The finding that jointly solving forward and inverse problems is helpful for the forward problem but not for the inverse problem is interesting. 
4.	(Originality) The proposed unifying framework helps to quickly understand the different existing approaches and will be helpful for future works to expand the framework or propose improvements within it.
5.	(Clarity) The paper is overall well written and easy to follow also for researches not actively involved in the field of subsurface imaging.
summary: The paper proposes a new training scheme combined with two new model architectures for the problem of learning the forward and inverse map in deep learning for subsurface imaging. 
The training scheme consists of jointly learning the encoders and decoders for the two domains (velocity maps and seismic waveforms) with the translation network transforming the latent space representations from one domain to another.
The two architectures differ in the characteristics of the translation network with one consisting of two separate unidirectional U-nets (latent U-net) and the other of a single bidirectional invertible network (invertible X-net).
Both architectures achieve a what seems to be significant improvement in performance over existing works.
Finally the paper formulates a unifying framework in which existing approaches can be categorized and distinguished.
weaknesses: Major concerns

1. (Clarity) The experiment 5.1.3 addresses the question ‘should we train latent spaces solely for reconstruction?’, which I assume corresponds to the question 3) formulated in line 083 ‘what is the role of manifold learning?’.  
1.1 For increased clarity I’d suggest not changing the phrasing of the questions.  
1.2 The setup is unclear to me. What are training fractions? A fraction of the training set or a fraction of the training time? And why is it instructive to look at 5% and 10% training fraction?  
1.3 So, is the conclusion form the experiment that learning manifolds is not important? It is difficult to grasp the conclusion if the wording is not consistent. Is “influence the training of the latent spaces based on the translation objectives” identical to “not learning manifolds”?  
1.4 To better answer the question I think the following baseline is missing that combines the two alternatives (reconstruct then translate and directly translate) into directly learning the translation networks while having at the same time reconstruction losses (basically what is described in lines 243-245). Seeing how this combination performs would help answering the question if manifold learning helps.  
2. (Originality) The idea of the Latent U-net seems to be very similar to the InversionNet following the comparison in Table 1.  
2.1 The Modelling Mode is the same because the forward and inverse part of the Latent U-net are completely disjoint and hence correspond to an InversionNet and a ForwardNet (as long as no manifold learning is performed).  
2.2 While the latent U-net allows manifold learning, the experimental results pertain to the case without manifold learning same as the InversionNet.  
2.3 Latent space  translation is said to be identity for the InversionNet and U-net for the Latent U-net, but couldn’t we also just claim some middle layers of the InversionNet to perform latent space translation?  
2.4 The size of the latent space is low in both cases.  
--> So to my understanding the difference between the two boils solely down to the architectural design and size of the network but not to any conceptual differences. If that is the case, it should be stated more directly in the paper and the exact architectural differences should be explained in more detail. 
3. (Clarity) From Section 5.1.4 it is not entirely clear what the answer to the question raised in the paragraph (Is it useful to jointly solve forward and inverse problem) is. Is the answer that no it is not useful for learning the inverse map but yes it is useful for learning the forward map (both answers based on the findings in Figure 3)?

Minor concerns

4. (Clarity) Please provide links to the exact Section in the Appendix rather than just referencing the entire appendix (e.g. lines 308, 311, 314, 377, 388). 
5. (Significance) Include quantitative scores in the qualitative examples (e.g. in Figures 4 and 8) to give a better understanding of how differences in scores translate to significantly perceptible differences.

---

confidence: 5
contribution: 2
initial_rating: 5
presentation: 3
questions: Invertible Neural Networks (INNs) often have limited flexibility in adjusting the dimensions of input and output, the number of layers, filters, and convolutional kernel sizes. These constraints can impede performance for complex tasks.

When mapping between velocity maps and seismic waveforms—data from two distinct domains (one spatial and the other spatio-temporal), why Invertible X-Net could demonstrate superior performance?
soundness: 3
strengths: 1. This work is the first to investigate key factors influencing mapping performance.
2. Extensive experiments were conducted to conclude the influence of these factors.
3. Two novel networks are introduced for mapping between velocity maps and seismic waveforms, achieving SOTA performance.
summary: This work summarizes recent advances in deep learning-based mapping between velocity maps and seismic waveforms, introducing a unified framework (GFI) to systematically characterize prior research. Using GFI, key factors affecting mapping performance are investigated, alongside the proposal of two novel model architectures. Comprehensive experiments reveal the impact of these factors, with the proposed models achieving state-of-the-art (SOTA) performance in both forward and inverse problems.
weaknesses: 1. This paper resembles a survey and might be more suitable for journal submission.
2. A significant portion of the content is placed in the appendices, making it less reader-friendly.
3. Aside from extensive experiments on factor influences, the work offers limited novelty and lacks theoretical analysis.
4. Additional factors should be investigated, such as the robustness of data noise across different network architectures.

Meta-review: This paper proposes a joint forward--inverse solver for wave imaging problems using two autoencoders whose latent spaces (of possibly different dimensions) are paired via learned "latent space translations". The reviewers praised clarity of presentation, and extensive experiments. Initially there was some concern about the contribution and vagueness, but this has been in good part eliminated in the rebuttal phase. My opinion is that this is a nice paper very slightly above the acceptance bar. It is clear, the prior work section is very nice, and the experiments are complete and informative. I will recommend acceptance but I would not mind if it ends up being rejected. The reason for hesitation is that there is a lot of very related work which has not been acknowledged or compared to, for example

https://arxiv.org/pdf/2405.13220
https://arxiv.org/pdf/2305.13314

but also that I think that as a community we should finally go beyond OpenFWI benchmarks with 2D 70x70 pixel, highly stylized, layered velocity models which have a handful of degrees of freedom.

Decision: Accept (Poster)