Efficient Numerical Wave Propagation Enhanced By An End-to-End Deep Learning Model

An efficient numerical solver $\mathcal{G}_{\Delta t}\mathfrak{u}\equiv\mathcal{G}_{\Delta t}[\mathfrak{u},c]$ is used to propagate a wave $\mathfrak{u}(x,t)=(u,\partial_{t}u)$ for a time step $t+\Delta t$ on a medium described by the piecewise smooth wave speed $c(x)$ for $x\in[-1,1]^{2}$. This method is computationally cheap since the advancements are computed on a coarse grid using a large time step within the limitation of numerical stability; however, it is consistently less accurate than an expensive fine solver $\mathcal{F}_{\Delta t}\mathfrak{u}\equiv\mathcal{F}_{\Delta t}[\mathfrak{u},c]$. Consequently, the solutions $\mathcal{G}_{\Delta t}\mathfrak{u}$ exhibit numerical dispersion errors and miss high-fidelity details. In a supervised learning framework, we aim to reduce this discrepancy using the outputs from $\mathcal{F}_{\Delta t}$ as the examples.

We define a restriction operator $\mathcal{R}$ which transforms functions from a fine grid to a coarse grid. Additionally, for mapping coarse grid functions to a fine grid, we integrate a neural network $\mathcal{I}^{\theta}$ to augment the under-resolved wave field. We can now define a neural propagator $\Psi_{\Delta t}[\mathfrak{u},c,\theta]\equiv\Psi_{\Delta t}^{\theta}$ that takes a wave field $\mathfrak{u}$ defined on the fine grid, propagates it on a coarser grid, and returns the enhanced wave field on the fine grid,

The models are parameterized by the family of initial wave fields $\mathfrak{F}_{\mathfrak{u}_{0}}$ and wave speeds $\mathfrak{F}_{c}$.

As the two solvers operate on different Cartesian grids, with $\delta x<\Delta x$ and $\delta t<\Delta t$, we define the restriction operator $\mathcal{R}$, which transforms functions from a fine to a coarse grid, and the prolongation operator $\mathcal{I}$, which maps the inverse relation. The enhanced variants consist of (a) bilinear interpolations denoted as $\mathcal{R}$ and $\mathcal{I}^{0}$, while $\mathcal{I}^{0}\mathcal{R}\mathfrak{u}\neq\mathfrak{u}$, and (b) neural network components denoted as $\mathcal{I}^{\theta}\equiv\Lambda^{{\dagger}}\mathcal{I}_{\Delta t^{\star}}^{~{}\theta}\Lambda$, while the lower index indicates that the neural networks are trained when the step size $\Delta t^{\star}$ is used. For improved neural network inference, we use the transition operator $\Lambda$ to transform physical wave fields $(u,u_{t})$ to energy component representations $(\nabla u,c^{-2}u_{t})$, with $\Lambda^{{\dagger}}$ as the pseudo-inverse (see also nguyen2022numerical ; energy1 ). Figure 1 provides a schema visualizing the wave argument transitions.

A simple model with bilinear interpolation (E2E-V, $\mathcal{I}^{0}\Psi_{\Delta t^{\star}}^{\theta}\mathcal{R}$) is used as a baseline. Each variant changes the baseline by exactly one aspect. This allows us to isolate the effect of each architecture modification. The four investigated end-to-end models $\Psi_{\Delta t^{\star}}^{\theta}\coloneqq\mathcal{I}_{\Delta t^{\star}}^{~{}\theta}\mathcal{G}_{\Delta t^{\star}}\mathcal{R}$ are:

E2E-JNet3:

E2E 3-level JNet (Figure 1)

E2E-JNet5:

E2E 5-level JNet

E2E-Tira:

Tiramisu JNet tiramisu

E2E-Trans:

Transformer JNet petit2021unet

The second baseline is taken from nguyen2022numerical and denoted as the modular, not end-to-end 3-level JNet (NE2E-JNet3), $\mathcal{I}_{\Delta t^{\star}}^{~{}\theta}(\mathcal{G}_{\Delta t^{\star}}\mathcal{R}$), while results of $\mathcal{G}_{\Delta t^{\star}}$ are used to separately train the E2E-JNet3 upsampling component.

