Adaptive Fourier Neural Operators: Efficient Token Mixers for Transformers

In plain words

Transformers for images cut a picture into many small patches and let every patch look at every other patch. That "looking" step, self-attention, costs more as the image gets larger: doubling the number of patches roughly quadruples the work. High-resolution images make this too slow.

The authors borrow an idea from the study of physical systems such as fluid flow. A global mixing step can be done cheaply by moving the data into frequencies, changing it there, and moving it back. This is a Fourier transform. Their module, called AFNO, does exactly that and then learns which frequencies to keep and how to combine the channels inside each frequency. Frequencies that are weak are set to zero, which both saves work and acts as a regulariser.

The result is a drop-in replacement for the attention step. It mixes information across the whole image in a cost that grows almost linearly with the number of patches. In the paper's experiments it does about as well as attention on inpainting and beats other cheap alternatives on high-resolution segmentation. These results come from the authors' own experiments.

TL;DR

• Replaces the token-mixing (self-attention) step of a vision transformer with a learned operator in the 2D Fourier domain §4.
• Three changes to the basic Fourier neural operator: block-diagonal channel mixing, a two-layer MLP whose weights are shared across tokens, and soft-thresholding of frequencies §4 eq. 5 eq. 6 eq. 8.
• Cost is quasi-linear in token count and linear in memory [Tab. 1]; reported 65k-token Cityscapes runs §1 §5.3.
• Author-reported: 80.9 mIoU on Cityscapes at 1024x1024 with SegFormer-B3, against 79.7 for efficient self-attention and 80.5 for Long-Short [Tab. 4]. As read, the official code does not perform the paper's hard-truncation path (see Pitfalls).

Problem

Self-attention mixes tokens by computing an N x N score matrix, where N is the number of tokens. Its cost therefore grows quadratically with N §1 §3.1. The paper's target is high-resolution vision, where N is large §1. Common remedies (sparse, low-rank, kernel and clustering approximations; MLP mixers) reduce the cost but are reported to lose accuracy or scale poorly with resolution §2. Fourier-based mixers (FNet, GFN) are cheap, but GFN has no channel mixing and its parameter count grows with N §1 §2.

Method

Setup. The image is a grid of h x w patches, each a d-dimensional token, so N = hw §3. The paper views token mixing as a kernel integral. Under a shift-invariant kernel this becomes a global convolution, which the FFT computes efficiently §3.1 eq. 3 eq. 4.

FNO as token mixer. An FNO layer runs three steps per frequency: (1) token mixing by a discrete Fourier transform; (2) channel mixing by a complex d x d matrix per frequency; (3) token demixing by the inverse transform §3.2. The paper notes that GFN is the special case where that matrix is diagonal §3.2. Plain FNO has O(N d^2) parameters and O(N d^2 + N d log N) cost §3.2 [Tab. 1].

AFNO changes, as the paper describes them §4:

1. Block-diagonal channel mixing. The d channels are split into k blocks of size d/k, and each block is mixed independently §4 eq. 5. Each block plays a role like an attention head §4.
2. Adaptive weight sharing. The per-frequency matrices are replaced by one two-layer MLP with a nonlinearity, shared across all tokens §4 eq. 6. The paper says this lets the mixing depend on the input §4.
3. Soft-thresholding. The per-frequency output is passed through the soft- threshold operator S_lambda, which is the solution of a LASSO-style sparse objective §4 eq. 7 eq. 8. Small frequency coefficients become zero.

A residual (bias) term is added to the output §3.2 fig. 3. The module is drawn in fig. 2 and its pseudocode is given in fig. 3 §4.

# AFNO token mixer  [fig. 3] [§4]
# Input:  tokens (B, h, w, d)
# Output: mixed tokens (B, h, w, d)

