[Paper Review] Adaptive Fourier Neural Operators: Efficient Token Mixers for Transformers

John Guibas, Morteza Mardani+4 more|arXiv (Cornell University)|Nov 24, 2021

CCD and CMOS Imaging Sensors78 citations

TL;DR

This paper introduces Adaptive Fourier Neural Operators (AFNO) as an efficient token mixer for vision transformers, combining Fourier-based global convolution with adaptive, sparse, and block-structured channel mixing to achieve quasi-linear complexity and improved downstream performance, especially in high-resolution and few-shot segmentation tasks.

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.

Motivation & Objective

• Motivate efficient token mixing for high-resolution vision transformers beyond quadratic self-attention complexity.
• Develop AFNO by adapting Fourier Neural Operators to handle images with discontinuities and high resolutions.
• Introduce architectural modifications (block-diagonal weights, adaptive weight sharing, sparsity) to improve efficiency and generalization.
• Demonstrate AFNO’s effectiveness across upsampling/inpainting pretraining, few-shot segmentation, and Cityscapes/high-resolution segmentation.
• Provide code and empirical comparisons to state-of-the-art mixers.

Proposed method

• Frame token mixing as a continuous global convolution via operator learning.
• Adopt Fourier Neural Operators (FNO) with discretized FFT-based implementation for token mixing.
• Impose block-diagonal structure on channel-mixing weights to reduce parameter count.
• Introduce adaptive weight sharing via a shared two-layer MLP across tokens.
• Apply soft-thresholding (sparsity) to frequency modes to focus on informative components.
• Allow adaptive weight masking and residual local terms to handle non-periodic boundaries and local features.

Experimental results

Research questions

• RQ1 Can AFNO achieve competitive or superior accuracy to self-attention with substantially reduced computational cost on vision transformers?
• RQ2 Does enforcing a block-diagonal channel structure, weight sharing, and sparsity improve generalization and efficiency for high-resolution image tasks?
• RQ3 How does AFNO perform in pretraining tasks (inpainting), downstream few-shot segmentation, and high-resolution Cityscapes segmentation compared to existing mixers?
• RQ4 Is AFNO resolution-invariant and scalable to very large token sequences without prohibitive memory consumption?

Key findings

• AFNO achieves competitive PSNR/SSIM in image inpainting while using ~30% fewer GFLOPs than self-attention.
• In few-shot segmentation, AFNO surpasses self-attention on 2/3 datasets and is on par on the others with lower FLOPs.
• On Cityscapes segmentation with SegFormer-B3, AFNO-100% attains a mIoU of 80.9, outperforming Efficient Self-Attention, LS, and GFN variants under comparable settings.
• ImageNet-1K classification with ViT-S/4 shows AFNO achieving top-1 accuracy of 80.89% (vs 80.87% for self-attention), with similar parameter count and FLOPs.
• Ablation studies show that sparsity (lambda ~0.01), 8 blocks, and adaptive weights collectively contribute to AFNO’s performance gains over FNO and non-adaptive variants.
• AFNO demonstrates resolution-invariance and efficient channel mixing, offering a principled alternative to traditional self-attention for high-resolution vision tasks.