$c\sim\mathfrak{F}_{c}$ are sampled from randomly chosen subregions of two synthetic geological structures, Marmousi marmousi and BP bp , that are mapped onto the spatial grid $h\mathbb{Z}^{2}\cap[-1,1]^{2}$ (see Figure 2). Four manually modified media (cf. nguyen2022numerical ) are added during testing to examine rapid variations in wave speed.

The average training time of each variant is approximately $73$ CPU core hours. Due to resource constraints, we therefore limit our main analysis to one end-to-end variant. Here, we selected the most promising approach from four deep learning architectures trained on $\mathcal{D}$.

We train the chosen end-to-end variant from experiment 1 on $\mathcal{D}^{m}$ using an equal number of training points as in $\mathcal{D}$. The test set is consistent with $\mathcal{D}$ to enable comparison with other experiments.

The setup follows experiment 2, while the models are trained on $\mathcal{D}^{w,m}$.

We explore improvements to our variants using the Parareal scheme in two datasets:
A. Comprehensive Training ($\mathcal{D}^{p}_{\text{train}}$): The models are trained according to the Parareal scheme in Eq. 7 and Eq. 8 with $K=4$ using a random sample that constitutes a quarter of the original dataset $\mathcal{D}$ for fair comparisons. The gradients are determined by summing the losses of a Parareal iteration.
B. Fine-tuning ($\mathcal{D}^{p}_{\text{refine}}$): Rather than employing an un-trained model, we deploy variants that were pre-trained on a random subset containing half of $\mathcal{D}$. Then, for another subset that constitutes an eighth of $\mathcal{D}$, $\Psi_{\Delta t}^{\theta}$ is applied with Parareal.

The first important observation based on Figure 3 is that integrating NE2E-JNet3 into a single, end-to-end system (E2E-JNet3) improved the average accuracy on the validation set by more than $46\%$, and on the test set by ca. $53\%$. The ability to include the loss of both the coarse solver and downsampling layer also caused a lower standard deviation and fewer outliers, since the mean is well above the median for NE2E-JNet3 compared to E2E-JNet3.

Introducing a multi-step training loss enhanced the benefits of an end-to-end architecture even further (cf. E2E-JNet3 ($\mathcal{D}^{m}$) in Figure 3) without increasing the number of model parameters. Figure 4 depicts how all end-to-end models had a much lower relative energy MSE (cf. nguyen2022numerical ) increase particularly for the first three time steps on the test set. Hence, we conclude that connecting wave states to incorporate temporal propagation dynamics in the training data appears to be especially important for the early stages of wave advancements. Additionally, by taking fewer steps through sampling from a normal distribution that is being shifted along the x-axis (cf. E2E-JNet3 ($\mathcal{D}^{w,m}$)), we successfully avoid high performance fluctuations when the model is only partially trained. Figure 5 visualizes the correction of the low-fidelity solution of $\mathcal{G}_{\Delta t^{\star}}$ by E2E-JNet3 ($\mathcal{D}^{w,m}$).

An overview of the upsampling architecture performances can be found in Table 1. As expected, the larger networks (E2E-Tira and E2E-Trans) performed slightly better compared to the 3-level JNet architecture, but for the ResNet architecture (E2E-JNet5), more weights did not increase accuracy by much. Consequently, we theorize that the ResNet design may be insufficient for capturing high-fidelity wave patterns, while especially highly-connected layers with an optimized feature and gradient flow (E2E-Tira) are better suited. Given that E2E-JNet3 ($\mathcal{D}^{w,m}$) had only a slightly worse average energy MSE on the test set, we generally advise against using the expensive models in our setup.

While models trained with $\mathcal{D}^{p}_{\text{train}}$ have an unstable training progress and diverging loss, applying E2E-JNet3 ($\mathcal{D}^{p}_{\text{refine}}$) within the Parareal scheme showed better accuracy than E2E-JNet3 with Parareal (cf. Figure 6). As this training method improved the stability of Parareal, sampling the causality of concurrently solving multiple time intervals is an efficient enhancement to our end-to-end structure.