def afno_mix(tokens, block_weights, block_biases, lam):
    # Keep the input for the residual term  [§3.2] [fig. 3]
    residual = tokens
    # Token mixing: real 2D FFT over the grid  [§3.2 step (1)]
    spectrum = rfft2(tokens)                     # (B, h, w//2+1, d)
    # Split the channels into k blocks of size d/k  [§4] [eq. 5]
    spectrum = reshape(spectrum, (B, h, w//2+1, k, d//k))
    # Two-layer complex MLP per block, shared across tokens  [eq. 6]
    hidden = relu(block_matmul(spectrum, block_weights_in) + block_biases_in)
    spectrum = block_matmul(hidden, block_weights_out) + block_biases_out
    # Soft-threshold each frequency coefficient  [§4] [eq. 8]
    spectrum = soft_shrink(spectrum, lam)
    # Back to the spatial grid  [§3.2 step (3)]
    tokens = irfft2(reshape(spectrum, (B, h, w//2+1, d)))
    return tokens + residual

Key equations:

\tilde z_{m,n} = W_{m,n}\, z_{m,n} \text{ with } W \text{ block-diagonal, } k \text{ blocks} eq. 5

\tilde z_{m,n} = W_2\, \sigma(W_1 z_{m,n}) + b eq. 6

\min \|\tilde z_{m,n} - W_{m,n} z_{m,n}\|^2 + \lambda \|\tilde z_{m,n}\|_1 eq. 7

S_{\lambda}(x) = \operatorname{sign}(x)\, \max\{|x| - \lambda,\, 0\} §4

Cost (per [Tab. 1], asymptotic, as the paper states it): AFNO is O(N d^2 / k + N d log N) in FLOPs and O((1 + 4/k) d^2 + 4d) in parameters, against O(N^2 d + 3 N d^2) and 3 d^2 for self-attention.

Hyperparameters

Values are from the appendix unless marked. Batch size, warmup and weight decay are "not stated" where the paper is silent.

Results

Author-reported. Each table lists its setup. Parameter counts and FLOPs are as the paper reports them.

Complexity tab. 1:

Setup: ImageNet-1K inpainting, ViT-B/4 §5.1 tab. 2:

Setup: few-shot segmentation after inpainting pretraining, mIoU §5.2 tab. 3:

Setup: Cityscapes at 1024x1024, SegFormer-B3, mIoU §5.3 tab. 4. Mixer GFLOPs are the attention-replacement part only; "Total" includes the MLPs.

Setup: ImageNet-1K classification, 224x224, ViT backbone, 16M params §5.4 tab. 5:

Setup: ablation on ViT-XS/4 (16M), inpainting pretraining on 10% of ImageNet-1K, CelebA-Faces mIoU §5.5 tab. 6:

Scaling: AFNO with the hierarchical SegFormer backbone is reported to beat the Swin backbones at matched parameter counts on Cityscapes §5.6 fig. 1.

Paper vs code

Code: official repository NVlabs/AFNO-transformer at commit 941e924ffdb34cb93897b5fd6e22e5965fdc1111 (code_license: CC BY-NC-SA 4.0, non-commercial). The paper's footnote points to github.com/jtguibas/AdaptiveFourierNeuralOperator, which returns 404 when cloned. The NVlabs repository is the official PyTorch implementation of the paper. Its README cites this paper and its authors. Code was read, not run.

• Mixer equation. The paper's eq. 6 uses dense weights W1, W2 shared across tokens eq. 6. The code and the listing use block-diagonal per-block weights of shape (2, k, d/k, d/k) afno2d.py fig. 3 eq. 5.
• Bias term. Appendix A.2 and A.5 describe a 1D convolution as the AFNO bias. The mixer used by the classification model adds only the identity residual afno2d.py afno2d.py. The conv bias exists only in an unused class afnonet.py and is switched on by --fno-bias, which is off by default main_afnonet.py.
• Sparsity. Appendix A.5 gives lambda = 0.1 for classification. The classification block hard-codes 0.01 afnonet.py. Appendix A.6 also uses 0.01.
• Number of blocks. Appendix A.5 says 16 blocks for classification. The command-line default is 1 main_afnonet.py. Which value the reported runs used is not stated in the repository.
• Hard frequency truncation. The paper keeps 32 of 56 frequencies per dimension app. A.6 and the Cityscapes AFNO-25% variant keeps 25% of low frequencies tab. 4. The code computes the kept count from the flattened size N//2+1 with N = h*w afno2d.py, but slices the width axis, which has length w//2+1 afno2d.py. With 56x56 tokens and fraction 0.35 the count is 549, and the axis has 29 entries, so nothing is removed. Reading the code, the truncation in the paper cannot be reproduced from it as written. Not run.
• Soft-threshold on complex values. The code applies the shrinkage to the real and imaginary parts separately afno2d.py. The paper defines S_lambda on a real argument §4 and does not say how it applies to complex coefficients.
• ReLU on complex values. The code applies ReLU to real and imaginary parts separately afno2d.py. The paper's listing shows one ReLU between two matmuls fig. 3 and does not say how it treats complex values.
• FFT normalisation. The code uses orthonormal FFTs in both directions afno2d.py afno2d.py. The paper does not state the normalisation.
• Square grids. Without an explicit spatial size, the code assumes a square grid: H = W = sqrt(N) afno2d.py. The paper's method is stated for h x w grids §3.
• Learning rate. Matches: the code scales 5e-4 by batch/512 main_afnonet.py main_afnonet.py, as appendix A.5 says.
• Model size. Matches on the defaults: 12 layers, hidden 384, 300 epochs, weight decay 0.05, 5 warmup epochs main_afnonet.py main_afnonet.py main_afnonet.py main_afnonet.py main_afnonet.py.

Pitfalls

• Figure and section references in the paper disagree with its own numbering. Counted from the LaTeX, the mixer diagram is Figure 2 and the listing is Figure 3, but the text cites "Fig 1" and "Fig 2" §4. The sparsity sweep is Figure 4, but the text cites "Figure 3" §5.5. The Figure 1 caption cites "Section 5.7", which does not exist (the comparison is §5.6). The few-shot text points to "B.1", but the appendix is labelled A §5.2 app. A.4. Trust the caption numbers and check against the PDF.
• Lambda 0.1 vs 0.01. Two settings in the paper disagree (see Paper vs code). Use the value named for the experiment you are reproducing, and expect the 0.01 setting to be what the released code uses.
• Hard truncation is not in the released code. A reimplementation that truncates as the paper says will not match the code's behaviour. Decide which one you follow and write it down.
• Shared MLP vs block-diagonal. Eq. 6 and eq. 5 describe different channel mixers. The listing and code use the block-diagonal form. Implement the block-diagonal form, since the released code does.
• Complex nonlinearities are unspecified. Both ReLU and soft-shrinkage act on real and imaginary parts separately in the code. Do the same, or the outputs will differ.
• Cost is not fully on the page. The mixer's own cost is small for GFN (2.6 GFLOPs) but for AFNO-100% it is 23.7 GFLOPs of a 440 GFLOP total; the MLPs dominate tab. 4 fig. 1. The O(N log N) claims are asymptotic and do not say what happens at small N.
• Few-shot numbers use the best validation epoch. The paper reports the maximum across training, measured every 100 epochs app. A.4. This is selection on the validation set; the reported numbers are optimistic relative to a fixed-epoch protocol.
• Half precision and LS. LS is omitted from Cityscapes results because of instability in half-precision training, and self-attention cannot be run at 1024x1024 because of memory §5.6 tab. 4.
• Resolution assumption. The code requires a square token grid by default afno2d.py. Pass the spatial size explicitly for h != w.
• Not stated. Warmup for inpainting, weight decay for few-shot (stated as none), drop-path in the paper's classification runs, and batch size in the inpainting and few-shot runs.

Lineage

• Builds on: Fourier Neural Operator (Li et al., 2020) for the global convolution view and the Fourier layer §3.2; Global Filter Networks (Rao et al., 2021) as the depthwise, channel-separable special case §2 §3.2.
• Compared against: self-attention and efficient self-attention (sequence reduction, Wang et al., 2021) §5.6; Long-Short transformer (Zhu et al., 2021) §5.1; MLP-mixers §2; Swin backbones §5.6.
• Superseded by: not known from content in this wiki.