MGCFNN: A Neural MultiGrid Solver with Novel Fourier Neural Network for High Wave Number Helmholtz Equations
Authors: Yan Xie
Minrui Lv
Chen-Song Zhang
Venue: ICLR 2025 Poster (2025)
Keywords: neural multigrid solver
Fourier transform
Helmholtz equations
high wave number
heterogeneous medium
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: Solving high wavenumber Helmholtz equations is notoriously challenging. Traditional solvers have yet to yield satisfactory results, and most neural network methods struggle to accurately solve cases with extremely high wavenumbers within heterogeneous media. This paper presents an advanced multigrid-hierarchical AI solver, tailored specifically for high wavenumber Helmholtz equations. We adapt the MGCNN architecture to align with the problem setting and incorporate a novel Fourier neural network (FNN) to match the characteristics of Helmholtz equations. FNN, mathematically akin to the convolutional neural network (CNN), enables faster propagation of source influence during the solve phase, making it particularly suitable for handling large size, high wavenumber problems. We conduct supervised learning tests against numerous neural operator learning methods to demonstrate the superior learning capabilities of our solvers. Additionally, we perform scalability tests using an unsupervised strategy to highlight our solvers' significant speedup over the most recent specialized AI solver and AI-enhanced traditional solver for high wavenumber Helmholtz equations. We also carry out an ablation study to underscore the effectiveness of the multigrid hierarchy and the benefits of introducing FNN. Notably, our solvers exhibit optimal convergence of $\mathcal{O}(k)$ up to $k \approx 2000$.

Published as a conference paper at ICLR 2025
MGCFNN: A NEURAL MULTIGRID SOLVER WITH NOVEL FOURIER NEURAL NETWORK FOR HIGH WAVENUMBER HELMHOLTZ EQUATIONS
Yan Xie, Minrui Lv& Chen-song Zhang
State Key Laboratory of Mathematical Sciences, Academy of Mathematics and Systems Science, Chinese Academy of Sciences, Beijing 100190, China; School of Mathematical Sciences, University of Chinese Academy of Sciences, Beijing 100049, China. (xi xiieyan2221, 1 lyminrui, zhangos) @ sec. cc ac cn
ABSTRACT
Solving high wavenumber Helmholtz equations is notoriously challenging. Tradi- tional solvers have yet to yield satisfactory results, and most neural network meth- ods struggle accurately solve cases with extremely high wavenumbers within heterogeneous media. This paper presents an advanced multigrid-hierarchical Al solver, tailored specifically for high wavenumber Helmholtz equations. We adapt the MGCNN architecture to align with the problem setting and incorporate novel Fourier neural network (FNN) match the characteristics of Helmholtz equations. FNN, mathematically akin to the convolutional neural network (CNN), enables faster propagation of source influence during the solve phase, making particu- larly suitable for handling large size, high wavenumber problems. We conduct supervised learning tests against numerous neural operator learning methods to demonstrate the superior learning capabilities of our solvers. Additionally, we perform scalability tests using an unsupervised strategy to highlight our solvers' significant speedup over the most recent specialized Al solver and Al-enhanced traditional solver for high wavenumber Helmholtz equations. We also carry out an ablation study underscore the effectiveness of the multigrid hierarchy and the benefits of introducing FNN. Notably, our solvers exhibit optimal convergence of Okk) up to k 2000.
INTRODUCTION
1.1 BACKGROUND
The Helmholtz equation, a fundamental partial differential equation (PDE), is pivotal in describing wave propagation phenomena across various disciplines, including acousties, electromagnetics, and seismology. Despite its importance, solving the Helmholtz equation presents significant challenges due to the high wavenumber and variable wave speed. These factors necessitate large number of grid points to accurately capture the wave shape and result in complex-valus highly ill-conditioned, indefinite system. Indeed, is widely acknowledged that the optimal number of iterations required for solver with complexity O(N) to solve high wavenumber Helmholtz equation is Olk). This due to the complex propagation pattern of the wave in heterogeneous media, with slow decay rate.
Numerous research efforts have been dedicated to developing efficient traditional solvers, such as the Born series method Osnabrugge al. (2016), the factorization method Ostiho]oOeeKKufffooo & Saad (2010), the shifted Laplacian method Gander et al. (2015) Sheikh et al. (2013); Calandra et al. (2013), domain
decomposition method Chen & Xiang 2013 3 Leng & Ju (2022 2019). and multigrid method Brandt
"Corresponding author.

---

Published as a conference paper at ICLR 2025
& Livshits 1997). However, these methods have not yet achieved satisfactory results Ernst & Gander (201 1).
Data-driven methods have recently emerged popular approach solving PDEs, either by optimizing certain solver components or by learning the entire solver. One significant research area Physics Informed Neural Networks (PINNs) Raissi et al. (2019) which use neural networks to approximate PDEs' solution functions. While recent works Song et al. (2022); Escapil-Inchauspé & Ruz (2023) have applied PINNs to Helmholtz equations, they do not address high wavenumber cases.
Furthermore, our focus is on inference rather than online training to gain the solution.
more related research area is neural operator learning, which learns the operator (mapping) between the infinite-ding dineiisoedi parameter and solution spaces of PDEs. This includes works such as the Convolutional Neural Operators (CNO) Raissi et (2019), Deep Operator Network (DeepONet) Lu et al. (202T), Fourier Neural Operator (FNO) Li et al. (022)) U-NO Rahman et al. (2023), and MgNO He et al. (2024). These methods solve PDEs during the inference phase, functioning more like solver for discretized PDEs. The MgNO He al. (2024) work has tested their method on Helmholtz
equations and compared with other neural operator learning methods, but did not consider high wavenumber cases.
1.2 RELATED WORKS
In practice, it is nearly impossible to directly solve high wavenumber Helmholtz equations within heterogeneous wave speed (or slowness) field a desired accuracy with just single application of the network (see Subsection 4.2). Actually, both traditional and AI solvers require more iterations when dealing with higher wavenumber. Recent studies Azulay Treister (2022); Cui et al. (2022); Stachenfeld et al.] (202T); [Stanziola et al.] (202T) have developed iterative Al solvers to tackle this challenging problem, capable of achieving higher accuracy through iteration. However, the wavenumbers tested in these studies remain limited. Our objective is to solve high wavenumber Helmholtz equations up to k = 2000, aiming to match the application scope of traditional solvers, but with significantly improved speed.
Several latest studies have made significant progress in this area. The study Lerer et al. (2024) introduced pure Al solver, termed encoder-solver, and the work Cui et al. (2024) presented a dedicated A Al-enhanced traditional multigrid solver framework, Wave-ADR-NG for Helmholtz equations. Both of them demonstrated superior performance over traditional shifted Laplacian solvers. However, the encoder-solver relies on an existing traditional solver for data generation and solve phase assistance. Meanwhile, Wave-ADR-NS requires extensive problem insights construct solver with good convergence, limiting its applicability to other wave equations.
Moreover, the MGCNN work Xie et al. (203)) laid out principles for building an efficient AI solver for discretized linear PDEs, guarantecing quick convergence and adaptability to different iterative frameworks for desired accuracy. Despite these advantages, MGCNN has not been applied solve the Helmholtz equation, as its architecture is not directly suitable for this specific problem.
1.3 OUR CONTRIBUTIONS
In this work, we propose a multigrid-hinn hhtirrid-hiirarchicaal AI solver that effectively addresses the challenges of solving high wavenumber Helmholtz equations. The innovations and contributions are listed as follows:
We significantly modify the spectral convolution operation in the FNO to create the Fourier Neural Network (FNN), which handles full-mode information, making suitable for high wavenumber problems and scalable for larger problem sizes.
The FNN, which functions mathematically equivalent to a CNN with an extended kernel size, facilitates wave propagation. We adapt the MGCNN address high wavenumber Helmholtz equations and incorporate the FNN into its solve phase. This integration forms the MGCFNN solver, which combines the complementary features of the novel FNN and the multigrid hierarchy, resulting in improved convergence and shorter solve time.
In contrast to recent specialized iterative Al solvers for high wavenumber Helmholtz equa- tions, our solvers do not necessitate specialized domain knowledge. This makes them easy
2

