Factorized Fourier Neural Operators

Alasdair Tran1 , Alexander Mathews1 , Lexing Xie1 , Cheng Soon Ong1,21Australian National University2
Data61, CSIRO{alasdair.tran,alex.mathews,lexing.xie,chengsoon.ong}@anu.edu.au

Abstract
The Fourier Neural Operator (FNO) is a learning-based method for efficiently
simulating partial differential equations. We propose the Factorized Fourier Neural
Operator (F-FNO) that allows much better generalization with deeper networks.
With a careful combination of the Fourier factorization, a shared kernel integral
operator across all layers, the Markov property, and residual connections, F-FNOsachieve a six-fold reduction in error on the most turbulent setting of the NavierStokes benchmark dataset. We show that our model maintains an error rate of 2%
while still running an order of magnitude faster than a numerical solver, even when
the problem setting is extended to include additional contexts such as viscosity andtime-varying forces. This enables the same pretrained neural network to model
vastly different conditions.

1

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 solutions requires computationally expensive numerical solvers, sometimes
consuming millions of core hours and terabytes of storage [2].
Deep learning architectures have come full circle, from initially replacing frequency-domain representations with learned features [1] to a resurgence of Fourier operators due to the O(n log n) timecomplexity of the fast Fourier transform (FFT) and its ability to represent long-range dependencies.
Two notable examples are FNet [5] for representing semantic relationships in text classificationand the Fourier Neural Operator (FNO) [7, 8] for flow simulation. FNOs, in particular, have beenshown to approximate any continuous operator [4]. However, the original FNO architecture has nottaken advantage of the separable nature of Fourier transforms in space, more parsimonious neuralrepresentations, and the neural network’s versatility to solve more general settings. We address theselimitations through the following three contributions:
1. We propose the Factorized Fourier Neural Operator (F-FNO) (Fig. 1), which uses a separable
Fourier representation and a shared kernel integral operator across all layers to reduce theparameter count from 100M to 1M and enable the network to grow to 24 layers on a single GPU(Fig. 3a).
2. For the Navier-Stokes equations on the torus, our approach outperforms the state of the art,
reducing the normalized mean squared error from 15.56% to 2.29% and maintaining an error
reduction of more than 80% even when predicting 10 steps ahead. (Fig. 2).
3. F-FNOs can take additional contexts as input (Fig. 4), allowing us to solve more general NavierStokes problems with different parameters without the need to retrain. On newly created datawith a range of viscosities and time-varying forces, F-FNOs maintain an error of 2% (Table 1).
Code and datasets are released on GitHub.1
1

https://github.com/alasdairtran/fourierflow

Fourth Workshop on Machine Learning and the Physical Sciences (NeurIPS 2021).

Figure 1: The architecture of the Factorized Fourier Neural Operator (F-FNO). The iterative process(Eq. (1)) is shown at the top, with a zoomed-in operator layer (Eq. (4)) at the bottom.

2

The Factorized Fourier Neural Operator

Preliminaries We are interested in learning neural operators between infinite-dimensional function
spaces. We consider the 2D Navier-Stokes equations as the guiding example, in which the function
space Ω(D; R) consists of vorticity fields ωt ∈ Ω over a bounded spatial domain, i.e. ωt : D → R
with D ⊂ R2 . The operator G : Ω → Ω would then learn to map this field ωt at some time step t to
the next time step t + 1. Motivated by the kernel formulation of the solution to linear PDEs using
Green’s functions, Li et al. [6] propose an iterative approach to map ωt to ωt+1 ,ωt+1 = G(ωt ) = (Q ◦ L(L) ◦ L(L−1) ◦ · · · ◦ L(1) ◦ P)(ωt ),(1)

where ◦ indicates function composition, L is the number of layers/iterations, P is the lifting operator(0)
that maps the input function ωt into the first latent representation zt , L(ℓ) is the ℓ’th non-linear(L)
operator layer, and Q is the projection operator that maps the last latent representation zt to the
output function ωt+1 . Fig. 1 (top) contains a schematic diagram of this iterative process. Originally,
Li et al. [7] formulate each operator layer as




(ℓ)
(ℓ)
(ℓ)
L(ℓ) zt= σ W (ℓ) zt + b(ℓ) + K(ℓ) (zt ) ,(2)(ℓ)

where σ : R → R is a point-wise non-linear activation function, W (ℓ) zt +b(ℓ) is an affine point-wise
map in the spatial domain, and K(ℓ) is a kernel integral operator using the Fourier transform,




