Fourier Neural Operator

This blog takes about 10 minutes to read. It introduces the Fourier neural operator that solves a family of PDEs from scratch. It the first work that can learn resolution-invariant solution operators on Navier-Stokes equation, achieving state-of-the-art accuracy among all existing deep learning methods and up to 1000x faster than traditional solvers. Also check out the paper, code, article, and project page.

Operator learning

Thinking in continuum gives us an advantage when dealing with PDE. We want to design mesh-indepedent, resolution-invariant operators.

Problems in science and engineering involve solving partial differential equations (PDE) systems. Unfortunately, these PDEs can be very hard. Traditional PDE solver such as finite element methods (FEM) and finite difference methods (FDM) rely on discretizing the space into a very fine mesh. And it can be slow and inefficient.

In the previous post, we introduced the neural operators that use neural networks to learn the solution operators for PDEs. That is, given the initial conditions or the boundary conditions, the neural network directly output the solution, kind of like an image-to-image mapping.

The neural operator is mesh-independent, different from the standard deep learning methods such as CNN. It can be trained on one mesh and evaluated on another. By parameterizing the model in function space, it learns the continuous function instead of discretized vectors.

Operator learning can be taken as an image-to-image problem. The Fourier layer can be viewed as a substitute for the convolution layer.

Framework of Neural Operators

Just like neural networks consist of linear transformations and non-linear activation functions, neural operators consist of linear operators and non-linear activation operators.

Let $v$ be the input vector, $u$ be the output vector. A standard deep neural network can be written in the form:

where $K$ are the linear layer or convolution layer, and $\sigma$ are the activation function such as ReLU.

The neural operator shares a similar framework. It’s just now $v$ and $u$ are functions with different discretizations (say, some inputs are $28 \times 28$, some are $256 \times 256$, and some are in triangular mesh). To deal with functions input, the linear transformation $K$ is formulated as an integral operator. Let $x, y$ be the points in the domain.

The map $K: v_{t} \mapsto v_{t+1}$ is parameterized as

Where $\kappa$ is a kernel function and $W$ is the bias term.

For the Fourier neural operator, we formulate $K$ as a convolution and implement it by Fourier transformation.

Fourier Layer

The real-world images have lots of edges and shapes, so CNN can capture them well with local convolution kernel. On the other hand, the inputs and outputs of PDEs are continuous functions. It is more efficient to represent them in Fourier space and do global convolution.

There are two main motivations to use Fourier transformation. First, it’s fast. A full standard integration of $n$ points has complexity $O(n^2)$, while convolution via Fourier transform is quasilinear. Second, it’s efficient. The inputs and outputs of PDEs are continuous functions. So it’s usually more efficient to represent them in Fourier space.

The convolution in the spatial domain is equivalent to the pointwise multiplication in the Fourier domain. To implement the (global) convolution operator, we first do a Fourier transform, then a linear transform, and an inverse Fourier transform, As shown in the top part of the figure:

The Fourier layer just consists of three steps:

1. Fourier transform $\mathcal{F}$
2. Linear transform on the lower Fourier modes $R$
3. Inverse Fourier transform $\mathcal{F}^{-1}$

We then add the output of the Fourier layer with the bias term $W v$ (a linear transformation) and apply the activation function $\sigma$. Simple as it is.

In practice, it’s usually sufficient to only take the lower frequency modes and truncate out these higher frequency modes. Therefore, we apply the linear transformation on the lower frequency modes and set the higher modes to zeros.

Notice the activation functions shall be applied on the spatial domain. They help to recover the Higher frequency modes and non-periodic boundary which are left out in the Fourier layers. Therefore it’s necessary to the Fourier transform and its inverse at each layer.

Filters in convolution neural networks are usually local. They are good to capture local patterns such as edges and shapes. Fourier filters are global sinusoidal functions. They are better for representing continuous functions.

Higher frequency modes and non-periodic boundary

The Fourier layer on its own loses higher frequency modes and works only with periodic boundary conditions. However, the Fourier neural operator as a whole does not have these limitations (examples shown in the experiments). The encoder-decoder structure helps to recover the higher Fourier modes. And the bias term $W$ helps to recover the non-periodic boundary.

Complexity

The Fourier layer has a quasilinear complexity. Denote the number of points (pixels) $n$ and truncating at $k_{max}$ frequency modes. The multiplication has complexity $O(k_{max}) < O(n)$ . The majority of the computational cost lies in computing the Fourier transform and its inverse. General Fourier transforms have complexity $O(n^2)$, however, since we truncate the series the complexity is in fact $O(n k_{max})$, while the FFT has complexity $O(n \log n)$.

Resolution-invariance

The Fourier layers are discretization-invariant, because they can learn from and evaluate functions which are discretized in an arbitrary way. Since parameters are learned directly in Fourier space, resolving the functions in physical space simply amounts to projecting on the basis of wave functions which are well-defined everywhere on the space. This allows us to transfer among discretization. If implemented with standard FFT, then it will be restricted to uniform mesh, but still resolution-invariant.

Experiments

1. Burgers Equation

The 1-d Burgers’ equation is a non-linear PDE with various applications including modeling the one-dimensional flow of a viscous fluid. It takes the form