---

Published as a conference paper at ICLR 2025
train and utilize, with the potential to directly tackle broader range of wave equations. Notably, the proposed MGCENN achieves at least 4.8x speedup compared with other AI iterative solvers under the same problem settings.
Although our solvers are designed for unsupervised learning, we also conduct supervised experiments. Our solver demonstrates superior learnability compared to other operator learning networks, achieving lower error with shorter training time and fewer parameters.
2 PRELIMINARIES
2.1
FORMULATION AND DISCRETIZATION
The Helmholtz equation expressed as
2 c(x)
- -u((s)
(1 (1- (i))))(()==[(a x [0, 1J3
where u represents the unknown field, c is the wave speed, is the angular frequency, and f is the source term. In line with related works/Azulay & Treister (2022); [erer et al. (2024); Cui et al. (2024). we define <(x)) с[œ] as the wave slowness and refer to k = WK the wavenumber. i v the imaginary unit, and is the damping factor, indicating the medium's
absorption property. This work primarily focuses on the hardest case of ạ = O, while also comparing with work Lerer et al. (2024) under 7 = 0.01 in experiments. As for boundary conditions, employ absorbing boundary conditions Engquist & Majda 1979); Erlangga et al. (2006) to reduce wave reflection at the boundary, simulating radiation in an infinite domain. Specifically, set y increase quadratically from 0 to 1 at certain distance from the boundary.
Applying second-order finite difference discretization to the Helmholtz equation I on a uniform grid with spacing h yields the linear system
A Sh,
where A₁ represents the discretized Helmholtz operator, expressed stencil form as
-1
0
A =
- 1 4 -
-1
(3)
h2
and fr. is the discretized source term. To accurately capture the wave shape, the spacing h must be sufficiently small, typically requiring at least 10 grid points per wavelength. This implies that increase the grid size 1/h in our experiments to accommodate the increased wavenumber If multiplying both sides of the equation by the factor 1 and denoting wkh along with as the problem coefficient coef, we find that coe f is scale-invariant, serving as normalized input to the network for different scales.
2xr wkh To Consequently,
2.2 FOURIER TRANSFORM
The Fourier transform is mathematical operation that decomposes function into harmonic wave components with different frequencies, making it particularly well-suited for wave propagation problems. In the frequency domain, globally propagated wave with specific frequency can be simply represented within a circle. From another perspective of view, the convolution theorem Smith (1997) states that the convolution of two functions in the spatial domain corresponds the pointwise multplicationn of their Fourier transforms in the frequency domain, i.c.,
(4) where F denotes the Fourier transform. Consequently, the Fourier transform can be utilized to perform a convolution operation with an extended kernel size, enabling the rapid propagation of wave source influence during the solve phase. As we use Fast Fourier Transform in our experiments, we denote FFT as the Fourier transform operation and IFFT as the inverse Fourier transform operation.
F(f g) = F())F(9).
'To illustrate high wavenumber cases and keep the comparison experiments feasible, may break the rule high wavenumber with relatively few grid points.
3

---

Published as a conference paper at ICLR 2025
3 NETWORK ARCHITECTURE
The architecture comprises a setup phase network for managing problem coefficients and a solve phase network that utilizes the handled coefficient information to process the source term, also known as the right-hand-side term (see Appendix[C) [C), to propose a solution. We use multigrid hierarchy to build the two phases, and design novel FNN kernel to handle the high frequency wave propagation in the coarse levels of the solve phase.
3.1 FOURIER NEURAL NETWORK
While the Fourier Neural Operator (FNO) Li al. (021) demonstrates impressive learnability across range of PDEs, it encounters limitations when solving high wavenumber Helmholtz equations. We make significant modifications to FNO's s spectral convolution operation.
We separate the processing of the coefficient and source term, ensuring the network linearly operates the source terms, as will be shown in the level status process equation
We assume that the weights form smooth function in the frequency domain (see Fig. Ia). and thus use smaller weight tensor to interpolate weights across the entire frequency do- main. This approach ensures high-frequency modes are captured within limited parameter size. This step is denoted case in Fig.
To apply the FNN to problems with larger size and higher wavenumber, we first perform IFFT to obtain the convolutional kernel in the space domain, then pad it to the size of the larger target problem domain, and finally apply FFT to obtain the weights in the frequency domain. This step is denoted case 2 in Fig. 2
Moreover, the inpul tensor is interpreted the concatenation of real and imaginary parts of a complex tensor, halving the number of channels prior to the operation. Conversely, the outpur tensor channels are doubled from complex tensor. These decrease the computational cost the Fourier transform while maintaining the network's 's predictive performance.
(a) weights Irequency
(b) equivalent space domain
Figure 1: The setuped weight tensor of FNN.
gedd
interp
weightts
Interp:
Input
emple pad
case
Space
Frequency Domain
Domain
Frear mag
satuped weights
Outpul
real
unpad
Figure 2: The network architecture of FNN. Case I: Interpolated weights of the size the target problem. Case 2: The target problem is larger than the interpolated weights. propectional number channels, and box height is proportional to the number of grid points
4

---

Published as a conference paper at ICLR 2025
Remark 3.1. For higher wavenumber problems, initially planned leverage the invariance of the function in the frequency domain (see Appendix(A) and apply the interpolation technique again to obtain the weight tensor. However, a straightfrward interpolation method performs poorly for substantially larger wavenumber. The mathematical equivalence in equation aids us in proposing a more elegant and efficient approach.
Remark 3.2. FNN is mathematically connected to the inverse of the damped Helmholtz operator, as detailed in Appendix AA] Additionally, the role it plays in solve phase is mathematically related to the traditional Born series method Osnabrugge et (0166). which solves the equation with an iterative nheudo-propggation process. A detailed discussion on this connection is provided in Appendix B
In summary, the network architecture of the FNN is depicted in Fig. 2 and its detailed pseudocode provided in Appendix D It's important to note that for given problem size, FNN first performs interpolation or even additional IFFT and FFT to set up the weight tensor. This is one-time cost for all same-sized problems.
Effective Kernel Size. Despite FNN's utilization of all Fourier transform modes, its effective influence length is limited, as illustrated in Fig. [bb] On one hand, fixed-length padding, which corresponds to the effective kernel size of the equivalent convolution operation, is adequate. On the other hand, a multigrid hierarchy is still crucial to handle problems on multiple levels to accelerate influence transmission when the multiple grids resolve the high frequency wave, shown in Subsection [4.4] We then introduce the multigrid-hierarcheee setup and solve phase network to efficiently address the challenge in conjunction with FNN.
3.2 SETUP PHASE
The setup phase processes the coefficients of the PDEs, including the wave speed c(x) (or slowness x(x)) and the damping factor (x)), preparing necessary information for each level in solve phase. We streamline the setup phase of MGCNN from Xie et al. (2023) by integrating the downsampling and nonlinear ResCNN processing into single workflow (see Alg. [T. The setup phase only needs to be executed once for fixed PDE coefficients and multiple right-hand-side tensors, as is the case for the updated residual terms in an iterative solving framework (see Appendix C.
Algorithm 1 Setup Phase Network
INPUT: coef problem coefficient tensor, level is number of levels for multigrid hierarchy.
2: OUTPUT: setup outs is a list of setup tensors for each level.
3: procedure SETUP(COe lewel)
4: selup_ouly <- C
5: for ( in 1,2...., level do
6: setup_oud out;
Nonlinear processing
7: if < level then setup_outiz
8:
RestridCCN
> Downsampling
9: end
end for
11: return selup-ouls + [sulup_oul], selup_oulz. selupeoulievel
end procedure
3.3
SOLVE PHASE
The solve phase framework of MGCFNN closely follows that of MGCNN described in Xie et al.) (2023). The key distinction lies in the replacement of the CNN components in ResNets at certain coarse levels with our custom-designed FNN. During both the down and up cycles, each level involves three primary steps:
1.
Utilize ResNet (either ResCNN or ResFNN) conjunction with the setup tensor at the current level to process the status tensor. Specifically, we have
xt ResNNel(setup oul,, mn) Nel(setup oul x1) + 31.
(5)
5

---

Architecture Techniques. In response to these challenges, besides our efforts of developing FNN to transmit wave influence more globally, we use fixed number of levels of the multigrid hierarchy in our experiments and increase the network's computational complexity when moving to coarser level. Specifically, number of levels are determined by number points per wavelength and we double the solvers coarser level. The detailed Published as a conference paper at ICLR 2025
2. During the down cycle, after processing, employ the restriction operator to transfer the status tensor to the next coarser level.
3. During the up cycle, before processing, add the prolonged status tensor from the coarser level.
opef
CNN:
Multiplication:
Skip Addition:
Update:
Residual CNN:
Resirual FNN:
Figure 3: The network architecture of MGCENN. Its coarsest two levels use FNN kernels.
By combining the aforementioned setup and solve phase networks, we arrive at our novel solver, MGCFNN, as shown in Fig.
Architecture Naming. We use the following naming conventions for our hhitgrrid-hierrachcccall networks: MGCFNN refers to networks that utilize FNN in some coarse levels of the solve phase.
MGFNN refers to networks that use FNN at all levels. MGCNN refers to networks that use CNN at all levels.
Multigrid Challenges. As demonstrated in previous literature Xie et al. (2023), MGCNN can function as a global operator, using more levels with weight sharing across levels to handle problems of larger sizes. However, applying this strategy to higher wavenumber cases in our scalability test settings presents two challenges. First, the grid must be sufficiently large to accommodate the high wavenumber. Actually, it is common practice to maintain fixed few number of levels of the multigrid hierarchy in traditional methods Calandra al. (2013) for higher wavenumber. Second, the coarse grid problem faces the same wavenumber with fewer grid points, resulting in dissimilar properties across different levels.
number of sweeps (layers) for all of our architecture settings are provided in Appendix H
Remark 3.3. The settings of network are closely related to frequency. MGCFNN replaces the CNN with FNN at levels where the grid size is nearly insufficient to resolve high-frequency waves but still too large for CNN handle. From perspective, using FNN all levels is unnecessary, shown in Subsection 4.4 From another perspective, there is no need to use FNN for low-freyiency waves, where more coarse grids still resolve waves. Therefore, the optimal of FNN is problem-dependent, which is a common feature for iterative solvers.
In summary, combine the complementary features of the novel FNN and the specialized multigrid hierarchy to develop the MGCENN solver, which effectively handles high wavenumber Helmholtz equations.

---

Published as a conference paper at ICLR 2025
4 NUMERICAL EXPERIMENTS
4.1 BASIC SETTINGS
Comparison Methods. We first compare solvers with well-known operator learning methods to demonstrate our networks' learnability. While there are some Graph Neural Network (GNN) based methods Trilonov al. (2024): : Li et al. (2023), they are designed for unstructured problems, do not utilize structural features, and are typically computationally slower than CNNs. Operator learning methods are not directly suitable for iterative applications to improve accuracy, so we only perform single inference to compare results in Subsection 4.2 Some strategies Azulay & Treister (2022);
Lerer et al. (2024): Rudikov et al. (2024) Kopaničáková & Karniadakis 2025) exist to improve the accuracy of these methods, such as using flexible Krylov methods for nonlinear preconditioning and generating suitable right-hand-side (source) terms for training. In Subsection 4.3
we compare our
work with AI, Al-enhanced traditional solvers and sparse direct solver.
Loss function & Data. Both supervised and unsupervised training strategies are applicable. However, recommend unsupervised learning as it eliminates the need for labeled data generation. Supervised learning is only employed in Subsection 4.2 facilitate a fair comparison with other operator learning methods under identical settings. The loss functions are the averaged square of the La norm of error and residual terms of the linear equations respectively:
Lerr E N Lrea= = N Ei=1 ||rhs - 1,,ol:12 unsupervised learning,
llu sol: supervised learning,
where N is the data count, u, is the true solution, rhs; is the input source term, sol, is the output guessed solution, and A; is the discretized linear operator. The operator is determined by the coefficient tensor coefi, which includes the normalized slowness field x and y. We generate rhs; as a white noise tensor to cover all frequency ranges. The k field is generated using the CIFAR- 10 Ko/Krizheysky et al. (2009) dataset, with images interpolated to the target problem size and linearly transformed to range Kmiin Kmaa]] In practice, fo should be generated fit specific application scenarios. We also present results from an ultrasonic CT scenario in Subsection 4.4, which are more representative of real-world applications. Furthermore, Appendix G shows generalization results on STL-10 Coaates et al. (2011) dataset and the benchmark Marmousi Brougois et al. (1990) data, both of which are of higher resolution.
Although the solve phase network can be applied several times when solving in an iterative framework, we only perform the solve phase once for each coeff during training.
Default Settings. Unless stated otherwise, experiments follow these defaults: slowness bounds Kmin and Kmax are 0.25 and 1.0 respectively. Solve test results represent the median of 50 unseen data. Training is unsupervised, executed on RTX 4090 GPU using PyTorch with CUDA 12.4. Both training and inference use float32 precision, while the solver employs a mixed precision strategy C achieve rtol=1E-7 tolerance for the norm of relative residual. We refer to the solver within stationary iterative framework as a standalone solver and also show the results of using GMRES to speed up the convergence. For scalability tests, we only train our models on 511 511 grid and test on grid size up to 4095 x 4095.
Note that larger models are used in small size problems fully utilize GPU resources, while smaller models are employed in scalability tests to fit within GPU memory constraints. For detailed hyperparameters, refer Appendix H
4.2 COMPARISON DIFFERENT SOLVERS
To provide a fair illustration of our solvers' learnability, we employ supervised learning, a training strategy commonly used in other operator learning methods. As most operator learning methods are not designed for use in an iterative framework, we compare the results after single network inference. We limit these tests to a grid size of 256 x 256 and use a relatively high angular frequency = 800 to set high wavenumber k. We compare our methods with recent well-known operator learning methods, including Mg et al. (2024), FNO2d Li et al. (202T), U-NO Rahman et al. Rahmaa
(202)),
et al. (2021) LSM Wu al. 2023).
MWT Cupta et al. (202)) and
7