(ℓ)
K(ℓ) zt= IFFT R(ℓ) · FFT(zt )(3)Note that the frequency-domain weight matrices {R(ℓ) | ℓ ∈ {1, 2, . . . , L}} take up most of themodel size, requiring O(LH 2 M 2 ) parameters, where H is the hidden size and M is the number of
top modes being kept.
F-FNO We consider a more general setting in which the operator G maps the vorticity field ωt alongwith some contexts such as the viscosity ν and the forcing function ft , i.e., ωt+1 = G(ωt | ν, ft ).
Furthermore, we design the following operator layer:


h
i
(ℓ)
(ℓ)
(ℓ)
(ℓ)
(ℓ) 
(ℓ) 
(ℓ)
L(ℓ) zt= zt + σ W2 σ W1 K zt + b1 + b2
(4)
Note that we apply the residual connection after the non-linearity to allow an easier flow of gradients.
Our kernel integral operator is shared across all layers. We also factorize the Fourier transforms over
the two spatial dimensions, giving us



 

(ℓ)
(ℓ)
(ℓ)
K zt
= IFFT Rx · FFTx (zt ) + Ry · FFTy (zt )(5)
Our approach uses two weight matrices, Rx and Ry , that together have only O(H 2 M ) parameters,
irrespective of the number of layers. Fig. 1 (bottom) gives the overview of an F-FNO operator layer.
We will show in the next section that the factorization and weight sharing in K, along with the residual
connection in L and other deep learning techniques, allow F-FNOs to vastly outperform the state of
the art.
2

Normalized MSE (%)

6

20

20

5

15

4

10

10

3

5
0

4

8 12 16 20 24
Number of Layers

FNO (proposed by Li et al. [2021a])
FNO-TF (with teacher forcing)
FNO-M (with Markov assumption)
FNO++ (with a bag of tricks)
F-FNO-NW (without weight sharing)
F-FNO (our full model)

15
5
0

4

8 12 16 20 24
Number of Layers(a)

0 1 2 3 4 5 6 7 8 9Inference Step

(b)

(c)

Figure 2: Performance (lower is better) of six model variants on TorusConst, with error bars
showing the min and max values over three trials. In (a), we improve the original FNO [7] by using
teacher forcing, the Markov property, among other techniques, resulting in FNO++. In (b), we show
the effect of Fourier factorization and weight sharing. In (c), we show (for a four-layer network) theerror as we infer more time steps into the future.

Training Time (h)Parameter Count

107
106
0 4 8 12 16 20 24
Number of Layers

Inference Time (s)

8

108

6
4
2
0

102

100

5 10 15 20
Normalized MSE (%)(a)

Crank Nicolson numerical simulator
FNO (proposed by Li et al. [2021a])
FNO++ (with a bag of tricks)
F-FNO-NW (without weight sharing)
F-FNO (our full model)

101

0

5 10 15 20
Normalized MSE (%)

(b)

(c)

Figure 3: The resource usage of four model variants, in terms of the parameter count (a), training
time (b), and inference time (c). Error bars show the min and max values over three trials.

3

Experiments on Navier-Stokes Equations

Datasets We test the F-FNO on three fluid dynamics datasets. TorusConst [7] is used to benchmark our model against the state of the art. TorusV and TorusVF are new datasets we generate to
test the generalization of F-FNOs across Navier-Stokes tasks with different viscosities and forcing
functions. The training split of each dataset consists of 1000 simulations of the 2D Navier-Stokes
equations on a unit torus, and each simulation contains 20 steps of a 64 × 64 grid of the vorticity
field. The viscosity is fixed at 10−5 (Re = 2000) in TorusConst, while it varies between 10−4 and10−5 in TorusV and TorusVF. The forcing function on the torus is defined asf (t, x, y) = 0.1

2 X
1 X
1 h
X

i
αpij sin 2πp(ix + jy) + δt + βpij cos 2πp(ix + jy) + δt

(6)

p=1 i=0 j=0

where the amplitudes αpij and βpij are sampled from the uniform distribution in TorusV and
TorusVF. Furthermore, δ is set to 0 in TorusV, making the forcing function constant across time;
while it is set to 0.2 in TorusVF, giving us a time-varying force. Finally in TorusConst, the forcing
function is fixed at f (x, y) = 0.1[sin(2π(x + y)) + cos(2π(x + y))] across all samples and time
steps.
Training details For the experiments involving the original FNO, FNO-TF (with teaching forcing),
and FNO-M (with Markov assumption), we use the same training procedure as Li et al. [7]. For allother experiments, we train for 100,000 steps, warming up the learning rate to 2.5 × 10−3 for the
first 500 steps and then decaying it using the cosine function [9]. We use ReLU as our non-linear
activation function, clip the gradient value at 0.1, and use the Adam optimizer [3] with β1 = 0.9,
β2 = 0.999, ϵ = 10−8 . The weight decay factor is set to 10−4 and is decoupled from the learning
rate [10]. Models are implemented in PyTorch [11] and trained on a single Titan V GPU. Finally, thePB
i −ω∥2
loss function is the normalized mean squared error, defined as N-MSE = B1 i=1 ∥ω̂∥ω∥, where2∥·∥2 is the 2-norm, B is the batch size, and ω̂ is the prediction of the ground truth ω.

3

Table 1: Performance of F-FNO on TorusV and
TorusVF, with ablation studies to test the contribution of each context. Each result is accompanied by the standard deviation from three trials.
N-MSE (%)

Figure 4: F-FNO solves the Navier-Stokes
equations with a range of parameters without
needing to retrain, by taking viscosity ν andforcing function ft as input.

TorusV
Without context
With force
With force and viscosity

43.8 ± 0.9
4.48 ± 0.09
2.06 ± 0.02

TorusVF
Without context
With force
With force and viscosity

33.0 ± 0.4
4.2 ± 0.1
2.01 ± 0.05

Results against state of the art The performance on TorusConst is plotted in Figs. 2 and 3, withmore details in Appendix Table 3. First, note that in FNO, Li et al. [7] use the previous 10 steps as
input and predict the next 10 steps incrementally (by using each predicted value as the input to thenext step). We find that the teacher forcing strategy (FNO-TF, orange line), in which we always use
the ground truth as input, helps in shallow networks. Enforcing the Markov condition (FNO-M, dottedgreen line), where only one step of history is used, further improves the performance. Including two
or more steps of history makes no difference to the results.
However, all of these three models (FNO, FNO-TF, and FNO-M) do not scale with network depth,
even failing to converge at 24 layers. The real improvement in deeper networks does not come until
we move the residual connection outside the non-linear activation to ease the gradient flow, andcouple it with other deep learning techniques such as normalizing the inputs, adding Gaussian noise,
and using a cosine learning rate scheduler with warmup (FNO++, cyan line). Furthermore, if we
use Fourier factorization (F-FNO-NW, yellow dashed line), the error drops by an additional 64%,
from 3.65% to 2.33%, at 24 layers (Fig. 2b), while the parameter count is reduced by an order ofmagnitude (Fig. 3a). Sharing the weights in the Fourier domain (F-FNO, red line) helps the most withthe deepest network, pushing the error down to 2.29% and reducing the parameter count by anotherorder of magnitude to 1M. Finally, even though our models do take longer to do inference, the deepest
network is still an order of magnitude faster than the Crank-Nicholson numerical simulator (Fig. 3c).
Results on more general settings The same pretrained F-FNO can handle Navier-Stokes equations
with different viscosities (in TorusV) and time-varying forcing functions (in TorusVF). Fig. 4 shows
an example where our model takes these two contexts as additional input channels. On both datasets,
the F-FNO is able to maintain an error of 2% (Table 1). We also conduct ablation studies to show theimportance of the contexts. In particular, if we remove the viscosity information, the error doubles. If
we further remove the forcing function from the input, the error increases by an order of magnitude,
showing the substantial influence the force has on the vorticity field.

4Conclusion

The Fourier transform is a powerful tool to learn neural operators that can handle long-range
dependencies efficiently. By factorizing the transform, sharing the kernel integral operator acrosslayers, enforcing the Markov condition, and carefully constructing the training process, our proposed
F-FNO outperforms the state of the art by 85% and handles a range of Navier-Stokes parameters
without retraining. For future work, we plan to compare our method against baseline spectral solvers
at different resolutions to illustrate the full Pareto frontier of accuracy vs. inference speed. We arealso interested in extending the Fourier layers to irregular structures such as graphs and meshes.
Finally we would like to examine equilibrium properties of generalized Fourier layers in the limit of
an infinite number of layers.

4

Broader Impact
This work focuses on fundamental technology that is typically used as a component in simulation andcontrol systems. Such larger systems could bring significant social benefits (e.g., weather forecast,
airplane design) or harm (e.g., missiles). The advantage of making this work available outweighs thepotential harm it could cause. Since this work does not pose an immediate risk, making the work
publicly available could help to ensure that any negative effects are anticipated and mitigated inadvance since keeping a technology secret is not a valid or effective preventative measure of harm.
Acknowledgments
This research is supported in part by the Australian Research Council Project DP180101985 and
AOARD project 20IOA064. We thank NVIDIA for providing us with Titan V GPUs for experimentation.

References[1] Dario Amodei, Sundaram Ananthanarayanan, Rishita Anubhai, Jingliang Bai, Eric Battenberg, Carl
Case, Jared Casper, Bryan Catanzaro, Qiang Cheng, Guoliang Chen, Jie Chen, Jingdong Chen, Zhijie
Chen, Mike Chrzanowski, Adam Coates, Greg Diamos, Ke Ding, Niandong Du, Erich Elsen, Jesse
Engel, Weiwei Fang, Linxi Fan, Christopher Fougner, Liang Gao, Caixia Gong, Awni Hannun, TonyHan, Lappi Johannes, Bing Jiang, Cai Ju, Billy Jun, Patrick LeGresley, Libby Lin, Junjie Liu, Yang
Liu, Weigao Li, Xiangang Li, Dongpeng Ma, Sharan Narang, Andrew Ng, Sherjil Ozair, Yiping Peng,
Ryan Prenger, Sheng Qian, Zongfeng Quan, Jonathan Raiman, Vinay Rao, Sanjeev Satheesh, David
Seetapun, Shubho Sengupta, Kavya Srinet, Anuroop Sriram, Haiyuan Tang, Liliang Tang, Chong Wang,
Jidong Wang, Kaifu Wang, Yi Wang, Zhijian Wang, Zhiqian Wang, Shuang Wu, Likai Wei, Bo Xiao,
Wen Xie, Yan Xie, Dani Yogatama, Bin Yuan, Jun Zhan, and Zhenyao Zhu. Deep speech 2 : End-to-end
speech recognition in english and mandarin. In Maria Florina Balcan and Kilian Q. Weinberger, editors,
Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of
Machine Learning Research, pages 173–182, New York, New York, USA, 20–22 Jun 2016. PMLR. URL
https://proceedings.mlr.press/v48/amodei16.html.
[2] S.M. Hosseini, R. Vinuesa, P. Schlatter, A. Hanifi, and D.S. Henningson. Direct numerical simulation of
the flow around a wing section at moderate reynolds number. International Journal of Heat and Fluid
Flow, 61:117–128, 2016. ISSN 0142-727X. doi: https://doi.org/10.1016/j.ijheatfluidflow.2016.02.001.
URL https://www.sciencedirect.com/science/article/pii/S0142727X16300169. SI TSFP9
special issue.
[3] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International
Conference on Learning Representations, 2015.
[4] Nikola B. Kovachki, S. Lanthaler, and Siddhartha Mishra. On universal approximation and error bounds
for fourier neural operators. ArXiv, abs/2107.07562, 2021.
[5] J. Lee-Thorp, Joshua Ainslie, Ilya Eckstein, and Santiago Ontañón. Fnet: Mixing tokens with fouriertransforms. ArXiv, abs/2105.03824, 2021.
[6] Zong-Yi Li, Nikola B. Kovachki, K. Azizzadenesheli, Burigede Liu, K. Bhattacharya, Andrew Stuart, andAnima Anandkumar. Neural operator: Graph kernel network for partial differential equations. ArXiv,
abs/2003.03485, 2020.
[7] Zongyi Li, Nikola B. Kovachki, K. Azizzadenesheli, Burigede Liu, K. Bhattacharya, Andrew Stuart, andAnima Anandkumar. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations, 2021. URL https://openreview.net/forum?id=
c8P9NQVtmnO.
[8] Zongyi Li, Nikola B. Kovachki, K. Azizzadenesheli, Burigede Liu, K. Bhattacharya, Andrew Stuart, andAnima Anandkumar. Markov neural operators for learning chaotic systems. ArXiv, abs/2106.06898, 2021.
[9] Ilya 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 https://openreview.net/forum?id=Skq89Scxx.
[10] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on
Learning Representations, 2019.
[11] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming
Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in PyTorch. In NIPS
Autodiff Workshop, 2017.

