PaperlayerView exact revision on arXiv

Source-linked AI summary

Adaptive Fourier Neural Operators: Efficient Token Mixers for Transformers

John Guibas, Morteza Mardani, Zongyi Li, Andrew Tao, Anima Anandkumar, Bryan Catanzaro

arXiv:2111.13587v2cs.CVcs.LG

TL;DR

High-resolution vision needs token mixing that preserves long-range relationships without self-attention's quadratic sequence scaling. The paper introduces AFNO, an adaptive Fourier-domain mixer based on continuous global convolution and modified FNO components. AFNO delivers favorable efficiency-accuracy trade-offs for few-shot segmentation and strong high-resolution segmentation results, including Cityscapes sequences of 65k tokens.

  • Problem

    Self-attention mixes tokens effectively but scales quadratically with sequence size, limiting efficient training on high-resolution images.

  • Method

    AFNO adapts Fourier Neural Operators by learning continuous global convolution in the Fourier domain with block-diagonal weights, adaptive sharing, and frequency sparsity.

  • Results

    AFNO outperforms existing mixers for few-shot segmentation in efficiency and accuracy and achieves state-of-the-art Cityscapes segmentation with Segformer-B3 at sequence size 65k.

  • Takeaways & Limitations

    AFNO provides a highly parallel token mixer with quasi-linear sequence complexity and favorable accuracy-efficiency trade-offs for high-resolution segmentation.

  • Takeaways & Limitations

    Standard FNO has O(Nd^2) parameters because each token has separate channel-mixing weights, motivating AFNO's modifications for image-resolution scaling.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper develops AFNO as a Fourier-domain token mixer to address the quadratic scaling of self-attention on high-resolution inputs. It adapts FNO with structural modifications for efficient, expressive, and generalizable visual token mixing, achieving strong segmentation results.

  • Motivation: High-resolution vision requires token mixers that capture long-range dependencies while scaling with sequence size and transferring across tasks.Self-attention provides effective mixing but has quadratic sequence complexity, motivating alternatives with better computational and memory trade-offs.
  • Approach: AFNO frames token mixing as continuous global convolution through operator learning and solves the mixing efficiently in the Fourier domain.The approach adapts Fourier Neural Operators, originally developed for operator learning in PDEs, to visual token mixing.
  • Approach: AFNO improves FNO with block-diagonal channel weights, adaptive weight sharing across tokens, and frequency sparsity through soft-thresholding.These modifications target channel complexity, parameter efficiency, generalization, and high-resolution image discontinuities.
  • Results: AFNO outperforms existing mixers for few-shot segmentation in efficiency and accuracy, including 30% higher efficiency than self-attention with slightly better accuracy.The reported comparison uses a ViT-B backbone and includes GFN, LS, and self-attention.
  • Efficiency: AFNO has quasi-linear complexity in sequence length and is designed as an efficient alternative to quadratic-complexity self-attention.The paper compares AFNO with FNO, GFN, and self-attention in terms of complexity, memory, parameters, and interpretation.

2 RELATED WORKS

Related work improves transformer efficiency through graph-based, MLP-based, and Fourier-based token mixers, while operator learning provides a function-to-function perspective for visual computation.

  • Graph-Based Mixers: Graph-based mixers approximate self-attention using sparse attention, low-rank attention, or related structural constraints.These methods seek efficient surrogates for self-attention's token-similarity graph.
  • MLP-Based Mixers: MLP-based mixers spatially mix tokens with MLP projections but remain limited by quadratic projection complexity and parameter inefficiency for high-resolution images.MLP-Mixer, ResMLP, and gMLP relax graph-similarity constraints while retaining scalability challenges.
  • Fourier-Based Mixers: Fourier-based mixers use Fourier transforms for spatial token mixing, but FNet lacks adaptive filtering and GFN lacks channel mixing and filter adaptivity.AFNO is positioned as a Fourier-based alternative that addresses these limitations through adaptive architectural design.
  • Operator Learning: Operator learning maps functions to functions, allowing images to be treated as RGB-valued functions on a two-dimensional domain.The paper adopts Fourier Neural Operators, which implement global convolution through FFT and have been successful for nonlinear and chaotic PDEs.

3 PRELIMINARIES AND PROBLEM STATEMENT

