Neural Spectral Methods: Self-supervised learning in the spectral domain
Authors: Yiheng Du
Nithin Chalapathi
Aditi S. Krishnapriyan
Venue: ICLR 2024 poster (2024)
Keywords: Machine learning for PDEs
spectral methods
neural network differentiation
spectral loss
PDEs
neural operators
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We present Neural Spectral Methods to solve parametric PDEs in the spectral domain.

Abstract: We present Neural Spectral Methods, a technique to solve parametric Partial Differential Equations (PDEs), grounded in classical spectral methods. Our method uses orthogonal bases to learn PDE solutions as mappings between spectral coefficients, instantiating a spectral-based neural operator. In contrast to current machine learning approaches which enforce PDE constraints by minimizing the numerical quadrature of the residuals in the spatiotemporal domain, we leverage Parseval's identity and introduce a new training strategy through a spectral loss. Our spectral loss enables more efficient differentiation through the neural network, and substantially reduces training complexity. At inference time, the computational cost of our method remains constant, regardless of the spatiotemporal resolution of the domain.  Our experimental results demonstrate that our method significantly outperforms previous machine learning approaches in terms of speed and accuracy by one to two orders of magnitude on multiple different problems, including reaction-diffusion, and forced and unforced Navier-Stokes equations. When compared to numerical solvers of the same accuracy, our method demonstrates a $10\times$ increase in performance speed. Our source code is publicly available at https://github.com/ASK-Berkeley/Neural-Spectral-Methods.

Published as a conference paper at ICLR 2024
NEURAL SPECTRAL METHODS:
SELF-SUPERVISED LEARNING IN THE SPECTRAL DOMAIN
Yiheng Du, Nithin Chalapathi, Aditi S. Krishnapriyan {yihengdu, nithinc, Saitikk)@berkeley.eee
University of California, Berkeley
ABSTRACT
We present Neural Spectral Methods, a technique to solve parametric Partial Dif- ferential Equations (PDEs). grounded in classical spectral methods. Our method uses orthogonal bases to learn PDE solutions as mappings between spectral co- efficients, instantiating spectral-basee -bared neural operator. In contrast to current machine learning approaches which enforce PDE constraints by minimizing the numerical guadrature of the residuals in the spatiotemporal domain, we leverage Parseval's identity and introduce a new training strategy through a spectral loss. Our spectral loss enables more efficient differentiation through the neural net- work, and substantially reduces training complexity. At inference time, the com- putational cost of our method remains constant, regardless of the spatiotemporal resolution of the domain. Our experimental results demonstrate that our method significantly outperforms previous machine learning approaches in ferms of speed and accuracy by one to two orders of magnitude on multiple different problems. including ceaction-diffusion. and forced and unforced Navier-Stokes equations. When compared to numerical solvers of the same accuracy, our method demon- strates 10x increase in performance speed. Our source code is publicly available at hhttps://glthue wtt.s/hu.coooocccooooooooocoeeeh ASK Baw.-Bokkoey//eeeeeeoooo
Spect -Methos
INTRODUCTION
Partial differential equations (PDEs) are fundamental for describing complex systems like turbulent flow (Temam| 2001). diffusive processes (Friedman 2008). and thermodynamics (Van Kampen 1992). Due to their complexity, these systems frequently lack closed-form analytical solutions, prompting the use of numerical methods. These numerical techniques discretize the spatiotemporal domain of interest and solve set of discrete equations to approximate the system's behavior.
Spectral methods are one such class of numerical techniques, and are widely recognized for their effectiveness (Boyd [200T] Gottlieb & Orszag 1977). These methods approximate PDE solutions as a sum of basis functions and transform the equations into the spectral domain. Spectral methods are known for their fast convergence and computational efficiency, especially for problems with smooth solutions. They are notably impactful in fields like computational fluid dynamics (Peyret 2002).
Numerical methods can be computationally expensive because fine discretization of the physical domain and large number of time-stepping iterations are often required to achieve high accuracy. Additionally, many engineering applications require solving systems under different parameters or initial conditions, necessitating multiple iterations to repeatedly solve such systems. The aforemen- tioned spectral methods also rely on time-stepping schemes, and present similar challenges to other numerical methods. These factors underscore the need for more efficient computational strategies.
Recent advances machine learning (ML) highlight neural networks (NNs) as potential alternatives or enhancements to traditional numerical solvers. Consider PDEs on regular domain Ω Rd:
Fo(u(x)) (u(x)) = O on an,
in Z2
where is the classical solution, Fg is the potentially nonlinear differential operator, and Bg the boundary condition(s). Both operators are parameterized by o. which could correspond to initial conditions or parameters associated with Ω such as diffusion coefficients. A common MI. approach

---

Published as a conference paper at ICLR 2024
is to train NNs datasets of numerical solutions spanning various parameters, and then learn the mapping Go: Φ ++ Ug from parameters to solutions. This approach is typically done via supervised learning, where the NN is trained by minimizing the error between the predicted solution and the numerical solution. At inference time, the goal is to generalize to previously unseen parameters and directly predict solutions. By amortizing computational costs during training and having rapid infer- ence, these data-driven approaches have the potential to be more efficient than traditional numerical solvers.
When there is information about the governing physical laws, another common ML approach is train the NN through loss function that imposes a soft constraint on these laws (Raissi et al.) 2019). In this case, the NN predicts a solution ug. and then the corresponding residual function, R(x) = ()), is minimized on set of sampled spatial and/or temporal points. Specifically, an additional loss term is defined as the numerical quadrature to the residual norm:
LPINN == N
1
Σ
R(xn)² ((20))))))))) In D i...d. U(2).
NE(N]
This is often called a Physics-Informan Neural Network (PINN) loss function, which we denote as (PINN(u)). In practice, this approach does not require solution data the interior of the domain, and the NN can be trained by only satisfying the PDE(s) of interest. This loss function can also be combined with data fitting loss and trained together. However, this often be impractical, as it requires knowing both the underlying governing physical laws and having solution data (either through a numerical solver or through observational measurements).
Most current ML methods to solve PDEs (Kochkov et al. 2021 Li et al. 2020) are grid-based ap- proaches. To model parametric solutions, these models create a mesh, and perform the parameters- to-solutions mapping through non-local transformations on function values at points in the mesh. One such example are neural operators (Kovachki et al., 2021), which parameterize the mapping using iterative kernel integrals, and have been applied to wide range of engineering problems (Zhang et al., 2022] Kurth et al., 2023). Neural operators can be trained through supervised learn- ing procedures (a loss function that matches solution data), self-supervised methods (such as the aforementioned PINN loss function), and a combination of both.
Data-fitting and PINN loss approaches both have several limitations:
Data availability. The effectiveness of data-driven methods generally depends on the avail- ability of large datasets consisting of PDE parameters and corresponding solutions. For complex problems, solution data can only be generated from expensive solvers (or through observations), and also includes inherent numerical errors. When solving new systems, new solution data often needs to be generated, which can be time-consuming.
Empirical evidence has suggested that minimizing the PINN loss often en- counters convergence issues, particularly for complex problems, resulting in subpar accu- racy. This is likely attributed to the ill- pll-posed nature of the optimization problem that arises from incorporating PDE constraints into the model (Krishnapriyan et al. 2021 Wang al., 2021a).
Optimization.
Computation cost. Computing the PINN loss involves evaluating the differential opera- tor F at sampled points, which requires computing higher-order derivatives of ug. As the complexity of tig increases, the computation cost of back proppaatpattiion scales significantly. Moreover, accurate estimation of the residual norm requires substantial amount of sam- pled points to enforce the PDE. For neural operators such as the commonly used Fourier Neural Operator (FNO) Li et al. (2020), differentiation costs scale quadratically with the number of points. This is due to the use of Fourier transform, and this makes it intractable take exact derivatives for large numbers of sampled points (see discussions in $B)
To address these issues with accuracy and efficiency, our work explores the incorporation of ML. with spectral methods. We focus on a data-constrained setting to learn the solution operator of parameterized PDEs, where we assume that we have no solution data on the interior of the spa- tiotemporal domain and only train our model by minimizing the PDE residual. Given the form of differential operator Fs, the model learns map the parameter function á to the corresponding solution ug. Our key insights are to learn the solution as a series of orthogonal basis functions, and leverage Parseval's 's Identity to obtain loss function in the spectral domain. While prior approaches
2

---