5

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope? [Yes](b) Did you describe the limitations of your work? [Yes] See Section 4.
(c) Did you discuss any potential negative societal impacts of your work? [Yes] See theBroader Impact section after Section 4.
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem? [Yes]2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A](b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)? [Yes] See Section 3.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they
were chosen)? [Yes] See Section 3.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [Yes] See Fig. 2 and Fig. 3.
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)? [Yes] See Fig. 3b and Fig. 3c.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] We cited Li et al. [7]for the TorusConst dataset.
(b) Did you mention the license of the assets? [Yes] See the Appendix B.
(c) Did you include any new assets either in the supplemental material or as a URL? [Yes]See the Appendix B.
(d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating? [N/A](e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable? [N/A](b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable? [N/A](c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation? [N/A]6

A

The Navier-Stokes Equations

Our experiments are conducted on simulations of the incompressible 2D Navier-Stokes equations on
the unit torus using the vorticity form:∂ωt (x, y) + vt (x, y) · ∇ωt (x, y) = ν∆ωt (x, y) + ft (x, y)∂t∇ · vt (x, y) = 0
ωt (x, y) = ω0 (x, y)

(7)
(8)
(9)

where x, y ∈ (0, 1) are the positions on the unit torus, t ∈ (0, T ] is the time step, ν is the viscosity,vt (x, y) is the velocity field, ω := ∇ × v is the vorticity field, and ft (x, y) is the forcing function.
Given the forcing function
f (x, y) = 0.1[sin(2π(x + y)) + cos(2π(x + y))],(10)

the Reynolds number is approximately
√

Re ≈

B

0.1
ν(2π)3/2(11)

Code and Datasets

The code, pretrained models, and datasets are licensed under the MIT License. They can be found athttps://github.com/alasdairtran/fourierflow.

C

Further Results

Detailed results, including the normalized mean squared error, along with complexity measures suchas the parameter count and inference time, are shown in Table 3.

7

Dataset
TorusConst [7]TorusV
TorusVF

Size

Viscosity

1,200
1,400
1,400

−5

Forcing function varyingν = 10ν ∈ [10−5 , 10−4 )ν ∈ [10−5 , 10−4 )across samples

across time

✗
✓
✓

✗
✗
✓

Table 2: An overview of the three datasets used in our experiments, including key differences between
them. Observe that our datasets, TorusV and TorusVF, contain simulation data with a more varietyof viscosities and forces than TorusConst. Note also that Li et al. [7] do not include a validation set,
splitting the data into 1,000 training examples and 200 test examples. In our datasets, we generate
200 additional simulations to be used as validation.

No. ofLayersNo. ofParameters

N-MSE
(%)

Training
Time (h)Inference
Time (s)ResNet [7]TF-Net [7]U-Net [7]FNO [7]

4

266,641
7,451,724
7,451,724
414,517

27.53
22.68
19.82
15.56

-

-

FNO (reproduced)

4
8
12
16
20

926,357
1,849,637
2,772,917
3,696,197
4,619,477

13.81 ± 0.04
15.44 ± 0.09
16.84 ± 0.32
17.62 ± 0.20
18.37 ± 0.38

2.1
3.4
4.8
6.0
7.4

0.4
0.9
1.4
1.9
2.4

FNO-TF (with teacher forcing)

4
8
12
16
20

926,357
1,849,637
2,772,917
3,696,197
4,619,477

12.76 ± 0.20
11.03 ± 0.07
11.91 ± 0.10
13.82 ± 0.47
18.17 ± 2.64

2.0
3.3
4.6
6.0
7.3

0.4
0.9
1.4
1.9
2.4

FNO-M (with Markov assumption)

4
8
12
16
20

926,177
1,849,457
2,772,737
3,696,017
4,619,297

10.76 ± 0.14
8.89 ± 0.22
8.88 ± 0.31
9.95 ± 0.41
10.68 ± 0.45

1.5
2.4
3.3
4.2
5.0

0.4
0.9
1.4
1.9
2.4

4
8
12
16
20
24

16,919,746
33,830,594
50,741,442
67,652,290
84,563,138
101,473,986

4.98 ± 0.03
4.11 ± 0.02
3.89 ± 0.02
3.74 ± 0.01
3.66 ± 0.02
3.65 ± 0.04

1.4
2.2
3.1
4.0
4.8
5.7

2.5
4.9
7.2
9.6
11.9
14.3

4
8
12
16
20
24

1,191,106
2,373,314
3,555,522
4,737,730
5,919,938
7,102,146

5.06 ± 0.08
3.16 ± 0.02
2.71 ± 0.01
2.49 ± 0.02
2.40 ± 0.01
2.33 ± 0.01

1.5
2.4
3.6
4.7
5.8
6.8

3.2
6.3
9.4
12.5
15.5
18.6

4
8
12
16
20
24

404,674
538,306
671,938
805,570939,202
1,072,834

5.70 ± 0.09
3.42 ± 0.02
2.86 ± 0.02
2.61 ± 0.01
2.41 ± 0.03
2.29 ± 0.04

1.4
2.4
3.3
4.5
5.7
6.5

3.2
6.3
9.4
12.5
15.6
18.6

FNO++ (with bags of tricks)

F-FNO-NW (without weight sharing)

F-FNO (our full model)

Table 3: Performance on TorusConst. When possible, the N-MSE is accompanied by the standard
deviation from three trials. A hyphen indicates that the data is not available.

8