Factorized Fourier Neural Operators
Authors: Alasdair Tran
Alexander Mathews
Lexing Xie
Cheng Soon Ong
Venue: ICLR 2023 poster (2023)
Keywords: fourier transform
fourier operators
pde
navier stokes

Abstract: We propose the Factorized Fourier Neural Operator (F-FNO), a learning-based approach for simulating partial differential equations (PDEs). Starting from a recently proposed Fourier representation of flow fields, the F-FNO bridges the performance gap between pure machine learning approaches to that of the best numerical or hybrid solvers. This is achieved with new representations – separable spectral layers and improved residual connections – and a combination of training strategies such as the Markov assumption, Gaussian noise, and cosine learning rate decay. On several challenging benchmark PDEs on regular grids, structured meshes, and point clouds, the F-FNO can scale to deeper networks and outperform both the FNO and the geo-FNO, reducing the error by 83% on the Navier-Stokes problem, 31% on the elasticity problem, 57% on the airfoil flow problem, and 60% on the plastic forging problem. Compared to the state-of-the-art pseudo-spectral method, the F-FNO can take a step size that is an order of magnitude larger in time and achieve an order of magnitude speedup to produce the same solution quality.

Published as a conference paper at ICLR 2023
FACTORIZED FOURIER NEURAL OPERATORS
Alasdair Tran"
Alexander Mathews 1
Lexing Xie
Cheng Soon Ong 1,2
1 Australian National University
Data6l, CSIRO
ABSTRACT
We propose the Factorized Fourier Neural Operator (F-FNO), a learning-based approach for simulating partial differential equations (PDEs). Starting from a recently proposed Fourier representation of flow fields, the F-FNO bridges the performance gap between pure machine learning approaches to that of the best numerical or hybrid solvers. This is achieved with new representations separable spectral layers and improved residual connections and a combination of training strategies such the Markov assumption, Gaussian noise, and cosine learning rate decay. On several challenging benchmark PDEs on regular grids, structured meshes, and point clouds, the F-FNO scale deeper networks and outperform both the FNO and the geo-FNO, reducing the error by 83% on the Navier-Stoker problem, 31% on the elasticity problem, 57% on the airfoil flow problem, and 60% on the plastic forging problem. Compared to the state-of- ofte-offthhe-rt pseudo-spectral method, the F-FNO can take step size that is an order of magnitude larger in time and achieve an order of magnitude speedup to produce the same solution quality.
INTRODUCTION
From modeling population dynamics to understanding the formation of stars, partial differential equations (PDEs) permeate the world of science and engineering. For most real-world problems, the
lack of closed-form solution requires using computationally expensive numerical solvers, sometimes
consuming millions of core hours and terabytes of storage (Hosseini et al. 2016). Recently, machine
learning methods have been proposed to replace part (Kochkov et al. 2021 or all (Li et al. 2021al of a numerical solver.
Of particular interest are Fourier Neural Operators (FNOs) (Li et al. 2021a), which are neural networks that can be trained end-to-end learn mapping between infnite-dimensionl function spaces. The FNO can take a step size much bigger than is allowed in numerical methods, can perform super-resolution, and can be trained on many PDEs with the same underlying architecture. more recent variant, dubbed geo-FNO (Li al., 2022), can handle irregular geometries such structured meshes and point clouds. However, this first generation issues. Lu et al. (2022) find that the performance of the FNO deteriorates significantly geometries and noisy data. In our own experiments, we observe that both the FNO and the geo-FNO perform worse
neural operators suffers from stability
complex
increase the network depth, eventually failing converge at 24 layers. Even at 4 layers, the error between the FNO and a numerical solver remains large (14% error on the Kolmogorov flow).
In this paper, we propose the Factorized Fourier Neural Operator (F-FNO) which contains an improved representation layer for the operator, and better of training approaches. By learning features in the Fourier space in each dimension independently, a process called Fourier factorization, we are able to reduce the model complexity by an order of magnitude and learn siggee-diimesional problems such as the 3D plastic forging problem. The F-FNO places residual connections after activation, enabling our neural operator to benefit from a deeply stacked network. Coupled with training techniques such as teacher forcing, enforcing the Markov constraints, adding Gaussian noise to inputs, and using a cosine learning rate scheduler, we are able to outperform the state of the art by a large margin three different PDE systems and four different geometries. On the Navier-Stoked (Kolmogorov flow) simulations on the torus, the F-FNO reduces the error by 83% compared to the FNO, while still achieving an order of magnitude speedup method (Figs. 3 and 4 On point clouds and structured meshes, the F-FNO outperforms the geo-FNO on both structural mechanies and fluid dynamics PDEs, reducing the error by up to 60%
the state-of-the of-the--rt pseudo-spectral
(Table 2).

---

Published as a conference paper at ICLR 2023
Overall, we make the following three key contributions:
1. We propose new representation, the F-FNO, which consists of separable Fourier represen- tation and improved residual connections, reducing the model complexity and allowing it to scale deeper networks (Fig. 2 and Eqs. (7) and (8).
2. We show the importance of incorporating training techniques from the existing literature, such as Markov assumption, Gaussian noise, and cosine learning rate decay (Fg.. 3) and investigate how well the operator can handle different input representations (Fig.
3. We demonstrate F-FNO's strong performance in variety of geometries and PDEs Table 2). Code, datasets, and pre-trained models are available
Fig.
2 RELATED WORK
Classical methods to solve PDE systems include finite element methods, finite difference methods, fi- nite volume methods, and pseudo-spectral methods such Crank-Nicholson and Carpenter-Kennee In these methods, space is discretized, and a more accurate simulation requires a finer discretization which increases the computational cost. Traditionally, we would use simplified models for specific PDEs, such as Reynolds averaged Navier-Stokes (Alfonsi, (2009) and large eddy simulation (Lesieur & Métais 1996), to reduce this cost. More recently, machine learning offers an alternative approach to accelerate the simulations. There are two main clusters of work: hybrid approaches and pure machine learning approaches. Hybrid approaches replace parts of traditional numerical solvers with learned alternatives but keep the components that impose physical constraints such as conservation laws; while pure machine learning approaches learn the time evolution of PDEs from data only.
Hybrid methods typically aim to speed up traditional numerical solvers by using lower resolution grids (Bar-Sinai et al. [2119] Um et al. [2020; Kochkov et al. 2021), or by replacing computationally expensive parts of the solver with learned alternatives Tomisson et al. 2017), Obiols-Sales al. (2020).
Bar-Sinai et al. (219)) develop a data driven method for discretizing PDE solutions, allowing coarser grids to be used without sacrificing detail. Kochkov et al.) (202T) design a technique specifically for the Navier-Stokes equations that uses neural network-based interpolation to calculate velocities between grid points rather than using the more traditional polynomial interpolation. Their method leads to more accurate simulations while at the same time achieving an 86-fold speed improvement over Direct Numerical Simulation (DNS). Similarly, [Tmpson et al. (2017) employ a numerical solver and decomposition specific to the Navier-Stokes equations, but introduce a convolutional neural network to infer the pressure map at each time step. While these hybrid methods are effective when designed for specific equations, they are not easily adaptable other PDE tasks.
An alternative approach, less specialized than most hybrid methods but also less general than pure machine learning methods, is learned correction (Um et al. 2020; Kochkov al. 2021) which involves learning a residual term to the output of a numerical step. That is, the time derivative is now u, = ut + LC(u,), ). where ut is the velocity field provided by a standard numerical solver on coarse grid, and LC((u;) ) is neural network that plays the role of super-resolution of missing details.
Pure machine learning approaches eschew the numerical solver altogether and learn the field directly, i.c., u, ((,-1) 1). where G is dubbed a neural operator. The operator can include graph neural networks Li et al. 2020a b). low-rank dccomposition Kovachki et al. (2021b), or Fourier transforms (Li ct al. 202 b). Pure machine Icarning models can also incorporate physical constraints, for example, by carefully designing loss functions based on conservation laws (Wandel et al. 2020). They can even be based on existing simulation methods such the operator designed by Wang et al. (2020) that uses learned filters in both Reynolds-averaged Navier-Stokes and Large Eddy Simulation before combining the predictions using U-Net. However, machine learning methods need not incorporate such constraints - for example, Kim et al. (2019) use a generative CNN model to represent velocity fields in low-dimensional latent space and feedforward neural network to advance the latent space to the next time point. Similarly, [Bhattacharya et al. (2020) use PCA to map from an infinite dimensional input space into a latent space, on which a neural network operates before being transformed to the output space. Our work is most closely related to the Fourier
https:///iiii htttps://github..comm/ waw.om/aaaaaaaaaa htwtairrrrraaoooorrrrraia
2

---

Published as a conference paper at ICLR 2023
evolve vorticty
compute Tlow velocity
(3) Torus
(c) Airfoil
compute stress
depllccment (8 over time
(b) Elasticity
(d) Plasticity
An illustration of the input and output different PDE problems. See the accompanying Table for details. On the torus datasets (a), the operator learns evolve the vorticity Elasticity (b), the operator learns to predict the stress value on each point on a point cloud. On Airfoil
Figure 1:
time. On
(c). the operator learns to predict the flow velocity on each mesh point. On Plasticity (d), the operator learns the displacement of each mesh point given an initial boundary condition.
transform bansform-bbaedd approaches (Li al. 202 2022) which can efficiently model PDEs with zero-shot super-resolution but is not specific to the Navier-Stokes equations.
Fourier representations are popular deep-learning due to the efficiency of convolution operators in the frequency space, the On log n) time complexity of the fast Fourier transform (FFT), and the ability to capture long-range dependencies. Two notable examples of deep learning models that employ Fourier representations are FNet (Lee-Thorp et al. (021) for encoding semantic relationships text classification and FNO (Li et al., 2021a) for flow simulation. learning mappings between function spaces, the FNO outperforms graph-based neural operators and other miite-dimensional operators such as U-Net. In modeling chaotic systems, the FNO has been shown to capture invariant
properties of chaotic systems (Li et al. 2021b). More generally, Kovachki al. (2021a) prove that the FNO can approximate any continuous operator.
THE FACTORIZED FOURIER NEURAL OPERATOR
Solving PDEs with neural operators An operator G A - LI is mapping between infinite- dimensional function spaces AA and 4. Exactly what these function spaces represent depends on the problem. In general, solving PDE involves finding solution u E l given some input parameter E A. and we would train a neural operator to learn the mapping + u. Consider the vorticity formulation of the 2D Navier-Stokes equations,
dw
+ Va = AV
Ot where the velocity field, is the vorticity, and is the external forcing function. These the governing equations for the torus datasets (Fig. lab. The neural operator would learn to evolve this field from one time step the next: Wt + W141. Or consider the equation for solid body in structural mechanics,
8uu
06
where is the mass density, u is the displacement vector and is the stress tensor. Elasticity (Fig. 1b and Plasticity (Fig. Id) are both governed by this equation. Plasticity, we would learn map [0, L] + IR to the grid position and displacement of each grid
the initial boundary condition point over time: sa + (x, u, t). Elasticity. we are instead interested in predicting the stress value for each point:
Finally consider the Euler equations to model the airflow around aircraft
wing (Fig. Ic):
0p
0pu
OE a1.
+ V (pu)
+ + (pu u u
+ (EE p)u) = 0 (3)
01.
al.

---

Published as a conference paper at ICLR 2023
9 Debormatioon
a(x.a)
P Lifting
(2.
C Layer
A Defirmation
Layer
Layer
Projection
FFT,
IFFT,
WMAA
Fondbrrazz g righ
(xt, Ro
97" ganur
FFT,
IFFT,
facter
neidhal connection
somboe oftup
gpd demanatoo ggs iter yamamiii
dostain
Xoa-liaes dostain
Figure 2: The architecture of the Factorized Fourier Neural Operator (F-FNO) for a 2D problem. The iterative process (Fq. (4)) is shown the top, which the input function aii j) is first deformed from an irregular space into a uniform space a(x, y), and is then fed through series of operator layers in order to produce the output function u(i,j). A zoomed-in operator layer (Eq. (7) is shown at the bottom which shows how process each spatial dimension independently in the Fourier space, before merging them together again the physical space.
where p the fluid mass density, p the pressure, is the velocity vector, and the energy. Here the operator would learn to map each grid point to the velocity field at equilibrium:
Original FNO and geo-FNO architectures Motivated by the kernel formulation of the solution to linear PDEs using Green's 's functions, [i et al. (202bb) 2022) propose an iterative approach to map input function a to output function u,
u = ((a) = (6 Q C(L) .... ((1) o Po P (a))
(4)
where indicates function composition, T. is the number layers/iterations, that maps the input the first latent representation 100), ((6) is the l'th non-linear operator layer, and Q is the projection operator that maps the last latent representation 2(4) the output. On irregular geometries such as point clouds, we additionally define a coordinate map é. parameterized by a small neural network and learned end-to-end, that deforms the physical space of irregular geometry into regular computational space. The architecture without this coordinate map is called FNO, while the one with the coordinate map is called geo-FNO, Fig. (top) contains a schematic diagram of this iterative process.
the lifting operator
Originally, Li al. (021a) formulate each operator layer as
( <0 (:(0) ( (w(0)) + (()) + xx(0
(5)
where a : + R point-wise non-linear activation function, 20(0)(00 +b() is an affine point-wise map in the physical space, and x(()) is a kernel integral operator using the Fourier transform.
KC (2(0) IFFT R(()) FFT(2)
The Fourier- domain weight matrices {R( (1, L}}
take up most of the model size, requiring DLHHMM)) parameters, where H is the hidden size, M the number of top Fourier modes being kept, and / is the problem dimension. Furthermore, the constant value for M and the affine pise map allow the FNO to be resolution-inddeenndnn.
Our improved F-FNO architecture
We propose changing the operator layer in Eq. to: (h(
c( (200 20 + ==0)
(w( + (1)) [
Note that we apply the residual connection (2 term) after the non- con-linearity to preserve more of the layer input. We also use a two-layer feedforward, inspired by the feedforward design used in transformers (Vaswani al. 2017). More importantly, we factorize the Fourier transforms over the problem dimensions, modify Siddfyingl Eq. (6)
(1)) x00)
Σ dED
IFFT (()) FFFa(=(1))
(8)
4

---

Published as a conference paper at ICLR 2023
Table I: An overview of the datasets and the corresponding task.
Dataset
Geometry
Dim.
Problem
Input
Output
TorusLi
regular grid
2D
Kolmogorov flow
We
W++1
TorusKochkov
regular grid
2D
Kolmogorov flow
Wr
W[+1
Torus Vis
regular grid
2D
Kolmogorov flow
wy and v
W+++
TorusVisForce
regular grid
2D
Kolmogorov flow
wy and and ft
W++1
Elasticity
point cloud
2D
hyper-elastic material
point cloud
stress
Airfoil
structured mesh
2D
transonic flow
mesh grid
velocity
Plasticity
structured mesh
3D
plastic forging
boundary condition
displacement
The seemingly small change from R0 to R2 in the Fourier operator reduces the number of parame- ters to O(LHHMM D). This is particularly useful when solving higher-dimensional problems such as 3D plastic forging (Fig. ld). The combination of the factorized transforms and residual connections allows the operator to converge in deep networks while continuing to improve performance (Fig. 3). is also possible to share the weight matrices Ra between the layers, which further reduces the parameters O MD). Fig. 2 (bottom) provides an overview of an F-FNO operator layer.
Furthermore, the F-FNO is highly flexible in its input representation. which means anything that is relevant to the evolution of the field can be an input, such as viscosity external forcing functions for the torus. This flexibility also allows the F-FNO to be easily generalized to different PDEs.
Training techniques to learn neural operators We find that combination of deep learning techniques are very important for the FNO to perform well, most of which were overlooked in Li et al. (2021a)'s original implementation. The first is enforcing the first-order Markov property. We find et al. (2021a)'s use of the last 10 time steps as inputs to the neural operator to be unnecessary. Instead, is sufficient feed information only from the current step. just like a numerical solver. Unlike prior works (Li et al. 202 Kochkov et al. (021), we do not unroll the model during training but instead use the reacher forcing technique which is often seen in time series and language modeling. In teacher forcing, we use the ground truth as the input to the neural operator. Finally during training, we find it useful normalize the inputs and add a small amount of Gaussian noise, similar to how Gonzalez et al.] (2020) train their graph networks. Coupled with cosine learning rate decay, we are able to make the training process of neural operators more stable. Ablation studies for the new representation and training techniques can be found in Fig.
4 DATASETS AND EVALUATION SETTINGS
PDEs on regular grids The four Torus datasets on regular grids (TorusLi, TorusKochkov, TorusVis, and Torus VisForce, summarized in Table I) are simulations based on Kolmogorov flows which have been extensively studied in the literature (Chandler & Kerswell 2013). In particular, they model turbulent flows the surface of 3D torus (i.e., 2D grid with periodic boundary conditions). TorusLi is publicly released by Li et al. (2021a) and is used to benchmark our model against the original FNO. The ground truths are assumed be simulations generated by the pseudo-spectral Crank- Nicholson second-order method on 64x64 grids. All trajectories have constant viscosity = 10-5 (Re 2000), use the same constant forcing function, (a, y) = 0.1 (21 (x++y)) cos (x+y and differ only in the initial field.
Using the same and Torus VisForce, to test the generalization of the F-FNO across Navier-Stokes tasks with different viscosities and forcing functions. particular, for each trajectory, we vary the viscosity between 10-1 and 10-5 and set the forcing function to
numerical solver, we generate two further datasets, called Torus Vis
f(t,,,)) x) = 0.1 E.
Omg sin (2τp(ir + ju) + 5t) + Bmn cos (2mp(iz + +Jv) + st)] (9)
where the amplitudes Opij and Spiy are sampled from the standard uniform distribution. Furthermore, is set to 0 Torus Vis, making the forcing function constant across time; while it is set to 0.2 in Torus Fis giving us time-varying force.
5

---

Published as a conference paper at ICLR 2023
Finally, we regenerate TorusKochkov (Fig. lap using the same settings provided by Kochkov et al. (2021) but with different initial conditions from the original paper (since the authors did not release the full dataset). Here the ground truths obtained from simulations 2048x2048 grids using the pseudo-spectral Carpenter-Kenn downsampled smaller grid sizes, allowing to study the Pareto frontier of the speed vs accuracy space (see Fig. 4ah. TorusKochkov uses a fixed viscosity of 0.001 and a constant forcing function = 4 cos 4y) ã 0.lu, but on the bigger domain of [0, 27.. Furthermore, we generate only 32 training trajectories to test how well the F-FNO can learn on low-data regime.
fourth-order method. The full-scale simulations are then
PDEs on irregular geometries The Elasticity, Airfoil, and Plasticity datasets (final three rows in Table 1) are taken from Li et al. (2022). Elasticity is point cloud dataset modeling the incompressible Rivlin-Saunders material (Pascon 2019). Each sample is unit cell with a void in the center of arbitrary shape (Fig. lb). The task to map each cloud point to its stress value. Airfoil models the transonic flow over an airfoil, shown as the white center in Fig. Ic The neural operator would then learn to map each mesh location to its Mach number. Finally, Plasticity models the plastic forging problem, in which a die, parameterized by an arbitrary function and traveling at a constant speed, hits a block material from above (Fig- lch. Here the task is to map the shape of the die to the 101 x 31 structured mesh over 20 time steps. Note that Plasticity expects a 3D output, with two spatial dimensions and one time dimension.
Training details For experiments involving the original FNO, FNO-TF (with teaching forcing), FNO-M (with the Markov assumption), and FNO-N (with improved residuals), we use the same training procedure as Li et al. 2021a). For own models, we train for 100,000 steps the regular grid datasets and for 200 epochs for the irregular geometry datasets, warming up the learning rate to 2.5 x 10-3 for the first 500 steps and then decaying it using the cosine function (Loshchilov & Hutter, 2017). We use ReLU as our non- liner activation function, clip the gradient value at 0.1, and the Adam optimizer (Kingma& & Ba 2015) with 81 0.9, B2 0.999, c 10-8 The weight decay factor is set to 10-4 and is decoupled from the learning rate (Loshchilov & Hutter (001)). In each operator layer on the torus datasets, we always throw away half of the Fourier modes (e.g., on a 64x64 grid, we keep only the top 16 modes). Models are implemented in PyTorch (Paszke et al. 2017) and trained on a single Titan GPU.
Evaluation metrics
We use the normalized mean squared error as the loss function, defined
1 ف wll2
N-MSE
B
| 2
where |1-112 is the 2-norm, B is the batch size, and @ is the prediction of the ground truth w. In addition to comparing the N-MSE directly, for TorusKochkov, we also compute the vorticity correlation, defined as
p(w, ()
ΣΣ
Wei
A||2
and from which we measure the time until this correlation drops below 95%. To consistent with prior work, we use the N-MSE to compare the F-FNO against the FNO and geo-FNO Li et al. 2202 2022). and the vorticity correlation to compare against Kochkov et al. 2222 's work.
5 RESULTS FOR NAIVER-STOKES ON A TORUS
Comparison against FNO The performance on TorusLi is plotted in Fig. 3 with the raw numbers shown in Table A.3 We note that our method F-FNO is substantially more accurate than the FNO regardless of network depth, when judged by N-MSE. The F-FNO uses fewer parameters than the FNO, has similar training time, but generally has a longer inference time. Even so, the inference time for the F-FNO is still up to two orders of magnitude shorter than for the Crank-Nicolson numerical solver.
In contrast to our method, Li et al. (2021a) do not use teacher forcing during training. Instead they use the previous 10 steps input to predict the next 10 steps incrementally (by using each predicted

---

Published as a conference paper at ICLR 2023
5.5
jpropesed by [2021a]]
Homalbed NSE
15.0
4.5
FNO TF (FNO with teacher forcing)
FND-M (PNO-TF with Narkov assumption)
improved FND-- (FNG-R with a bag of tric
10.0
7.5
3.0
F-FND IFNC+ with Fourier factorization)
F-FNO EFD weight sharing)
5.0
16 layers
20
24
12 16
20
24
Number layers
(a)
(b)
Figure 3: Performance (lower is better) on TorusLi, with error bars showing the min and max values over three trials. We show the original FNO (Li et al. 202 along with variants that use: teacher forcing, Markov assumption, improved residuals, a bag of tricks, Fourier factorization, and weight sharing. Note that F-FNO and F-FNO-WE are presented on separate plot (b) to make visualizing the improvement easier (if shown in (a), F-FNO and F-FNO-WS would just be straight line).
Tie until comelationn 95%
1024x10266
Time until comelatioon 95%
4-
128x128
512x512
256x256
64x64
(CarpertanKknnedly 4tb order)
128x128
2-
F-FND-WS (-IND with weight sharing)
66484
1-
32x32
10-2
10-1
10
10 -1
10-7
107
Runtime per time unit (s)
Step size
(a)
(b)
Figure 4: Performance of F-FNO on TorusKochkov. In (a), we plot the time until the correlation with the ground truths in the test set drops below 95% on the y-axis, against the time it takes to run second of simulation on the x-axis. In (b), show how, the validation set of Tocuskoc given a fixed spatial resolution of 64x64, changing the step size has no effect on the numerical solver; however there is an optimal step size for the F-FNO at around 0.2.
value the input to the next step). We find that the teacher forcing strategy (FNO-TF, orange line), in which we always use the ground truth from the previous time step as input during training, leads to smaller N-MSE when the number of layers less than 24. Furthermore, enforcing the first-order Markov property (FNO-M, dotted green line), where only one step of history is used, further improves the performance over FNO and FNO-TF. Including two more steps of history does not improve the results.
The models FNO, FNO-TF, and FNO-M do not scale with network depth, as seen by the increase in the N-MSE with network depth. These models even diverge during training when 24 layers are used. FNO-R, with the residual connections placed after the non-linear,ty, does not suffer from this problem and can finally converge at 24 layers. FNO++ further improves the performance, as result of careful combination of: normalizing the inputs, adding Gaussian noise to the training inputs, and using cosine learning rate decay. In particular, we find that adding small amount of Gaussian noise to the normalized inputs helps to stabilize training. Without the noise, the validation loss the early stage of training can explode.
Finally, if we use Fourier factorization (F-FNO, yellow dashed line), the error drops by an additional 35% (3. 73% + 2.41%) at 24 layers (Fig. 3b), while the parameter count is reduced by order of magnitude. Sharing the weights in the Fourier domain (F-FNO-WS. red line) makes little difference to the performance especially at deep layers, but it does reduce the parameter count by another order of magnitude to IM (see Fig. A. and Table A.3
Trade-off between speed and accuracy From Fig. 4a we observe that our method F-FNO only needs 64x64 input grids reach similar performance to 128x 128 grid solved with DNS. At the
7

---

Published as a conference paper at ICLR 2023
1.0
Tonuvits TorusVisorcce
40-
0.8
10 NMSSE (55
0.4
0.2
Vertidty
10
0.0
Vorbidty Cosecinates
Verlidty Conrciattes Velncity
contest
viscosity
(a) Performance of F-FNO on different input features: having only vorticity as an input with no further con- text (first group); having vorticity and the force field inputs (sccond group); and having vorticity, the force field, and viscosity as inputs (third group). The error bars are the standard deviation from three trials.
(b) Effect having the coordinates velocity as additional input channels on TorusKochkov. A higher model that can correlate with the ground-trun truth vorticity for longer. Error bands corre- spond to min and max values from three trials.
line corresponds
Figure 5: Performance of F-FNO different contexts and input representations.
same time, the F-FNO also achieves an order of magnitude speedup. While the highly specialized hybrid method introduced by Kochkov al. (2021) can achieve a speedup closer two orders of magnitude over DNS, the F-FNO takes much more flexible approach and thus can be more easily adapted other PDES and geometries.
The improved accuracy of the F-FNO over DNS when both methods are using the same spatial resolution can be seen graphically Fig. A.4 In this example, the F-FNO on 128x128 grid produces vorticity field that is visually closer to the ground truth than DNS running the same
grid size. This is also supported by comparing the time until correlation falls below 95% in Fig. 4a
Optimal step size The F-FNO includes step size parameter which specifies how many seconds of simulation time one application of the operator will advance. A large step size sees the model predicting far into the future possibly leading large errors, while small step size means small errors have many more steps to compound. We thus try different step sizes in Fig. 4b
In numerical solvers, there is close relationship between the step size in the time dimension and the spatial resolution. Specifically, the Dourant-Friedrichss step size given space discretization: At Cmax Aa/ max. This means that if we double the grid size, the solver should take step that is twice as small (we follow this approach to obtain DNS's Pareto frontier in Fig. 4ah. Furthermore, having step sizes smaller than what is specified by the CFL condition would not provide any further benefit unless also reduce the distance between two grid points (see purple line Fig. 4b). On the other hand, a step size that is too big (e.g., bigger than 0.05 on a 64x64 grid) will lead to stability issues in the numerical solver.
(CFL) condition provides an optimal
For the F-FNO, we find that we can take step size that is at least an order of magnitude bigger than the stable step size for numerical solver. This is the key contribution to the efficiency of neural methods. Furthermore, there is sweet spot for the step size around 0.2 on TorusKochkov - and unlike its numerical counterpart, we find that there is no need reduce the step size as train the F-FNO on a higher spatial resolution.
Flexible input representations The F-FNO can be trained to handle Navier-Stoke equations with viscosities (in TorusVis) and time-varying forcing functions (in Torus VisForce) provided at inference time. Our model, when given both the force and viscosity, in addition to the vorticity, is able to achieve an error of 2% (Fig. 5ap. If we remove the viscosity information, the error doubles. Removing the forcing function from the input further increases the error by an order of magnitude. This shows that the force has substantial impact on the future vorticity field, and that the F-FNO can use information about the forcing function to make accurate predictions. More generally, different datasets benefit from having different input features Table A.7 shows the minimum set of features to reach optimal performance on each of them. We also find that having redundant features does not significantly hurt the model, so there need do aggressive feature pruning in practice.
8

---

Published as a conference paper at ICLR 2023
Table 2: Performance (N-MSE, expressed as percentage. where lower is better) on point clouds (Elasticity) and structured meshes (Airfoil and Plasticity) between our F-FNO and the previous state-of-the-art geo-FNO et et al., 2022). Cells with a dash correspond to models which do not converge. The N-MSE is accompanied by the standard deviation from three trials. More detailed results are shown
Tables A.4 to A.6
Elasticity
Airfoil
Plasticity
No. of layers
geo-FNO
F-FNO
geo-FNO
F-FNO
geo-FNO
F-FNO
4 layer
2.5 0.1
3.16 + 1.29
1.9 0.4
0.79 + 0.02
0.74 0.01
0.48 + 0.02
8 layer
3.3 1.3
2.05
1.4
0.64 0.01
0.57 = 0.04
0.32 0.01
12 layer
16.8
1
4.1
0.62 0.03
0.45 = 0.03
0,25 0.01
16 layer
16.3 + 0.4
1.86
0.61 0.01
0.22 0.00
20 layer
16.0 + 0.7
1.002
0.57
0.20 + 0.02
24 layer
15.9 + 0.5
1.04
0.58
0.18
Our experiments with different input representations also reveal an interesting performance gain from the double encoding of information (Fig. (b). All datasets benefit from coordinate encoding - i.e. having the (x, coordinates two additional input channels even if the positional information is already contained in the absolute position of grid points (indices of the input array). We hypothesize that these two positional representations
used by different parts of the F-FNO. The Fourier transform uses the absolute position of the grid points and thus the Fourier layer should have no need for the (ar, positional features. However, the feedforward layer in the physical space is pointwise operator and thus needs to rely on the raw coordinate values, since it would otherwise be independent of the absolute position of grid points.
RESULTS FOR PDES ON POINT CLOUDS AND MESHES
As shown in Table the geo-FNO (Li et al. 2022). similar to the original FNO, also suffers from network scaling. It appears to be stuck in local minimum beyond 8 layers in the Elasticity problem and it completely fails converge beyond 12 layers Airfoil and Plasticity. Plasticity is the only task in which the geo-FNO gets better as we from 4 to 12 layers (0.74% + 0.45%). In addition to the poor scaling with network depth, we also find during our experiments that the geo-FNO can perform worse as we increase the hidden size H. This indicates that there might not be enough regularization in the model as we increase the model complexity.
Our F-FNO, on the other hand, continues to gain performance with deeper networks and bigger hidden size, reducing the prediction error by 31% % on the Elasticity point clouds (2.51% - 1.74%) and by 57% on the 2D transonic flow over airfoil problem (1.35% - 0.58%). Our Fourier factorization particularly shines in the plastic forging problem. in which the neural operator needs to output a 3D array, i.e., the displacement of each point 2D mesh over 20 time steps. As shown in Table A.6 our 24-layer F-FNO with 11M parameters outperforms the 12-layer geo-FNO with 57M parameters by 60% (0.45% + 0.18%).
CONCLUSION
The Fourier transform powerful tool to learn neural operators that can handle long-range spatial dependencies. By factorizing the transform, using better residual connections, and improving the training setup, our proposed F-FNO outperforms the state of the art on PDEs variety of geometries and domains. For future work, we are interested in examining equilibrium properties of generalized Fourier operators with an infinite number of layers and checking if the universal approximation property (Kovachki et al.) 2021a) still holds under Fourier factorization.

---

Published as a conference paper at ICLR 2023
REFERENCES
Giancarlo Alfonsi. Reynolds-averaged navier-stokes equations for turbulence modeling. Applied Mechanics Reviews, 62:040802, 2009.
Yohai Bar-Sinai, Stephan Hoyer, Jason Hickey, and Michael Brenner. Learning data-driven discretizations for partial differential equations. Proceedings of the National Academy of Sciences, 11663))11334-15549, 2019. ISSN 0027-8424. doi: 10.10733pnas.1811 181.ma8..11148881111 16. URL, https 7/www. www.ppnas... wttpr///oooneeeee/////1///
Kaushik Bhattacharya, Bamdad Hosseini, Nikola B Kovachki, and Andrew M Stuart. Model reduction and neural networks for parametric PDEs. arXiv preprint arXiv:2005.03/14 2020.
Gary J. Chandler and Rich R. Kerswell. Invariant recurrent solutions embedded a turbulent two-dimensional kolmogorov flow. Journal of Fluid Mechanics, 722:554 595, 2013. doi: 10.1017/ 2fm..013.122.
S.M. Hosseini, R. Vinuesa, P. Schlatter, A. Hanifi, and D.S. Henningson. Direct numerical simulation of the flow around wing section at moderate reynolds number. International Journal of Heat and Fluid Flow, 61:117-128 2016. ISSN 0142-727X. doi: http:///ii.or/11.10111111 ofgh ijheatfluidflow
2016.02.000 02.00. S014272 6 6300169
URL https: / /twww/wwww......//////////// www. scwencedireceeeo com/ scwence/ attp///cee/atiiccepppp../. SI TSFP9 special issue.
Byungsoo Kim, Vinicius C Azevedo, Nils Thuerey, Theodore Kim, Markus Gross, and Barbara Solenthaler. Deep fluids: A generative network for parameterized fluid simulations. In Computer Graphics Forum, volume 38, pp. 59-70. Wiley Online Library, 2019.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015.
Dmitrii Kochkov, Jamie A. Smith, Ayya Alieva, Qing Wang. Michael P. Brenner, and Stephan computational fluid dynamies. Proceedings of the National
Hoyer. Machine laaming--acclerren Academy of Sciences, 1 118(11), 2021. ISSN 0027-8424. doi: 20.1073/paas.20018811 URL
https: www pnas /ttp:///eeeeet/////eeeeeeee/
Nikola Kovachki, Samuel Lanthaler, and Siddhartha Mishra. On universal approximation and error bounds for fourier neural operators. Journal of Machine Learning Research, 22(290):1-76, 2021a.
URL http:http///mt fmmrr wttppppers//pppepppee 220p:/000222220066. html
Nikola B. Kovachki, Zong-Y Yi Li, Burigede Liu, Kamyar Arizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Animashree Anandkumar. Neural operator: Learning maps between function spaces. ArXiv, abs/2108.08433 088811 2021b.
J. Lee- Lherp. Joshua Ainslie, llya Eckstein, and Santiago Ontañón. Fnet: Mixing tokens with fourier transforms. ArXiv, abs/2105.03824 2021.
Marcel R. Lesieur and Olivier Métais. New trends in large-eddy simulations of turbulence. Annual Review of Fluid Mechanics, 28:45 82, 1996.
Zong-Yi Yi Li, Daniel Z. Huang, Burigede Liu, and Anima Anandkumar. Fourier neural operator with learned deformations for PDEs on general geometries. ArXiv, 0bs/2207.05209, 2022.
Zongyi Li, Nikola B. Kovachki, Azizzadenesheli, Burigede Liu, Bhattacharya, Andrew Stuart, and Anima Anandkumar. Multipole graph neural operator for parametric partial differential equations. ArXiv, abs/2200600935, 2020a.
Zongyi Li, Nikola B. Kovachki, Azizzadenesheli, Burigede Liu, K. Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Graph kernel network for partial differential equations. ArXiv, abs/2003.03485, 2020b.
Zongyi Li, Nikola B. Kovachki, K. Azizzadenesheli, Burigede Liu, K. Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations, 2021a. btRpstss: tetetps://openee //operrevieww net/ wttp:/ummmicc88pm00m d=
10

---

Published as a conference paper at ICLR 2023
Zongyi Li, Nikola B. Kovachki, K. Azizzadenesheli, Burigede Liu, K. Bhattacharya, Andrew Stu- art, and Anima Anandkumar. Markov neural operators for learning chaotic systems. ArXiv, abs/2106.066 06898, 2021b.
Hya Loshchilov and Frank Hutter. SGDR: stochastic gradient descent with warm restarts. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26. 2017, Conference Track Proceedings. OpenReview.net, 2017. URL /ttppsss//ppeeeeieee net
Iyya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Reprexentations, 2019.
Lu Lu, Xuhui Meng, Shengze Cai, Zhiping Mao, Somdatta Goswami, Zhongqiang Zhang, and George Em Karniadakis. A comprehensive and fair comparison of two neural operators (with practical extensions) based on fair data. Computer Methods in Applied Mechanics and Engineering, 2022.
Octavi Obiols-Sales, Abhinav Vishnu, Nicholas Malaya, and Aparna Chandramowliswhaaan. Cfd- net: A deep learning-based accelerator for fluid simulations. In Proceedings of the 34th ACM International Conference on Supercomputing, pp. 1-12, 2020.
João Paulo Pascon. Large deformation analysis of plane-stress hyperelastic problems via triangular membrane finite elements. International Journal of Advanced Structural Engineering. 2019.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary De Veto,o, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in PyTorch. In NIPS Autodiff Workshop, 2017.
Michael Poli, Stefano Massaroli, Federico Berto, Jinkyoo Park, Tri Dao, Christopher Re, and Stefano Ermon. Transform once: Efficient operator learning in frequency domain. In /CMI. 2022 2nd AI for net/
Science Workshop, 2022.
j41N
Alvaro Sanchez-Gonzalee Jonathan Godwin, Tobias Pfaff, Rex Ying. Jure Leskovec, and Peter W. Battaglia. Learning simulate complex physies with graph networks. In Proceedings of the 37th International Conference on Machine Learning, ICML'20. JMIR.org. 2020.
Jonathan Tompson, Kristofer Schlachter, Pablo Sprechmann, and Ken Perlin. Accelerating eulerian fluid simulation with convolutional networks. In International Conference on Machine Learning. pp. 1424-3433 PMLR, 2017.
Kiwon Um, Robert Brand, Yun Fei, Philipp Holl, and Nils Thuerey, Solver-in-the-Loop Learning from Differentiable Physics to Interact with Iterative PDE-Solvers. Advances in Neural Information Processing Systems, 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar. Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garett (eds.), Advances in Neural Information Processing Systems. volume 30. Curran Asso- ciates, Inc., 2017. URL 3f5eee2354 9 (ww.:0555ccc444400 c4 Paper .ddf]
newrips. wttp://ppppppp///////LLL
Nils Wandel, Michael Weinmann, and Reinhard Klein. Learning incompressible fluid dynamics from scratch-towards fast, differentiable fluid models that generalize. arXiv preprint arXiv:2006.08762 2020.
Rui Wang, Karthik Kashinath, Mustafa Mustafa, Adrian Albert, and Rose Yu. Towards physics- informed deep learning for turbulent flow prediction. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 1457-1466, 2020.
11

---

Published as a conference paper at ICLR 2023
A
APPENDIX
Table A.1: An overview of the four fluid dynamics datasets on regular grids. Our newly generated datasets, Torus Vis and Torus VisPorce, contain simulation data with more variety of viscosities and forces than TorusLi (Li et al. 2021a) and TorusKochkov (Kochkov et al. 2021). Note that Li et al. (2021a) did not generate validation set.
Train valid test split
Trajectory length
Force varying across
Dataset
Domain
Viscosity
samples time
TorusLi
1000
20
[0,
10
TorusKochkov
32/ 4/4
34
[x.
= 10
Torus Vis
1000 200 200
20
00, 1]
e [10-5 10
Torus VisForce
1000 / 200 200
20
[0,1]
E [10-5,10 1-1,,01 1011)
V
Table A.2: An overview of the three PDE datasets on irregular geometries. generated by Li al. 2022).
These datasets were
Dataset
Train
Valid
Test
Governing equation
Problem dimension
Elasticity
1000
200
200
Equation of solid body
point cloud on 2D unit cell
Airfoil
1000
200
200
Euler's equation
2D structured mesh
Plasticity
827
80
80
Equation of solid body
2D structured mesh over 1D time
10a
Paramater count
Purtimm unit
10-1
DNS (Cank-NNcolion 2nd-arder)
20'-
FND (proposed el al, 12001al)
(with tricks)
10"3_
F-FNO (without weight sha ring)
F-FND (with weight sharing)
108
10
1.0 Number of layers
12
10
20
Normalized MSE (56)
(a)
(b)
Figure A.I: The resource usage of four model variants, in terms of (a) the parameter count and (b) inference time (the time it takes to run one second of simulation). Error bars, when applicable, show the min and max values over three trials. (b), we move along line, we increase the number of layers. We observe that only our model variants (F-FNO) have the desired slope, that is, as we use more resources (increasing the inference time), we obtain better predictions.
12

---

Published as a conference paper at ICLR 2023
6X
F-TNO (2561256)
0.8
spectram
DAS (128x128)
Voriety conelaion
DNS (250x258)
DNS (512x512)
11024x1024)
DNS (2048X204B)
0.2
F FNO [32x82)
FNO (64x64)
FMF
128x1281
Isuper resolution 256X236)
10
4. 6x
10'
2x 3x100'
Simulatio
(a) Zhro-shot super-resolution performance F-FNO.
(h) Energy spectra F-FNO and DNS on various grid sizes. The spectra are computed by averaging the ki-
We train the model on 32x32 and 64x64 grids of
TorusKochkov, and evaluate on the larger 128x128 and
netic energy for each wavenumber between / 12 and
256x256 grids. We observe correlation with the ground truths on unseon grid sizes.
degradation the
= 34. when the predictions from all methods have decorrelated with the ground truths.
Figure A.2: Performance of F-FNO on zero-shot superresolution and its ability to capture the energy spectrum DNS on TorusKochkov.
Zero-shot super-resolution In Fig. A.2al we train the F-FNO once on 32x32 and 64x64 grids from TorusKochkov, and then perform inference and evaluation on 28xx128 and 256x256 grids. This extends the super-resolution setting presented by Li et al. (2021a) they only worked on simple PDEs such as the ID Burger's equation and the 2D Darcy flow. We find that although the F-FNO can do zero-shot super-resolution unlike traditional CNN which by design cannot even accept inputs of variable size - its performance does degrade on grid sizes not seen during training. This is seen by the lower vorticity correlation of the super-resolution F-FNO settings Fig. A.2a We posit that the super-resolution performance could be improved by training a variety of grid sizes (c.g., by downsampling each training example random size). We leave such exploration for future work.
Capturing the energy spectrum In addition to having high vorticity correlation, good model should also produce predictions with an energy spectrum similar to the most accurate numerical methods. Given the Fourier transform of a velocity field ú (u), wavenumber k; the corresponding kinetic energy as Fi(k) = 2 spectrum of both the F-FNO and DNS at different resolutions. These multiple DNS resolutions are included both as reference solution in the case of DNS 2048x2048, and to demonstrate that increasing the resolution of DNS further is not likely substantially change the energy spectrum.
compute, for each 1 | 12 Fig. A.2b shows the energy
We observe that compared to DNS on 2048x2048, the F-FNO trained on 64x64 grids produces an energy spectrum that has substantially lower energy at high wavenumbers. This is expected at this spatial resolution we only select the top 16 modes in each Fourier layer. Even so, the F-FNO can still capture the long term trend much better than running DNS on a grid four times its size (see (sceFiig. 4ap. As we select more Fourier modes on bigger grids (top 32 modes on 128x1 28 grids and top 64 modes on 256x256 grids), the energy spectrum produced converges towards that of the reference solution (DNS on 2048x2048). This gives some indication that the F-FNO is able to accurately predict both high and low frequency details.
Effect of using cosine transforms As an alternative the Fourier transform, Poli et al. (2022) proposed using the cosine transform, which has the advantage of being real-valued, thus halving the number of parameters. Let the Factorized Cosine Neural Operator (F-CNO) be the operator where the Fourier transform is replaced with the cosine transform. In Fig. A.3) we observe that on Airfoil, the F-CNO outperforms the F-FNO especially at deeper layers. On Plasticity, the F-CNO performs comparably to the F-FNO on the same depth, while using fewer parameters. We have not had much success in training the F-CNO on torus datasets such as Kochkoo. We leave the investigation of how stable the cosine transform is on different domains to future work.
13

---

Published as a conference paper at ICLR 2023
Table A.3: Detailed performance TorusLi. These results used to generate Fig. 3 in the main paper. We run three trials for each experiment, each with a different random seed. We report the mean N-MSE from the three trials, along with the min and value. A dash indicates that the data is not available.
No. of layers
No. of parameters
N-MSE (%)
Training time (h)
Mean
Min
Max
ResNet Li et al.
2021a
266,641
27.53
TF-Net et.
2021a
7,451,724
22.68
U-Net
7,451,724
19.82
FNO
(L. al.
2021a)
414,517
15.56
926,357
13.80
13.75
1,849,637
15.45
15.40
15.50
FNO (reproduced)
12
2,772,917
16.86
16.55
17.27
16
3,696,197
17.59
17.41
17.87
20
4,619,477
18.44
17.91
19.24
does not converge
926,357
12.82
12.50
13.13
1,849,637
11.05
10.98
11.10
FNO-TF (FNO with teacher forcing)
12
2,772,917
11.83
11.74
11.91
5
16
3,696,197
14.10
13.55
14.80
6
20
4,619,477
17.17
15.40
19.64
does not converge
926,177
10.94
10.47
11.35
1,849,457
8.48
9.67
FNO-M (FNO-TF with Markov assumption)
12
2,772,737
9.15
9.02
9.25
16
3,696,017
10.39
9.61
10.81
20
4,619,297
11.73
11.50
12.12
4
does not converge
4
926,177
10.37
10.08
10.69
1,849,457
8.36
8.22
8.46
FNO-R (FNO-M with improved residuals)
12
2,772,737
7,88
7.78
7.95
16
3,696,017
7.90
7.48
8.25
20
4,619,297
7.38
7.34
7.45
5
24
5,542,577
7.33
7.31
7.36
4
16,919,746
5.05
5.02
5.06
33,830,594
4.16
4.11
4.19
12
50,741,442
3.90
3.85
3.93
FNO+ (FNO-R with bags of tricks)
16
67,652,290
3.82
3.77
3.84
20 84,563,138
3.72
3.65
3.79
5
24
101,473,986
3.73
3.70
3.76
4
1,191,106
5.05
5.00
5.09
2,373,314
3.22
3.17
3.29
12
3,555,522
2.75
2.70
2.77
F-FNO (FNO+ with Fourier factorization)
16
4,737,730
2.58
2.57
2.59
20
5,919,938
2.39
2.37
2.42
5
24
7,102,146
2.41
2.37
2.47
4
404,674
5.74
5.69
5,79
538,306
3.43
3.42
3.44
671,938
2.87
2.84
2.90
F-FNO-WS (F-FNO with weight sharing)
805,570
2.56
2.54
2.57
20
939,202
2.42
2.38
2.45
5
24
1,072,834
2.37
2.31
2.45
14

---

Published as a conference paper at ICLR 2023
Table A.4: Detailed performance Airfoil. These results are more detailed version of Table in the main paper. We run three trials for each experiment, each with different random seed. We report the mean N-MSE from the three trials, along with the min and max value.
No. of layers
No. of parameters
N-MSE (%)
Training time (h)
Mean
Min
Max
2,368,033
1.87
1.40
2.27
geo-ENO (reproduced)
4,731,553
1.35
1.02
2.00
12
7,095,073
4.11
0.92
10.29
1,715,458
0.79
0.76
0.82
8
3,421,954
0.64
0.63
0.65
4
F-FNO
12
5,128,450
0.62
0.59
0.67
16
6,834,946
0.61
0.59
0.62
8,541,442
0.57
0.56
0.58
24
10,247,938
0.58
0.56
0.64
535,810
0.98
0.90
1.03
0.4
8
669,442
0.72
0.70
0.75
0.7
E-FNO-WS (F-FNO with weight sharing)
12
803.074
0.68
0.66
0.70
16
936,706
0.67
0.63
0.70
1,070,338
0.64
0.63
0.66
24
1,203,970
0.66
0.60
0.70
Table A.5: Detailed performance on Elasticity. These results are more detailed version of Table 2 in the main paper. We run three trials for each experiment, each with a different random seed. We report the mean N-MSE from the three trials, along with the min and max value. Note that for given layer, our F-FNO (whether with weight sharing or without) has slightly more parameters than the geo-FNO. This is due the F-FNO using bigger hidden size H. We find that the geo-FNO, increasing its hidden size does not necessarily translate to a better performance.
No. of layers
No. of parameters
N-MSE (%)
Training time (h)
Mean
Min
Max
4
1,546,403
2.51
2.43
2.59
0.4
geo-FNO (reproduced)
8
2,730,659
3.30
2.32
5.10
0.5
12
3,914,915
16.76
16.17
17.72
0.7
4
3,205,763
3.16
2.23
4.98
8
4,338,051
2.05
2.04
2.06
5,470,339
1.93
1.98
F-FNO
6,602,627
1.86
1.83
1.88
20
7,734.915
1.84
1.82
1.86
2
24
8,867,203
1.74
1.70
1.78
4
2,681,775
3.55
2.36
5.84
0.2
8
2,765,187
2.23
2.18
2.29
0.2
2,848,899
2.12
2.10
2.16
0.3
F-FNO-WS (F-FNO with weight sharing)
2,932,611
2.08
2.06
2.10
3,016,323
2.04
1.99
2.07
0.5
24
3,100,035
1.97
1.94
2.01
0.6
15

---

Published as a conference paper at ICLR 2023
Table A.6: Detailed performance on Plasticity. These results more detailed version of Table 2 in the main paper. We run three trials for each experiment, each with a different random seed. We report the mean N-MSE from the three trials, along with the min and max value.
No. of layers
No. of parameters
N-MSE (%)
Training time (h)
Mean
Min
Max
18,883,492
0.74
0.73
0.75
geo-FNO (reproduced)
37.762.084
0.57
0.55
0.63
12
56,640,676
0.45
0.41
0.49
1,846,920
0.48
0.47
0.51
3,684,488
0.32
0.31
0.34
8
12
5,522,056
0.25
0.24
0.26
12
F-FNO
16
7,359.624
0.22
0.21
0.22
16
20
9,197.192
0.20
0.18
0.22
20
24
11,034,760
0.18
0.17
0.18
24
568.968
0.58
0.57
0.60
702,600
0.50
0.46
0.52
12
836,232
0.44
0.42
0.48
12
E-FNO-WS (F-FNO with weight sharing)
16
969.864
0.40
0.36
0.44
16
20
1,103,496
0.34
0.31
0.37
19
24
1,237,128
0.30
0.28
0.35
21
Table A.7: The F-FNO flexible in input representation. We find that different datasets benefit from having different features. Shown here is the optimal input combination for each dataset on the torus.
Coordinates
Vorticity
Velocity
Viscosity
Forcing
Dataset
TorusLi Kochsko Torus Vis Torus VisForce
0.80
F-ON
F-C00
6.75
box (6)
0.70
Tea box
0.65
0.30
0.25
2M
6M
BM
13M
4M
5M1
Paramater count
Paramaterr count
(a) Airfoil
(b) Plasticity
Figure A.3: Effect of the cosine transform on Airfoil and Plasticity. We plot the test loss (y-axis) against the model parameter count (x-axis). Error bars show the min-max values from three trials. As we move long each line, we make the network deeper, which increases the number of parameters. On Airfoil (a), the F-CNO outperforms the F-FNO deeper layers. On Plasticity (b), the performance between the two is mostly similar for the same depth. Since cosine transforms are real-valued, the F-CNO requires only half many parameters as the F-FNO.
16

---

Published as a conference paper at ICLR 2023
time 0.2805
time 3.085
time 6.171
10.0
SNG =
2002x8007
7.5
5.0
= model
2.5
F-FNO
0.0
128x122
-255
-5.0
POOO 1aB SNN 128x128
-7.55
10.0
2
6
4
6
2
4
Figure A.4: Similar to Kochkov et al. 2021), visualize how the correlation with the ground truths varies between different models. The heatmaps represent the surface torus mapped onto a 2D grid, with color representing the vorticity (the spinning motion) of the fluid. We observe that the vorticity fields predicted by the F-FNO trained on 128x 28 grids (middle row) correlates with the ground truths (top row) for longer than if This is especially evident after seconds of simulation time (compare the green boxes). In other words, for the same desired accuracy, the F-FNO requires smaller grid input than numerical solver. This observation is also backed up by Fig. 4a
DNS the
spatial resolution (bottom row).
17


Reviews:
[{"review_id": "ro8Leumpc7O", "reviewer": "Reviewer_KRSZ", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "In this work, the authors proposed a novel neural operator architecture that factorizes the convolution on Fourier space into separate dimensions. Consequentially, the F-FNO model can scale up to a higher number of layers and achieve smaller errors. The paper has a comprehensive numerical study on multiple types of partial differential equations, considering chaotic systems and complex geometries. It also has a careful comparison with numerical solvers on the trade off between speed and accuracy.", "summary_of_the_review": "Overall I find this paper interesting. It has a concrete contribution to the community. I recommend acceptance. ", "strength_and_weaknesses": "Strength:\n- The work shows significant improvement on previous methods.\n- The paper has a comprehensive study on many pde problems.\n- It has a careful cost-accuracy study with the numerical solver. I especially love figure 4.\n\nweak:\n- The author improve FNO++ with many tricks. I assume F-FNO is also equipped with these tricks. It could be better to clearly list what these tricks are, and to what aspect they contribute to the overall improvement.\n- It would be better to provide some intuition why the factorized structure help F-FNO scale with more layers.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "nvR_jarSqTD", "reviewer": "Reviewer_Uuaf", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "In this paper, the author proposee Factorized Fourier Neural Operator (F-FNO), a learning-based approach for simulating partial differential equations (PDEs). This work is seen as an improvement of the work by [Li2021] in which the Fourier Neural Operator was introduced. In this work, the author aims at improving the stability of the algorithm under complex geometries and noisy data by introducing more regularizations in the neural operator. The proposed F-FNO method is able to reduce the size of parameter by an order of magnitude while improving the performance over FNO significantly.\n\n\n[Li2021] Zongyi Li, Nikola B. Kovachki, K. Azizzadenesheli, Burigede Liu, K. Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations, 2021a.\n", "summary_of_the_review": "This paper has proposed an important update over existing Fourier Neural Operator and has brought it closer to practical use. With additional factorization, the F-FNO framework reduces the parameter size while attaining a great improvement in stability and performance when the layer of the neural network goes deep. This justifies my decision to accept this paper.   ", "strength_and_weaknesses": "This paper has several aspects of strength:\n\n1. Compared to [Li2021], the F-FNO framework has several changes: 1) a residual connection after the feedforward connection; 2) a spatial factorization; and 3) during the training, instead of unrolling the model, the F-FNO use teacher forcing and make online update based on Markov assumption;\n\n2. The performance improvement of F-FNO over FNO is significant. Under several synthesis dataset  and using deep architectures, the reduction on MSE and time until correlation are both significant. \n\n3. The introduction of forcing function in the F-FNO allows the model to take additional contextual information that helps to solve the equation. \n\nThis paper has several aspects for improvement:\n\n1. While claiming to improve over the state-of-the-art algorithm, the experiments in the main paper only compare with FNO while the comparison with other methods are left in appendix. This is a missed opportunities since it is important to provide a comprehensive review on the related methods in both the parameter size, the MSE and time until correlation, as in [Li2021].  It is suggested to put some of the comprehensive study in the early main page.\n\n2. The spatial factorization is a significant simplification for the model while it may over introduce additional bias since it implicitly assumes that the solution also factorizes spatially. Does it make sense physically in the scenario when the Navier-Stokes equations are applied ? What is the tradeoff for this factorization ? \n\n3. As a PDE solver in the spectral domain, it is expected to discuss some tradeoff on the choice of spatial and spectral resolution (i.e. the design of the wavelength number vs. the dimension of input). It is known from Sampling Theorem that this choice of sampling rate would affect the performance of the system when the complex geometries require more detailed information to be kept after the discretization. This practical issue seems not discussed but it is interesting to have some thoughts on that.   \n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "SCz1oBjRPE", "reviewer": "Reviewer_jqw2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper presents the factorized Fourier neural operator (F-FNO), which brings a set of techniques on FNO. The first technique is the separable Fourier representation to improve model stability and performance. Second, the improved residual connections. Third, the training strategies are carefully designed. Impressive results are found on a variety of benchmarks.", "summary_of_the_review": "Incremental contributions; Clearly below the acceptance bar; results are not convincing.\n\n——————————————————————————————————————————————————————\nAfter discussing with other reviewers via the Zoom meeting, some of my questions are addressed, and it would be better if the authors clarify the mentioned points in the next revision. Also, some senior reviewers with great expertise in this domain insist that this paper has a clear empirical advantage over the baselines. Therefore, I agree to accept this paper right now. I don't think a simple method should not be accepted, whereas the AC also thinks so. I think the authors should clarify their central contribution (eq.8) in this case and not hide their contribution behind a set of tricks. I think a revision would be beneficial to the community and also to the authors.", "strength_and_weaknesses": "# Strengths\n1. The results seem to be impressive. \n2. This paper clearly presented their method.\n3. Some of the visualizations might inspire other researchers in this domain.\n\n# Weaknesses\n\n## 1. The proposed contributions are incremental and are not even well-defined.\nSupporting details:\n\n(1) The residual connections are not novel, and a naive application is not enough for ICLR.\n\n(2) In the abstract, the authors say, \"carefully designed training strategies.\" This claim is too vague.\n\n(3) In the introduction, the authors claim, \"In our exploration of various regularization techniques, we find that it is possible to reduce the parameter count by up to an order of magnitude while still outperforming state of the art by a significant margin.\" This is too incremental. Applying a kind of regularization cannot be a major contribution. Also, how and why the regularization work remains unclear.\n\n(4) The core of the introduction and the method are too few. The major body of the methods is about the formulation of the problem and the baseline methods.\n\n(5) In the introduction, the authors claim three key contributions. The first is about the factorized Fourier neural operator. However, why is it called factorized? I think the only change is the residual connection. This should not be called the factorized Fourier neural operator, and this contribution is still weak.\n\n(6) In the introduction, the authors list two of the experimental findings as their contributions. However, the experimental findings are not that important because the technical contributions are weak, and the experiments are not convincing.\n\n## 2. The experimental results are not convincing. In general, I do not believe the baselines and the improved models are really working well.\n\n(1)  All the error numbers are too large. F-FNO achieves 3.16 errors in elasticity. However, the original FNO paper reports numbers at ~0.01. I understand the settings might be different. If this is the case, why do not the authors put all the methods under the same setting?\n\n(2) Can the authors provide a visualization comparing the GT with the prediction? Why does the modification work under this setting?\n\n(3) Is the baseline strong? Why don't the authors conduct experiments on the official settings of FNO?\n\n(4) In figure captions, the briefing settings should be presented instead of saying, \"our models work well, ... \". It is always better for the readers to discover the improvements from the figures rather than reading general over-stated claims.\n\n(5) In Figure 2, how to understand the figure? What are the inputs to each model? What are the outputs? Why are there big holes in the figure? What does the color mean? How to understand the differences between different modalities? Can the baselines achieve this? If the baselines cannot support varying geometries, why? What are the corresponding gt? What do geometries mean here? Are the geometries important in this domain?\n\n(6) The error bars should be shown in each figure.\n\n(7) It seems that the changes are not significant. The improvements should not be so significant.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "DLwG8U4Q1s7", "reviewer": "Reviewer_WWi7", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper uses the Fourier transform to learn neural operators that can handle long-range spatial dependencies. By factorizing the transform, using better residual connections, and improving the training setup, the proposed F-FNO outperforms the state of the art on PDEs on a variety of geometries and domains. ", "summary_of_the_review": "I would suggest to accept the current paper. ", "strength_and_weaknesses": "Pros:  The proposed method outperforms the state of art on PDEs with geometies and domains. \nCons: The literature reviews on the classical and recent PDE solvers are not very complete. I would like to see some comparison with the classical numerical method, such as the vortex model in approximating 2D Navier-Stokes/Euler with FFO. \n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "ABTcEA0iLDJ", "reviewer": "Reviewer_n5hB", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper introduces Factorized-FNO, where they consider separable Fourier representation (by taking the fourier transform of each dimension separately and independently) the authors achieve more “stable” models, whose performance increases as the networks are made deeper (something that does not happen with baseline FNOs), for 2D navier stokes equations.\n\nThe authors also improve the performance of original FNO baseline by introducing techniques like teacher forcing, a markov property. \n\nThey also test F-FNO on irregular geometries.", "summary_of_the_review": "The authors have tested F-FNO on a single family of PDEs (and only in 2D) and it is hard to judge if it will always improve performance. The paper also lacks comparison with baselines (like FNO) for their abilities to hand irregular geometries. ", "strength_and_weaknesses": "[Strenghts]\n\n- The main idea introduced by the authors, of factorizing the input dimension and treating the fourier transform of each dimension separately seems to work well for the various 2D Navier Stokes equations considered.\n- The method seems to get better with size of depth, and requires less parameters to reach similar performance as normal FNO.\n\n[Weakness]\n\n- The main idea (of treating each dimension independently) is only testing for ONE family of PDEs and that too in just two dimensions. It is hard to know if factorizing will *always* perform better. For example, one scenario where it might not is for PDEs that have cross terms involving derivatives w.r.t different dimensions.\n- Most of the contribution made by the authors seem to be utilizing training techniques (like teacher forcing) that have are shown to give better performance in other domains and test them for FNOs. This does not seem to be a very novel.\n- The authors test F-FNO on irregular geometries, is it established that normal FNOs definitely do not do well of these irregular geometries considered by the authors. That comparison seems to be lacking in the paper.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster