---
title: "Neural Operators with Accurate Jacobian for High-Fidelity Image-Domain Seismic Inversion"
header-includes:
  - \usepackage{amsmath}
  - \usepackage{amssymb}
author:
  - name: Jeongjin Park
    affiliations:
      - name: Georgia Institute of Technology
  - name: Huseyin Tuna Erdinc
    affiliations:
      - name: Georgia Institute of Technology
  - name: Grant Bruer
    affiliations: 
      - name: Georgia Institute of Technology
  - name: Richard Rex
    affiliations:
      - name: Terra AI
  - name: Nisha Chandramoorthy
    affiliations:
      - name: University of Chicago
  - name: Felix J. Herrmann
    affiliations:
      - name: Georgia Institute of Technology
bibliography: abstract.bib
crossref:
  fig-prefix: Figure
  eq-prefix: Equation
filters:
    - SLIM
    - critic-markup
number-sections: true
format:
  html:
    page-layout: full
    sidebar: false
    lightbox: true
    crossrefs-hover: true
  pdf:
    template: IMAGEAbstractTemplate.latex
    # csl: apa.csl
    cite-method: natbib
    keep-tex: true
    pdf-engine: pdflatex
abstract: Neural operators (NOs) provide efficient surrogates for wave-equation-based imaging operators, but accurate forward prediction alone does not guarantee reliable inversion. As inversion updates depend on the adjoint Jacobian, errors in learned sensitivities lead to unphysical updates and degraded quality in recovered models. To this end, we propose a Jacobian-informed training strategy that enforces accuracy in Jacobian in the directions where gradient descent occurs during Least-Squares. Numerical experiments with Fourier Neural Operators demonstrate that supervising only a small fraction ($0.1 \sim 0.01\%$) of Jacobian information significantly improves gradient quality and leads to more accurate subsurface velocity reconstruction. In particular, the proposed approach reduces mid to deep layer artifacts and improves structural continuity compared to models trained only with forward misfit term. These results highlight that learning inversion-relevant sensitivities, rather than forward mappings alone, is critical for reliable surrogate-assisted seismic inversion.
---

# Introduction