---

Published as a conference paper at ICLR 2025
U-Net Ronneebrger et al. (2015). We utilize the official implementations of these solvers, adjusting the tunable parameters to suit the problem and balance the train error, train time and parameter size. All solvers are trained under identical settings (refer to Appendix To assess the model's '' ability to generalize, we record the point Where the Best Test Error is achieved. This includes the test error, the corresponding training error, and the epoch index within the total of 120 epochs.
We decrease the value of Emin to aggravate the heterogeneity of the media, and separately evaluate models for Kiin 0.75, 0.5, 0.25. As demonstrated in Table and Table 2 our proposed models, MGCFNN and MGFNN, outperform other solvers in terms of train and test error, time efficiency, and parameter size, which represent the state-of- of-the-art in this field. Notably, while models like FNO2d, MWT2d. and U-NO can achieve certain low training errors, they are prone to high test errors and early overfitting. This further underscores the superior generalization and robustness our models.
Table 1: Supervised learning grid 256 x 256 with 80m and Kmin 0.75. We record information Where the Best Test Error reaches.
Mmin = 0.75
Where Best Test Error
Model
Train Error
Test Error
Train Error
Epech
Train Timo(s/opoch)
Paramoters(MBB
MGENN
0.035
0.061
8.9
MGCENN
0.046
0.070
0.046
120
67.5
5.3
MGNO
0.079
4.6
FNO2D
0.085
0
0.496
103.6
46.1
0.119
0.527
0.475
26.0
U-NO
0.408
0.850
0,870
4
101.6
86.7
U-NET 0.534
0.803
0.758
31
89.9
31.0
DIL-RESNET
0.605
0.006
0.605
116
140.0
0.6
LSM
0.722
0.783
0.739
66
230.7
4.9
Table 2: Supervised learning on grid 256 256 with w = 80w and Kmin = 0.50, 0.25. We record
information Where the Best Test Error reaches.
Main 0.50
Kmin
0.25
Who here Best Test Error
Where Best Test Error
Model
Train Error
Error
Epoch
Train
Emorr
Epoch
MGFNN
0,075
0,187
0,122
25
0.126
0,431
0,253
18
MGCENN
0.101
0.206
0.104
0.190
0.488
0.396
MGNO
0.209
0.333
0.221
81
0.361
0.633
0.614
14
FNO2D
0.122
0,749
0.718
3
0.154
0.851
0.818
MWT20
0.169
0.728
0.186
0.845
0.821
U-NO
0.389
0.871
0.864
3
0.362
0.935
0.930
U-NET
0.634
0.774
0.716
0.575
0.806
0.778
26
DES-RE
0.627
0.629
0.627
117
0.699
0.702
0.699
119
LSM
0.852
0,860
0.852
117
0.867
0.878
0.874
38
These results, especially in Table suggest that applying a neural network just once to achieve the desired accuracy for high wavenumber and heterogeneous Helmholtz equations may not be feasible. In subsequent experiments, we will concentrate solely on iterative solving. Moreover, in scalability tests, solvers expected to handle higher angular frequency up to 640m z= 2000.
4.3 SCALABILITY COMPARISON
In this paper, we place greater emphasis on the scalability to larger problems with higher wavenumber. The solver should be capable of iteratively achieving any desired level of accuracy. The latest note- worthy iterative Al solvers to handle the challenging high wavenumber Helmholtz equations include pure AI solver Lerer et al. (2024). termed encoder-solvers. and the Wave-ADR-NV solver, an Al-enhanced traditional multignd solver et al.] (2024). Both methods have shown superior conver- gence and solve time performance compared to the traditional shifted Laplacian solver. However, we are unable to replicate the encoder-solver results, and the Wave-ADR-NS while effective, is complex to implement and not currently open source. Therefore, we use their published results for comparison under the same data and hardware settings. Furthermore, while direct solvers are potentially less scalable to 3D problems than iterative methods, they are a viable option for challenging 2D problems. We compare our results with the sparse LU solver CHOLMOD Davis & Hager (209)); Chen et al.) (2008); Davis et al, (2004); Amestoy et al. (1996) 2004) in the SuiteSparse Davis (2024) software. Its GPU acceleration Rennich al. (2014) is not beneficial in problems, run CHOLMOD
8