The paper frames token mixing as continuous global convolution and connects it to Fourier neural operators, offering resolution-independent mixing with quasi-linear sequence complexity. It identifies FNO’s parameter scaling and static weights as limitations motivating adaptive visual architectures.

  • Problem Statement: High-resolution token mixing must preserve global relationships while avoiding self-attention’s quadratic scaling with sequence size.Transformers require effective token mixing for rich representations, but self-attention impedes high-resolution training.
  • Kernel Integration: Self-attention can be represented as a kernel summation and extended to a continuous kernel integral over the image domain.Tokens become samples of a spatial function, allowing the network to be viewed as an operator on functions.
  • Global Convolution: When the kernel depends only on relative displacement, the kernel integral becomes global convolution, which captures global interactions with beneficial regularization.Global convolution is a smaller complexity class than integration and can be implemented efficiently using FFT.
  • Fourier Neural Operator: FNO performs token mixing with a DFT, per-frequency channel mixing, and inverse transformation on discrete image grids.The complex-valued tensor W parameterizes the kernel, while a residual term can compensate for local features and non-periodic boundaries.
  • Fourier Neural Operator: FNO is invariant to grid discretization and can be evaluated at another resolution after training on one resolution.Higher-frequency information is encoded in channels, so truncating frequency modes can still produce the full spectrum.
  • Fourier Neural Operator: FNO has O(N log(N)d^2) sequence complexity but O(Nd^2) parameters because each token has separate channel-mixing weights.Its weights are also static, which can negatively impact generalization; GFN is a diagonal, channel-separable special case.

4 ADAPTIVE FOURIER NEURAL OPERATORS FOR TRANSFORMERS

AFNO modifies FNO for visual inputs by improving scalability, adaptivity, and sparsity. Its design uses block-diagonal channel weights, shared adaptive token processing, and soft-thresholding-based frequency sparsification.

  • Design Motivation: AFNO addresses FNO’s image-specific shortcomings to improve scalability and robustness.The section introduces modifications motivated by FNO’s parameter count and static-weight limitations.
  • Block-Diagonal Structure: Block-diagonal channel weights divide each d × d matrix into k independent blocks, reducing parameterization and enabling parallel computation.Each block acts like a multi-head subspace, while block size one recovers GFN kernels.
  • Weight Sharing: Adaptive weight sharing uses a two-layer perceptron so tokens can interact while deciding which frequency modes to pass.The perceptron weights are shared across tokens, significantly reducing parameter count.
  • Soft-Thresholding and Shrinkage: Soft-thresholding and shrinkage sparsify Fourier-domain representations according to token importance and use λ to control sparsity.The nonlinear LASSO channel mixing replaces the linear combination used in FNO and can regularize the network.
  • AFNO Module: AFNO combines these modifications in a mixer evaluated against FNO, GFN, and self-attention for interpretation, memory, and complexity.The comparison is summarized in Table 1 and the overall module is accompanied by pseudocode.

5 EXPERIMENTS

Experiments evaluate AFNO across inpainting, classification, few-shot segmentation, high-resolution Cityscapes segmentation, and ablations. AFNO generally matches or exceeds competing mixers while reducing computational cost and retaining performance under frequency sparsification.

  • 5.1 IMAGENET-1K INPAINTING: AFNO matches self-attention on ImageNet-1K inpainting while using significantly fewer FLOPs, and outperforms LS and GFN on PSNR and SSIM.Inpainting uses PSNR and SSIM between reconstructed and ground-truth images.
  • 5.2 FEW SHOT SEGMENTATION: AFNO performs on par with self-attention for few-shot segmentation and slightly outperforms it on the out-of-domain ADE-Cars and LSUN-Cats datasets.Performance is measured by validation mIoU after inpainting pretraining.
  • 5.3 CITYSCAPES SEGMENTATION: AFNO-100% achieves the highest Cityscapes mIoU, while AFNO-25% loses only 0.05 mIoU after truncating 75% of high-frequency modes.The evaluation uses 1024 × 1024 semantic segmentation, and MLP layers account for most computations.
  • 5.4 IMAGENET-1K CLASSIFICATION: AFNO outperforms GFN by more than 2% top-1 accuracy on ImageNet-1K classification and is competitive with LS.The reported comparison attributes the GFN improvement to adaptive weight sharing, which permits a larger channel size.
  • 5.5 ABLATION STUDIES: The sparsity threshold λ = 0.01 produces peak inpainting PSNR and segmentation mIoU, while eight blocks achieve the best inpainting PSNR under matched parameter counts.The ablations compare λ values from 0 to 10 and block counts from 1 to 64.
  • 5.5 ABLATION STUDIES: Adaptive weights are crucial to AFNO’s performance, and AFNO’s modifications are crucial relative to FNO for both inpainting pretraining and few-shot segmentation.The comparisons are performed at the same parameter count where specified by the ablation table.
  • 5.6 COMPARISON WITH DIFFERENT TRUNKS AT DIFFERENT SCALES: Across Segformer backbones and parameter sizes, AFNO outperforms other mixers, and hierarchical Segformer with AFNO consistently outperforms Swin for semantic segmentation.The comparison includes AFNO, GFN, efficient self-attention, Swin, ResNet, and MobileNetV2 at multiple scales.