In seismic inversion, Image-Domain Least-Squares Migration (IDLSM) constructs subsurface images by solving a linearized inverse problem based on the Born modeling operator and its adjoint [@symes2008migration]. While producing high-resolution subsurface images, these methods remain computationally expensive due to repeated large-scale wave-equation solves, with cost scaling with acquisition size [@tarantola1984inverse; @virieux2009overview; @louboutin2017full]. To mitigate these costs, recent work has explored Machine Learning approaches that accelerate components of the seismic imaging workflow [@geng2022deep; @orozco2024machine; @yin2024wise; @yin2025wiser; @erdinc2025power]. In this work, we train a neural surrogate for a nonlinear, image-domain operator induced by migration. Specifically, we approximate an imaging operator that maps subsurface and background-velocity models to an RTM image:
$$
\mathcal{R}: (\mathbf{m}, \mathbf{m}_0) \mapsto \delta \mathbf{m}_{\mathrm{RTM}},
$$ {#eq-forward}
where $\mathbf{m}$, $\mathbf{m}_0$ and $\delta \mathbf{m}_{\mathrm{RTM}}$ are the subsurface velocity model, smooth background-velocity model and the migrated image, respectively. We approximate this mapping using neural operators, which learn mappings between function spaces governed by PDEs [@kovachki2023neural; @FNO_one]. Once trained, these surrogates enable fast foward prediction and gradient computation via automatic differentiation [@siahkoohi2022velocity; @ma2025velocity], replacing the numerical simulator and potentially providing efficient Gauss–Newton approximations for preconditioning.

However, accurate forward prediction alone does not ensure that the learned operator captures the derivative information required for inversion [@park2024dynamical; @rosofsky2023applications; @li2024physics; @o2024derivative; @park2025reduced]. Since gradient-based updates are governed by the adjoint Jacobian [@tarantola1984inverse; @virieux2009overview], errors in the learned sensitivities lead to inaccurate gradients, unphysical updates, and degraded reconstruction quality (see rightmost column of @fig-recovered). This highlights the need to incorporate inversion-relevant sensitivity information during training.

![Schematic of the forward problem. The neural operator learns $\mathcal{R}_{\mathrm{nn}}: (\mathbf{m}, \mathbf{m}_0) \mapsto \delta \mathbf{m}_{\mathrm{RTM}}$.](figs/forward_obj2.png){#fig-forwardobs width=99%}

![Inversion workflow using $\mathcal{R}_{\mathrm{nn}}$. Starting from an initial guess, $\mathcal{R}_{\mathrm{nn}}$ predicts RTM images used to compute a misfit, and the velocity model is updated iteratively.](figs/inverse_obj2.png){#fig-inverseobs width=99%}

To address this limitation, we introduce a Jacobian-informed training strategy that incorporates sensitivity information from Least-Squares Migration. Specifically, we supervise the neural operator along a small number of well-illuminated directions in model space, obtained from the dominant eigenvectors of the Gauss–Newton Hessian averaged over the training distribution. This yields a set of globally informative directions that capture the dominant sensitivity structure of the inverse problem. We further construct a dataset with variability in both subsurface and background-velocity models to enable inversion-oriented learning. Numerical experiments with Fourier Neural Operators (FNOs) [@FNO_one] show that the proposed approach produces more physically-aligned gradients, leading to improved recovery of subsurface velocity structures and reduced artifacts compared to standard training algorithm.

# Theory: Informative Directions for Inversion

In the following subsections, we define the forward problem and identify directions that drive gradient-based inversion.

## Linearized Modeling and RTM

Let $\mathbf{m}$ define the discretized subsurface velocity model. When the forward operator, $\mathcal{F}$, is linearized around a background model $\mathbf{m}_0$ and $\mathbf{d}_\text{obs}$ is seismic data, the RTM image at $\mathbf{m}_0$ is
$$\delta \mathbf{m}_{\mathrm{RTM}} = \mathbf{J}(\mathbf{m}_0)^\top (\mathcal{F}(\mathbf{m}_0) - \mathbf{d}_\text{obs}),$$
where $\mathbf{J}(\mathbf{m}_0)$ denotes the Born modeling operator.


## Direction of gradient descents during inversion

During gradient-based seismic inversion, not all directions in the model space are equally constrained by the data [@tarantola1984inverse; @virieux2009overview], due to limited acquisition geometry and band-limited sources. As a result, when performing gradient descent without additional prior information (e.g., maximum likelihood estimation), only a subset of model perturbations, which coincides with as well-illuminated directions, can be reliably recovered. That is, in Least-Squares Migration (LSM), the objective function and its gradient are given by
$$
\begin{aligned}
\mathcal{L}(\mathbf{m}) &= \frac{1}{2}\|\mathcal{R}(\mathbf{m}) - \delta \mathbf{m}_\mathrm{RTM} \|^2_2 \\
\nabla_\mathbf{m} \mathcal{L}(\mathbf{m}) &= \Bigl(\frac{\partial \mathcal{R}(\mathbf{m})}{\partial \mathbf{m}} \Bigr)^\top \mathbf{r}(\mathbf{m}) 
\end{aligned}
$$
where $\mathbf{r}(\mathbf{m}) = \mathcal{R}(\mathbf{m}) - \delta \mathbf{m}_\mathrm{RTM}.$ The curvature of this objective function, which characterizes how perturbations in the model affect the data misfit and therefore which components of the model can be reliably recovered, is characterized by the Gauss–Newton Hessian (GNH),
$$ \mathbf{H}_\text{GN}(\mathbf{m}) = \mathbf{J}(\mathbf{m})^\top\mathbf{J}(\mathbf{m}).$$ To identify the shared directions that primarily drive inversion updates across varying models, we compute the dominant eigenvectors of the Gauss-Newton Hessian,
$$
\begin{aligned}
&\bar{\mathbf{H}}_\text{GN} = \mathbb{E}_\mu \Bigl[ \mathbf{J}(\mathbf{m})^\top\mathbf{J}(\mathbf{m}) \Bigr] \\ &\bar{\mathbf{H}}_\text{GN}\bar{\mathbf{v}}_j = \lambda_j \bar{\mathbf{v}}_j,
\end{aligned}
$$ {#eq-GlobalGNH}
where $j$ indexes the eigenvector and the expectation is taken over a distribution $\mu$ of groundtruth and background models. These eigenvectors correspond to well-illuminated directions in model space that dominate inversion updates. We refer to them as globally informative directions, and are used to augment the dataset. By computing a shared set of globally informative directions, we significantly reduce computational cost while retaining the key sensitivity structure relevant for inversion.

<!-- (Claerbout, 1985; Nemeth et al., 1999; Schuster, 2017). -->


## Jacobian action along informative directions

To supervise the Jacobian of the neural operator during training, we compute the action of the Jacobian along the $\bar{\mathbf{v}}_j$. Using $\mathcal{R}(\mathbf{m}) = \mathbf{J}^\top \mathbf{r}(\mathbf{m})$, the differential satisfies
$$d\mathcal{R}(\mathbf{m}) \approx - \mathbf{J}^\top\mathbf{J} d\mathbf{m},$$
which corresponds to the action of the GNH. By probing along $\bar{\mathbf{v}}_j$, we obtain derivative information that reflects the most informative directions that observation can give on the model. 

# Methodology

![Background models $\mathbf{m}_0$ generated by smoothing randomly in time or depth, leading to shifts in RTM images and enabling generalization over different $\mathbf{m}_0$.](figs/dataset.png){#fig-background2 width=90%}


## Jacobian-informed dataset with variability in $\mathbf{m}_0$ {#sec-data}

<!-- RTM is based on a linearization of the forward model around a background-velocity model $\mathbf{m}_0$; thus, the migrated image strongly depends on $\mathbf{m}_0$.  -->

![Vertical velocity traces comparing groundtruth and varying background models](figs/trace_comparison_817.png){#fig-data-trace width=80%}

To ensure that the learned surrogate generalizes across different models, we construct a dataset that captures variability in both the groundtruth and background model. We draw pairs $(\mathbf{m}, \mathbf{m}_0) \sim \mu$ where  $\mathbf{m}_0$ is generated by applying randomized smoothing operations to the slowness field in either the depth or time domain (@fig-background2). These perturbations introduce kinematic variability, resulting in shifts and distortions of migrated events (@fig-background2). In addition to RTM images, $\delta\mathbf{m}_{\mathrm{RTM}}$, we include action of Jacobian along a small set of informative directions shared across all models $\bar{\mathbf{V}}_r = [\bar{\mathbf{v}}_1, ..., \bar{\mathbf{v}}_r]$, where $r$ denotes the number of directions used. The resulting dataset is given by
$$
\mathcal{D} =
\left\{
\left(
\mathbf{m}^{(i)}, \mathbf{m}_0^{(i)}, \delta \mathbf{m}_{\mathrm{RTM}}^{(i)}, \bar{\mathbf{V}}_r,
\left(\frac{\partial \mathcal{R}}{\partial \mathbf{m}}\bigg|_{\mathbf{m}^{(i)}}\right)\bar{\mathbf{V}}_r
\right)
\right\}_{i=1}^{N}
$$
where $i$ indexes training samples. This dataset enables training of a neural operator that is amortized across different groundtruth and background models, while learning accurate Jacobian along the directions that govern inversion updates.

![Velocity models recovered from a test sample through LS inversion using different neural surrogates. ](figs/obj_final.png){#fig-recovered fig-env="figure*" fig-pos="t"  width=99%} 

## Jacobian-informed training

With this training dataset, we train $\mathcal{R}_{\mathrm{nn}}(\mathbf{m}, \mathbf{m}_0)$ to approximate the generalized RTM mapping as defined in @eq-forward. In addition to matching the RTM images, we enforce accuracy in the learned Jacobian along the $\bar{\mathbf{V}}_r$. This encourages the neural operator to learn sensitivities that are consistent with the physical inversion process. The training loss is:
$$
\begin{aligned}
\mathcal{L}_\text{train}(\mathbf{m}, \mathbf{m}_0) = &\|\mathcal{R}_{\mathrm{nn}}(\mathbf{m}, \mathbf{m}_0) - \delta \mathbf{m}_{\mathrm{RTM}}\|^2_2 \\ 
&+ \lambda \|\Bigl(\frac{\partial \mathcal{R}_{\mathrm{nn}}}{\partial \mathbf{m}} \Bigr) (\mathbf{m}, \mathbf{m}_0)\bar{\mathbf{V}}_r - \Bigl(\frac{ \partial \mathcal{R}}{ \partial \mathbf{m}} \Bigr) (\mathbf{m}, \mathbf{m}_0)\bar{\mathbf{V}}_r\|^2_2.
\end{aligned}
$$
This regularization term encourages the surrogate to learn inversion-relevant sensitivities, leading to more accurate gradient directions during inversion.


## Least-Squares inversion with $\mathcal{R}_{\mathrm{nn}}$

After training, the neural surrogate, $\mathcal{R}_{\mathrm{nn}}$, is used within a least-squares inversion framework. Given an observed RTM image $\delta \mathbf{m}_{\mathrm{RTM}}$, we estimate the velocity model by solving
$$\min_\mathbf{m} \|\mathcal{R}_{\mathrm{nn}}(\mathbf{m}, \mathbf{m}_0) - \delta \mathbf{m}_{\mathrm{RTM}}\|^2_2 \: \: \text{s.t.} \: \|\nabla \mathbf{m}\|_1 \leq \tau,$$
where the total variation (TV) constraint promotes piecewise-smooth structure while preserving sharp interfaces and $\tau$ controls the strength of the TV constraint [@esser2018total; @Bas_TV].

# Numerical Experiments

We compare standard neural operators trained only with forward prediction loss against the proposed Jacobian-informed training. 

<!-- Our goal is to assess whether a learned surrogate $\mathcal{R}_{\mathrm{nn}}$ can reliably replace the physical operator $\mathcal{R}$ in least-squares inversion, and especially, whether Jacobian supervision improves gradient quality and reconstruction. -->

## Training Setup

To create the dataset, we use Julia package, JUDI.jl [@witteJUDI2019; @mathias_louboutin_2023_7785440]. Velocity models and RTM images are discretized to 256 x 256 [@BG]. In this work, we use 202 training samples and 22 test samples. 

We evaluate our training algorithm using FNO [@FNO_one], which models operators via global convolutions in Fourier space. We train two variants:
*(i)* a standard FNO trained with forward loss, and
*(ii)* a Jacobian-informed FNO (JI-FNO) with additional Jacobian supervision. To understand the impact of the number of directions in Jacobian we supervise, we train JI-FNO with 8 and 64 directions, which is 0.01% and 0.1% of the model of size, 256 x 256. Detailed schematic for training objective is shown in @fig-forwardobs.

## Inversion setup

![Data residuals for different models, showing improved data fitting with Jacobian-informed training.](figs/metric_data_relL2 (3).png){#fig-data_residual width=99%}

After training, $\mathcal{R}_{\mathrm{nn}}$	​replaces the physical operator within the inversion loop (@fig-inverseobs). For each test sample, inversion is performed by updating the velocity model while keeping the background model fixed [@ma2025velocity]. The initial model is a smoothed version of the true velocity. Because the surrogate is differentiable, gradients can be computed efficiently via automatic differentiation, making each iteration significantly cheaper than wave-equation-based approaches. This allows us to run 1200 iterations of gradient-based optimization using Adam [@kingma2014adam].

# Result

::: {.content-visible when-format="html"}
| Metric | FNO | JI-FNO | JI-FNO |
|---|---:|---:|---:|
| $r$ | N.A. | 8 | 64 |
| Model Error<br>(Rel. $L_2$) | $0.29$ | **0.21** | **0.21** |
| Data Residual<br>(Rel. $L_2$) | $0.28$ | $0.19$ | **0.12** |
| Recovered Model<br>SSIM | $0.50$ | $0.68$ | **0.76** |

: Inversion metrics evaluating velocity model fidelity and data residuals across different training algorithms on a test sample. The number of directions used for the JI-FNO framework is denoted as $r$. {#tbl-inversion-metrics}
:::

::: {.content-visible when-format="pdf"}
\begin{table}
\centering
\caption{Inversion metrics evaluating velocity model fidelity and data residuals across different training algorithms on a test sample. The number of directions used for JI-FNO framework is denoted as $r$}
\label{tbl-inversion-metrics}
\begin{tabular}{l | l l l l}
\hline
Metric & FNO & JI-FNO & JI-FNO \\
\hline
$r$ & N.A. & 8 & 64 \\
\hline
\begin{tabular}[c]{@{}l@{}}Model Error \\ (Rel. $L_2$)\end{tabular} & $0.29$ & $\mathbf{0.21}$ & $\mathbf{0.21}$ \\
\begin{tabular}[c]{@{}l@{}}Data Residual \\ (Rel. $L_2$)\end{tabular} & $0.28$ & $0.19$ & $\mathbf{0.12}$ \\
\begin{tabular}[c]{@{}l@{}}Recovered Model \\ SSIM\end{tabular} & $0.50$ & $0.68$ & $\mathbf{0.76}$ \\
\hline
\end{tabular}
\end{table}
:::

@tbl-inversion-metrics summarizes the inversion performance across all neural surrogates. Incorporating Jacobian supervision consistently improves both data fitting and model reconstruction quality. The JI-FNO reduces the relative model error by 8% and achieves significantly lower data residuals (decreases by 16%) compared to the standard FNO. In addition, structural similarity (SSIM) improves substantially (increases by 26%), indicating better recovery of subsurface features. These results show that aligning the surrogate with inversion-relevant sensitivities leads to more accurate gradient updates and improved convergence (as shown in @fig-data_residual). 

![Vertical velocity trace comparison at lateral position 2,800 m comparing the recovered models against the ground truth.](figs/recovered_model_traces_224.png){#fig-trace width=99%}

@fig-recovered compares velocity models recovered using neural surrogates trained with different training strategies. The standard FNO (the right most column) produces reconstructions with noisy artifacts from mid to deep layers and exhibits distortions in the deeper high-velocity layers. In contrast, the Jacobian-informed models yield significantly cleaner images, with improved layer continuity and more accurate recovery at depth. Notably, even when using only 8 eigenvectors, the Jacobian-informed model already leads to substantial improvements over the baseline. Increasing the number of eigenvectors to 64 further enhances layer sharpness and continuity. These improvements are more pronounced in deeper regions (@fig-trace).

We further analyze the gradients produced by each model (second row of @fig-recovered). Gradients from the standard FNO are dominated by high-frequency noise and lack coherent spatial structure. In contrast, the Jacobian-informed models produce gradients that are more structured and aligned with reflector geometry, particularly in the shallow to mid-depth regions. This difference explains the improved inversion performance: by learning gradients that follow the physically meaningful directions of the inverse problem, the proposed method produces more stable and accurate model updates.


# Conclusion and discussion

We showed that accurate forward prediction alone is insufficient for surrogate-assisted seismic inversion, as inversion updates depend on the adjoint Jacobian. If a neural operator learns incorrect sensitivity, gradient-based updates may introduce unphysical artifacts and fail to recover deeper structures. By incorporating Jacobian supervision during training, the proposed approach ensures that the learned operator produces gradients consistent with the physical inversion geometry. Numerical experiments demonstrate that the proposed approach improves gradient quality, reduces artifacts, and yields more accurate recovery of deeper velocity structures compared to models trained with forward prediction loss alone. These results highlight that reliable surrogate modeling for seismic inversion requires learning inversion-relevant geometry, not just forward mappings. 

Also, we observed as we add more directions during Jacobian-informed training, the surrogate Jacobian during inversion becomes increasingly accurate in model recovery. We therefore expect the benefits of the proposed approach to become even more pronounced in more challenging inversion scenarios, such as cases with poorer background models or limited training data, which we will continue to explore in the future work. 

Lastly, the proposed approach has direct implications for Gauss–Newton FWI. Improving the accuracy of the learned Jacobian enables more accurate curvature information and search directions. This opens the possibility of using neural surrogates as data-driven nonlinear preconditioners, reducing the number of Gauss-Newton iterations while preserving inversion accuracy.

# Acknowledgement

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

# References {.unnumbered}

::: {#refs}
:::