---

Published as a conference paper at ICLR 2025
with single precision CPU (Intel(R) Xeon(R) Gold 6458Q). The results are presented in Table 3 and our proposed MGCFNN significantly outperforms other solvers.
Table 3: Scalability comparison with other solvers.
mtl-11E7
0.01
(RTX338D)|
2024
0.0 (A100 Cui al.
2024
7 (CPU)
time
lime ltens
grid
MGCFNN ENCGDER-SOLLVER
speedup
MGCFXN
WAVE-ADR-NS
speedup
CHOLMOD
speedup
511 511
0.129)
0.6543)
5.5
0,16
(3.07(25)
94,8
0.49
3.1
1023
0.19(11)
0.30/22)
29.5
320m
2047
2047
0.58(14)
3,40185)
1.15(40)
91.63(122)
79.6
38.91
33.8
6408
4005
4096
2.77(18)
13,34(117)
48
8,55083)
286. 14(247))
33.5
183.6
21.5
4.4 ABLATION STUDY
Firstly, we examined the FNN, MGENN, and MGCENN models on two datasets with fixed grid sizes. The first dataset has grid size of 256 x 256 and an angular frequency of 800.. The second dataset, from an ultrasonic CT competition has grid size of 480 x 480 and an angular frequency of w = 150m. Table + presents the results for random sources, while Fig. 4b and Fig. 4a show single source results for illustration. As shown Table the multigrid hierarchy allows MGFNN to converge faster than FNN, but its use of FNN all levels only reduces few iterations compared to MGCFNN. MGCFNN solves problems faster in terms of time and is expected to be comparably faster for larger problems, Moreover, MGENN uses more memory than MGCENN, making it less suitable for large-scale problems.
Table 4: Model architecture ablation study.
2566
480,
1506
mol1FE7
standalane
GNIRES
standalane
GMRES
MGFNN
PNN
MGCFNN 13
FNN
MICCFNN
PNN
MGCFNN
Biess
22
19
35
13
24
12
11
lime(s)
0.136
0.135
0.270
0.111
0.116
0.151
0,093
0.101
0.724
0.105
0.124
B.365
300
16
12
30
Single source wave with 80s, in slowness field of cell.
C6
3
380
64
20
18
16
25
15 180 151
(b) Single source wave with 150x, a slowness field generated from CIFAR-10 data.
Figure 4: From left to right: slowness field s(r), real value of wave field and wave field in frequency domain, iterative solving history. The cross symbol in the first column denotes the position of source.
Secondly, we assess the scalability of MGCFNN and MGCNN across varying grid sizes and angular frequencies. Detailed results are Table Both models show optimal O(k) convergence for high
AT4S Cup- Prediction of Wavefield in Ultrasonic Computed Tomography

---

Published as a conference paper at ICLR 2025
wavenumber k, whether standalone solvers preconditioners in the GMRES solver. Notably, MGCENN performs effectively as a standalone solver. In contrast, MGCNN requires 2 to times more iterations and relies GMRES to significantly reduce its iteration count. This confirms MGCFNN's effective learning of the Helmholtz equation's s inverse operator. Despite the unexpected computation inefficiency of linear map in frequency domain (see Appendix F). .GCENN still outperforms MGCNN in solve time.
Table 5: Scalability comparison between MGCENN and MGCNN.
1t-7=1F-7
standalone
GMRES
MGCENN
MGCNN
MGCENN
MGCNN
size
Iters
time(s)
iters
time(s)
Iters
time(s)
Iters
time(s)
80m
511 511
0.17
35
0.27
12
0.17
21
0.20
160
1023 x 1023
0.31
0.55
20
0.31
36
0.37
320
2047 x 2047
41
1.18
115
2.33
35
1.13
71
1.65
640
4095 X 4095
83
8.18
231
18.66
72
8.48
146
14.56
Remark 4.1. The solve time for the largest problem is approximately 8.2 seconds. This represents significant improvement over recent parallel traditional solvers, such as source transfer Ju (2019) and trace transfer Leng & Ju (2022) methods, which typically require over fifty seconds to handle similar problems.
5 CONCLUSIONS
We modify MGCNN architecture and further introduce a novel network component, the Fourier Neural Network (FNN), improve performance for high wavenumber Helmholtz equations. FNN manages the information of full modes and scales to larger size and higher wavenumber problems. In standard supervised learning scenarios, MGCFNN surpasses other operator learning methods, showcasing superior learnability. Focusing on iterative solving for any desired accuracy and scala- bility to larger, higher wavenumber problems, MGCFNN significantly outperforms the recent specialized Al solver and AI Al-enhanced traditional solver for Helmholtz equations and sparse direct solver CHOLMOD. The ablation studies further demonstrate the benefits of multigrid hierarchy and the hybrid model pattern, verifying the complementary features of FNN and multigrid hierarchy. Both MGCENN and modified MGCNN exhibit an optimal convergence property of O(k) for high wavenumber & up to approximately 2000, and MGCFNN excels terms of number of iterations and solve time. All the experiments establish MGCENN as the satte-of-the-are Al solver for high wavenumber Helmholtz equations. Moreover, it reaches high wavenumber scope comparable the advanced parallel traditional solvers, while significantly reducing solve time.
In the present study, the interpolation operation in FNN remains simplistic, and the properties and best practices of FNN for solving Helmholtz equations are not yet fully explored. Furthermore, we have addressed various boundary conditions, including reflection boundary conditions, which may conflict with the periodic boundary effect of the Fourier transform or the padding technique. These limitations highlight areas for future exploration and improvement.
Reproducibility Statement. The code and configuration scripts used in our experiments are available ul/ttps:///itte com www..//iwwwooooooooooooeoe /ttpi/inno///////////////aa/a Our models can be trained using unsu- pervised learning, and we recommend this approach. To reproduce our results, you need some open datasets for speed (or slowness) fields. The speed data from the Ultrasonic CT area sourced from the A14S Cup and we provide the Marmousi data tested in Appendix G Other datasets are directly accessible through the provided PyTorch code.
Acknowledgements. This research is supported by the Strategic Priority Research Program of the Chinese Academy of Sciences, Grant No. XDB0640000, XDB0640300.
REFERENCES
Patrick Amestoy, Timothy A Davis, and Iain S Duff. An approximate minimum degree ordering algorithm. SIAM Journal Matrix Analysis and Applications, 17(4):886990) 1996.
10