Published as a conference paper at ICLR 2024
minimize the approximated residual function by computing higher-order derivatives the sampled points, our method exploits properties of the spectral representation to obtain the exact residual via algebraic operations on the prediction.
Our contributions are summarized as follows:
We propose Neural Spectral Methods (NSM) Iearn PDE solutions in the spectral domain. Our model parameterizes spectral transformations with NNs, and learns by minimizing the norm of the residual function in the spectral domain. Since solution data can be expensive to generate for every new problem, we focus on scenarios with no solution data on the interior of the domain.
We introduce a spectral-base bassd neural operator that can learn transformations in a spectral basis. By operating on fixed collocation points, our proposed spectral-based neural operator avoids aliasing error and avoids scaling the computational cost with grid resolution.
We introduce a spectral loss to train models entirely in the spectral domain. By utilizing the spectral representation of the prediction and Parseval's identity, the residual norm is computed by exact operations on the spectral coefficients. This approach avoids sampling a large number of points and the numerical quadrature used by the PINN loss, thereby simplifying computation complexity.
We provide experimental results three PDEs: Poisson equation
Deaction-Diiffusion equation $4.2, and Navier-Stokes equations $4.3 Our approach consistently achieves a minimum speedup of 100x during training and 500 during inference. It also surpasses the accuracy of grid-based approaches trained with the PINN loss by over 10x. When tested on different grid resolutions, our method maintains constant computational cost and solution eтror. In comparison to iterative numerical solvers that achieve equivalent accuracy, our method an order of magnitude faster.
2 RELATED WORKS
ML methods for solving PDEs. Using NNs to solve PDEs has become an active research focus in
scientific computing (Han et al. 2018 Raissi et al. 2019
Lu et al. 2021) He et al. (2018); Mitusch
eet al. (202T) explore finite element methods in NNs. Yu et al. (018); Sirignano & Spiliopoulos
(2018) ; Ainsworth & Dong (202]); Bruna et al. (0022 recast PDEs into variational forms and ap-
ply the Galerkin projection minimization via sampling-based losses. (hhamaa & Shankar (2022)
accelerate discretized PDE residual computation. Wang et al. (2021b) focus on Fourier features and data-fitting; Dresdner et al.) (2022) learn to correct numerical solvers: Xiia et al. (2023): Lütjens et al. (2021) use spectral representations in spatial domains; and Meuris al. (2023) extract basis functions from trained NNs for downstream tasks. These studies differ in problem settings and deviate from our method in both architecture and training.
eigenvalue problems. In the context of spectral methods, |ange al. (2021) focus
Neural operators. Neural operators (Kovachki et al. 2021: Li et al. 2020; Gupta et al. 2021) learn mappings between functions, such as PDE parameters or initial conditions to the PDE solu- tions. They are typically trained in a supervised learning manner on and aim to learn the functional mappings between them (i.c., the solution operator). However, the supervised learning setting poses a challenge in data generation for new or complex problems, es- pecially when the data is scarce the numerical solvers generating it are inefficient.
datasets,
One of the most common neural operators is the Fourier Neural Operator (FNO) eLi al. 2020 Kurth et al. 2023 Zhang et al., 20022). The training process for FNO consists of performing convo- lutions through Fourier layers and learning in the frequency domain. The Spectral Neural Operator (SNO) (Fanaskov & Oscledets 2022) was proposed to reduce aliasing error in general neural op- erators by utilizing ferd-forward network to map between spectral coefficients. Similarly, the TransformOnce (TI) (Poli et al. 2022) model looks learning transformations in the frequency domain with an improved model architecture. However, both models have number of architecture and training differences from NSM, and as we will show, have much poorer accuracy. They also only look at the supervised learning setting.
Physics-Informel Neural Networks (PINNs). The physics-informed neural networks (PINNs) framework (Raissi et al.] 2019) adds the governing physical laws (i.e., PDE residual function), esti- mated on sampled points, into the NN's loss function. This additional term, which refer to
3

---

Published as a conference paper at ICLR 2024
PINN loss (Eq. 2), acts as a soft constraint to regularize the model's prediction of the PDE solution, and can be considered a self-supervised loss function. This approach can also be used in a operator learning setting across various architectures, where the base architecture is neural operator and the PINN loss is used to train the model Li et al. (021); Tripura et al. (2033: : Rosofsky et al. (2023) Goswami al. (2022).
However, the PINN approach requires evaluating the PDE residual on a large number of sampled points in the interior domain. In scenarios with higher-order derivatives in the PDE residual, multi- ple differentiations through the NN are required. When using the PINN loss with grid-based neural operators such as FNO, the Fast Fourier Transform (FFT) in the forward pass escalates to quadratic in batched differentiation with respect to the number of sampled points, making exact residual com- putation through aute-differentiation computationally expensive (see discussions in $B for more details). As we will show, these grid-based methods are inaccurate and often overfit to the training grid size because of aliasing error.
NEURAL SPECTRAL METHODS
We introduce Neural Spectral Methods (NSM), general class of techniques to incorporate spectral methods with NNs. NSM learns the mapping from PDE parameters to PDE solutions, i.c. Ge : + ue and is shown in Fig. NSM consists of two key components: base NN architecture (Fig. Ta that maps the spectral representation the parameters ø to that of its solutions ug, and a spectral training procedure (Fig. lb) that minimizes the spectral norm of the PDE residual function.
I
In $32 we describe our core NN architecture, which incorporates spectral methods within a neural operator framework. In $3.3 we introduce our spectral training procedure. Finally, in $4 we demon- strate the strong empirical performance of NSM for learning solutions to different PDE classes
3.1 BACKGROUND
Notation. A set of functions Ω is denoted as {. (x)) Qm where I is countable index set. We denote integer indices by [n)] = 21,2..., n. The nith component of a vector & is denoted as Tn- Given set of basis functions, the spectral coefficients of function u are denoted as ū. For an integer k > 0, the kith order Sobolev space (Evans 2022) on domain Ω C Rd is denoted as 7t (S2).
Orthogonal basis. Orthogonal basis functions are fundamental components of spectral methods. For completeness, the definition of orthogonality is provided in S The choice of the basis functions problem-specific and they must have desirable spectral properties. In this paper, we focus on two commonly used orthogonal bases for periodic and other types of boundary conditions, respectively:
Example (Fourier basis). {sinn (mx): : E BrTmee w.r.t Lebesgue measure.
Example 2 (Chebyshev polynomials). (T)( E []1, MEN W.E. de
1/vI- 2², where Tm (x) == cos cos 1 (x)) = (-1) sin(m sin (x)).
du
For multi-dimensional problems, we can verify that the product of basis in each dimension preserves completeness and orthogonality (see Proposition []. Spectral representations are known for their efficiency in representing smooth functions. Specifically, Fourier and Chebyshev interpolations pos- sess the well-known spectral decay for sufficiently smooth functions (Mason & Handscomb] 2002):
Fact For any I E HT*),,, its Fourier series coefficients [ = O(1/mp).
Fact For any f E [P[[[1, 1J its Chebyshev expansion Ifml = O(1/m²).
Neural operators.
Neural operators parameterize the mapping Go + up
Go Q - (W)())) + xx(1)) o.... ((W(1) x P.
where a is a non-lincarity. The operator iteratively updates v/" + R4, where dj is the hidden dimension of layer 1. The input layer ( P()) and output layer uo (() are parameterized by wise NNs. The affine operator With) and the kernel integral operator k00) are defined as,
(w)))) = w(e(()) + 10) (x)))) = S. k (.. ((..)))())d(()))
(4)
4

---

