Multigrid-augmented deep learning preconditioners for the Helmholtz equation
Abstract
We present a data-driven approach to iteratively solve the discrete heterogeneous
Helmholtz equation at high wavenumbers.
We combine multigrid ingredients with convolutional neural networks (CNNs) to form a preconditioner which is applied within a Krylov solver.
Two types of preconditioners are proposed 1) U-Net as a coarse grid solver, and 2) U-Net as a deﬂation operator with shifted Laplacian
V-cycles.
The resulting CNN preconditioner can generalize over residuals and a relatively general set of wave slowness models.
On top of that, we offer an encodersolver framework where an “encoder” network generalizes over the medium and sends context vectors to another “solver” network, which generalizes over the right-hand-sides.
We show that this option is more efﬁcient than the stand-alone variant.
Lastly, we suggest a mini-retraining procedure, to improve the solver after the model is known.
We demonstrate the efﬁciency and generalization abilities of our approach on a variety of 2D problems.
Introduction
The efﬁcient numerical solution of the Helmholtz equation with a high and spatially-dependent wavenumber is a difﬁcult task.
The linear system that results from the discretization of the Helmholtz equation involves a large, complex-valued, sparse, and indeﬁnite matrix.
Over the years, various iterative methods were developed for solving the problem, e.g., [18, 7, 8, 25, 5, 27, 6, 16, 21], yet it remains expensive to solve.
One of the most common methods to solve it is the Shifted Laplacian (SL) multigrid method [6, 28], which is based on a multigrid preconditioner for a damped version of the equation.
That is, a complex shift is added to the matrix to form a preconditioner that is efﬁciently solved by multigrid components.
The preconditioning is applied with a Krylov solver, such as GMRES [23].
However, the SL preconditioner is not scalable due to the high shift that is required to solve the shifted system, and further investigation is needed.
In recent years, Deep Learning (DL) methods have revolutionized many ﬁelds, like computer vision and natural language processing [14].
Similarly, solving partial differential equations (PDEs) with a deep neural network has been recently considered in two main forms.
One common class of
DL-for-PDEs approaches, which are often referred to as physics-informed neural networks (PINNs), includes both PDE solvers [19, 20, 15, 9, 1], and PDE discovery [2].
The general idea is to implement a neural network that implicitly represents the PDE solution.
That is, the solution is approximated by the neural network as an analytical function at a given point in space.
This function is obtained using stochastic optimization for learning the network’s weights as an unsupervised learning task, given a large training set of points.
The loss function that is minimized forces the network to meet both the system of equations and the initial and/or boundary conditions (BC).
Following the learning process, the PDE coefﬁcients and BC are assumed to be embedded into the neural network weights.
Another family of approaches involves solvers in which DL is used to approximate the solution of the PDE in an end-to-end fashion given the problem properties (coefﬁcients, initial conditions or sources).
That is, the network needs to generalize over some of the PDE properties, which is the more common scenario in computer vision applications.
Some works employ Convolutional Neural
Networks (CNNs) which are typically used for applications involving data on regular grids, like images and videos, or, discretized PDEs on a structured grid.
Khoo and Ying [11] introduced a CNN named ’SwitchNet’, for approximating forward and inverse maps arising from the time-harmonic wave equation, assuming plane-wave sources and 2-4 Gaussian scatterers in a constant medium.
However, no neural network has been presented in the literature that effectively solves the Helmholtz equation at high wavenumber for rather general heterogeneous models and right-hand-sides.
In this work we utilize CNNs as preconditioners for the discretized Helmholtz equation.
We wish to deﬁne a solver that generalizes over right-hand-sides and slowness models.
We propose to use a CNN together with a classical multigrid approach—the Shifted Laplacian—which is efﬁcient at removing only part of the error.
We build on the similarity between a V-cycle and a U-Net [22] which is a multiscale CNN used for image-to-image mappings such as semantic segmentation, image denoising etc.
Indeed, a geometric V-cycle can be naturally implemented using DL frameworks, suggesting that a U-Net architecture, like a V-cycle, can be used as a preconditioner for solving
PDEs on a regular grid.
To precondition the Helmholtz equation, we augment the U-Net with a classical approach, to complement the network and stabilize the solution process.
We examine two preconditioners: one is a U-Net that acts as a non-linear coarse grid correction and is applied with a
Jacobi iteration for pre-and post-smoothing.
In the second formulation we apply alternating U-Net and SL multigrid V-cycle, which provides a more powerful smoothing and somewhat resembles the deﬂation preconditioner of [24, 5].
Both preconditioners are applied within a Krylov subspace method—ﬂexible GMRES.
We deﬁne our U-Net to generalize over a random right-hand-side and a random piece-wise smooth slowness model.
We also suggest two upgrades to the scheme above: (1) an encoder-solver framework and (2) a mini-retrain approach.
Both upgrades are described in details later, and can improve the standard U-Net tremendously.
Preliminaries and background
The heterogeneous Helmholtz problem is give by
[Equation]
The unknown u ( (cid:126)x, ω ) represents the pressure wave function in the frequency domain, ω =
πf denotes the angular frequency, ∆ is the Laplacian operator and κ ( (cid:126)x ) is the heterogeneous wave slowness - the inverse of its velocity.
The source term is the right-hand-side of the system, g .
The parameter γ indicates the fraction of attenuation (we focus on the hardest case γ = 0 ), and i =
√− 1 .
As boundary conditions, we consider an absorbing layer [6], but PML [26] or [17] can be viable options as well.
We consider the second order ﬁnite difference discretization of (2.1) on a uniform mesh of width h in both.
This leads to a system of linear equations
[Equation]
When discretizing the Helmholtz equation, we must obey the rule of thumb that at least 10 grid nodes per wavelength are used.
This typically requires a very ﬁne mesh for high wavenumbers, signiﬁcantly increasing the number of unknowns and making the system (2.2) ill-conditioned, in addition to being indeﬁnite.
Hence, the solution of the problem can be challenging and require creative iterative solvers.
2.1 Shifted Laplacian multigrid
Eq.
(2.2) is commonly solved by a preconditioned Krylov subspace method.
The SL operator
[Equation]
is used to accelerate the convergence of a Krylov subspace method for solving (2.2).
The preconditioning matrix M h is obtained from the discretization of (2.3), similarly to (2.1) yielding A h .
In this paper we focus on the pair α = 1 and β = 0 .
5 , which in [6] is shown to lead to a good compromise between approximating (2.2) and our ability to solve the (2.3) using multigrid tools.
The SL multigrid method is consistent and robust for non-uniform models, but is slow and for large wavenumbers.
For details on the geometric multigrid method used here to approximately invert (2.3) please see Sec.
A in our supplementary material.
Multigrid-augmented deep learning preconditioners
We utilize a U-Net architecture to accelerate the convergence of the Helmholtz solution.
In order for the computational complexity to be cost-effective we deﬁne the architecture to be relatively small in DL standards (see Fig. 1), and yet demonstrate a signiﬁcant acceleration to the solution.
Our
DL-based preconditioner is applied within the ﬂexible GMRES Krylov method [23].
The ﬂexible variant is important as our preconditioner is non-linear, hence also non-stationary.
The U-Net architecture [22] is a CNN for imageto-image tasks such as semantic segmentation.
In our case, given a slowness model and a righthand-side sampled on a grid (viewed as images), we wish to predict a solution for (2.2) on that grid.
On its way down the hierarchy, the architecture interchangeably applies downsampling convolutions and ResNet convolution blocks.
Then, going up the hierarchy it applies up-sampling convolutions.
The down- and upsampling are applied using strided convolutions, similarly to multigrid transfer operators.
U-Net as a preconditioner accelerator Since the medium can be very heterogeneous, the solution, even for a point source or a plane wave, can be highly complex at high wavenumber due to reﬂections and interference.
Therefore, we do not pursue a solution via a single application of a network, but aim for a network to be applied several times in the form of a preconditioner.
Speciﬁcally, we aim for the U-Net to complement the shifted Laplacian multigrid method, which is efﬁcient at handling error modes with rather large absolute eigenvalues.
We aim that the U-Net will focus on the these error directly, and it will be accompanied by classical smoothing via SL.
This guides our training procedure, which is described in Sec.
B in the supplementary material.
Furthermore, the smooting is important because the Krylov method is sensitive to high residuals.
Below we list two preconditioning schemes:
U-Net as a coarse grid solver in a two-level cycle In our ﬁrst version, denoted as M JU , we apply one pre- and post-
Jacobi relaxations with a U-Net in between.
Since the U-Net starts and ends with strided convolutions, this ends up being similar to a two-level cycle.
U-Net as a deﬂation operator
A slightly more evolved option would be to apply U-Net, and afterwords to smooth the output using a SL V-cycle.
More explicitly:
[Equation]
This option somewhat resembles the concept of deﬂation operator studied in [25].
Encoder-Solver: an encoder network as a preconditioner setup
The stand-alone solver described so far needs to generalizes on both the residual and the slowness.
Obviously, the less we have to generalize on, the solution will be easier for the network.
This is veriﬁed in our experiments—the network performs better when the slowness model is known during training (see Sec. C.2).
Hence, we wish to exploit the following: when we solve an equation using several iterations, the slowness model remains the same during those iterations—the U-Net is repeatedly applied with the same model for different residuals.
To exploit this, we propose to split the generalization over the residual and slowness model into two “encoder” and “solver” networks.
The encoder network receives only the slowness model as input, and prepares feature vectors that include the slowness model information.
Those vectors are computed only once at the beginning of the solution process, and are fed into the solver network at every iteration in the solution time.
The solver network is applied at every iteration, with the readily available feature vectors inside it.
This idea resembles a preconditioner setup that is applied once before the solution process, and is adequate for any right hand side.
Mini re- training Our last approach aims to further improve the solver U-Net by a mini retraining procedure at solve time for the given slowness model.
That is, given a slowness model κ 2 we apply a few training iterations over a randomly generated data set of residuals to improve the U-Net weights.
This is similar to the process obtained in PINNs, and is motivated by our relatively better performance for known models compared to the case where we generalize over the models.
This option would ﬁt cases where we need to solve multiple linear systems with the same slowness model but for many right-hand-sides, like when solving inverse problems, where the medium needs to be estimated.
Numerical Results
In this section we evaluate the performance of the preconditioner together with ﬂexible GMRES for solving the Helmholtz equation for random right-hand-sides.
In particular, we use a restarted and ﬂexible Block-GMRES(10), with a subspace of size 10, solving 10 different right-hand-sides together.
We consider it to be the preconditioner test.
Our code is written in the Julia language [3], utilizing the Flux DL framework [10].
Unless stated otherwise, the test cases that we present are for (2.1) deﬁned on a 2D domain of [0 , 1] 2 , discretized on a nodal regular grid of 128 × 128 cells, and a wave frequency deﬁned by ω = 20 π , which leads to about 12.5 grid points per wavelength.
We use an absorbing boundary layer of width 10.
We compare three preconditioners: The SL V-cycle M V , M JU and M VU for Eq.
We evaluate the methods as preconditioners by their average convergence factor denoted by ρ , calculated over the number of iterations required to reduce the residual norm by a factor of 10 6 .
For the SL V-cycles we use a shift of β = 0 .
5 , 3 levels, 1 pre- and post- damped
Jacobi smoothing with damping of 0.8, and a coarsest grid solution of GMRES(10) as in [4].
We now present the ability of the proposed methods to generalize for different slowness models and right-hand-sides.
In these experiments, the training phase was performed on a set of 20 ,pairs of ( r , e ) produced as described in Sec.
B.1, and κ 2 were randomly generated from a collection of the 50 , 000 images in the CIFAR10 data-set.
We train the network on a grid of 128 × 128 for κ 2 ∈ [0 .
25 , 1] , but also examine the performance on larger grids and frequencies.
In particular, we solve (2.2) for random right-hand-sides on grids of size 256 × 256 and 512 × 512 with ω = 40 π, 80 π , respectively , using a network trained on residuals and slowness models of size of 128 × 128 only.
In supplementary Sec.
C.3 we also show generalization for κ 2 ∈ [0 .
Figure 2 summarizes the results for all our three methods: stand-alone preconditioner, encoder-solver framework, and the mini-retrain approach.
Overall, it is clear our CNN can generalize over the right-hand-sides (residuals), and over the slowness models.
The stand-alone solver signiﬁcantly out-perform the SL solver, with M VU typically outperforming M JU .
Our encoder-solver framework, where two CNNs are used, indeed upgrades the stand-alone U-Net in all cases, and particularly when we test at a larger problem size than the network is trained on.
In addition, our mini-retraining phase consistently yields another upgrade to the performance.
While our preconditioner is not necessarily superior than traditional ones in terms of FLOPs, it can be applied more efﬁciently on hardware.
Furthermore, CNN models can be signiﬁcantly compressed without harming performance. [Figure]
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results? [Table]
/A] (b) Did you include complete proofs of all theoretical results?
[N/A]
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[No] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[No] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[No] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[N/A] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
Supplementary Material
A Background A.1 Geometric multigrid
In this section we describe the geometric multigrid approach that we use for solving (2.2).
Generally, multigrid methods are used to iteratively solve linear systems like (2.2) deﬁned on a ﬁne grid Ω h using a hierarchy of grids.
In a nutshell, in a two level setting, we project the ﬁne level problem onto a coarser grid, solve the problem on that coarser grid, and then use the coarse solution to correct the solution on the original ﬁne grid.
This is repeated recursively to form a V-cycle, which is applied iteratively to solve the problem.
Geometric multigrid can be implemented using CNN constructs quite naturally, resulting in a structure that resembles the well-known U-Nets.
More explicitly, multigrid methods are based on two complementary processes: relaxation and coarse grid correction.
The relaxation is obtained by a standard method like Jacobi or Gauss-Seidel, which is only effective at reducing part of the error 2 .
The remaining error, called “algebraically smooth”, is typically spanned by the eigenvectors of A h corresponding to small eigenvalues (in magnitude), i.e., vectors e h s.t.
[Equation]
To reduce such errors, multigrid methods use a coarse grid correction, where the error e h for some iterate u h is estimated by solving a coarser system and interpolating its solution:
[Equation]
Here, the matrix A H is an approximation of A h on a coarser mesh Ω H , obtained for width H = 2 h .
To interpolate the solution from coarse to ﬁne grids we choose the bi-linear interpolation operator
I hH which is suitable for the problem because the Laplacian operator in (2.3) is homogenous.
The restriction operator I Hh , dubbed “full-weighting”, projects the ﬁne-level residual to the coarse grid.
Since the coarse problem is still too large to solve, the process of relaxation and coarse grid correction is applied recursively resulting in the V-cycle algorithm given in Alg. 1.
Algorithm
Relax v 1 times on A h v h
= f h with v h as an initial guess ; if deepest level then v h ← Solve the system
A h v h
= f h directly; else f 2 h ← I 2 hh ( f h −
A h v h ) ; v 2 h ← V-cycle ( v 2 h ← 0 , f 2 h ) ; v h ← v h
+ I h
2 h v 2 h ; end
Relax v 2 times on A h v h
= f h with v h as an initial guess;
In Alg.
1 one has to choose the number of levels in the V-cycle.
Unlike other scenarios, the smooth error modes of the Helmholtz operator are oscillatory at high wavenumber, and cannot be represented well on very coarse grids.
Hence, the performance of the solver deteriorates as we use more levels.
For example, the results in [4] show that the best performance is achieved using three levels only, and the authors suggest to use GMRES as a coarsest grid solver.
A.2 Geometric multigrid using CNN modules
In our work, both the CNN and the multigrid preconditioner are used interchangeably in the same framework, using a GPU back-end.
We implement the multigrid method using CNN components, in order to enable a good integration of the V-cycle with the CNN in a uniﬁed framework within the preconditioner component.
The idea is based on the close connection between U-Nets and multigrid
V-cycles.
The Helmholtz operator needed for the residual calculation and the Jacobi relaxation is executed using a convolution operator according to the discretized Laplacian operator
[Equation]
[Equation] [Equation]
[Equation]
using a ﬁxed-weights 3 × 3 kernel plus an element-wise multiplication with the mass term in (2.3) or (2.1), i.e., whether with an artiﬁcial shift or not.
The high order discretization operators in [28] can also be obtained similarly.
The geometric multigrid transfer operators, the restriction (coarsening of the mesh) and prolongation (interpolation), are realized by convolution operators with the ﬁxed kernels
[Equation]
that correspond to the full-weighting and bi-linear operators, respectively.
Both kernels are applied with a stride of 2.
For the restriction, this results in a coarse mesh Ω H .
For the prolongation, we use a transposed strided convolution with K I hH , so that the mesh is reﬁned.
The solution for the coarsest mesh problem is done using GMRES with diagonal Jacobi preconditioner [4].
In addition, all these convolutions are applied with zero padding of size 1, which is consistent with the Dirichlet boundary condition (the ABC is obtained through the mass term).
Neumann BC can be obtained by using reﬂection or replication padding in CNN frameworks, which result in ﬁrst and second order Neumann
BC, respectively.
B Method: training, data generation and augmentation B.1 Training and data augmentation
In order for the U-Net to be efﬁcient at solve time as a preconditioner, it must practice or learn on
GMRES residual vectors that are as close to the solve time GMRES residuals as possible.
Let
[Equation]
be the forward application of the network for a given right-hand-size r and slowness model κ ( x ) 2 .
As before, θ denote the collection of network weights (convolutions, biases, etc.) that we wish to learn.
Since we wish to complement smoothing processes, in our supervised learning setup, we aim to minimize the error with respect to multiple residuals and models.
That is, given training data triplets { ( e truei , r i , κ 2 i ) } mi =1 , where A h e truei =
r i , we minimize
[Equation]
which is the mean (cid:96) 2 error norm, also known as the mean squared error (MSE).
The training phase is performed by a stochastic gradient descent optimizer, where in each epoch we sweep through all the training examples in batches and varying order.
Speciﬁcally, we use the ADAM optimizer [12] with a variable step size.
For supervised learning, we have to provide the network with the corresponding error solution for each residual vector and slowness model.
To have such triplets, in some cases one needs to invest a lot of time and resources in preparing the training-set, i.e., solve the PDE many times.
However, since we have a linear system here, we can simply generate a random e true i , and easily compute the residual.
But, we wish to create a data-set that is as close as possible to the situation in the real solution using
GMRES, and if we simply take the GMRES residual vectors, we will not have a corresponding error without solving the system.
To obtain that we need to generate data with rather smooth vectors in varying “smoothness levels”, since GMRES and V-cycles are smoothing operations.
Hence, we propose the following procedure for which, in our experience, the training error predicts the ﬁnal solution’s convergence factor quite accurately.
First, we generate a random normally distributed x i and compute b i =
A h x i .
Then, we apply a random number of GMRES iterations with SL V-cycle as preconditioner, starting from a zero vector, to get ˜ x i :
[Equation]
Following that, we compute a residual r i =
b i − A h
˜ x i and let r i be the network input.
The true error that we wish the network to recover is given by
[Equation]
Beyond the procedure above, we insert ˆ r i =
h 2 r i as input to the network.
This causes the network’s input and output to be of the same order of magnitude 3 , and because the system is linear, this change has no effect on the accuracy.
Furthermore, to avoid over-ﬁtting and achieve better generalization, we augmented the training examples using a linear combination of the errors and residuals to expand the set.
This summarizes the training data generation, given slowness models.
The training set and the validation set are manufactured in the same way, and we saw no difference in the loss between these two.
Overall, our network has 4 input channels of the size of the grid.
These are the real and imaginary parts of the residuals r i , the real (squared) slowness model κ i ( x ) 2 , and the real attenuation model γ ( x ) .
Although we do not generalize of γ ( x ) , we found it slightly beneﬁcial to include it as input, as it allows the network to have a different behaviour at the grid’s boundaries.
Training details The training process for all the U-Nets reported in this paper is performed by the
ADAM optimizer [12], and includes 80-120 epochs.
The initial learning rate is 10 − 4 and the initial mini-batch size is 20.
Along the epochs we reduce the learning rate and increase the mini-batch size.
The schedule of these changes as well as the amount of epochs we apply vary a bit between the different experiments and are adjusted by cross validation on the test set.
In the majority of the cases, after t 0 = 48 epochs we divide the initial learning rate by 10 and double the mini-batch size.
The next changes of the learning rate and mini-batch size occur after every t i =
t 0 2 i epochs.
That is, the changes occur at epocs 48, 72, 84, 90 etc, until the learning rate essentially reaches zero.
These settings can be used as default parameters for all the experiments that we show.
Remark
As an alternative to the error loss in (B.2) , one may use the residual norm (cid:107)
A h e neti − r i (cid:107) 2 , so that the training can be unsupervised [1].
This is a viable option, but it can be a bit tricky, since high errors can in fact have low residuals—these are “algebraically smooth errors”.
Since we can eliminate those errors with V-cycles or Jacobi iterations, we wish to let the network focus on those smooth modes and not on the whole spectrum.
We note that empirically, our experience in augmenting the error norms in (B.2) with residual norms generally yields inferior performance than focusing on the errors only during training.
B.2 Slowness models
Solving the Helmholtz equation for a uniform slowness model is not difﬁcult, as it can be performed using a Fourier transform.
When trying to solve this equation over a non-uniform media, the difﬁculty increases signiﬁcantly, especially if there are sharp jumps or large ratios between the smallest and highest slowness.
Here, we want the network to succeed in different types of models, smooth and non-smooth, that faithfully represent reality.
In certain applications (e.g., geophysical or medical imaging) one may have slowness models of a certain typical distribution coming from the application domain.
E.g., in geophysical imaging, one typically has a layered model with higher velocity values deeper in the earth.
In machine learning tasks, if we wish for an algorithm to generalize on an input, it is best if we let the algorithm learn on inputs from the same distribution as in the real life application.
That is a large part of the advantage of data-driven approaches.
Usually, however, to obtain a good generalization we need a large data set of inputs—a collection of slowness models in our case.
Here, we use a rather general slowness model collection to demonstrate our method.
The models that we use for the U-Net training, are created from a large set of natural images CIFAR-10 [13].
This data set is usually used as image classiﬁcation benchmark and contain 50K images of resolution 32 ×
[Figure]
pixels.
We use this data set and manipulate the images to become slowness models in order to have a signiﬁcant amount of training models to learn from.
In particular, we use a training set of 20 ,random sample images.
To transform the natural images into slowness models, we ﬁrst enlarge them to the grid size using a bi-linear interpolation, then smooth them with a single application of a Gaussian smoothing kernel, and normalize the values to the range of [0 . 25 , 1] .
Fig.
3 shows four different models produced from the CIFAR-10 image collection.
Fig.
4 presents how the smoothness level of the model, κ 2 , affects the difﬁculty of the solution.
It is evident that as the medium is less smooth, the solution is more complex and is harder to model by a neural network.
Similarly, the range of slowness values in the model also affects the the complexity of the solution.
Fig.
5 illustrates how the solution is again more complex as the contrast in the slowness model is larger.
C Additional Numerical Results
In addition, we present the relative error norm plot throughout the training period, which shows the value of the loss function (B.2) vs. the epoch number.
Ideally, the asymptotic value of the relative
[Figure] [Figure]
training error should indicate the real convergence factor at solve time using GMRES.
It is indeed more-or-less so thanks to the varying smoothing levels that we apply for our training residuals in Eq.
(B.3).
Using this plot, we can also estimate over-ﬁtting and adjust the hyper-parameters (learning rate and mini-batch size, for examples).
C.1 Performance of a homogeneous model
In the homogeneous test case, we solve the system (2.2) for a uniform model, known at training.
That is, the training is performed on random residual and error vectors only as described earlier.
Fig.describes the results of the proposed method for this case.
The average relative error, i.e. the minimum estimate value of the loss function (B.2), of the training and test sets at the end of the training process is 0 .
The GMRES method with U-Net-based preconditioner as described in Eq.
(3.1) converges within 10 iterations with convergence factor ρ VU = 0 .
231 compared to 0 .
88 using only V-cycles.
Our cheaper M JU preconditioner, which uses only our U-Net and Jacobi steps converges withiniterations with convergence factor ρ JU = 0 .
These are very favorable convergence factors given how poorly V-cycles behave for a homogeneous model.
On the other hand, homogeneous problems can be efﬁciently solved using fast Fourier transform.
C.2 Performance on a single heterogeneous model
In this section we examine our method for a single heterogeneous slowness model, known during the training phase.
In this scenario, the network can learn weights speciﬁc to that model.
As shown in
Fig.
7, for a model normalization of κ 2 ∈ [0 . 5 , 1] we obtained convergence factors of ρ VU = 0 .and
ρ JU = 0 .
383 , which are close to the performance on the uniform model.
For the more difﬁcult case of κ 2 ∈ [0 . 25 , 1] , we obtain ρ VU = 0 .
41 , ρ JU = 0 .
Furthermore, in plot (c) of Fig. 6 and in the bottom plots (c) and (e) of Fig. 7, we show the results of a single iteration of Eq.
(3.1) on a centralized point source, a problem that is not included in the training set.
It is evident that U-Net has indeed learned to produce a global solution, and the integration with the V-cycle is intended to “clean” or smooth the solution and ensure convergence.
C.3 Performance on generalized slowness models
As we showed in the previous section, our U-Net can learn to solve (2.2) quite efﬁciently for a given heterogeneous slowness model.
However, the training of the U-Net is quite expensive to apply at solve time.
In a more applicable approach for the realistic world, we wish the neural network to work and generalize on models it had not seen in the training phase.
Therefore, we also present the ability of the proposed method to generalize for different models.
Table 1 presents the results for the encoder-solver setting.
Regardless of the speciﬁc preconditioner, we see that if the networks are trained on a single model, they exhibit the best performance.
However, they fail to generalize on slightly different cases—even for the easier normalization of the same model (row (a) vs row (b) in
[Figure]
Single model
General model
Re-trained
M V
M VU
M JU
M VU
M JU
M VU
M JU ρ 0 .0 .0 .0 .0 .0 .0 .
T
Generalization for cases not seen in training ρ 0 .0 .fail
[Table]
(c) κ 2 ∈ [0 . 25 , 1]; n =
[Table] [Table]
[Table] [Table]
the single model column).
If we train the networks to generalize over slowness models, it can also generalize on the problem size and slowness contrast.
All cases are improved using mini-retraining.
C.3.1 The inﬂuence of the Krylov subspace and block sizes
In this set of experiments we examine the inﬂuence of the subspace sizes that the GMRES method uses.
We examine the inﬂuence of both the subspace size and the number of RHSs solved together (block size).
The experiment corresponds to the 256 × 256 experiment shown in Fig 2 (e).
Fig.summerizes the results.
It is obvious that whether by multiple RHSs or by subspace size, increasing the effective Krylov subspace can improve the performance of the solver signiﬁcantly.
That is in contrast to the SL V-cycle that is not improved by the same magnitude.
This shows that there is more potential to the U-Net based preconditioners that shown earlier, since in addition, we typically did not
[Figure]