---

Published as a conference paper at ICLR 2025
Patrick R Amestoy, Timothy A Davis, and lain S Duff. Algorithm 837: AMD, approximate minimum degree ordering algorithm. ACM Transactions on Mathematical Software (TOMS), 30 (3):881-388, 2004.
Yael Azulay and Eran Treister. Multigrid-augmentee deep learning preconditioners for the Helmholtz equation. SIAM Journal on Scientific Computing, 45(3):S127-S151, 2022.
Achi Brandt and Irene Livshits. Wave-ray multigrid method for standing wave equations. Electron. Trans. Numer. Anal, 6(162-181)991, 1997.
A Brougois, Marielle Bourget, Patriek Lailly. Michel Poulet, Patrice Ricarte, and Roelof Versteeg. Marmousi, model and data. In EAEG marrkhop-ppacttion aspects of seismic data inversion, pp. cp-108. European Association of Geoscientists & Engineers, 1990.
Henri Calandra, Senge Gratton, Xavier Pinel, and Xavier Vasseur. An improved two-grid pre- conditioner for the solution of dineesiinnni Helmholtz problems in heterogeneous media. Numerical Linear Algebra with Applications, 20(4))663-688 2013.
Yanqing Chen, Timothy A Davis, William W Hager, and Sivasankaran Rajamanickam. Algorithm 887: CHOLMOD, supernodal sparse Cholesky factorization and update/downdate. ACM Transactions on Mathematical Software (TOMS), 35(3):1-14, 2008.
Zhiming Chen and Xueshuang Xiang. A source transfer domain decomposition method for Helmholtz equations in unbounded domain. SIAM Journal on Numerical Analysis, (((4)2331-22566 2013.
Adam Coates, Andrew Ng. and Honglak Lee. An analysis of single-layer networks in unsupervised feature learning. In Proceedings of the fourteenth international conference on artificial intelligence and statistics, pp. 215-223 JMLR Workshop and Conference Proceedings, 2011.
Chen Cui, Kai Jiang, Yun Liu, and Shi Shu. Fourier neural solver for large sparse linear algebraic systems. Mathematics, 10(21):4014 2022.
Chen Cui, Kai Jiang, and Shi Shu. A neural multigrid solver for Helmholtz equations with high wavenumber and heterogeneous media. arXiv preprint arXiv:2404.002493, 2024.
Timothy A. Davis. Suitesparse: A suite of sparse matrix software. http://suttesparee.ccm 2024. Version 7.8.3.
Timothy A Davis and William W Hager. Dynamic supemodes in sparse Cholesky update/downdate and triangular solves. ACM Transactions on Mathematical Software (TOMS), 35(4): 1-:4:):12333 123, 2009.
Timothy A Davis, John R Gilbert, Stefan Larimore, and Esmond G Ng. Algorithm 836: COLAMD, column approximate minimum degree ordering algorithm. ACM Transactions Mathematical Software (TOMS), 30(3):377-380, 2004.
Björ Engquist and Andrew Majda. Radiation boundary conditions for acoustic and clastic wave calculations. Communications on Pure and Applied Mathematics, 32:313-357, 1979.
Yogi A Erlangga. Coreelis W Oosterlee, and Comells Vuik. A novel multigrid based preconditioner for heterogeneous Helmholtz problems. SIAM Journal on Scientific Computing, 27(4):1471-1492 2006.
Oliver G Ernst and Martin J Gander. Why it is difficult to solve Helmholtz problems with classical iterative methods. Numerical analysis of multiscale problems, pp. 925-363, 2011.
Paul Escapil-Inchauspé and Gonzalo A Ruz. Hyper-parameter tuning of chysics-informmd neural networks: Application Helmholtz problems. Neurocomputing, 561: 166::22626, 2023.
Jeffrey Galkowski, Shihua Gong, Ivan G Graham, David Lafontaine, and Euan A Spence. Conver- gence of overlapping domain decomposition methods with PML transmission conditions applied to nontrapping Helmholtz problems. arXiv preprint 200.v:2404.02/ 2024.
Martin Gander, Ivan G Graham, and Euan A Spence. Applying GMRES the Helmholtz equation with shifted Laplacian preconditioning: what is the largest shift for which eeeeuumbr--indepeneeee convergence is guaranteed? Numerische Mathematik, 131(3):567-614, 2015.
11

---

Published as a conference paper at ICLR 2025
Gene H Golub and Charles F Van Loan. Matrix computations. JHU press, 2013.
Ivan G Graham, Owen R Pembery, and Euan A Spence. The Helmholtz equation heterogeneous
media: a priori bounds, well-posedness, and resonances. Journal of Differential Equations, 266(6): 2869-2923, 2019.
Gaurav Gupta, Xiongye Xiao, and Paul Bogdan. Multiwavelet- operator learning for differential equations. In Proceedings of the 35th International Conference on Neural Information Processing Systems, pp. 24048-24066 2021.
Juncai He, Xinliang Liu, and Jinchao Xu. MgNO: Efficient parameterization linear operators via multigrid. In The Twelfth International Conference on Learning Representations, 2024.
Alena Kopaničáková and George Em Karniadakis. DeepONet based preconditioning strategies for solving parametric linear systems equations. SIAM Journal on Scientific Computing. 47(1): C15I-C181, 2025.
Alex Krizhevsky, Geoffrey Hinton, al. Learning multiple layers of features from tiny images. 2009.
Wei Leng and Lili Ju. An additive overlapping domain decomposition method for the Helmholtz equation. SIAM Journal on Scientific Computing, 4 A2): 1252-11227727 2019.
Wei Leng and Lili Ju. Trace transfer-be based diagonal sweeping domain decomposition method for the Helmholtz equation: Algorithms and convergence analysis. Journal of Computational Physics, 455:1 10090, 2022.
Bar Lerer, Ido Ben-Yair, and Eran Treister. Multigrridauugmeneed deep learning preconditioners for the Helmholtz equation using compact implicit layers. SIAM Journal Scientific Computing, pp. S123-S144, 2024.
Yichen Li, Peter Yichen Chen, Tao Du, and Wojeiech Matusik. Learning preconditioners for conjugate gradient PDE solvers. In International Conference on Machine Learning, pp. 19425-19439. PMLR, 2023.
Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli,, Kaushik Bhattacharya, Andrew Stuart, Anima Anandkumar, et al. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations, 2021.
Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators. Nature machine intelligence, 333):218-229, 2021.
Daniel Osei-Kuffuor and Yousef' Saad. Preconditioning Helmholtz linear systems. Applied Numerical Mathematics, 60(4):420 431, 2010.
Gerwin Osnabrugge, Saroch Lecdunronggwatthanakku and Ivo M. Vellekoop. A convergent Born series for solving the inhomogeneous Helmholtz equation arbitrarily large media. Journal of Computational Physics, 322: 1ee-11132244 2016. ISSN 0021-9991.
Md Ashiqur Rahman, Zachary Ross, and Kamyar Azizzadeneshelii U-NO: U-shaped neural operators. Transactions Machine Learning Research, 2023.
Maziar Raissi, Paris Perdikaris, and George Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 778:686-707, 2019.
James V Ralston. Trapped rays in spherically symmetric media and poles of the scattering matrix. Communications on Pure and Applied Mathematics, 24(4):571-582. 1971.
Steven C Rennich, Darko Stosic, and Timothy A Davis. Accelerating sparse Cholesky factorization on gpus. In 2014 4th Workshop on Irregular Applications: Architectures and Algoritims (IA 3), pp. 9-16. IEEE, 2014.
12