Published as a conference paper at ICLR 2024
a)
b)
(++
Ri
6%
Ma2
Rc
Parseval's Identity
Spectral Loss
La R(x))d((x)> ||
Ru
T
Spectral Physical
NSM
(xx
7-1
d) Autograd
((9)
ux(xs> w
(p(x))
Vur(x)) Vuy(())
PINN
(0xg)
(o(%))
Loss
(++1)
or FDM
d6xx)
Up
#p(Xx) Vay(Rt))
PINN
Figure 1: Schematic of NSM. We refer Neural Spectral Methods (NSM) general approach to learn PDE solution mappings in the spectral domain. Our method consists of two components: a) The parameters á are converted to spectral coefficients, Ø. In each NN layer 1, the spectral coeffi- cients 0g are transformed by linear operator K, with the activation then applied collocation points in the physical space. b) The prediction üg transformed by Fs. the spectral form of the differential operator, which gives the spectral coefficients R of the residual function. The exact residual norm is obtained by Parseval's Identity, giving the spectral loss [| We contrast our method against the commonly employed grid-based neural operators with PINN loss. c) General neural operators learn the PDE solutions as transformations of function values on x,,. consider the kernel integral in more general sense, with transformation T not restricted to Fourier basis. d) Autograd or finite difference methods are used to obtain the higher-order derivatives. The PINN loss is then obtained by approximating the norm of the residual function the sampled points.
(!()
where the input is + Rd41 and the outputs are in Rd,, The affine operator and kernel integral operator are used to capture local and non-local transformations, respectively. Given input grid [x;). general grid-based neural operators parameterize Go as mapping between function values:
[ ((I2)
x(xN)) + [u)(
ua(2)) uo(xN)]
(5)
3.2 SPECTRAL-BASHR NEURAL OPERATORS
In this work, we employ the neural operator architecture to model transformations between spectral coefficients. By fixing {fm (x) parameterized as the mapping between the coefficients the parameter functions & and the pre- dicted solutions ug. Suppose the series is truncated to M terms, then Go is parameterized in the spectral domain as:
Ω me the chosen basis functions, the solution operator
x(x)=
@m)m(()) uz(x) ME(M]
Σ
Σ ug (x), mE(M]
where &6 and ūg are the spectral expansion o and us under the basis f. For each layer I, the function ui (1) and kernel ((() are also parameterized under the same basis functions,
(1) (I)
Σ
(11) fma (x.v) Σ K fm. (x)fma())
(7)
ME(M]
ME[M]
where 00 R Mxd: and R(() E MxMxxxxd are coefficients in the spectral domain. Due to orthogonality, integral transformations K()) actually equivalent to tensor contractions a/(1)) ke(1)
(([-1) mymai Similarly, affine transformations Wil are equivalent to o0 w + (
Non-lineat activation functions. The activation is applied on the collocation points. We denote as the interpolation of function values at collocation points aligned with the basis, and T-1 as the function value evaluation on those collocation points. Then the spectral counterpart of the activation function a, is given by: T
(8)
5

---

Published as a conference paper at ICLR 2024
Aliasing error. General grid-based neural operators prone to aliasing error. When trained and tested on different grid resolutions, the interpolation error is inconsistent, leading to an increased error on the test grids that a different resolution from the training grid. In contrast, based approach circumvents aliasing error. By operating exclusively on fixed collocation points, the interpolation error remains consistent across different input resolutions. This also ensures that the model's 's computation cost and predictions are rssoltiin-indeeenden
spectral-
Kernel approximation. Within each layer, the computational cost of the kernel integral is quadratic. To mitigate this cost, the kernel can be confined to a low-rank or simplified form. Here, employ FNO's approach, which simplifies the kernel integral into convolution through Fourier transformation. More broadly, the kernel is restricted to diagonal form, K0()) Raxxxxd--1 and,
(x, 9)
Σ mEIM]
K fm))
3.3
SPECTRAL TRAINING
After our base NN architecture predicts the solution in Eq. 6 we train our model using a spectral training procedure. Here, we describe the details of our spectral training procedure.
Spectral form of the residual function. We can derive the exact residual function using the spec- tral representation of the prediction, denoted by u. The spectral representation has direct corre- spondence between operations performed on function values and on spectral coefficients. Addi- tionally, differentiation and integration are transformed to algebraic operations on the coefficients.
Given the PDE operator Fsr we convert it to its spectral correspondence F ügg + R, such that,
[x))
Σ (m)
(10)
mEI
where represents the spectral form of the PDE residual function. We describe in more detail how to obtain F& from Fg for typical nonlinear operators and bases composed of Fourier series and Chebyshev polynomials in A.
Spectral loss. After computing the residual function, we aim to minimize our spectral loss, ||| This method involves projecting the residual function onto a subspace spanned by truncated basis functions, known as the Galerkin projection. The orthogonality of the basis functions is crucial to this procedure, Leveraging Parseval's Identity, we can equate the spectral loss to the weighted norm of the residual function, outlined below: Theorem (Parseval's 's Identity). For R(x) =Emer fm (2)), we have,
J R(x))d((x)- = Σ RRn
(1I)
mEI
Contrast with the PINN loss. As previously discussed, the PINN loss is obtained by sampling a substantial number of points within the interior domain, followed by employing numerical quadra- ture to approximate the integral of R(x))² This requires differentiating through the NN, which can be computationally expensive, especially for higher-order derivatives, and when a large number of points are sampled to ensure accurate quadrature. For the grid-based PINN loss, which commonly uses finite difference methods to approximate the derivatives, we have provided an error analysis in [B.I From a theoretical perspective, we show that as long as the grid spacing is finite, the expected solution error can be non-zero, even if the grid-based PINN loss is minimized arbitrarily well.
We bypass this process by using the spectral representation of the solution, and apply spectral trans- formations represent the residual function in the same basis. This greatly simplifies the entire optimization procedure, and as we will demonstrate, significantly reduces training time. Note that even though the spectral loss is weighted norm, the corresponding PINN loss can also be readily constrained:
Corollary 1. For da
E (N), [ ((
du
This result follows directly by Hölder's inequality. Both Fourier series and Chebyshev polynomials fulfill this condition, ensuring that minimizing the spectral loss also minimizes the PINN loss.

---

Published as a conference paper at ICLR 2024
4 EXPERIMENTAL RESULTS
We compare NSM to different neural operators with diffecrent loss functions (PINN and spectral losses) several PDEs: 2D Poission ($4.1) 1D Reaction-Diffusie ($4.2), and 2D Navier-Stokes ($4.)) with both forced and unforced flow. NSM is consistently the most accurate method, and orders of magnitudes faster during both training and inference, especially on large grid sizes.
all experiments, we focus on the data-constrained setting, using no interior domain solution data during training (i.c., train only by minimizing the PDE residual). The task is to learn the mapping from PDE parameters @ to solutions u. During training, each model given Fo and the parameters 65, which independently sampled in every training iteration.
Problem setting.
Recall that the PINN loss (Eq. T2) and the spectral loss (Eq. 13) are used on grid-based and spectral- based models, respectively. For the PINN loss, higher-order derivatives computed using the finite difference method on a fixed rectangular grid [r1 22 n For our spectral loss, the M-term residual series directly transformed from the predicted solution ūg.
PINN Loss
Spectral Loss
L N
Σ RE[N]
(u, (zn..
(12)
Σ Σ mE[M]
Fo. (13)
[(.1
For each problem, the test set consists of = 128 PDE parameters, denoted by d;. Each Φs is sampled from the same distribution used at training time, and u, is the corresponding reference solution. For each prediction ug, we evaluate two metrics: L2 relative error ||wo, u/ / ///2// and the PDE residual error ([F)(uu,))1 Both metrics are computed on the test resolution, and averaged over the dataset. Additional details about the experimental setup are provided in $C
We include the following models for comparison:
FNO + PINN loss. A grid-based FNO architecture (Li et al. 2020). trained with a PINN loss. The model trained on different grid sizes, which are indicated by the corresponding labels (e.g., 647 means grid size of 64 x 64 is used to calculate the PINN loss).
SNO Spectral loss. Ablation model: A base SNO architecture (Fanaskov & Oseledets)
2022), trained with our spectral loss.
T1 + PINN loss / Spectral loss. Ablation model: A base TransformOnce architecture (Poli
et ai. 2022), trained with either a PINN loss or our spectral loss.
CNO + PINN loss (ours). Ablation model: The base architecture identical NSM, but trained with a PINN loss on the largest grid size used by FNO, i.e. 256 on each dimension.
NSM (ours). Our proposed spectral-based neural operator using Fourier and Chebyshev basis on periodic and non- nor-oodio dimensions, and it is trained with our spectral loss.
For fair comparison, the base architecture, number of parameters, and other hyperparameters kept exactly the same across neural operator-based models. Detailed parameters are provided in CC
4.1 POISSON EQUATION
We study the 2D Poisson equation, Auu( y) = s(x, y), y E T, with periodic boundary condi- tions. The source term sampled from random field with length scale of 0.2. The task is to learn the mapping from s to the potential field u. We evaluate the predictions using an La relative error with respect to the analytical solution. Additional details and results for Dirichlet boundary conditions are in C.I
Table 1: L2 relative error (%) for the periodic Poission equation.
SNO+Spectral
TINGP+PINN
TI+Spectral
PNOx642+PINN
FNOx 1N6x1286PPINN
FNOX256'-PINN
NSM (ours)
0.59 0.12
3.22 + 0.49
0.302 + 0.071
4.24 + 0.13
2.01 + 0.03
1.75 0.00
.057 + .012
The results are summarized in Tab. Since the solution operator an inverse Laplacian, all models can theoretically express it with one layer (see discussions in C.I). but the FNO + PINN models exhibits high error, even when trained with large grids and for a longer training time. This simple example highlights the inherent optimization challenges with the grid-based PINN loss.
Results.

---

Published as a conference paper at ICLR 2024
@
(C)
512x lest res.
FNOx64 iPIN (trahing]
Numcrioal solver NEM (ours)
(%5 different test resciutions
FNCx1265-PINN (taiiing)
PDE resicual
32x32
FNO:G56"-PPNN (takiing)
10-1
512x512
64
10-2
128 128
RNO .PNN
RND:058".PPN
NSM (OLr))
1.0
32x32
1%6 G
PPNN
CNO PINNour)
512 test res.
64
L. test
128 128
32 x4
128 128
0.1
256 256x2285 256
258 x26
0.1
10"
10
10"
10 10" 10-3
10-3
10
Training time (seconds)
Inference time (seconds)
Figure 2: Reaction-Diffusion equation with 0.01. In a) and b), the L2 relative error and PDE residual on the test set are plotted over training each model. The grid-based methods (FNO trained with PINN loss) show improved accuracy the grid resolution increases, but significantly slower to train. When tested on different resolutions, significant aliasing errors occur on test grid resolutions that differ from the training grid resolution. In contrast, NSM has much low error and PDE residual, and achieves this lower error 100x faster than the grid-based methods. In c), when compared with iterative numerical solvers on different resolutions, NSM achieves the same level of accuracy with a 10x speedup. Notably, both the accuracy and computational cost of NSM remains constant, regardless of grid resolution.
4.2 REACTION-DIFFUSIN EQUATION
We study the 1D periodic Reaction Deaction-Diffusion system with different diffusion coefficients:
Uy 1l pu(1 u), x E w(x, 0) = l((x),
[0,T],
(14)
TET,
where v is the diffusion coefficient, and h(r) is the initial condition sampled from a random field with a length scale of 0.2. Given h(x),, the model learns to predict u(x, t) up to time = 1. The initial condition is enforced by transforming the prediction u(z, t) to u(x, t) . h(x).
Table 2: L2 relative error (%) and computation cost (GFLOP) for the Reaction-Diffusioo equation.
SNO+Spectral
FNOx PNNNN
PINN
FNOx256
CNO+PINN (ours)
NSM (ours)
y = 0.005
4.56 = 0.99
2.30 + 0.19
0.94
0.20 = 0.01
.075 + .016
V 0.01
5.41 1 4.43
2.64 1.64:00.0 0.97
2.27 + 1.19
0.57 1 0.19
0.48 1 0.16
.086 1 .019
y 0.05
87.76 +
11.82 5.4
3.25 4 1.29
1.06 0.28
0.78 = 0.01
.083 .006
0.1
152.8
13.03 1 6.4
4.90 2.400
4.07 1: 2.00
0.42
.077
Train/Test
0.91/0.03
7.30/15.8
31.3/15.8
138.5/15.8
200.3/25.0
15.0/0.32
Results.
The main results summarized in Tab. including the results for TransformOnce (Poli
et al. 2022) in C22 The training curves for the PDE residual and relative error on the test set for diffusion coefficient
0.01 are shown Fig. g The grid-based models using the PINN loss improve in relative envor with higher-resolution grid, but require significantly longer training time. In contrast, NSM consistently achieves high accuracy while maintaining low computation cost. As the diffusion cocfficient increases, NSM shows strong robustness and consistently achieves low solution error, while the other models all increase significantly in solution error.
We also compare the ML models to standard numerical solver (Simpson & Landman 2006). as shown in Fig. 2c During inference time, the accuracy and computational cost of NSM remains constant, regardless of the spatiotemporal resolution of the domain. NSM exhibits 10x increase in speed when compared to a numerical solver with a comparable level of accuracy. Solution error distributions and additional details are in
8

---

Published as a conference paper at ICLR 2024
(a)
Training (seconcs)
Iritaa varticity
T-1
T-2
Namerical solver
PDE residual en test set
0.04
PNOver--PR -PNN
NSM (ours) predietion
FNOves"-PPR -PNNN
NSM loun)
0.02
(6)
10F
10
10P
Lg rel. error (%) on test sel
FNO PINN prediction
ENO PRNN
Figure 3: Navier-Stokes equation with v 10-4. In a) and b), the relative error and PDE residual on the test set are plotted. NSM achieves low Le error and PDE residual 100x faster than FNO PINN methods, and is order of magnitude more accurate. c) NSM captures fine features of the vorticity evolution accurately, while grid-based approach fails to predict the overall shape.
4.3 NAVIER-STOKES EQUATION
We study the vorticity form of the 2D periodic Navier-Stokes (NS) equations:
d.w + Vu = rAw + f. TET t [0, where
V u 0,
(15)
w(x, 0) wn(x),
TET².
X
where v is the viscosity, the vorticity field, and f the forcing term. The initial vorticity wo is sampled from a random field with a length scale of 0.8. Given w, the velocity field u is determined by applying the inverse Laplacian operator. The model is trained learn the evolution of vorticity.
We first consider the unforced flow with different y values and T 3s. The solu- tion is diffusive for large and becomes more challenging to learn as the viscos- ity decreases, due to the sharp features in the solution. For FNO trained with the PINN loss, using grid resolution larger than 96 becomes intractable to train, due
Table 3: La relative error (%) for unforced NS equations.
FNOX61$+PINN FNOx96
NSM (ours)
== 10
8.18 + 2.83
7.90 + 0.57
0.71 = 0.02
== 10-3
14.81
0.67
11.99 + 0.86
1.65 0.26
10-4
17.88 1 2.67
16.20 + 0.61
3.53 0.53
Train
41h
13th
10h
to the cost of compute and memory. The results summarized in Tab. and the case for 10 shown in Fig. B. NSM significantly outperforms grid-based FNO with the PINN loss terms of both error and computational speed, achieving accurate results x00x faster (see Fig. 3a and Fig. (3b)
Next, we consider the long temporal transient flow under the forcing term f = (sin(((2 y)) cos(2 (m v))) and T = 50s, following the setting in Li et al. (2000) 2021). ). This is a significantly more challenging task, as it requires propagating the initial condition over an extended time interval. Nevertheless, as summarized
Table 4: Le relative error (%) for the long temporal transient NS equation.
FNO 96 +PINN
NSM (ours)
1/500
55.1 1 17.4
13.2 1 0.57
in Tab. NSM maintains high accuracy, while brsed FNO with the PINN loss collapses during training and fails entirely. Further details for both unforced and forced flow can be found in C3
Conclusion. We introduce an ML approach for solving PDEs, inspired by spectral methods. By utilizing orthogonal basis functions and their spectral properties, demonstrate numerous advan- tages for learning PDEs in the spectral domain. Our method is evaluated on different PDEs, and achieves significantly lower emrr and increased efficiency when compared current MI methods.

---

Published as a conference paper at ICLR 2024
Acknowledgements. This work was supported by the U.S. Department of Energy. Office of Sci- ence, Office of Advanced Scientific Computing Research, Scientific Discovery through Advanced Computing (SciDAC) program under contract No. DE-AC02-05CHII23 It was also supported in part by the Office of Naval Research (ONR) under grant 100014-23-1-2587 We also acknowledge generous support from Google Cloud and AWS Cloud Credit for Research. We thank Rasmus Malik Høegh Lindrup and Sanjeev Raja for helpful discussions and feedback.
REFERENCES
Mark Ainsworth and Justin Dong. Galerkin neural networks: A framework for approximating varia- tional equations with error control. SIAM Journal on Scientific Computing, 4247442244,, 2021.
John Boyd. Chebyshev and Fourier spectral methods. Courier Corporation, 2001.
James Bradbury. Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Miin Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
//github. ctmpubbleom//gole/e/ee/
Joan Bruna, Benjamin Peherstorfer, and Eric Vanden-Eijnden. Neural galerkin scheme with active learning for high-dimensional evolution equations. arXiv preprint arXiv:2203.01336 2022.
Gideon Dresdner, Dmitrii Kochkov, Peter Norgaard, Leonardo Zépeda-Núñez, Jamie Smith, Michael Brenner, and Stephan Hoyer. Learning correct spectral methods for simulating turbulent flows. arXiv preprint arXiv:2207.00556 2022.
Lawrence C Evans. Partial differential equations, volume 19. American Mathematical Society, 2022.
Vladimir Fanaskov and Ivan Oseledets. Spectral neural operators. arXiv preprint arXiv:2205.1057: 2022.
Avner Friedman. Partial differential equations of parabolic type. Courier Dover Publications, 2008
Somdatta Goswami, Aniruddha Bora, Yue Yu, and George Em Karniadakis, Physies-informed neu- ral operators. arXiv preprint arXiv:2207.05748. 2022.
David Gottlieb and Steven A Orszag. Numerical analysis of spectral methods: theory and applica- fions. SIAM, 1977.
Gaurav Gupta, Xiongye Xiao, and Paul Bogdan. Multiwavelet-based operator learning for differen- tial equations. Advances in neural information processing systems, 34:24048-244 -24622 2021.
Jiequn Han, Arnulf Jentzen, and Weinan E. Solving high-dimensional partial differential equations using deep learning. Proceedings of the National Academy of Sciences, 1 115(34):8505-8510, 2018.
Juncai He, Lin Li, Jinchao Xu, and Chunyue Zheng. Relu deep neural networks and linear finite elements. arXiv preprint 1TXiv:1807.03973, 2018.
Dmitrii Kochkov, Jamie A Smith, Ayya Alieva, Qing Wang, Michael P Brenner, and Stephan Hoyer. Machine saaming-accelerateed computationa fluid dynamics. Proceedings of the National Academy of Sciences, 118021))e110178411,, 2021.
Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces. arXiv preprint 100.::2:000000) 08888) 2021.
Aditi Krishnapriyan, Amir Gholami, Shandian Zhe, Robert Kirby, and Michael W Mahoney. Char- acterizing possible failure modes in physics-inniine informed neural networks. Advances Neural In- formation Processing Systems, 34:26548-26566 2021.
10

