Solving PDEs with AI: How Fourier Neural Operators Work

This article explains how Fourier Neural Operators (FNOs) leverage artificial intelligence to solve complex partial differential equations (PDEs) by parameterizing integral kernels directly in the frequency space. It covers the fundamental shift from traditional neural networks to operator learning, the application of the Convolution Theorem via Fast Fourier Transforms, and the exact mechanics of the Fourier layer that enable fast, zero-shot super-resolution predictions for physical systems.

The Shift from Finite Spaces to Operator Learning

Traditional deep learning models map between finite-dimensional spaces, meaning they are tied to a specific grid resolution. To solve PDEs efficiently, AI must map between infinite-dimensional function spaces—taking an initial condition function and outputting the solution function. Neural operators achieve this by framing the neural network layer as an iterative update that combines a local linear transformation with a non-local integral operator.

In a standard neural operator, the integral operator is defined as a convolution of the input function with a learnable kernel function. Calculating this integral directly in physical space is computationally expensive, scaling quadratically with the number of grid points. FNOs bypass this bottleneck entirely by moving the computation into the frequency domain.

Parameterizing Kernels via the Convolution Theorem

The core innovation of the FNO lies in the Convolution Theorem, which states that a global convolution in the physical domain is equivalent to a point-wise multiplication in the frequency domain. Instead of defining the kernel function in physical space and integrating it, FNOs parameterize the kernel directly as a complex-valued weight tensor in Fourier space.

Inside a single Fourier layer, the parameterization process follows four direct steps:

  1. Forward Fourier Transform: The input spatial function is converted into the frequency domain using the Fast Fourier Transform (FFT). This decomposes the physical state into its constituent wave frequencies.
  2. Mode Truncation: High-frequency modes are filtered out, and only a fixed number of lower-frequency modes are retained. This step reduces the number of learnable parameters, prevents overfitting to grid noise, and ensures the model remains invariant to the spatial resolution of the input data.
  3. Frequency Multiplication: The retained Fourier modes are multiplied point-wise by the learnable complex weight tensor. This tensor is the parameterized integral kernel. The AI learns these weights during standard backpropagation, effectively learning how different spatial frequencies interact to evolve the physical system.
  4. Inverse Fourier Transform: The modified frequency modes are padded with zeros to restore the original dimensions and passed through an Inverse Fast Fourier Transform (IFFT) to return the data to the physical spatial domain.

Completing the Fourier Layer

Once the data returns to the physical domain via the IFFT, it has captured the global, long-range dependencies of the PDE. However, because high-frequency modes were truncated, some local details are lost, and non-periodic boundary conditions must still be accounted for.

To solve this, the FNO runs a standard, local linear transformation (such as a pointwise 1x1 convolution) in parallel with the Fourier space operation. The output of the IFFT and the output of the local linear transformation are added together. Finally, this combined result is passed through a non-linear activation function, such as GELU or ReLU. Stacking several of these layers allows the FNO to approximate highly non-linear differential operators.

Advantages for Solving PDEs

By parameterizing the integral kernel in frequency space, FNOs reduce the computational complexity of global convolutions to quasi-linear time, thanks to the efficiency of the FFT algorithm. Furthermore, because the learned weights are tied to fixed frequency modes rather than spatial grid points, FNOs are resolution-invariant. They can be trained on low-resolution data to save compute time and deployed immediately on high-resolution grids to accurately solve complex equations like the Navier-Stokes equations for fluid dynamics.