N ON - EQUISPACED F OURIER N EURAL S OLVERS FOR PDE S
Anonymous authors
Paper under double-blind review
A BSTRACT
Solving partial differential equations is difﬁcult.
Recently proposed neural resolution-invariant models, despite their effectiveness and efﬁciency, usually require equispaced spatial points of data.
However, sampling in spatial domain is sometimes inevitably non-equispaced in real-world systems, limiting their applicability.
In this paper, we propose a Non-equispaced Fourier PDE Solver (NFS) with adaptive interpolation on resampled equispaced points and a variant of Fourier
Neural Operators as its components.
Experimental results on complex PDEs demonstrate its advantages in accuracy and efﬁciency.
Compared with the spatiallyequispaced benchmark methods, it achieves superior performance with 42 .
85% improvements on MAE, and is able to handle non-equispaced data with a tiny loss of accuracy.
Besides, to our best knowledge, NFS is the ﬁrst ML-based method with mesh invariant inference ability to successfully model turbulent ﬂows in non-equispaced scenarios, with a minor deviation of the error on unseen spatial points.
I NTRODUCTION
Solving the partial differential equations (PDEs) holds the key to revealing the underlying mechanisms and forecasting the future evolution of the systems.
However, classical numerical PDE solvers require ﬁne discretization in spatial domain to capture the patterns and assure convergence.
Besides, they also suffer from computational inefﬁciency.
Recently, data-driven neural PDE solvers revolutionize this ﬁeld by providing fast and accurate solutions for PDEs.
Unlike approaches designed to model one speciﬁc instance of PDE (E & Yu, 2017; Bar & Sochen, 2019; Smith et al., 2020; Pan & Duraisamy, 2020; Raissi et al., 2020), neural operators (Guo et al., 2016; Sirignano & Spiliopoulos, 2018;
Bhatnagar et al., 2019; KHOO et al., 2020; Li et al., 2020b;c; Bhattacharya et al., 2021; Brandstetter et al., 2022; Lin et al., 2022) directly learn the mapping between inﬁnite-dimensional spaces of functions.
They remedy the mesh-dependent nature of the ﬁnite-dimensional operators by producing a single set of network parameters that may be used with different discretizations.
However, two problems still exist – discretization-invariant modeling for non-equispaced data and computational inefﬁciency compared with convolutional neural networks in the ﬁnite-dimensional setting.
To alleviate the ﬁrst problem, MPPDE (Brandstetter et al., 2022) lends basic modules in
MPNN (Gilmer et al., 2017) to model the dynamics for spatially non-equispaced data, but even intensiﬁes the time complexity due to the pushforward trick and suffers from unsatisfactory accuracy in complex systems (See Fig. 2(a)).
FNO (Li et al., 2020c) has achieved success in tackling the second problem of inefﬁciency and inaccuracy, while the spatial points must be equispaced due to its harnessing the fast Fourier transform (FFT).
To sum up, two properties should be available in neural PDE solvers: (1) discretization-invariance and (2) equispace-unnecessity.
Property (1) is shared by inﬁnite-dimensional neural operators, in which the learned pattern can be generalized to unseen meshes.
By contrast, classical vision models and graph spatio-temporal models are not discretization-invariant.
Property (2) means that the model can handle irregularly-sampled spatial points.
For example, graph spatio-temporal models do not require the data to be equispaced, but vision models are equispace-necessary, and limited to handling images as 2-d regular grids.
And recently proposed methods can be classiﬁed into four types according to the two properties, as shown in Fig.
As discussed, although the equispace-necessary methods enjoy fast parallel computation and low prediction error, they lack the ability to handle the spatially non-equispaced data.
For these reasons, this paper aims to design a mesh-invariant model (deﬁned in
Fig.
[Figure]
and high accuracy, by lending the powerful expressivity of FNO and vision models to efﬁciently solve the complex PDE systems.
Our paper including leading contributions is organized as follows:
• In Sec. 2, we ﬁrst give some preliminaries on neural operators as related work, with a brief introduction to Vision Mixers, to build a bridge between Fourier Neural Operator and Vision Mixers.
Thus, we illustrate our motivation for the work: To establish a mesh-invariant neural operator, by harnessing the network structure of Vision Mixers.
• In Sec.
3, we proposed a Non-equispaced Fourier Solver (NFS), with adaptive interpolation operators and a variant of Fourier Neural Operators as the components.
Approximation theorems that guarantee the expressiveness of the proposed interpolation operators are developed.
Further discussion gives insights into the relation between NFS, patchwise embedding and multipole graph models.
• In Sec. 4, extensive experiments on different types of PDEs are conducted to demonstrate the superiority of our methods.
Detailed ablation studies show that both the proposed interpolation kernel and the architecture of Vision Mixers contribute to the improvements in performance.
B ACKGROUND AND R ELATED W ORK 2. P ROBLEM S TATEMENT
Let
D ∈ R
d be the bounded and open spatial domain where n s -point discretization of the domain D written as X =
{ x i =
( x (1) i , . .
, x ( d ) i ) : 1 ≤ i ≤ n s } are sampled.
The observation of input function a ∈ A ( D ; R d a ) and output u ∈ U ( D ; R d u ) on the n s points are denoted by { a ( x i ) , u ( x i ) } n s i =1 , where
A ( D ; R d a ) and U ( D ; R d u ) are separable Banach spaces of function taking values in R d a and R d u respectively.
Suppose x ∼ µ is i.i.d.
sampled from the probability measure µ supported on D .
An inﬁnite-dimensional neural operator G θ : A ( D ; R d a ) → U ( D ; R d u ) parameterized by θ ∈ Θ , aims to build an approximation so that G θ ( a ) ≈ u .
A cost functional C : U ( D ; R d u ) × U
( D ; R d u ) → R is deﬁned to optimize the parameter θ of the operator by the objective
[Equation]
To establish a mesh-invariant operator, X can be non-equispaced, and the learned G θ should be transferred to an arbitary discretization X (cid:48) ∈ D , where x ∈ X (cid:48) can be not necessarily contained in
X .
Because we focus on spatially non-equispaced points, when the PDE system is time-dependent, we assume that timestamps { t j } are uniformly sampled, which means we do not focus on temporally irregular sampling or continuous time problem (Rubanova et al., 2019; Chen et al., 2019; Ça˘gatay
Yıldız et al., 2019; Iakovlev et al., 2020).
2. D ISCRETE F OURIER T RANSFORM
Let k l =
( k (1) l , . .
k ( d ) l ) the l -th frequency corresponding to X , with k l ∈ Z d .
The discrete Fourier transform of f : D → R d f is denoted by F ( f )( k ) ∈
C d f , with F − 1 as its inverse, then
F ( f ) ( j ) ( k l ) = n s (cid:88)
i =
f ( j ) ( x i )
e − 2 iπ< x i , k l > ,
[Equation]
(2) where j means the j -th dimension of f .
General Fourier transforms have complexity O ( n 2 s ) .
When the spatial points are distributed uniformly on equispaced grids, fast Fourier transform (FFT) and its inverse (IFFT) (Rader & Brenner, 1976) can be implemented to reduce the complexity to O ( n s log n s ) .
2. F OURIER N EURAL O PERATOR
Neural Operators.
To model one speciﬁc instance of PDEs, a line of neural solvers have been designed, with prior physical knowledge as constraints.
Different from these methods, neural operators (Lu et al., 2021; Nelsen & Stuart, 2021) require no knowledge of underlying PDEs, and only data.
Finite-dimensional operator methods (Guo et al., 2016; Sirignano & Spiliopoulos, 2018;
Bhatnagar et al., 2019; KHOO et al., 2020) are discretization-variant, meaning that the model can only learn the patterns of the spatial points which have been fed to the model in the training process.
By contrast, inﬁnite-dimensional operator methods (Li et al.
, 2020b;c; Bhattacharya et al., 2021;
Brandstetter et al., 2022) are proposed to be discretization-invariant, enabling the learned models to generalize well to
Kernel integral operator method (Li et al., 2020a) is a family of inﬁnite-dimensional operators, in which ( G θ ( a ))
( x ) =
Q ◦ v T ◦ · · · ◦ v 1 ◦ P ( a )( x ) is formulated as an iterative architecture.
A higher-dimensional representation function is ﬁrst obtained by v 0 =
P ( a ) ∈ U ( D ; R d v ) , where P is a shallow fully-connected network.
It is updated by
[Equation]
where K φ : A → L ( U ) is a kernel integral operator mapping, mapping a to bounded linear operators, with parameters φ .
W is a linear transform and σ is a non-linear activation function.
After the ﬁnal iteration, Q projects v T ( x ) back to U ( D ; R d u ) .
Fourier Neural Operator (FNO) (Li et al., 2020c) as a member in kernel integral operator methods, updates the representation by applying the convolution theorem as:
[Equation]
where R φ as the Fourier transform of a periodic kernel function κ φ , is directly learned as the parameters in the updating process.
To be resolution-invariant, FNO picks a ﬁnite-dimensional parameterization by truncating the Fourier series of both F ( v ) and R φ as a maximal number of modes k ( l )max for 1 ≤ l ≤ d .
Because the sampled spatial points are equispaced in FNO, it can conduct
FFT and IFFT to get the Fourier series, which can be very efﬁcient.
2. V ISION M IXER AND G RAPH S PATIO -T EMPORAL M ODEL
Vision Mixers (Tolstikhin et al., 2021;
Rao et al., 2021; Guibas et al., 2021) are a line of models with a stack of (token mixing) - (channel mixing) - (token mixing) as their network structure for vision tasks.
They are based on the assumption that the key component for the effectiveness of Vision
Transformers (V I T) (Dosovitskiy et al., 2020) is attributed to the proper mixing of tokens.
The deﬁned tokens are equivalent to equispaced spatial points in the former deﬁnition, and the research on the mixing of them can be an analogy to modeling the proper interaction or message-passing patterns among spatial points.
In speciﬁc, V I T uses a non-Mercer kernel function (Wright & Gonzalez, 2021) κ φ to adaptively learn the pattern of message-passing through the iterative updating process
[Equation]
where W is a linear transform called channel mixing layer because it transforms the input on the channel of an image whose dimension is equivalent to function dimension d f .
Note that we omit the residual connection in Eq.
(3) for simplicity.
Remark.
The FNO can be regarded as a member of the family of Vision Mixers.
The reason is that a component in an iteration in Eq.
(4) can be written as ( R φ · F ( v ))
( x ) =
R φ · (cid:80) i e − 2 πi< x , x i > v ( x i ) , because in the equispaced scenarios, x i can be regarded as lying on the same grids as k after scaling.
The kernel κ φ is parameterized by κ φ ( x , x i , v ( x ) , v ( x i ))
=
e − 2 πi< x , x i > , and the matrix multiplication of R φ also performs mixing on channels.
Besides, the inverse Fourier transform can also be regarded as token mixing layers, or so-called token demixing layers (Guibas et al., 2021).
However, the powerful ﬁtting ability and efﬁciency of Vision Mixers are limited to being applied to non-equispaced spatial points.
Another option for non-equispaced data is graph spatio-temporal models, in which interaction patterns among spatial points are modeled in a graph message-passing way (Gilmer et al., 2017; Atwood & Towsley, 2016; Defferrard et al., 2017).
The mechanism is similar to the token mixing in Vision Mixers by means of the summation in Eq.
(5) conducted in the predeﬁned neighborhood of each point.
Unfortunately, the graph spatio-temporal models (Seo et al., 2016; Li et al., 2018; Bai et al., 2020; Lin et al., 2021) suffer from high computational complexity and unsatisfactory accuracy in solving complex dynamical systems (such as turbulent ﬂows).
2. M OTIVATION
Since FNO belongs to Vision Mixers, this ﬁrstly raised a question to us: Do models employing
Vision Mixers’s architecture have the potential to model complex PDE systems?
Thus, experiments are conducted to give an intuitive explanation of our motivation as shown in Fig.
The data are generated by Navier-Stokes equations.
It is noted that graph spatio-temporal methods can also handle the equispaced data.
Detailed setup is given in Sec. 4.
    
OOM ORJ506(    
OOM O
RJ  VHF RQG
VITMLPMIXERGFNFNODCRNNAGCRNGCGRUMPPDENFS 2XUV
[Figure]
    
OOM ORJ506(    
OOM O
RJ  VHF RQG
DCRNNAGCRNGCGRUMPPDENFS 2XUV
[Figure]

     (SRFK  
 
  O
RJ  5
6
(
[Figure]
(T
[Figure]
We ﬁnd that (1) All of the evaluated Vision Mixers are able to model the dynamical systems effectively, in spite of FNO as the only discretization-invariant model; (2) The complex dynamics of the systems are hardly captured by graph spatio-temporal models, whose performance on both accuracy and computational efﬁciency is very unsatisfactory in either equispaced or non-equispaced scenarios.
Fig. 2(c) shows the problem of infeasibility of graph spatio-temporal models through the loss curves on the validation set, compared with Vision Mixers.
However, Vision Mixers fail to handle the non-equispaced data.
Therefore, we aim to (1) establish a mesh-invariant model, by harnessing the network structure of Vision Mixers to achieve competitive efﬁciency and effectiveness in equispaced scenarios, as shown in Fig. 2(a); (2) Besides, it should allow applicability and comparable accuracy in non-equispaced scenarios for solving PDEs, as shown in Fig. 2(b).
P ROPOSED M ETHOD 3. N ON - EQUISPACED F OURIER T RANSFORM
Nonuniformly signals are unavoidable in certain real-world physics scenarios, such as signals obtained by meteorological stations on the earth surface, which urge the fast Fourier transform (FFT) to be extended to non-equispaced data with efﬁcient implementation of FFT.
Non-equispaced FFTs usually rely on a mixture of interpolation and the judicious use of FFT, where the calculations of interpolation are no more than O ( n s log n s ) operations (Kalamkar et al., 2012; Cheema et al., 2017).
For example,
Lagrange interpolation is used to approximate the signal values on m s resampled equispaced points
[Figure]
{ x j } 1 ≤ j ≤ m s , and then implement FFT on the interpolated points.
A low rank approximation with complexity of O ( n s log(1 /(cid:15) )) is used to replace the interpolation with complexity of O ( n 2 s ) with (cid:15) as the precision of computations (Dutt & Rokhlin, 1995).
Another example is commonly used
Gaussian-based interpolation (Kestur et al., 2010).
Denote F as equispaced FFT in particular, and H as the interpolation operator, and the proposed non-equispaced FFT can be written as
[Equation]
H ( f )( x j ) =
(cid:80) n s i =
1 f ( x i )
h τ ( x i − x j ) interpolates values on resampled points via convolution with the periodic heat kernel h τ ( x − y )
= (cid:80) l ∈ Z d e − ( x − l ) 2 / 4 τ , with τ as a constant.
Multiplication of the inperploation matrix ( H τ ) i,j = h τ ( x i − x j ) and the signal vector ( f ) i =
f
( x i ) includes O ( n s m s )
≈
O ( n 2 s ) operations.
For the kernel h τ , it is a summation of Gaussian kernel, and convolution with a single Gaussian in each points x i ’s neighborhood N ( x i ) can yeid a tiny error depending on τ , so the interpolation operator can be approximate via H ( f )( x j ) =
(cid:80) x
j ∈N ( x i ) f ( x i ) h τ ( x i − x j ) .
Restricting the neighbor number to |N ( x i ) | ≤ log n s leads the complexity to reduce to O ( n s log n s ) .
3. N ON - EQUISPACED F OURIER N EURAL PDE S OLVER
Non-equispaced interpolation.
To harness the effectiveness of FNO, we use non-equispaced
Fourier token mixing instead of the equispaced one.
It generalizes the equispaced FFT in Eq.
(4) as
[Equation]
We denote H η : A → L ( U ) as the interpolation operator mapping, which maps parametric function to a bounded interpolation operator.
H η ( a ) gets the inerploated values on m s resampled equispaced points via the convolution with kernel h η as
[Equation]
where x j lies on resampled equispaced grids.
Another H (cid:48)
ζ interpolates back on the n s non-equispaced ones in the same way via the convolution with kernel h ζ .
To reduce the operations to no more than
O ( n s log n s ) , the summation is restricted in the neighborhood of x i and x j , such that |N ( x i ) |
≈ |N
( x j ) | ≤ c log n s with c as a predeﬁned constant determining the neighborhood size of spatial points.
We formulate the kernel with a shallow feed-forward neural networks.
Thanks to the universal approximation of neural networks, the following theorem assures that the interpolation operator can approximate the representation function v arbitrarily well.
(For detailed proof, see Appendix. A.3.)
Theorem 3.1 (Approximation Theorem of the Adaptive Interpolation) .
Assume the setting of Theorem
A2 in Appendix.
A.3 is satisﬁed.
µ is the probability measure supported on D .
For v ∈ U , suppose
U =
L p ( D ; R d v ) , for any 1 < p < ∞ .
Then, given (cid:15) > 0 , there exist a neural network h
η
R d × R d × R d a → R d v , such that || ˆ v − v || U ≤ (cid:15) , where ˆ v ( x ) =
(cid:82) D h η ( x − y , x , a ( y )) v ( y ) dµ ( y ) .
Applicability of Layer-Norm.
As shown in Fig.
3, besides the comparison of the proposed interpolation operator with the traditional ones, a notable difference between the original FNO and FNO layers in our Vision Mixer architecture is the applicability of normalization layers (Layer-Norm) which is usually used in Vision Mixers’ architecture.
FNO cannot adapt Layer-Norm layers, because the change of resolution will make the trained normalization parameters and spatial points disagree with each other.
In comparison, the resampled equispaced points are ﬁxed in our architecture, no matter how the discretization of the input changes.
Therefore, the normalization layers can be added, in a similar way to Vision Mixers, bringing considerable improvements (See Sec. 4.3).
Mesh invariance.
In the intermediate layers, which adopt equispaced FNO, the resampled points are ﬁxed in both training and inference process, invariant to the input meshes.
In the interpolation layers, the operator H η ( a ) is discretization-invariant because the kernel can be inductively obtained by the newly observed signals a ( x ) , its coordinate x and resampled spatial points’ relative coordinates x j − x .
In the same way, H (cid:48) ζ ( a ) is also mesh-invariant.
This allows the NFS to achieve zero-shot mesh-invariant inference ability, which is demonstrated in Sec. 4.2.
Complexity analysis.
The complexity of FNO is O ( n s log n s + n s k max ) .
In the interpolation layers, because the interpolated values of resampled points are determined by their neighbors, we set the size of each resampled point’s neighborhood in G and observed non-equispaced points’s neighborhood in G (cid:48) as |N ( x i ) |
≈ |N
( x j ) | ≤ c log( n s ) , for 1 ≤ i ≤ n s , 1 ≤ j ≤ m s .
And in this way, the sparsity of the interpolation matrix reduces the complexity of the two interpolation layers to
O ( c · n s log n s + c · m s log n s ) .
If we set the resampled points number as n s , the overall complexity is O (2 c · n s log n s + n s log n s + n s k max ) ∼ O ( n s log n s + n s k max ) .
3. F URTHER D ISCUSSION
Relation to Vision Mixer.
The interpolation can be compared to patchwise embedding in Vision
Mixers.
For example, MLPM IXER learns the token mixing patterns adaptively with a feed-forward network, but the high resolution of input images does not permit the global mixing of tokens due to the complexity of O ( n 2 s ) .
Therefore, the input images are ﬁrstly rearranged into patches, with each patch containing n p pixels.
In this way, the complexity is reduced to O ( n 2 s /n 2 p ) , enabling feasible token mixing.
The patchwise embedding is very similar to interpolating the values on resampled points, as the former one ﬁrst chooses patches’ centers as n 2 s /n 2 p resampled points, and ‘interpolates’ the resampled points by lifting the embedding dimension and the rearranging of their neighbors’ values as the interpolated values, rather than using a kernel.
Relation to multipole graph model.
The adaptively learned interpolation layer in NFS has a similar formulation of multipole graph models (Li et al., 2020b).
In multipole graph models, the high-level nodes aggregate messages from their low-level neighbors as v High ( x j )
=
|N
( x j ) | (cid:80) x i ∈N ( x j ) v Low ( x i )
h η ( x j , x i , a ( x j ) , a ( x i )) .
Compared to multipole graph models, the values of high-level resampled equispaced nodes are approximated with low-level non-equispaced nodes’ values in NFS, but nodes’ values of low levels are given in multipole graphs.
This causes differences in multipole graph models’ message-passing and NFS’s interpolation:
In the former one, messages ﬂow circularly among different levels of nodes, while in NFS, messages only exchange twice between the nodes of two levels – one is from low-level non-equispaced nodes to high-level resampled equispaced nodes, and the other is the opposite.
E XPERIMENTS 4. E XPERIEMNTAL S ETUP
Benchmarks for comparison.
For ﬁnite-dimensional operators, we choose Vision Mixers including
V I T (Dosovitskiy et al., 2020), GFN (Rao et al., 2021) and MLPM IXER (Tolstikhin et al., 2021)
as equispaced problem solvers, and graph spatio-temporal models including DCRNN (Li et al., 2018),
AGCRN (Bai et al., 2020) and GCGRU (Seo et al., 2016) as non-equispaced problem solvers.
For inﬁnite-dimensional operators, the state-of-the-art FNO (Li et al., 2020c) for equispaced problems and MPPDE (Brandstetter et al., 2022) for non-equispaced problems are chosen.
A brief introduction for these models is shown in Appendix.
B.1.
In Vision Mixers, the different timestamps in temporal axis are also regarded as ‘tokens’ in that timestamps are uniformly sampled.
Protocol.
The widely-used metrics - Mean Absolute Error (MAE) and Root Mean Square Error (RMSE) are deployed to measure the performance.
The reported mean and standard deviation of metrics are obtained through 5 independent experiments.
All the models for comparison are trained with target function of MSE, i.e. C ( u, v )( x ) =
|| u ( x ) − v ( x ) ||
2 corresponding to Eq.
(1), and optimized by Adam optimizer in 500 epochs.
The hyper-parameters are chosen through a carefully tuning on the validation set.
Every trial is implemented on a single Nvidia-V100 (32510MB).
4. N UMERICAL E XPERIMENTS
Data.
We choose four equations for numerical experiments, three of which are time-dependent (KdV,
Burgers’ and NS), while the other one is not (Darcy Flows).
For 1-d problem, we consider Korteweg de Vries (KdV) (given in Appendix.
B.2.)
and Burgers’ equation:
[Equation]
where x ∈ [0 , 1] , t ∈ [0 , 1] .
u ( x, 0)
=
u 0 ( x ) is the initial condition and ν ∈ R + is the viscosity coefﬁcient.
For 2-d PDEs, we consider Darcy Flow (given in Appendix.
B.2.) and Navier-Stokes (NS) equation for a viscous, incompressible ﬂuid in vorticity form on the unit torus:
[Equation]
where x ∈ [0 , 1] 2 , t ∈ [0 , 1] .
u is the velocity ﬁeld, w = ∇ × u is the vorticity, w 0 is the initial vorticity, ν ∈ R
+ is the viscosity coefﬁcient, and f is the forcing function.
The total number of instances is 1200, with percentages of 0.7, 0.1 and 0.2 for training, validating and testing, respectively.
The original simulated resolutions of PDE signals are 128 × 128 in NS equation.
For others, see Appendix.
B.2.
When evaluating their performance in equispaced scenarios of different resolutions, we can downsample the resolution for training to low-resolution data, e
[Table]
g.
64 × 64 in NS equation.
To evaluate their performance in non-equispaced scenarios of different meshes, we randomly choose n s spatial points for training
[Table]
[Table]
[Table]
Performance comparison.
In this part, for time-dependent PDEs, our target is to map the observed physical quantities from initial condition u ( X , T ) ∈ R n s × n t , where T = { t i : t i < T } 1 ≤ i ≤ n t , to quantities at some later time u
( X , T (cid:48) )
∈ R n s × n (cid:48) t , where T (cid:48)
= { t i : T < t
i < T (cid:48) } 1 ≤ i ≤ n (cid:48) t .
We set the input timestamp number n t as 1 (initial state to future dynamics) and 10 (sequence to sequence), and prediction horizon n (cid:48) t as 10, 20 and 40 as short-, mid- and long-term settings.
For Darcy Flows, which are independent of time, we directly build an operator to map a to u .
In equispaced scenarios, the resolution is denoted by r d = n s , where d is the spatial dimension. In non-equispaced scenarios, the spatial points number is denoted by n s .
The comparison of benchmarks with or without equispaceunnecessity are shown in Table.
1 and 2 respectively, and detailed results including KdV equations with RMSE and standard deviations are given in Appendix.
B.3.
It can be concluded that (1) In equispaced scenarios, the proposed NFS obtains the lowest error in most 1-d PDE settings, and in solving 2-d PDEs, its superiority over other Vision Mixers are signiﬁcant, with 42 .
85% improvements on MAE according to the trials of NS ( r = 64 , n t = 10 , n (cid:48) t = 40 ).
(2) In non-equispaced scenarios, the evaluated graph spatio-temporal models’ performance is unsatisfactory, especially in NS equations.
In comparison, NFS achieves comparable high accuracy to the equispaced scenarios, for instance, according to columns of NS ( r = 64 , n t = 10 , n (cid:48) t = 40 ) with ( n s = 4096 , n t = 10 , n (cid:48) t = 40 ).
In some trials such as Burgers’ ( n t = 1 ) in Table.
B4 in Appendix.
B.3, Vison Mixers including
FNO also suffer from non-convergence of loss; while NFS can still generate accurate predictions.
The explanation of the phenomenon will be our future work.
[Equation]
[Table]
Mesh-invariance evaluation.
We use ( u ( X , T ) , u ( X , T (cid:48) ))
as the training set, and evaluate the model’s performance of mesh-invariant inference ability on X (cid:48) , where | X (cid:48)
| = n (cid:48) s .
X (cid:48) is a different mesh with X ⊆ X (cid:48) .
The visualization results of NS ( n s = 4096 , n (cid:48) t = 40 ) are shown in Fig.
For a ﬁxed n (cid:48) s , we randomly sampled different X (cid:48) for 100 times, to get the mean errors and standard deviations (given in Appendix.
B.5) of different spatial meshes.
We can conclude from Table.
3 that (1) The errors on unseen meshes are larger than the errors on seen meshes, showing the overﬁtting effects.
However, the errors on unseen meshes are acceptable, since they are even lower than other models’ prediction error on seen meshes.
(2) Larger n (cid:48) s leads to higher prediction error because a large number of unseen points are likely to disturb the learned token mixing patterns.
On the other hand, NS ( n s = 1024 , n (cid:48) t = 10 ) implies that small spatial point numbers of training meshes ( n s ) hinder model’s generalizing ability on unseen meshes, due to excessive loss of spatial information.
4. A RCHITECTURE A NALYSIS
Two modules in NFS differ from FNO.
The ﬁrst is the interpolation layers at the beginning and the end of the architecture.
The second is the extra Layer-Norm in the FNO layers, which can be applicable in NFS thanks to its ﬁxed resampled equispaced points, but inapplicable to FNO for preserving its resolution-invariance.
We aim to ﬁgure out what makes NFS outperform FNO.
Effects of neighborhood sizes.
It is widely believed that modeling the long-range dependency among tokens brings improvements (Naseer et al., 2021; Tuli et al., 2021; Mao et al., 2021).
By contrast, some local kernel methods demonstrate their superiority (Yang et al., 2019;
Liu et al., 2021; Chu et al., 2021; Park & Kim, 2022).
For this reason, we ﬁrst conjecture that the large neighborhood sizes in the interpolation layer are conducive to predictive performance.
  
   ORJ0HDQ1HLJKERUKRRG1XPEHU       0 $ (   2XWRI0HPRU\ 0$(RI3)12 0$(RI1)  
    5 0 6(  
 506(RI3)12 506(RI1)
Besides, as demonstrated in Sec.
3.3, the patchwise embedding in Vision Mixers can be an analogy to the resampling and interpolating, so we further establish a patchwise FNO (PFNO), with patch size equaling to 4 and [4 , 4] in 1-d and 2-d PDE problems, equivalent to each resampled points aggregatingand 16 points in spatial domains in 1-d and 2-d situations respectively.
Layer-Norm is stacked in the
FNO layers in PFNO, for a fair comparison.
Results of Fig.
5 show that the long-range dependency may even compromise the performance, as larger mean neighborhood sizes often cause higher errors.
However, no matter how large is the neighbor size, the NFS outperforms PFNO.
More details are given in Appendix.
B.6.
Therefore, we rule out the possibility of performance gains brought form large neighborhood sizes and suppose that proposed kernel interpolation layers are the key, and is superior to the simple patchwise embedding methods.
Beneﬁts from learned interpolation kernel.
Since the kernel interpolation is likely to hold the key to improvements, we investigate the performance gains brought from adaptively learned interpolation kernels over the predeﬁned one (See Fig. 3).
We use an inﬂexible Gaussian kernel h
[Equation]
as a predeﬁned one as discussed in Sec. 3.1, where Γ = diag( γ (1) , . . .
, γ ( d ) ) , and µ ∈ R d , γ (1) , . .
, γ ( d ) , β ∈ R + are learnable parameters.
By setting all the other modules and the interpolation neighborhood sizes as the same, we compare performance on different meshes of the two interpolation kernels in Table.
3 (Gaus +
LN), where the adaptively learned kernels achieve better accuracy.
Beneﬁts from normalization layers.
Previous works demonstrated the normalization is necessary for network architecture, for fast convergence and stable training (Dong et al., 2021;
Ba et al., 2016).
A notable difference between NFS and FNO is that the Layer-Norm can be implemented in NFS’s layers without disabling its discretization-invariance.
The improvements brought from the normalization layers are given in Table.
3 (Flex + (cid:8)(cid:8) LN), where the performance gap is obvious on unseen meshes.
4. N ON - EQUISPACED V ISION M IXERS
Since NFS can be regarded as a combination of our interpolation layers with the revised FNO, our interpolation layers can also be implemented in the other Vision Mixers, so that these methods are equipped with the ability to handle non-equispaced data.
Details are given in Appendix.
B.7.
We ﬁnd that (1) From Table.
B10 and Table.
1, the degeneration of performance is obvious in other Vision
Mixers.
In comparison, FNO as intermediate equispaced layers, truncates the high frequency in its channel mixing and retains the low frequency shared by both resampled and original signals, so the loss of accuracy in non-equispaced scenarios is tiny in our NFS; (2)
Although the performance on unseen meshes is more stable in these non-equispaced Vision Mixers, the performance gap is still large, according to Table.
B10 and Table.
4. C OMPLEXITY COMPARISON
The discussed Fig.
2 shows time complexity of each method.
Vision Mixers are the fastest, while graph spatio-temporal models are far much slower.
NFS falls in between because the interpolation layer can be an analogy to a graph-message-passing layer (See Sec. 3.3), and the intermediate are equispaced token-channel mixing layers of Vision Mixers’ structure.
Memory usage shows storing each point’s neighbors in the interpolation layer is very memory-consuming.
C ONCLUSION AND F UTURE W ORK
A simple method called NFS is established based on FNO, with theoretically enough expressiveness, allowing non-equispaced scenarios and fast computation for solving PDEs, and achieves state-ofthe-art performance in equispaced and non-equispaced scenarios.
Problems still exist, including relatively high error on unseen meshes and high memory cost in the interpolation layers .
A M ETHODS A. N OTATION
Symbol
Used for
X
A discretization of the domain D , which is used to train the model.
X (cid:48)
A discretization of the domain D , which is used to evaluate the model.
x
Coordinate of spatial point in D .
[Table]
Number of spatial points of seen meshes for training, as | X | = n s .
[Table]
(cid:48)
s
Number of spatial points of unseen meshes for inference, as | X (cid:48) | = n (cid:48) s .
n
[Table]
Number of input timestamps as the number of input time-dependent
PDE’s initial states.
[Table]
Number of output timestamps as the number of output time-dependent
PDE’s future states.
a
Input function, where a ∈ A ( D ; R d a ) means for x ∈ D , a ( x ) ∈ R d a .
In time-dependent PDEs, a = u ( · , T ) , where T = { t i } n t i =1 .
[Table]
Target function for approximation, where u ∈ U ( D ; R d u ) means for x ∈ D , u ( x ) ∈ R d u .
[Table]
Representation function, where v ∈ U ( D ; R d v ) means for x ∈ D , v ( x ) ∈ R d v , which is a function obtained by a lifter which project a into a higher dimensional space.
[Table]
Approximation operator, where G θ ( a ) ≈ u .
[Table]
The probability measure for sampling Spatial points x , which is supported on D .
[Table]
Cost functional as the minimum optimization target.
[Table]
Discrete Fourier transform for equispaced spatial points.
[Table]
Discrete Inverse Fourier transform for equispaced spatial points.
[Table]
Project operator, with P ( a )( x ) ∈ R d v .
[Table]
Project operator, with Q ( v )( x ) ∈ R d u .
[Table]
t -th iterative representation function after kernel operators’ update.
K φ
Kernel integral operator mapping, which maps a to a bounded linear operator, with parameter φ
W Linear transform on the d v dimension (channel) of v ( x ) ∈ d v .
R
[Table]
Fourier transform of a periodic kernel function, which is learnable parameters in a single iterative process in FNO.
[Table]
Channel-mixing operator as a linear transform in the dimensoion of channel ( d v ).
[Table]
Token-mixing operator as a linear transform in the dimensoion of spatial points ( n s ).
[Table]
Proposed non-equispaced Fourier transform, where ˜ F = ( F ◦ H ( a )) .
[Table]
Spatial points’ number on resampled equi-spaced points.
H Interpolation operator to interpolate the non-equispaced spatial points on equispaced spatial grids.
[Table]
Parameter in Gaussian interpolation kernels controlling smoothness of the kernel.
[Table]
Parameter in Gaussian interpolation kernels, as the mean of Gaussian kernels.
[Table]
Interpolation operator mapping, where H η ( a ) is an interpolation operator used to map signals on the non-equispaced spatial points to equispaced spatial grids.
[Table]
Interpolation operator mapping, where H (cid:48) ζ ( a ) is a interpolation operator used to map signals on the equispaced spatial points to nonequispaced spatial grids.
[Table]
Neighborhood of spatial point x .
A. G RAPH CONSTRUCTION
Neighborhood construction.
Instead of using K-nearest neighborhood method, the neighborhood system in the interpolation layer is constructed by (cid:15) -ball, because in equispace scenarios, there will be multiple points as K-th nearest neighbor at the same time.
For point x , its neighbor is deﬁned according to
[Equation]
For given c deﬁned in Sec. 3.2, we can restrict (cid:15) so that E x ∼ µ [ |N ( x ) | ] < c log( n s ) .
A. P ROOF OF T HEOREM 3.1.
Our proof is mostly based on Chen & Chen (1995) and Kovachki et al.
For notation simplicity, in the proof, we directly write H η ( a ) as H η as the linear operator.
Lemma A1.
Let X be a Banach space, and U ⊆
X a compact set, and K ⊂ X a dense set.
Then, for any (cid:15) > 0 , there exists a number n ∈ N , and a series of continuous, linear functionals
G 1 , G 2 , . .
. , G n ∈ C ( U ; R ) , and elements ϕ 1 , . .
. , ϕ n ∈ K , such that
[Equation]
[Equation]
The proof is given in Lemma 7.
in Kovachki et al.
(2021), and Theorem 3.
and 4. for reference .
Theorem A2.
Let D
⊆
R
d be compact domain.
Let U be a separable Banach space of real-valued functions on D , such that C ( D, R ) ⊆
U is dense.
Suppose U = L p ( D ; R ) for any 1 < p < ∞ .
ν is a probability measure supported on U and assume that, E v ∼ ν
||
v
|| U < ∞ for any v ∈ U .
µ is a probabilistic measure supported on D , which deﬁnes the inner product of Hilbert space U as < f, g > U =
(cid:82) D f ( x ) g ( x ) dµ ( x ) .
Then, there exists a neural network h
η : R d × R d → R whose activation functions are of the Tauber-Wiener class, such that
[Equation]
where H ( v )( x ) = (cid:82) D h η ( x , y ) v ( y ) dµ ( y ) .
Proof.
Since U is a Polish space, we can ﬁnd a compact set K , such that ν
( U \ K ) ≤ (cid:15) .
Therefore,
Lemma A1 can be applied, to ﬁnd a number n ∈ N , a series of continuous linear functionals
G j ∈ C ( U ; R ) and functions ϕ j ∈ C
( D ; R ) such that
[Equation]
Denote ˆ H n ( v ) = (cid:80) n j =1 G j ( v ) ϕ j , and let 1 < q < ∞ be the Hölder conjugate of p .
Since
U =
L p ( D ; R ) , by Reisz Representation Theorem, there exists functions g j ∈ L q ( D ; R ) , such that G j ( v ) =
(cid:82) D v ( x ) g j ( x ) dµ ( x ) for j = 1 , . .
. , n and v ∈
L p ( D ; R ) .
By density of C ( D ; R ) in
L q ( D ; R ) , we can ﬁnd functions ψ 1 , . . .
, ψ n ∈ C ( D ; R ) , such that
[Equation]
Then, we deﬁne ˜
H n : L p ( D ; R ) →
C ( D ; R ) by
[Equation]
For the universal approximation (density) (Hornik et al., 1989) of neural networks, we can ﬁnd a Multi-layer Feedforward network h
η : R d × R d → R whose activation functions are of the
Tauber-Wiener class, such that
[Equation]
Let H η ( x ) =
(cid:82) D h η ( x , y ) v ( y ) dµ ( y ) .
Then, there exists a constant C 1 > 0 , such that || ˆ H n ( v ) − H ( v ) ||
L p ( D ; R ) ≤
[Equation]
1 ( || ˆ H n ( v )
[Equation]
˜ H n ( v ) ||
L p ( D ; R ) + ||
˜ H n ( v ) − H ( v ) ||
L p ( D ; R ) ) .
For the ﬁrst term, there is a constant C 2 > 0 , such that
[Equation]
(cid:88)
[Equation]
[Equation]
||
[Equation]
( y )
[Equation]
|
[Equation]
D
[Equation]
R ) , for some C 3 > 0 .
And for the second term, |
[Equation]
Therefore, there is a constant C > 0 , such that (cid:90) [Equation]
Because of the assumption that E v ∼ ν
||
v
|| U < ∞ , and (cid:15) is arbitrary, then
[Equation]
the proof is complete.
Corollary A3.
Deﬁne H η
( v ) =
(cid:82) D h η ( x − y , x , a ( y )) v ( y ) dµ ( y ) , the interpolation operator can also approximate v to any precision (cid:15) .
Proof.
We use a one-layer neural network h
η : D × D
→ R as an example, which is deﬁned as h η ( x , y , a ( y ) =
σ ( (cid:80) di =1 w x,i x ( i ) + w y,i y ( i ) + b ) .
We can rewrite it as
[Equation]
+
[Equation]
where w a,j = 0 .
Corollary A4.
The Theorem A2 and Corollary A3 can be extended for v : D → R d v , where d v > 1 .
Proof.
As v = ( v (1) , v (2) , . .
, v ( d v ) ) , for each v ( j ) , a single neural network can be used for approximation.
Moreover, in implementation, we make h η fully-connected, to improve the expressivity.
Remark.
As (cid:80) x i ∈
X v ( x i )
h η ( x − x i , x i , a ( x i )) is the unbiased estimation of E y ∼ µ ( h η ( x , y ) v ( y )) , we use the Equation.
(8) for the approximation.
B E XPERIMENTS B. B ENCHMARK M ETHOD D ESCRIPTION
Vision Mixers.
We provide a framework for vision mixers as PDE solvers, including VIT, MLPM IXER , FN ET , GFN, FNO, PFNO and our NFS.
The intermediate architecture of mixing layers is shown in Fig.
B1.
The code of our framework will be released soon.
And the resampling and back-sampling methods are stacked before ‘Equispaced Input’ and ‘Equispaced Output’.
In this way, the description of the Vision Mixers included in our framework can be described by different modules, as shown in Table.
B1.
All the trials on Vision Mixers set embedding size as 32 , batch size as 4 , layer number of the intermediate equispaced mixing layers as 2 .
In FNO and PFNO, the truncated K max is set as 16 .
The patch size of Vision Mixers with patchwise embedding are set as [4 , 2] in 1-d PDEs and [4 , 4 , 2] in 2-d PDEs.
The interpolation layers in NFS are composed of one layer of feed-forward network whose perceptron unit is equal to 4 × embedding size of the model.
TokenMixing
ChannelMixing
TokenDemixing
ChannelMixing
Norm
ResidualTransform
Activation EquispacedInput
EquispacedOutput
[Table]
Graph Spatio-Temporal Models.
The evaluated graph spatio-temporal neural networks are based on recurrent neural networks for dynamics modeling, where the spatial dependency is modeled by graph neural networks.
The spatial and temporal modules for AGCRN, DCRNN and GCGRU are shown in Table.
B2.
MPPDE used different architecture, with the pushforward trick used for taining, with rolling equaling 1 and time window equaling to 10 .
All the trials on these graph spatio-tempral models set embedding size as 64 , except MPPDE as 128 .
Batch size is set as 4.
When the graph convolution needs multi-hop message-passing, we set the hop as 2 .
For MPPDE, the layer number of
GNNs is 6.
The embedding dimension in AGCRN is set as 2 .
[Table]
B. D ATA G ENERATION
Burgers’ Equation.
The initial condition u 0 ( x ) is generated according to u 0 ∼ N (0 . 625( − ∆ + 25 I ) − 2 ) with periodic boundary conditions.
ν is set as 0 .
x ∈ [0 , 1] and t ∈ [0 , 1] .
The spatial resolution is 1024, and time resolution is 200.
The dataset generation follows FNO’s protocol, which can be downloaded from its source code on ofﬁcial Github.
KdV Equation.
The equation is written as
[Equation]
where x ∈ [0 , 1] .
The initial condition u 0 ( x ) is calculated as
[Equation]
where c i ∼ N (0 , σ i ) , and a i , b i > 0 .
The spatial resolution is 1024 .
The dataset is generated by scipy package, with fftpack.diff used as pesudo-differential method and odeint used as forward Euler method.
Darcy Flow.
The equation is written as
[Equation]
The original resolution is 256 × 256 .
a ( x ) is generated by Gaussian random ﬁeld, and we directly establish the operator to learn the mapping of a to u .
NS Equation.
Our generation of NS Equation is based on FNO’s Appendix.
A.3.3, with the forcing is kept ﬁxed.
The original spatial resolution is 128 × 128 , and time resolution is 200 .
B. C OMPLETE R ESULTS ON M ODEL C OMPARISON
Here we give complete results on the four Equations.
Table.
B3 give the performance comparison on
Darcy ﬂow of both equispaced and non-equispaced scenarios.
Table.
B4 and B5 gives performance comparison in equispaced scenarios on the other three time-dependent problems.
Table.
B6 and
B7 gives performance comparison in non-equispaced scenarios on the other three time-dependent problems.
In all the tasks except Darcy Flow , the depth of layer is set as 2 , and k max = 16 in both
NFS and FNO.
However, we ﬁnd in Darcy Flow , k max should be set much larger, or the loss will not decrease.
In the reported results, k max = 32 , 64 , 128 in Darcy Flow .
[Table]
NFS fails to model the non-equispaced Burgers’ Equation when n t is set as 1, in which the performance is far from it can achieve in equispaced scenarios.
Such problem will be our future work.
Validation loss on Burgers’ ( n t = 1) of VIT, GFN, and FNO does not converge.
The results show that the early-stopping occurs in the begining of training.
[Table]
[Table]
[Table]
[Table]
B. M ORE V ISUALIZATION
Here we provide more visualization results on the three equations.
See Fig.
B2, Fig.
B3 and Fig.
B
[Figure]
[Figure]
B. M ESH - INVARIANT E VALUATION
[Table] [Table]
The mesh-invariant evaluation on Burgers’ and KDV Equations of NFS are given in Table. B8.
In Table.
B8, when the spatial resolution is just 512 , inference performance on unseen meshes deteriorates.
This result also validate our conclusion (2) in the third paragraph in Sec. 4.
Besides, we give a full evaluation on mesh-invairance of NFS in NS equation, with its variants as a detailed results corresponding to Table.
B9.
B. N EIGHBORHOOD S IZE ’ S E FFECTS
The effects of mean neighborhood size on the predictive performance on Burgers’ ( n s = 512 , n t = 10 , n (cid:48) t = 40) and KDV ( n s = 512 , n t = 10 , n (cid:48) t = 40) are shown in Fig. B5.
      ORJ0HDQ1HLJKERUKRRG1XPEHU   
[Figure] ( 
[Figure]
  
[Figure]

[Figure]
   
[Figure]
    
[Figure]

[Figure]
      
[Figure]

[Figure]
B. I NTERPOLATION WITH O THER V ISION M IXERS
We conduct experiments on non-equispaced NS equations with the combination of our interpolation layers and other Vision Mixers to ﬁgure out if they can achieve camparable performance.
[Table]
[Table]
C OMPLEXITY C OMPARISON
We here ﬁrst give Table.
B11 to show the complexity of time and memory of all the evaluated methods on NS ( r = 64 , n t = 10 , n (cid:48) t = 40 ) [Table]
[Table]
[Table]
It demonstrates that our method has comparable efﬁciency to Vision Mixers.
For the graph spatialtemporal models, they suffer from the recurrent network structures and thus are extremely timeconsuming while the parameter number is small, limiting their ﬂexibility.
Time.
However, once we compare the used time in PFNO and NFS, we will ﬁnd that the interpolation layers are considerably time-consuming.
Another module that cost time complexity is the normalization layer, as the original FNO does not include Layer-Norm in its architecture, but it is stacked in PFNO.
Theoretically, PFNO handles down-sampled grids in a low resolution, because of the patchwise embedding.
However, it takes more time than FNO.
Therefore, we conclude that the time complexity brought from Layer-Norm is very signiﬁcant, but it is affordable because of the performance improvements.
Memory.
Besides, the operation of searching for each spatial point’s neighborhood and calculating weighted summation in Eq.
(9) and Eq.
(10) are very memory-consuming.
We test it on the same experiment, and give the memory usage of different models in forward process, as shown in Table.
B12.
The memory cost in backward process is 6902MB. C E MPIRICAL OBSERVATION FOR T HEOREM 3.
In Sec. 3.2, Theorem 3.1 is proved to assure the expressivity of NFS.
However, no further evidence gives the assurance of the convergence of the kernel interpolation.
Here we conduct empirical study to give some clues.
We conduct experiments on NS equation with n s = 4096 , n t = 10 , n (cid:48) t = 40 .
In a single trial, NFS is trained with ﬁxed meshes.
We repeated the trials 10 times with different meshes, and then give the one-v.s.-all deviations of the representation states calculated by
Diﬀ = 10 (cid:88) i (cid:54) =
j,i,j = |
H
i − H
j | | H i | + | H j | where H i is the representation states of the shape [ √ m s , √
m s , t ] .
If the Diﬀ is small in the beginning and end, it can be inferred that the interpolation kernel function converges to a similar mapping since the ﬁnal predictions are close to ground truth in these experiments, and the inputs are sampled from the same instance of PDEs.
We give the Diﬀ before the ﬁrst FNO and after the ﬁnal of FNO layers in Table.
C1.
The small values indicate that the trained model usually has similar representation states.
Figure.
C1 and C2 give visualizations of representation states obtained by one instance of NS equation in two different trials.
It indicates that the differences are getting smaller during the training.
[Table]
[Figure]
[Figure]
As a result
[Figure]
we present the one-v
[Figure]
s.-all differences of different epochs in the training process
[Figure]
to validate the convergence, as shown in Figure.
C3.
[Figure]