6 CONCLUSIONS

The paper proposes AFNO, an efficient Fourier-domain token mixer for images, and reports favorable accuracy-efficiency trade-offs for few-shot and high-resolution segmentation.

  • AFNO adapts Fourier Neural Operators for image token mixing through block-diagonal weights, adaptive weight sharing, and frequency sparsification.
  • AFNO has quasi-linear complexity in sequence length and is designed to improve efficiency relative to self-attention.
  • Experiments show favorable accuracy-efficiency trade-offs for few-shot segmentation and competitive high-resolution segmentation against state-of-the-art methods.

A APPENDIX

This appendix section presents visualizations of AFNO and provides experimental details.

  • The section includes visualizations of AFNO representations and weights.
  • The section includes details of the experiments.
  • The visualizations and experimental details support analysis of AFNO’s behavior and evaluation.

A.1 VISUALIZATION OF AFNO

The appendix visualizes AFNO’s token clustering and Fourier-domain sparsity, comparing representations across transformer layers and mixers.

  • Token clustering: Token clustering is visualized after each layer of a 10-layer transformer pretrained on 10% of ImageNet-1k.Spectral clustering uses a k-NN kernel with k = 10 and four clusters.
  • Token clustering: AFNO clusters appear as good as self-attention clusters and more aligned with image objects than GFN clusters.
  • Sparsity masks: Soft-thresholding and shrinkage produce sparsity masks showing that input images are sparse in the Fourier domain.The visualized threshold is λ = 0.1, with layers 1 to 5 shown from left to right.
  • Sparsity masks: The observed Fourier sparsity suggests that higher frequencies can be aggressively truncated while maintaining performance.

A.2 INPAINTING

The inpainting experiments use a ViT-B/4 backbone and compare AFNO with self-attention, Long-Short, and Global Filter Network mixers under specified training settings.

  • Backbone and mixers: The ViT-B/4 backbone has 12 layers and uses 4×4 patches to model long sequence sizes.
  • Backbone and mixers: The compared mixers use specified hidden sizes and architectural settings, including AFNO’s one block, 750-dimensional hidden state, and threshold 0.1.
  • Training procedure: Training masks image pixels through a transition process, reconstructs the image with a linear projection, and uses masked-pixel mean squared error.
  • Training procedure: Models train for 100 epochs with Adam, cosine-decayed learning rates, gradient clipping at 1.0, and weight decay of 0.01.

A.3 CITYSCAPES SEGMENTATION

Cityscapes experiments use SegFormer-B3 with consistent four-stage depth while comparing token mixers under specified backbone configurations and training procedures.

  • SegFormer-B3 reduces sequence size while increasing hidden size across its four stages, whose layer counts remain [3, 4, 18, 3].
  • Efficient Self-Attention uses hidden sizes [64, 128, 320, 512] and head counts [1, 2, 5, 8] across the four stages.
  • Global Filter Network uses hidden sizes [128, 256, 440, 512] and bilinear interpolation because it is not resolution invariant.
  • Long-Short uses hidden sizes [128, 256, 360, 512] and [1, 2, 5, 8] attention heads across the four stages.
  • AFNO uses hidden sizes [208, 288, 440, 512] and [1, 2, 5, 8] blocks across the four stages to match other networks' parameter counts.
  • The backbone is pretrained on ImageNet-1K for 300 epochs, then finetuned on Cityscapes for 450 epochs using random and evaluation crops of 1024x1024.

A.4 FEW-SHOT SEGMENTATION

Few-shot segmentation and related experiments use specified mixer configurations, pretrained models, and optimization procedures across the reported datasets and ablations.

  • Few-shot models use inpainting-pretrained weights and are finetuned on CelebA-Faces, ADE-Cars, and LSUN-Cats at 224x224 resolution.
  • Validation performance is measured every 100 epochs, with the maximum across 2000 training epochs reported for each mixer.
  • Classification models use 4x4 patches, 12 transformer layers, and global average pooling to produce softmax probabilities for 1,000 ImageNet-1K classes.
  • Self-Attention uses 12 heads and hidden size 324, while Long-Short uses window size 4, dynamic projection rank 8, and hidden size 312.
  • GFN uses hidden size 245, reduced to keep model parameter counts equal across classification models.
  • AFNO classification uses 16 blocks, hidden dimension 384, sparsity threshold 0.1, and a 1D convolution layer as bias.
  • Classification models train for 300 epochs with Adam, cross-entropy, warmup, cosine decay, gradient clipping at 1.0, and weight decay 0.05.
  • ViT-XS ablations use five-layer, 16M-parameter backbones with attention-specific hyperparameters adjusted accordingly.
Loading 2111.13587v2…