with periodic boundary conditions where $u_0 \in L^2_{\text{per}}((0,1);\R)$ is the initial condition and $\nu \in \R_+$ is the viscosity coefficient. We aim to learn the operator mapping the initial condition to the solution at time one, defined by $u_0 \mapsto u(\cdot, 1)$ for any $r > 0$.

2. Darcy Flow

We consider the steady-state of the 2-d Darcy Flow equation on the unit box which is the second order, linear, elliptic PDE

with a Dirichlet boundary where $a \in L^\infty((0,1)^2;\R_+)$ is the diffusion coefficient and $f \in L^2((0,1)^2;\R)$ is the forcing function. This PDE has numerous applications including modeling the pressure of the subsurface flow, the deformation of linearly elastic materials, and the electric potential in conductive materials. We are interested in learning the operator mapping the diffusion coefficient to the solution, defined by $a \mapsto u$. Note that although the PDE is linear, the solution operator is not.

Benchmarks for time-independent problems (Burgers and Darcy):

• NN: a simple point-wise feedforward neural network.
• RBM: the classical Reduced Basis Method (using a POD basis).
• FCN: a the-state-of-the-art neural network architecture based on Fully Convolution Networks.
• PCANN: an operator method using PCA as an autoencoder on both the input and output data and interpolating the latent spaces with a neural network.
• GNO: the original graph neural operator.
• MGNO: the multipole graph neural operator.
• LNO: a neural operator method based on the low-rank decomposition of the kernel.
• FNO: the newly purposed Fourier neural operator.

3. Navier-Stokes Equation

We consider the 2-d Navier-Stokes equation for a viscous, incompressible fluid in vorticity form on the unit torus:

where $u$ is the velocity field, $w = \nabla \times u$ is the vorticity, $w_0$ is the initial vorticity,
$\nu$ is the viscosity coefficient, and $f$ is the forcing function. We are interested in learning the operator mapping the vorticity up to time 10 to the vorticity up to some later time $T > 10$, defined by $w|_{(0,1)^2 \times [0,10]} \mapsto w|_{(0,1)^2 \times (10,T]}$. We experiment with the viscosities $\nu = 1\mathrm{e}{-3}, 1\mathrm{e}{-4}, 1\mathrm{e}{-5}$, decreasing the final time $T$ as the dynamic becomes chaotic.

Benchmarks for time-dependent problems (Navier-Stokes):

• ResNet: 18 layers of 2-d convolution with residual connections.
• U-Net: A popular choice for image-to-image regression tasks consisting of four blocks with 2-d convolutions and deconvolutions.
• TF-Net: A network designed for learning turbulent flows based on a combination of spatial and temporal convolutions.
• FNO-2d: 2-d Fourier neural operator with an RNN structure in time.
• FNO-3d: 3-d Fourier neural operator that directly convolves in space-time.

The FNO-3D has the best performance when there is sufficient data ($\nu=1\mathrm{e}{-3}, N=1000$ and $\nu=1\mathrm{e}{-4}, N=10000$). For the configurations where the amount of data is insufficient ($\nu=1\mathrm{e}{-4}, N=1000$ and $\nu=1\mathrm{e}{-5}, N=1000$), all methods have $>15\%$ error with FNO-2D achieving the lowest. Note that we only present results for spatial resolution $64 \times 64$ since all benchmarks we compare against are designed for this resolution. Increasing it degrades their performance while FNO achieves the same errors.

FNO-2D, U-Net, TF-Net, and ResNet all use 2D-convolution in the spatial domain and recurrently propagate in the time domain (2D+RNN). On the other hand, FNO-3D performs convolution in space-time.

4. Bayesian Inverse Problem

In this experiment, we use a function space Markov chain Monte Carlo (MCMC) method to draw samples from the posterior distribution of the initial vorticity in Navier-Stokes given sparse, noisy observations at time $T=50$. We compare the Fourier neural operator acting as a surrogate model with the traditional solvers used to generate our train-test data (both run on GPU). We generate 25,000 samples from the posterior (with a 5,000 sample burn-in period), requiring 30,000 evaluations of the forward operator.

The top left panel shows the true initial vorticity while the bottom left panel shows the true observed vorticity at $T=50$ with black dots indicating the locations of the observation points placed on a $7 \times 7$ grid. The top middle panel shows the posterior mean of the initial vorticity given the noisy observations estimated with MCMC using the traditional solver, while the top right panel shows the same thing but using FNO as a surrogate model. The bottom middle and right panels show the vorticity at $T=50$ when the respective approximate posterior means are used as initial conditions.

Conclusion

• We propose a neural operator based on Fourier Transformation. It is the first work that learns the resolution-invariant solution operator for the family of Navier-Stokes equation in the turbulent regime, where previous graph-based neural operators do not converge.
• By construction, the method shares the same learned network parameters irrespective of the dis- cretization used on the input and output spaces. It can do zero-shot super-resolution: trained on a lower resolution directly evaluated on a higher resolution.
• The proposed method consistently outperforms all existing deep learning methods for parametric PDEs. It achieves error rates that are $30\%$ lower on Burgers’ Equation, $60\%$ lower on Darcy Flow, and $30\%$ lower on Navier Stokes (turbulent regime with Reynolds number $10000$).
• On a $256 \times 256$ grid, the Fourier neural operator has an inference time of only $0.005$ compared to the $2.2s$ of the pseudo-spectral method used to solve Navier-Stokes.