Fourier neural operator

Overview

A Fourier layer moves a grid of features into frequency space with a discrete transform, applies a learned map to each retained frequency, and moves back. Multiplying in frequency space is a global convolution in the original space, so one layer mixes every position with every other position 2111.13587 §3.1. The shared invariant is that the mixing is global and its weights are indexed by frequency, not by position. The AFNO card gives the three steps of the plain layer 2111.13587 §3.2.

The original FNO paper has no card in this wiki, so it is not cited on this page. Its results appear here only as baselines in the MWT card.

Variants

Plain Fourier layer

Each retained frequency gets a complex d x d channel-mixing matrix. Token mixing is a discrete Fourier transform, and demixing is the inverse transform. The parameter count is O(N d^2), where N is the number of tokens and d the channel width, and the cost is O(N d^2 + N d log N) 2111.13587 §3.2 2111.13587 tab. 1. Because the parameter count grows with N, the layer is tied to the grid it was sized for. The MWT card uses this layer as its baseline on several PDE benchmarks 2109.13459 tab. 1.

Adaptive Fourier token mixer (AFNO)

A replacement for the self-attention step in a vision transformer 2111.13587 §4. Three changes are made to the plain layer:

• Block-diagonal channel mixing. The channels are split into k blocks and each block is mixed on its own 2111.13587 §4 2111.13587 eq. 5.
• Shared two-layer map. One small nonlinear MLP, shared across tokens, is applied per frequency, so the mixing depends on the input 2111.13587 §4 2111.13587 eq. 6.
• Soft-thresholding. Weak frequency coefficients are set to zero by a shrinkage operator, which is the solution of a sparse objective 2111.13587 §4 2111.13587 eq. 7.

A residual connection is added to the output 2111.13587 §3.2. The parameter count does not depend on N, and the cost is O(N d^2 / k + N d log N) 2111.13587 tab. 1. Author-reported results: 80.9 mIoU on Cityscapes at 1024x1024 with SegFormer-B3, against 79.7 for efficient self-attention and 80.5 for Long-Short 2111.13587 tab. 4. On ImageNet classification with ViT-S/4 it scores 80.89 top-1, against 80.87 for Long-Short and 78.77 for GFN 2111.13587 tab. 5. The authors' ablation finds the input-dependent mixer ahead of a static-weight version (49.49 against 44.04 mIoU) and a hard 35% frequency cut behind both (34.17 mIoU) 2111.13587 tab. 6. Card: Adaptive Fourier Neural Operators: Efficient Token Mixers for Transformers.

Spectral kernels inside a multiwavelet operator (MWT)

The input is split into scales with a fixed multiwavelet decomposition, a small learned kernel acts on each scale, and the result is reconstructed. The kernels are described as convolutions in the paper, but the authors' code uses Fourier-domain maps in the 1D model. The card's Paper vs code section lists which kernels are spectral in each model 2109.13459 §2.3. On the card's benchmarks the Legendre variant has lower relative L2 error than FNO. KdV at s=64 gives 0.00338 against 0.0125 2109.13459 tab. 1. Darcy flow at s=32 gives 0.0152 against 0.0177 2109.13459 tab. 2. Burgers at s=256 gives 0.00199 against 0.00332 2109.13459 tab. 7. The margin is 1.2 to 3.7 times on the rows checked. Card: Multiwavelet-based Operator Learning for Differential Equations.

When to use which

Implementation: neither card allows copying code. The AFNO repository is CC BY-NC-SA, which is non-commercial, and the MWT repository has no license. Write the layer from the description and the pseudocode below.

Implementation sketch

The common core is the plain layer. The AFNO variant replaces the channel map and adds soft-thresholding. Pseudocode in the style of the paper cards.

# Input: grid x of shape (h, w, d). Output: same shape.
def fourier_layer(x, weights, lam):
    # Keep the input for a residual connection (AFNO) [arxiv:2111.13587 §3.2]
    skip = x
    # Step 1: move to frequency space, 2D real FFT  [arxiv:2111.13587 §3.2]
    s = rfft2(x)                          # (h, w//2+1, d), complex
    # Keep the lowest modes; the MWT paper does not state how many
    s = s[:, :modes, :]
    # Step 2: per-frequency channel map. Plain: one d x d matrix per
    # frequency. AFNO: block-diagonal blocks and a shared MLP [arxiv:2111.13587 §4]
    s = channel_map(s, weights)
    # AFNO only: soft-threshold weak frequencies to zero [arxiv:2111.13587 §4]
    s = soft_shrink(s, lam)
    # Step 3: back to the grid with the inverse transform  [arxiv:2111.13587 §3.2]
    x = irfft2(pad(s), shape=(h, w))
    # Add the residual  [arxiv:2111.13587 §3.2]
    return x + skip

Where the variants differ:

• Mode truncation. Keeping a fixed number of low modes is an FNO-family choice, and the MWT card treats the number as a hyperparameter that the paper does not state 2109.13459 §3. The AFNO paper's hard cut is a separate, weaker path 2111.13587 tab. 6.
• Complex values. The paper does not say how ReLU and the shrinkage act on complex numbers. The AFNO code applies both to the real and imaginary parts separately, so match that or the outputs will differ 2111.13587 §4.
• Grid shape. The AFNO method is stated for h x w grids, but its code assumes a square grid unless the size is passed in 2111.13587 §3.

Contradictions

• AFNO, channel mixer. Equation 6 describes one dense shared MLP, while equation 5 describes block-diagonal weights 2111.13587 eq. 5 2111.13587 eq. 6. The paper's listing uses the block-diagonal form.
• AFNO, sparsity. The classification appendix gives lambda = 0.1 2111.13587 app. A.5. The sparsity sweep peaks at 0.01 2111.13587 §5.5.
• MWT, kernels. The text calls the kernels single-layer convolutions 2109.13459 §2.3. The authors' 1D code uses Fourier-domain maps, so the two are not the same operator.
• Cross-paper comparison. The MWT card reports error margins over FNO on its own tasks 2109.13459 tab. 1, but its Burgers and Navier-Stokes FNO baselines are taken from published results, not re-run 2109.13459 §3.3. The AFNO card has no PDE benchmark, so the two cards share no task and their numbers are not directly comparable.

Papers

• Multiwavelet-based Operator Learning for Differential Equations: a multiwavelet operator with learned kernels per scale. It is compared against FNO and other neural operators on KdV, Burgers, Darcy and Navier-Stokes.
• Adaptive Fourier Neural Operators: Efficient Token Mixers for Transformers: a Fourier-domain token mixer for vision transformers, with block-diagonal channel mixing and soft-thresholding.