---

Published as a conference paper at ICLR 2025
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical image computing and computer- assisted intervention-Mee MMCCCAI 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part IIl 18, pp. 234-241. Springer, 2015.
Alexander Rudikov, Vladimir Fanaskov, Ekaterina Muravleva, Yuri M Laevsky, and Ivan Oseledets. Neural operators meet conjugate gradients: The FCG-NO method for efficient PDE solving. In International Conference Machine Learning, pp- 42766-42782. PMLR, 2024.
Abdul Hanan Sheikh, Domenico Lahaye, and Cornelis Vuik. On the convergence of shifted Laplace preconditioner combined with multilevel deflation. Numerical Linear Algebra with Applications, 20(4):645-662, 2013.
Steven W Smith. The scientist and engineer's guide to digital signal processing. California Technical Pub, 1997.
Chao Song, Tariq Alkhalifah, and Umair Bin Waheed. versatile framework solve the Helmholtz equation using ohysics-informed neural networks. Geophysical Journal International, 228(3): 1750-1762, 2022.
Kim Stachenfeld, Drummond Buschman Fielding, Dmitrii Kochkov, Miles Cranmer, Tobias Pfaff, Jonathan Godwin, Can Cui, Shirley Ho, Peter Battaglia, and Alvaro Sanchez-Gonzaler Learned simulators for turbulence. In International conference on learning representations, 2021.
Antonio Stanziola, Simon R Arridge, Ben T Cox, and Bradley E Treeby. A Helmholtz equation solver using unsupervised learning: Application to transcranial ultrasound. Journal of computational physics, 441:1 10430, 2021.
Vladislav Trifonov, Alexander Rudikov, Oleg Iliev, Ivan Oseledets, and Ekaterina Muravleva. Learn- ing from linear algebra: A graph neural network approach to preconditioner design for conjugate gradient solvers. arXiv preprint arXiv:2405.15557 2024.
Haixu Wu, Tengge Hu, Huakun Luo, Jianmin Wang, and Mingsheng Long. Solving high-dimensional PDEs with latent spectral models. In International Conference Machine Learning, pp. 37417- 37438. PMLR, 2023.
Yan Xie, Minrui Lv, and Chensong Zhang. MGCNN: learnable multigrid solver for linear PDEs on structured grids. arXiv preprint 11X1v:2312.11093, 2023.
A FNN WEIGHTS FREQUENCY AND SPACE DOMAINS
Consider the weights of the FNN in the frequency domain and their equivalent kernels in space domain at each level of a MGFNN, as depicted in Fig. 5 and Fig. 6 The weights in the frequency domain appear smooth, and in the space domain, they exhibit larger kernel size than common CNNs. These observations align with our design assumptions. Moreover, the weights in the space domain behave like damped wave propagation.
(a) Level
(b) Level 2
(c) Level
Figure 5: The weights of FNN in frequency domain.
13

---

Published as a conference paper at ICLR 2025
4050
540
1.0
(a) Level
(h) Level 2
(c) Level
Figure 6: The equivalent kernels of the weights of FNN in space domain.
(a)
(b)
Figure 7: The inverse of the damped Helmholtz operator in frequency domain and its corresponding weights in space domain. = 0.5, wkh 0.5 and N 1000.
To further elucidate this, we examine the expression of the inverse of a damped Helmholtz operator in the frequency domain. The one-dimensional stencil of the damped Helmholtz operator can be expressed as follows:
[-1,2 - - 1]]
where k2,, ==>>22(11->)) 71). For mode -2man, with frequency m, where n is the grid index and N is the grid size, the inverse of the damped Helmholtz operator on this mode is:
2(1 - cos( 2m N T))
2-2 Kiy
If ki.. remains invariant with the increase of N, and increases the same rate, this operation result also remains invariant. Given large damping factor 0.5, damped Helmholtz operator in the frequency domain in Fig. 7a and its corresponding weights the space domain in Fig. 7b] Similar patterns Thus, we can interpret FNN as learning a damped Helmholtz operator in the frequency domain.
view the inverse of the
observed in both the frequency and space domains.
Remark A.1. We observe that the weights in the frequency domain are not smooth enough, which is expected since wave with certain frequency will exhibit centralized energy in the frequency domain. As shown in Fig. 5 the circles appear suboptimal with some serration resulting from simple interpolation operation. This is potential area for improvement future work.
B CONNECTION WITH BORN SERIES METHOD
The Born series method Osnabrugge et 2016) uses the Green's function theorem solve the Helmholtz equation in heterogeneous medium. Specifically, for a Helmholtz equation
Vu((x) (()) fxx), xE Rdd
(9)
14

---

