…in which 𝜎(𝑊𝑖 + 𝒦𝑖 ) is the spectral convolution layer 𝑖 with the point-wise linear transform 𝑊𝑖

and activation function 𝜎(⋅). 𝒫 is the point-wise lifting network that projects the input into a
higher-dimensional latent space, 𝒫 ∶ ℝ𝑑𝑖 𝑛 → ℝ𝑘 .
Similarly, 𝒬 is the point-wise fully-connected decoding network, 𝒫 ∶ ℝ𝑘 → ℝ𝑑𝑜 𝑢𝑡 . Since all fully-
connected components of FNO are point-wise operations, the model is invariant to the dimension-
ality of the input.
For more details, please refer the Modulus User Documentation.

1.2 Adaptive Fourier Neural Operators (AFNO)

In contrast with the Fourier Neural Operator which has a convolutional architecture, the AFNO
leverages contemporary transformer architectures in the computer vision domain. Vision trans-
formers have delivered tremendous success in computer vision. This is primarily due to effective
self-attention mechanisms. To cope with this challenge, Guibas et al. proposed Adaptive Fourier
Neural Operator (AFNO) as an efficient attention mechanism in the Fourier Domain. AFNO is
based on principled foundation of operator learning which allows us to frame attention as a con-
tinuous global convolution efficiently in the Fourier domain. To handle challenges in vision such as
discontinuities in images and high-resolution inputs, AFNO proposes principled architectural mod-
ifications to FNO which results in memory and computational efficiency. This includes imposing a
block-diagonal structure on the channel mixing weights, adaptively sharing weights across tokens,
and sparsifying the frequency modes via soft-thresholding and shrinkage.
The AFNO model typically includes the following steps: 1. Dividing the input image into a regular
grid with ℎ × 𝑤 equal sized patches of size 𝑝 × 𝑝. 2. Embed the patch into a token of size 𝑑, the
embedding dimension resulting in a token tensor (𝑋ℎ×𝑤×𝑑 ) of size ℎ × 𝑤 × 𝑑. 3. Pass the token
through multiple layers of transformer architecture performing spatial and channel mixing. 4. At
the end of all the transformer layers, convert the feature tensor back to image space using a linear
decoder.
For each layer in the Step 3, the AFNO architecture implements the following operations:
The token tensor is first transformed to the Fourier domain with

𝑧𝑚,𝑛 = [DFT(𝑋)]𝑚,𝑛 , (3)

where 𝑚, 𝑛 is the index the patch location and DFT denotes a 2D discrete Fourier transform.
The model then applies token weighting in the Fourier domain and promotes sparsity with a Soft-
Thresholding and Shrinkage operation as

𝑧𝑚,𝑛
̃ = 𝑆𝜆 (MLP(𝑧𝑚,𝑛 )), (4)

where 𝑆𝜆 (𝑥) = sign(𝑥) max(|𝑥| − 𝜆, 0) with the sparsity controlling parameter 𝜆, and MLP(⋅) is
a 2-layer multi-layer perceptron with block-diagonal weight matrices which are shared across all
patches.
The last operation in an ANFO layer is an inverse Fourier to transform back to the patch domain
and add a residual connection as