---

Published as a conference paper at ICLR 2024
Thorsten Kurth, Shashank Subramanian, Peter Harrington, Jaideep Pathak, Morteza Mardani, David Hall, Andrea Miele, Karthik Kashinath, and Anima Anandkumar. Fourcastnet: Accelerating global high-resolution weather forecasting using adaptive fourier neural operators. In Proceedings of the Platform for Advanced Scientific Computing Conference, pp. 1-11, 2023.
Henning Lange, Steven L Brunton, and J Nathan Kutz. From fourier to koopman: Spectral methods for long-term time series prediction. The Journal of Machine Learning Research, 22(1):1881- 1918, 2021.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint arXiv:2010.08889 2020.
Zongyi Li, Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Azizzadenesheli, and Anima Anandkumar, Physics-informed neural operator for learning partial differential equations. arXiv preprint arXiv:2111.03794, 2021.
Lu Lu, Xuhui Meng, Zhiping Mao, and George Em Karniadakis. Deepxde: deep learning library for solving differential equations. SIAM review, 63(1):208-2288 2021.
Bjöm Lütjens, Catherine H Crawford, Mark Veillette, and Dava Newman. Spectral pinns: Fast un- certainty propagation with physics-informet neural networks. In The Symbiosis of Deep Learning and Differential Equations, 2021.
John Mason and David C Handscomb. Chebyshev polynomials. CRC press, 2002.
Brek Meuris, Saad Qadeer, and Panos Stinis. Machine-learning-benne based spectral methods for partial differential equations. Scientific Reports, 13(1):1739, 2023.
Sebastian K Mitusch, Simon W Funke, and Miroslav Kuchta. Hybrid fem-nn models: Combining artificial neural networks with the finite element method. Journal of Computational Physics, 446: 110651, 2021.
Roger Peyret. Spectral methods for incompressible viscous flow, volume 148. Springer, 2002.
Michael Poli, Stefano Massaroli, Federico Berto, Jinkyoo Park, Tri Dao, Christopher Ré, and Ste- fano Ermon. Transform once: Efficient operator learning in frequency domain. Advances in Neural Information Processing Systems, 35:7947-7959, 2022.
Maziar Raissi, Paris Perdikaris, and George E Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.
Shawn G Rosofsky, Hani Al Majed, and EA Huerta. Applications of physics informed neural oper- ators. Machine Learning: Science and Technology, 4(2):025022, 2023.
Ramansh Sharma and Varun Shankar. Accelerated training of physics-informed neural networks (pinns) using meshless discretizations. Advances in Neural Information Processing Systems, 35: 1034-1046, 2022.
Matthew J Simpson and Kerry A Landman. Characterizing and minimizing the operator split error for fisher's equation. Applied mathematics letters, 19(7):604-612, 2006.
Justin Sirigmano and Konstantinos Spiliopoulos. Dgm: A deep learning algorithm for solving partial differential equations. Journal of computational physics, 375:1339-1368 1364, 2018.
Roger Temam. Navier-Stokes equations: theory and numerical analysis, volume 343. American Mathematical Soc., 2001.
Tapas Tripura, Souvik Chakraborty, et al. Physics informed wno. arXiv preprint arXiv:2302.05925, 2023.
Nicolaas Godfried Van Kampen. Stochastic processes in physics and chemistry, volume 1. Elsevier, 1992.
11