Published as a conference paper at ICLR 2025
rewrite it in damped form
(2((2) + (Kg ic)u(x) = f(x) (Mx(2) ic)u(x:).
(10)
Suppose go is the Green's function of the damped Helmholtz operator, and denote G = F- goF. Further, denote V the diagonal matrix of k(z)² kū ie. We can rewrite Eq. 10
goFF
u = GVu Gf.
(11)
The Born series method solves Eq. 11 iteratively as series expansion of the operator 1/(1 GV), i.c.,
u =
GV + (v (.v))++..)G)))
(12)
This method can be iteratively written
u A+1 = GVu" + Gf,
(13)
which is similar to solve phase in FNN
2 = FNN (setuppout mi) + xy.
(14)
Recall that FNN performs like a damped Helmholtz operator, setup out, contains information from k<x) and xy the current state, starting from f initially. Thus, the FNN in our solve phase level can be interpreted as learnable Born series method for solving the Helmholtz equation.
C ITERATIVE FRAMEWORK AND MIXED PRECISION STRATEGY
In the context of a linear system A soi = rhs, where A is square matrix, sol is the solution vector, and rhs is the right-hand-side vector, the objective of linear solver find the solution sol that meets satisfactory tolerance which can be obtained iteratively enhancing the solution's 's accuracy.
One such iterative method is the stationary iterative method, which solves linear systems by iteratively addressing the error equation A sol where r = rhs A sol represents the residual vector. This method can be defined as follows:
+
sol + B nK = sol + B (rhs sot
(15)
where so the solution the k-th iteration and B serves as the solver, providing an approximate сттоr correction for each residual p² In our experiments, our neural network functions as the solver B in the iterative framework.
The Krylov subspace method Golub & Van Loan 2013) is another iterative framework where the solver B
preconditioner. This method often speeds up the iteration process, albeit with some overheads. In this category, we use the Generalized Minimal Residual (GMRES) method and apply a restart technique to prevent storage excessively large subspace vectors and to reduce the computations required for subspace orthogonalization. In our experiments, we set the restart number to 10, unless convergence is achieved rapidly, in which case it is set
To attain high accuracy, we employ straightforward mixed precision strategy for both iterative frameworks. The computationally intensive part, either stationary iteration's 's solver network B the entire GMRES solve process before restart, operates in float32 precision. To avoid error accumulation, the residual is updated in float64 precision.
D FNN PSEUDOCODE
FNN operates in two phases: a setup phase (Alg. 2 to initialize the weight tensor in the frequency domain, and solve phase (Alg. B) perform spectral convolution. Note that this setup phase should not be confused with the setup phase network described in Subsection 3.2 which takes the problem coefficients input.
15

---

Published as a conference paper at ICLR 2025
Algorithm 2 FNN setup phase
1: INPUT: size is the size of the input tensor of the FNN.
2: OUTPUT: weigght the tensor the FNN weights in frequency domain to do linear mapping.
3: procedure SETUPENN (EGGHT(size)
4: init_weight < loaded initialized
5: interp.sise size f- loaded or settled
6:
if size interp size then setup_weight + interpolate(miit.ueigghe
case 1
7:
8: else
case
9:
space_weight + F wiiit))eigh))
space_weight pad space_weight to size end
11:
setup weight + Fssaace weight)
13:
return sclup_weight
end procedure
Algorithm 3 FNN spectral convolution
1: INPUT: impul is the input tensor of the FNN, selup_weight is the tensor of the FNN weights in frequency domain.
2: OUTPUT: oulput is the output tensor of the FNN.
3: procedure FNNCoNv(input,, setup_weight
4: C' number of channels of the input tensor
5: input emplax t- imput(; C/2] imput :)
D real complex value
6: pad input_cmppl.
7:
input-freq f (inputt.cmpl output-freq +- linear tenearr_mpp(e
8:
setup.weight,
9:
t F (output-ffee)
10:
unpad output_cmpla
11:
output +- concat concat(output cmplæ. output emplæ.imag)
b complex to real value
12: return outpul
13:
end procedure
The linear map in Alg. uses setup_wee weight per pixel to transform the input channels to output channels per pixel, which can be expressed as follows in Einstein summation convention:
(16) where b denotes the batch index, is the output channel index, i is the input channel index, and x. are the spatial indices. Therefore, we observe that setup weight has a size of C²n², which restricts the application of FNN to large-scale problems. Fortunately, it is only required for certain coarse levels in the multigrid hierarchy.
OUTi y] 1N,.; uj (o y]
E A COMPUTATIONAL INEFFICIENCY ISSUE
drawback of the FNN is the unexpected inefficiency computing the linear map operation (see Eq. 16) in the frequency domain. The computational complexity of this operation is O(C2n²) for a n grid with channels, which should ideally be as swift as convolution with 1 1 kernel. However, our timing tests reveal that this operation takes even longer than the FFT2d in PyTorch. This issue less apparent in FNO-like methods, which consider significantly fewer modes. Consequently, future work should focus on developing more efficient implementation of the linear map in FNN.
F SCALABILITY TEST DETAILS
We conduct scalability tests on 50 unseen data with varying grid sizes and angular frequencies. The results, depicted in Fig. & reveal that the iterations required for convergence vary across
16

---

Published as a conference paper at ICLR 2025
different data, characteristic typical of the Helmholtz equation. Indeed, certain speed (or slowness) distributions can exacerbate the ill-conditioned nature of the Helmholtz equation, leading to trapping phenomena Graham al. (2019); Ralston (197I.
40
5333533
182251100
33
322385223
4080s41090
23
422504023
32.3
17.3
13.0
12.5
10.0-
7.5
25
20
100
500
0.2
0.5
2.5
12.
60.5
Figure 8: Scalability test over 50 unseen data. The upper is the results of MGCFNN, and the lower row is the results MGCNN.
G GENERALIZATION TO OTHER DISTRIBUTIONS
We further investigate the generalization of MGCFNN to other distributions. the experiments, we apply the model trained on the CIFAR-10 dataset to the STL-10 dataset, which has higher image resolution 96 x 96, and The results are presented in Table 6 where MGCFNN demonstrates good generalization performance with slight increase in number of solving iterations. This increase is expected, as more complex medium results in more optic rays, which influence the convergence of the iterative solver Galkow/ll et al. 2024). Figures Ya and 9b illustrate the single source results on a STL-10 data and Marmousi model, respectively.
benchmark Marmousi data, which a resolution of 362 x 1101.
Table 6: Generalization results number iterations on STL-10 dataset and Marmousi. The model trained on CIFAR-10 dataset.
rtol=1B-7, iters
standalone
GMRES
w
size
CIFAR-
STL-10
Marmousi
CIFAR-
STL-10
Marmousi
80t
511 x 511
14
39
26
12
32
21
160m
1023
1023
22
41
35
20
33
29
320m
2047 x 2047
41
59
52
51
42
640m
4095 x 4095
83
109
92
72
94
79
H MODEL ARCHITECTURE SETTINGS
Before discussing the architecture settings, it is important to note that the number of points per wavelength critical factor in determining the model architecture. For problems smaller than 511 x 511, there are about 6 points per wavelength, while for larger problems in scalability tests, there are approximately 12 points per wavelength. We set the number of multigrid levels such that the coarsest level has around points per wavelength, and then add more levels to optimize the results. For MGCENN, the level with around points per wavelength is also the starting level where the FNN used in the solve phase network. We find that using FNN in the coarsest levels is beneficial. Additionally, to fully utilize GPU resources, we set larger number of channels for smaller problems.
Detailed model architecture settings for the solve phase presented in Tables and g
17

---

Published as a conference paper at ICLR 2025
300
(a) Single source wave with w 80m, in slowness field generated from a STL-10 data.
Wiail
80
(b) Single source wave with 160T, a slowness field from Marmousi.
Figure 9: From left to right: slowness field x(x), real value of wave field and wave field in frequency domain, iterative solving history. The cross symbol in the first column denotes the position of source.
Table 7: Solve phase settings for around points in a wavelength.
Model
Settings
MGCENN, MGENN
32 channels, 1, 2, 4 sweeps for level 1, 2,
FNN
32 channels, 8 sweeps for one level
Table Solve phase settings for around points wavelength.
Model
Settings
MGCFNN MGCNN
12 channels, 1, 2, 4, 8 sweeps for level 2, 3, 4
12 channels, 4, 8, 16 sweeps for level 1, 2, 3, 4,
For the kernel sizes in both the solve and setup phases, CNNs within ResNets 5, while the restriction and prolongation CNNs use a kernel size of All skip add operations are performed with kernel size of 1. In the setup phase, we set two sweeps for each level, and the number channels is consistent with the solve phase to ensure network compatibility.
kernel size of
I TRAINING SETTINGS
Table Training settings.
Parameter
Value
Description
epochs
120
num_data
10000
number data used in training
Ir
0.001
initial learning rate
optimizer
Adam
scheduler: step size=6, gamma=0.8
batch_size
10
18


Reviews:
[{"review_id": "LM2n2fQDDG", "reviewer": "Reviewer_T4nv", "summary": "Authors develop a hybrid iterative solver for Helmholtz equation. A proposed technique is a combination of multigrid-like architecture with a specially modified Fourier Neural Operator and a classical modified Richardson iteration. As I understand (although the description is not complete as I discuss below), authors apply a neural network to produce parameters of linear operators from PDE data and later apply these linear operators within a multigrid-like architecture. The resulting solver shows promising results on several high frequency problems.", "questions": "1. The description of the method is not entirely clear.\n   Authors provide descriptions of their approach in various places in the article. In my opinion these descriptions are confusing and incomplete. It is possible that other readers will find descriptions insufficient as well. Below is a list of questions on architectures authors used:\n   1. Appendix E contains parameters of MGCFNN and MGFNN. Are these parameters for the setup phase network or for the solve phase network (or both)?\n   2. Does the solve phase network have other parameters than the one produced by the setup phase network? If the answer is yes, what are those parameters?\n   3. Is it the case that authors explicitly used residual only on the fine level, i.e., in equation (9) network implement linear operator B?\n   4. Authors write about inefficiency of linear transformation in appendix D. From the number of operations authors provide it seems that this is a convolution with kernel size 1. Is that the case? Can the authors provide a number of floating point operations their architectures perform? This is highly relevant for comparison with classical iterative methods.\n1. Not all relevant baselines are covered.\n   1. Many well developed classical methods are available for the Helmholtz equation. It looks problematic that authors do not compare with any of them. For some reason all baselines are other ML approaches. I suggest the authors provide results for at least shifted-Laplace preconditioners and ideally with the optimized Schwartz method https://arxiv.org/abs/1610.02270. It is especially interesting to consider time needed to achieve desired accuracy since ML-based preconditioners are typically more computationally involved.\n   2. Since authors consider only 2D cases, sparse direct solvers are also relevant baseline (see section 4 in https://arxiv.org/abs/1610.02270). I suggest providing solution time for modern CUDA implementation of sparse LU, preferable using freely available SuiteSparse solver https://people.engr.tamu.edu/davis/suitesparse.html.\n   3. Sparse algebraic ML-based preconditioners (often GNN-based) form another class of relevant baselines, for example, https://arxiv.org/abs/2405.15557, https://arxiv.org/abs/2305.16432, etc\n   4. Completely nonlinear methods based on flexible Krylov methods are also a valid solution approach https://arxiv.org/abs/2401.02016v2, https://arxiv.org/abs/2402.05598.\n2. Dataset construction and generalization.\n   First, the dataset authors used is detached from practically relevant applications of the Helmholtz equation. Second, a main dataset is formed from CIFAR-10, which has only 32x32 resolution which is interpolated to access performance on grids with higher resolutions. The troubling part is that fields, produced in such a way, have fixed frequency content, which is already known at the lowest resolution. Given that, it does not allow one to test whether the preprocessing network used to construct linear operators can generalize well to the problems of other resolutions. I suggest authors provide generalization tests where $k(x)$ contains higher frequencies on fine grids.\n4. The fixed number of levels is not completely satisfactory.\n   Authors used multigrid with a fixed number of layers. This means, with the increase in resolution the performance of the solver will degrade. I can suggest authors to empirically demonstrate that this is not the case. A possible experiment is to train architecture with two levels (i.e., the smallest possible) on a small grid and after that apply it to grids with increasing resolutions (much larger than the original grid). If one observes no degradation, the fixed number of levels is not a problem.\n5. Is a comparison with NO fair?\n   Neural operators have only access to input-output pairs, while the method in the current article explicitly uses discretisation. I think it is appropriate to explicitly discuss this distinction.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The problem authors consider is well motivated, since it is known that classical iterative methods suitable for elliptic problems perform poorly on the Helmholtz equation.\n2. A proposed method is a combination of multigrid and machine learning techniques, so, potentially, it scales well, is consistent and may be more robust than purely ML-based solvers.\n3. Benchmarks authors propose cover especially challenging problems with high wavenumbers.", "weaknesses": "An article has several weaknesses which I briefly describe below:\n1. The description of the method is not entirely clear.\n2. Not all relevant baselines are covered.\n3. The dataset is chosen in a very specific manner, making generalization hard to access.\nI elaborate on these and other issues in the next section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ogJI7Q973y", "reviewer": "Reviewer_pjJn", "summary": "The paper introduces MGCFNN, a multigrid-hierarchical neural solver designed to address high-wavenumber Helmholtz equations, which are challenging for both traditional and AI solvers. The authors propose a Fourier Neural Network (FNN), an enhancement of spectral convolution operations, specifically to manage full-mode information and accelerate wave propagation during solving. The authors evaluate MGCFNN on both supervised and unsupervised tasks, comparing its performance against various neural operator learning methods.", "questions": "1. Can you provide more comparison experiments especially with other types of equations?\n\n2. Can you provide details on your hyperparameter selection and cost of tuning and selection?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is well written and easy to follow\n2. The paper presents extensive experimental validation, including scalability tests and comparisons with recent operator learning methods.\n3. This is interesting idea for generalizing FNO to a more generic architecture.", "weaknesses": "1. Comparison is limited. The paper mainly focuses on a few high-wavenumber scenarios for Helmholtz equations. There are a broad range of equations worth to be evaluated. \n2. Hyperparameters: the paper lacks discussion on the choice and tuning cost of their hyperparameters, such as number of multigrid levels, number of Fourier modes, etc.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LdNL5ARNOK", "reviewer": "Reviewer_yYix", "summary": "This work presents an AI-driven PDE solver that specifically utilizes a Fourier neural network. By independently processing the coefficient and source term using a modified spectral convolution operation, it effectively handles high wavenumber Helmholtz equations.", "questions": "* In Figure 2, the visualization of the network architecture is difficult to understand", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The modification of spectral convolution successfully handles full-mode information\n* MGCFNN effectively handles the linearity of source terms\n* It demonstrates excellent scalability with high-wavenumber problems", "weaknesses": "* Does MGCFNN work well only for high-frequency problems?\n* Can MGCFNN be applied to other PDEs in operator learning?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7D0MTJ0TJ4", "reviewer": "Reviewer_q7vM", "summary": "This paper introduces MGCFNN, an AI-based solver tailored for high wavenumber Helmholtz equations, which are  challenging, especially in heterogeneous media. Building on the MGCNN framework, the authors incorporate a Fourier neural network (FNN) to improve solution efficiency and accuracy, leveraging FNN's CNN-like structure for faster propagation of source influence. Experimental results reveal MGCFNN's promising performance in supervised learning compared to existing neural operator methods, as well as notable scalability improvements in unsupervised settings. An ablation study underscores the value of the multigrid hierarchy and FNN, with the solver achieving optimal convergence up to $k\\approx 2000$.", "questions": "In addition to the points noted in the \"Weaknesses\" section, I have the following questions:\n1. Tables 1 and 2 present results with $\\omega=80\\pi$. How do the results compare at $\\omega=640\\pi$, given that the primary goal of this work is to address high-wavenumber cases?\n2. The numerical results in Table 4 do not show a significant difference between MGFCNN and MGFNN. Additional arguments and comparisons between these two methods would help clarify their respective advantages.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The framework presented in this work, with its division into setup and solve phases, is both well-structured and innovative.\n2. The paper introduces a novel Fourier neural network (FNN) that enhances the solver’s ability to manage high wavenumber problems efficiently, particularly within heterogeneous media, providing a significant advancement over traditional CNN-based approaches.\n3. The manuscript includes multiple numerical examples that demonstrate improved performance over existing neural operator architectures. Additionally, the results in Tables 3 and 5 are particularly compelling and highlight the model's effectiveness.", "weaknesses": "1. The paper would benefit from a more detailed discussion of its novelty relative to prior work on MGCNN, explicitly highlighting the advancements MGCFNN brings over previous frameworks.\n2. While the motivation appears intuitively reasonable, adding solid mathematical insights to support the design of MGCFNN would strengthen the rationale and deepen the reader’s understanding of the architectural choices.\n3. The explanation in Section 3.1 lacks clarity, particularly regarding its connection to Sections 3.2 and 3.3. Providing clearer linkages between these sections would help readers follow the logic and flow of the methodology.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)