Adaptive Fourier Neural Operators: Efficient Token Mixers for Transformers

TL;DR

The paper tackles the scalability challenge of token mixing in vision transformers for high-resolution inputs by introducing Adaptive Fourier Neural Operators (AFNO) as efficient, Fourier-domain token mixers. By incorporating a block-diagonal channel structure, adaptive weight sharing via a token-shared MLP, and soft-thresholding to sparsify frequency modes, AFNO achieves quasi-linear computational complexity and linear memory with strong performance on few-shot segmentation and Cityscapes segmentation, while remaining competitive on ImageNet-1K classification. The work situates AFNO within operator learning and Fourier-based mixers, showing it outperforms or matches existing approaches like GFNs and self-attention across diverse tasks, with notable efficiency gains. The findings suggest AFNO as a practical, scalable alternative for high-resolution vision transformers and point toward future extensions using alternative spectral bases such as wavelets.

Abstract

Vision transformers have delivered tremendous success in representation learning. This is primarily due to effective token mixing through self attention. However, this scales quadratically with the number of pixels, which becomes infeasible for high-resolution inputs. To cope with this challenge, we propose Adaptive Fourier Neural Operator (AFNO) as an efficient token mixer that learns to mix in the Fourier domain. AFNO is based on a principled foundation of operator learning which allows us to frame token mixing as a continuous global convolution without any dependence on the input resolution. This principle was previously used to design FNO, which solves global convolution efficiently in the Fourier domain and has shown promise in learning challenging PDEs. To handle challenges in visual representation learning such as discontinuities in images and high resolution inputs, we propose principled architectural modifications 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 resulting model is highly parallel with a quasi-linear complexity and has linear memory in the sequence size. AFNO outperforms self-attention mechanisms for few-shot segmentation in terms of both efficiency and accuracy. For Cityscapes segmentation with the Segformer-B3 backbone, AFNO can handle a sequence size of 65k and outperforms other efficient self-attention mechanisms.

Figures (6)

• Figure 1: Parameter count and mIoU for Segformer, Swin, and other models at different scales. AFNO consistently outperforms other mixers (see Section 5.7).
• Figure 2: The multi-layer transformer network with FNO, GFN, and AFNO mixers. GFNet performs element-wise matrix multiplication with separate weights across channels ($k$). FNO performs full matrix multiplication that mixes all the channels. AFNO performs block-wise channel mixing using MLP along with soft-thresholding. The symbols $h$, $w$, $d$, and $k$ refer to the height, width, channel size, and block count, respectively.
• Figure 3: Pseudocode for AFNO with adaptive weight sharing and adaptive masking.
• Figure 4: Ablations for the sparsity thresholds and block count measured by inpainting validation PSNR. The results suggest that soft thresholding and blocks are effective
• Figure 5: Spectral clustering of tokens for different token mixers. From top to bottom, it shows the input and the layers $2,4,6,8,10$ for the inpainting pretrained model.