---

Published as a conference paper at ICLR 2024
Sifan Wang, Yujun Teng, and Paris Perdikaris. Understanding and mitigating gradient flow patholo- neural networks. SIAM Journal on Scientific Computing, 43(55:
gies in
A3081, 2021a.
Sifan Wang, Hanwen Wang, and Paris Perdikaris. On the eigenvector bias of fourier feature net- works: From regression to solving multi-scale pdes with physics-informed neural networks. Com- puter Methods Applied Mechanics and Engineering. 384: 13938, 2021b.
Mingtao Xia, Lucas Böttcher, and Tom Chou. Spectrally adapted physics-informed neural networks for solving unbounded domain problems. Machine Learning: Science and Technology, 4(2): 025024, 2023.
Bing Yu et al. The deep ritz method: a deep learning-based numerical algorithm for solving varia- tional problems. Communications in Mathematics and Statistics, 6(1):1-12, 2018.
Kai Zhang, Yuande Zuo, Hanjun Zhao, Xiaopeng Ma, Jianwei Gu, Jian Wang. Yongfei Yang. Chuan- jin Yao, and Jun Yao. Fourier neural operator for solving subsurface oil/water two-phase flow partial differential equation. SPE Journal, 27(03):1815-1888 2022.
12

---

Published as a conference paper at ICLR 2024
A ORTHOGONAL BASIS
Definition (Orthogonal basis). A set of real-valued functions (mm(x): x E mmmm said to be orthogonal w.r.f probability measure μs . RR, for any two indices i, I we have,
J. f = p.òdy.
(16)
where constants P4 .R+ and diz Li=j. uf 1, they also orthonormal. Furthermore, when their finite linear combinations are dense in 222 p.) they are called orthogonal basis.
Given that set of orthogonal basis functions can be easily normalized, assume their orthonor- mality unless otherwise specified this paper. For multi-dimensional problems, we can use the cross product of basis functions each dimension to form another orthogonal basis:
Proposition Given orthogonal basis (m (x1) :2 Siimez, and { (22):x2 E S22mmeza in (911 (1) and (Sb2 μ2), the cross product [f (21) ff (22) : E Un x Sggmmez xI, is an orthongonal basis in the product space L (22 x S22 ji (2)-
A.1 SPECTRAL CORRESPONDENCE
This section details operations in common PDE operators and their spectral correspondence with Fourier and Chebyshev bases. While operations on Fourier bases are well- well-known, for an in-depth introduction of Chebyshev polynomials, see Mason & Handscomb (2002).
Let u, v be the 1-dimensional function in consideration. Denote u, E RM as their M-term expan- sion under Fourier Chebyshev basis. We have the following properties in Tab.
Table 5: Spectral correspondence of common operations.
Fourier basis
Chebyshev polynomials
f = cos (2rmm))))
fm (x) = cas(m cos (x))
mms = sin(2mmx)
u+
11
j j= u
im.coa
=
inn F' n,00s
fm =
Σ
Un -n
n nSm, Uhe sin m,nn. nssin
nSm.M
i
Σ Un tin n Mu mmmoos Jm..oo 2rmümm,sin
Σ Umtn
um
M
du f =
8m
2m 1))m++2242
dm
sm. 2
Note that sums convolution form Σn a&bb- - can be accelerated by Fast Fourier Transforms (FFT). For multi-dimensional problems, the above operations can be done on each dimension sequentially.
B GRID-BASED METHODS
Complexity of computing residuals. We continue the discussion in our introduction and Li et al. (2021), where we look at the complexity of computing exact derivatives for grid-based neural oper- ators. Consider differentiating a parameterized kernel IC:
u(x) = (K))) (xx,, v)(())d(())
Σ E/ WEZ2
(17)
For given set of sampled points {a suppose both u and are queried these points. Then, de((ar))
regardless of how the integral is implemented or approximated, generally we have au(sy) * for every i,j [N]. Therefore, the complexity obtain the exact derivatives is quadratic in N.
For implementation purposes, use shifted Chebyshev polynomial on [0. 1] throughout this work, i.c. consider functions T (z) (2x 1), where T., is the standard Chebyshev polynomial [-1,1]
13

---

Published as a conference paper at ICLR 2024
B.1 ERROR ANALYSIS FOR GRID BASED PINN LOSS
Grid-based methods estimate residuals on a discretized domain. As grid spacings (e.g., Az, At) approach zero, the estimated residuals tend towards being exact. However, the grid size is limited by the expensive backpropagation in the training time, causing discretization emrors. In this section, we demonstrate this with an error analysis of the grid-based PINN loss for elliptie PDEs (Evans) 2022
Problem setting. Consider an operator L on bounded open subset Q C IR:
Lu = V (AVu),
(18)
where A(x) E Rdd are coefficient functions. In order guarantee the uniqueness of the solution, we assume satisfies the following uniformly elliptic condition for almost every æ f2:
*llell S 5" ((x) S 2 for all E R.
(19)
where S,S > 0 are absolute constants. Given & E (2t),, consider the Dirichlet problem Lus for us (S2) satisfying u 0 on On. The task is to learn the transformation G - ué.
Error analysis. Given an input o, denote the model prediction Ug. Consider minimizing the grid-based PINN loss using finite difference methods with grid spacing h O, i.e., the discrete gradient operator Vn is used to approximate V in L. Suppose that the grid-hased PINN loss is minimized, that for some small 6, 0, the probability is,
P( (AV u + @llca())
2e)
(20)
Proposition 2. Under the assumptions of Eq. and Eq. error holds:
20 the estimation for the expected solution
s V/VVnus us 3 (2) + VllV S -S -
Bl - unl uolle>(0)] Z
(21)
VV (oV - 6)
Proof. Consider the inner product (V (AV (A), Us - up)n By Cauchy-Schward
( A(V)us) us 04)0 I V (AV, (²))) ||ū - us (((()))
(22)
We can also use integration by parts, and then this inner product equals:
/ (up vn (AVaug)de,
/ Vn(up- (T) (AVa(Ms ko)) + (āg Vo (AVnū.)dz
(23)
( Vn(ug- no) (AV -o))dd
Vkūg (AVaū,)dz VoVn (
By the uniformity condition Eq. 19 the first two terms can estimated by:
/ Vn(us - ūo) (AVn(ug - az))do 2 (u - 02)
(24)
/
n VAus (AVaūg)dr S|| 2
By assumption Eq. 20 the following holds with probability 5:
Us Vn (AV, ds - us V (AVug)da + < - Vus [22 +
(25)
Then, with high probability, the inner product in Eq. 22 can be estimated by P((V, (AVAus), Us
uain 2 sl|V,(u, T=) + s|l V S V u 1 (2 -e) > 1 8. Rearranging the
terms, the desired result follows by Markov's inequality.
14

---

Published as a conference paper at ICLR 2024
Discussion. This analysis implies that when the difference between the exact derivatives Vug. and the approximated derivatives V, Ng, is sufficiently large, there exists some prediction üg. such that the solution error for the prediction that minimizes the grid-based PINN loss can also be large. The spectral loss has no such issue, as it is able to always obtain exact residuals.
C EXPERIMENT SETUP AND DETAILS
Our code is publicly available at hctps: httee://ggitee //tgth/ggtbb//.//////e com/aSk Beckeley/Seuzall 90pp:// Mechods All ex- periments are implemented using the JAX framework (Bradbury et al. 2018).
Random fields. We use random fields for the distribution of PDE parameters, such as initial con- ditions coefficient fields. For random fields associated with periodic boundaries, employ
periodic kernel k(x,x')= exp sin" (31==x(1) to sample values on uniform grid and use
14/2
Fourier basis for interpolation. Otherwise, we employ RBF kernel k(z,x') exp (- -- 2F
to
sample values on Chebyshev collocation points and use Chebyshev polynomials for interpolation.
Training. During training, the PDE parameters sampled online from random fields, describe above. We use a batch size of 16. The learning rate initialized to 10-3, and exponentially decays to 10-6 throouhh the whole training. Experiments are run with different random seeds, and averaged over the seeds.
Testing. For each problem, the consists of 128 parameters sampled from the same distribution used for training (different from the training samples). The reference solutions are gen- erated specified SC.T. C.2 and Both L2 relative error and PDE residuals are averaged over the whole test set.
Profiling. During training or inference, the wall clock time is measured by the average elapsed time during 16 calls of just-in-time compiled gradient descent step forward pass. The com- putation cost of numerical solvers the Reaction-Diffusinn problem is measured similarly. The parameter countings for each model used in each experiment provided in Tab. 6
Table 6: Model size (number trainable parameters) in each experiment.
PDE
Poisson equation
equation
Navier-Stoker equation
Forced NS
Model
SNO
FNO NSM
SNO
CNO
FNO/ NSM
TI
FNO /NO/NSM
T1
FNO / NSM
Param.
7.5M
31.6M
33.6M
67.2M
67.2M
167.8M
236.2M
782.4M
82.7M
Boundary conditions. Similar to classic spectral methods, the imposition of boundary and initial conditions is straightforward. Periodic boundaries can be automatically satisfied by using Fourier basis, and other types of boundary conditions can be enforced by applying mollifier to the coeffi- cients. This procedure can be specific to the problem at hand, and we show examples in $4
C.I POISSON EQUATION
Besides the periodic Poisson equation, we also show results for Dirichlet boundary conditions.
Parameters. For periodic boundary conditions, all models use Fourier bases in æ and dimen- sions. For Dirichlet boundary conditions, both SNO and NSM use Chebyshev polynomials both dimensions. For each model, use 4 layers with 64 hidden dimensions. In each layer, use 31 modes in both dimensions. All models are trained for 30k steps, except for FNO with a grid size of 256, which requires 100k steps to converge. For the PINN loss, the discrete Laplacian is computed with a 5-point stencil, and with a periodic wrapping for the periodic boundary case. The analytic solutions evaluated resolution 256 x 256.
All models use ReLU activations, except those using a PINN loss which totally collapse during the training. Therefore, use tanh activations for ENO+PINN and T1+PINN, and report these results.
15

---

Published as a conference paper at ICLR 2024
Source term s
(b)
Solution
Lz rol. error (%) 256 x26 tost ros.
100
6.2
0.8
0.3
0.4
0.4
1.0
NSM
SNO
FMDxs
.PNN
0.1
.PNNN
FNO PPN
0.0
0.0
0.0
139
10P
Figure 4: (a) Source term (b) Solution of periodic Poisson equation. An additional constraint (0, 0) 0 is imposed ensure the uniqueness of the solution. (c) Relative error curves dur- ing training. NSM converges significantly faster and to a lower error than the other methods.
Table 7: L2 relative error (%) of the 2D Poission equation with Dirichlet boundary conditions.
Spoctplltral
TI+Spectral
6NON66P+PINN
FNOx 128"
FNOx
NSM (ours)
39.9 36.6
2.24 1.61
0.407 + 0.008
3.52 0.29
3.61 0.27
0.413 + 0.032
0.280 0.017
Results. The L2 relative error for both types of boundaries are summarized in Tab, and Tab, 7 and the case for periodic boundary conditions is shown in Fig. 4 For the periodic case, the inverse Laplacian operator the solution with one layer. However, FNOs with PINN loss fail to achieve an reasonable level of accuracy. As the grid size increases, error improves
convolution with g(x) oc ||a||2. Therefore,
models can exactly represent
requires more training steps to converge.
For Dirichlet boundary conditions, the solution operator is no longer diagonal under either Fourier or Chebyshev basis. Nevertheless, NSM still achieves the highest accuracy among each models.
C.2 REACTION-DIFFUSE EQUATION
Parameters. Both CNO and NSM use Chebyshev polynomials in the dimension, and Fourier bases in the dimension. All models ReLU activations. For each model, 64 hidden dimensions. In each layer, we use 32 modes in the time dimension and 64 modes in the spatial dimension. For v 0.005 and 0.01, all models are trained for 30k steps. Due to the increasing problem difficulty as v becomes larger, the training steps are increased to 100k for 0.05, and 200k for 0.1, respectively.
layers with
The reference solution is generated by standard operator splitting method (Simpson & Landman 2006 with resolution 4096 4096, and then downsampled to resolution 512 512.
P= 0.005
= 0.01
>=
1.0
1.0
1.0
1.0
1.0
0.8
0.8
0.8
0.0
0.0
0.8
0.6
0.6
0.8
0.4
0.4
0.4
0.4
0.2
0.2
0.0
0.2
0.4
0.8
0.8
0.0
0.2
0.4
0.8
0.8
0.0
0.2
0.4
0.6
0.8
0.0
0.2
0.4
0.6
0.8
Figure 5: Solutions of the ID Deaation-Diffusion equation with Liffferen@tiffussiio coefficients v. Following Krishnapriyan et al. (201)) we use reaction coefficient 5 and four different values for the diffusion coefficient. As increases, the problem becomes progressively harder to solve.
"The domain the z-axis is taken as [0, 1). unlike Krishnapriyan et al. 2021 which considered the interval [0, 2T). Diffusion coefficients rescaled accordingly for similar solution behaviors. For diffusion coefficients larger than 0.1, the solution becomes almost fully diffusive and constant the domain.
16

---

Published as a conference paper at ICLR 2024
Results. The Ly relative error for all models are summarized in Tab. and Tab and distributions of absolution eπror over the test set are shown in Fig. 6 As v increases, the error distribution of the grid-based methods significantly increases, while the error of NSM remains almost identical.
Table 8: L2 relative error (%) and computation cost (GFLOP) for the Deaction-Diffussee equation, compared to T1 Spectral loss and T1 PINN loss (continued from Tab. 2
T1+Spartaa
TIX64*+PINN
TIx256 T1X128* +PINN
CNO+PINN (ours)
NSM (ours)
0.005
0.68 +: 0.02
2.30 0.19
0.94 0.11
0.11
0.32 + 0.04
0.20 0.01
.075 + .016
0.01
0.84 0.08
2.64 0.97
2.27 + 1.19
0.57 4 0.19
0.48 1 0.16
.086 + .019
0.05
41.36
60.10 8.5
81.87 3.1
90.86 = 6.8
0.78 0.01
.083 + .006
0.1
196.5 = 0.6
94.91 + 23.7
101.7
94.89 = 5.4
1.28 0.42
.077 + .005
Train/Test 4.08/0.17
6.64/6.35
16.5/6.35
60.1/6.35
200.3/25.0
15.0/0.32
Both SNO and T1 perform non-linear activations directly on spectral coefficients, design choice that differs from FNO's approach of activating function values. This distinction leads to poor ro- bustness in SNO and TI for increasing values. In contrast, NSM shows notable robustness the increasing diffusivity, validating our decision to apply activations at the collocation points.
001
Acaaciile ertor dirritation on
DCC0
Abockue Cnee Con1Du901 for 0.01
8
Absokiu darliuution forv 0.00
[t)
Abaatible diatiaation
8
3.0
3.20
INGHS Hicoaa 10000 jam0
HOUS
HRUA HOVE 11o000" OMG
Figure 6: Reaction-Diffusion equation: Error distribution on the test set for each value of v.
C.2.1 ABLATION STUDY: FIXED-GRID TRAINING
An insightful observation in our study is the connection between our spectral training and the prac- tice of training grid-based FNOs on fxxe-resolution grids, followed by interpolation for inference on different grid sizes. In this section, we provide ablation study this fixed-grid training strategy.
Problem setup. We use exactly the same PDE problem settings as in $4.2 The ablation model is trained on its collocation points, i.e. fixed-resolution grid with 32 x 64 nodes. During inference time, the model prediction interpolated using the Fourier basis, and queried at the test resolution.
The L2 relative етrоr for this ab- lation model is summarized in Tab. g FNO trained with fixed-resolution grid fails to converge, showing the important distinction between our spectral training and the practice of fixing training resolution and interpolating.
Results.
Table 9: L2 relative error (%) for fixed-grid -giid training.
0.005
0.01
== 0.05
0.1
55.7 + 0.01
76.9 + 0.02
79.5 + 0.04
80.4 + 0.23
C.3 NAVIER-S STRERSSTOKESSEE EQUATION
Parameters for the unforced flow. In each layer, we use 12 modes the time dimension and 31 modes in both spatial dimensions.
For each model, 10 layers with 32 hidden dimensions.
Parameters for the forced flow. For each model, use layers with 32 hidden dimensions. In each layer, we use 18 modes in the time dimension and 31 modes in both spatial dimensions.
The reference solution is generated numerically using time-stepping scheme similar to Li et al. (2020), with a spatial resolution of 256 x 256. The time stepping interval At set to 10 for the unforced flow, and 5 10-3 for the forced flow. Solution is recorded at temporal resolution of 64.
17

---

Published as a conference paper at ICLR 2024
Numerioll sploar
(cunt)
FNC+ Prodiction
Figure 7: Visualization of unforced flow time T = 3 for viscosity 10-1, using various initial vorticity samples.
NSM uses a Fourier basis the x, y dimensions and Chebyshev polynomials in the dimension. All models use ReLU activations. All models' learning rate schedulers target 200k steps. For FNO with the PINN loss, the training is limited to 1 day for grid size 643 and 3 days for 963 respectively, after the loss plateaus.
Table 10: T2 relative error (%) for unforced NS equation, compared to T1 + Spectral loss and T1 + PINN loss (continued from Tab. 3).
FNOx61++PING
TIx64 T
TI+Spectral
NSM (ours)
10 17.88 2.67
34.6 + 0.21 9.22 + 0.05
3.53 : 0.53
Results. The La relative errors for both problems are summarized in Tab. 3 Tab. and Tab. To and a visualization of the predicted vorticity for the unforced flow is shown in Fig. FNO with PINN loss fails to predict both the overall shape and the fine features of the vorticity evolution.
D DISCUSSION: LIMITATIONS
NSM applies to wide range problems, but is not a one-size-fitsall solution. Similar those in numerical analysis, efficient spectral representations necessitate sufficient regularities of the so- lution. Therefore, predicting solutions that involve singularities requires more basis functions. Ad- ditionally, NSM is designed for low oow-dimensional problems, e.g., problems described in x, 2 (and coordinates. Extensions of our spectral training strategy to high-dimensional problems are left for future research.
18


Reviews:
[{"review_id": "wnDkt7JekU", "reviewer": "Reviewer_KbDK", "summary": "The paper introduces Neural Spectral Methods (NSM) for learning solutions to Partial Differential Equations (PDEs) in the spectral domain, with a model that parameterizes spectral transformations. The authors introduce a new class of spectral-based neural operators and leverage Parseval's identity to derive a spectral loss that does not require auto-grad or finite-differences to approximate derivatives. Adopting a *data-constrained* setting, they conducted a set of experiments to validate their approach on three different PDEs.\nTheir method demonstrates significant speedup and accuracy improvements over considered baselines (FNO + PINN loss, SNO + spectral loss , NSM + PINN loss).", "questions": "The exposed approach has been tested when $\\phi$ is either an initial condition or a forcing term in the equation. While this is encouraging, I wonder if the method would remain applicable if we changed the boundary conditions between samples ? Let's say on a circle or a square. Would there be an issue between the change of dimensions between the 1D input parameter function and the 2D output ?\n\nI am also curious to understand the inference of NSM. For a new input function, is it purely operator-based or do you also finetune the model to reduce the PDE loss ? What kind of guarantee do you provide on the PDE solution for a new parameter function ?\n\nThe operators $T$ and $T^{-1}$ are not detailed. What is their complexity with respect to the input length ? How do you choose the collocation points ? Does the number depend on the difficulty level of the PDE ?\n\nCould you elaborate on this point ? ```The Fast Fourier Transform (FFT) in the forward pass of FNO escalates to quadratic in batched differentiation```.\n\nWhat is $\\mathcal{H}$ throughout the paper ?", "limitations": "", "rating": 3, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is well written and the method section is easy to follow. The spectral loss seems like a promising direction for solving PDEs with Neural Networks. Experimentally, the method outperforms all considered baselines in terms of L2 relative error. Reaction-diffusion and  Navier-stokes experiments were done for different values of diffusion and viscosity coefficients, which highlights the robustness of the method. Figures 2 and 3 show that NSM converges faster during training and is also insensitive to the spatial resolution at inference.", "weaknesses": "The motivation of the paper is not transparent to me. The authors propose in this paper two novelties for PDE-based neural networks : a spectral loss and a general design for spectral-based neural operators. While I understand that the first is supposed to simplify the training of PINNs, the second one seems to be a new architecture for solving operator learning tasks like FNO or SNO. Therefore, their method is not a PINN-like solver with a new loss, but rather a deep surrogate model that can approximate the PDE solution from an initial condition or forcing term, without data in the domain. This positioning should be stated explicitly. \n\nAs a consequence, the selected baselines and chosen setting cannot lead to a fair comparison between the different methods. Neural operators such as FNO or SNO have been proposed to learn mappings between functions that can be accessed through point-wise evaluations. Therefore, they require data for training and should not be trained with a PDE loss only. I think the authors should focus their comparison with classical PINN methods, and show first that their model is capable of solving a single equation for diverse types of PDEs, and then compare it to PINN approaches that target generalization through meta-learning such as [1], [2].\n\nI also do not understand why NSM is claimed to be insensitive to the grid size. Is a grid ever used for NSM ?  The forcing terms or initial conditions seem to be queried at the collocation points. This avoids aliasing problems rather than preventing them.\n\nThere is overall a lack of details regarding the implementation of the proposed method. We do not know if the experiments were done with fourier or chebyschev basis, or both. The truncated number of basis functions considered is also not mentioned. \n\n[1] Meta-Auto-Decoder for Solving Parametric Partial Differential Equations. Huang et. al. Neurips 2022.\n\n[2] Learning differentiable solvers for systems with hard constraints. Geoffrey Négiar, Michael W. Mahoney, Aditi S. Krishnapriyan. ICLR 2023.", "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": "Z7AAVfP2uQ", "reviewer": "Reviewer_9377", "summary": "The paper proposed a novel spectral-based neural operator, Neural Spectral Method (NSM), trained by self-supervised loss, i.e., it needs no data. Previous methods including FNO usually are trained with data, though they can also be trained in self-supervised way with PINN loss. This work extends the idea of self-supervised neural operators and proposes a new loss by Parseval's identity. \n\nFor periodic and non-periodic boundary conditions, NSM uses Fourier basis and Chebyshev basis respectively. \n\nUnder self-supervised setting, the paper compared NSM with some baseline modes, e.g., SNO, FNO combining with PINN loss. It shows that the proposed method has advantages including faster training convergence, higher accuracy, especially at super-resolution inference.", "questions": "Regarding aliasing error, it seems to me the solution proposed in the paper is to fix the resolution of grids and do interpolation in higher resolution. If this is correct, it is essentially advising people not to use super resolution inference like FNO does. Hence, this is not an advantage of NSM, but an insight about super resolution, which can be also applied to FNO, i.e., fixing resolution grids + interpolation.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**Originality:** The paper has several novelties, including a new design of neural operator with fixed bases. A novel residual loss by Parseval's identity.\n\n**Quality:** The paper has carefully benchmarked the proposed method in training cost, inference cost and accuracy. It exhibits advantages over several baseline methods.\n\n**Clarity:** The paper is well organized and clearly presented.\n\n**Significance:** The proposed method can in inspiring to the community in both", "weaknesses": "More comparison can be done, for example, Transform Once\n\n>Poli, Michael, et al. \"Transform once: Efficient operator learning in frequency domain.\" Advances in Neural Information Processing Systems 35 (2022): 7947-7959.\n\nwhich can be combined with PINN loss or the loss proposed here.", "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": "WdU0JVA7Dh", "reviewer": "Reviewer_dP1d", "summary": "The paper introduces Neural Spectral Methods (NSM) as a novel approach to solving partial differential equations (PDEs). What makes NSM unique in comparison to the numerous ML methods out there is that it's structurally integrated with spectral representations, Parseval's Identity, and orthogonal basis functions, all of which are classic numerical receipes that were well studied. It's a solid step towards hybrid NN and classic numerical method.", "questions": "Since the proposed resesarch is embedded in the spectral method, it's clearly limited by the issues of many spectral method. For example, as discussed in \"Implicit Neural Spatial Representations for Time-dependent PDEs (ICML 2023), spectral methods also have a global support, just like neural networks. So I would be curious to understand more of these limitations. What are some disadvantages of spectral method, in comparison say grid-based neural methods. Also how does the spectral approach compare to Implicit Neural Spatial Representations for Time-dependent PDEs (ICML 2023) which avoids grid completely and only user a MLP architecture.\n\n\"Focus on scenarios with no solution data.\" This is understandable but the scope of the paper. Since the author's method is a very general spectral formulation, why not try it on some data? My biggest concern right now is that the examples shown are all 2D toy example. Would it make sense to try the method on large scale 3D problems where data is already available?\n\n\"Dresdner et al. (2022); Xia et al. (2023); Meuris et al. (2023) use spectral representations in the spatial domain. These studies differ in problem settings and deviate from our method in both architecture and training\" I looked in details of these papers and can confirm that the author's work is unique. But I think for a complete paper, the author should discuss their similarities and differences.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Like discussed in the summary, the strength of the paper is the elegance integrating NNs with existing spectral methods. In particular, spectral basis are used to represent the functions and NNs are only used to map from spectral coefficients to spectral coefficients (with neural operator architecture). In addition, by leveraging the Parseval's Identity, the proposed method is able to train without using the expensive PDE residual norm (PINN loss). The authors mathematically proved the equivalence of the spectral loss and the PINN loss. This approach avoids sampling a large number of points and the numerical quadrature.\n\nThe author backed up these theoretical advantages with experiments that demonstrate 100X training time speedup.\n\nThe intro is well written. The short summary on of the the limitations of existing methods (data availability, optimization, computation cost) are very good. And the paragraph above it clearly place the author's work in the literature.", "weaknesses": "The paper briefly mentions some limitations of existing ML methods for solving PDEs, but it could benefit from a more extensive discussion of the potential drawbacks and challenges specific to NSM.\n\nBasis are predefined manually. It's unclear how to choose these basis functions.", "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": "xg3xErMDaM", "reviewer": "Reviewer_4uLV", "summary": "The authors propose an improvement upon Spectral Neural Operators, essentially proposing to train these architectures in spectral space using Parseval's identity and input data in form of spectral coefficients rather than signals evaluated on grid points. As such, the approach has some novelty and is relevant to the ML+PDE community. The authors show that this approach can outperform existing methods on 3 simple PDE problems (Poisson, Reaction-Diffusion, 2D NS).", "questions": "Major:\n1. NSM is trained using Parseval's identity using a spectral loss. From the results I can see that NSM outperforms CNO (the same but PINN loss). Why is that? Mathematically speaking, the two loss functions are identitcal and so should be the gradients. The only difference I can spot is potentially an added spectral transform in the gradients due to the summation in spectral space. Numerically, this would make a difference but I have a hard time understanding why this should lead to better performance. Any ideas?\n2. Results specify resolutions for FNO but do not disclose the spectral resolution/param count for the NSM. I encourage the authors to add this information\n3. Figure 2b - why is there only a single point for all resolutions of NSM?\n\nMinor/Clarity:\n1. in eq. 5 what is phi? It doesn't seem to be defined\n2. The statement that the neural operator is a mapping between grid points is limiting. It is simply realised as such and trained at a fixed resolution but conceptually, nothing stops the user from evaluating it on another grid. I believe this is what the authors mean with \"implemented\" in the last sentence of page 4 but I find this sentence to be somewhat misleading. A clarification would help in my opinion.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The paper introduces some novel ideas and is well-written. These are:\n- novel approach for training with spectral inputs + spectral evaluation of the Loss via Parseval's identity\n- relevant baselines and ablations that show the effect of the individual components\n- good motivation for the overall approach", "weaknesses": "Some interesting claims are made, for which I would hope for more clarification and theory if possible. These are:\n- Why is training in spectral space better? analytically both approaches should be identitcal (due to the identity) and gradients should be the same. Some theoretical analysis on this would have made the paper better.\n- how aliasing effects are avoided by using spectral coefficients as input data rather than the signal itself. Ideally, this should come with an ablation study on its own\n- the baselines could be better explained - I would like to see parameter counts and spectral resolutions of all approaches.", "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)