Adaptive Fourier Neural Operators: Efficient Token Mixers for Transformers (2111.13587v2)

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.

Efficient Token Mixers for Vision Transformers: Adaptive Fourier Neural Operators

The proposed work on "Adaptive Fourier Neural Operators" (AFNO) presents a new token mixing mechanism designed to enhance the efficiency and accuracy of vision transformers. This effort addresses a significant challenge in representation learning, namely the quadratic scaling complexity of conventional self-attention mechanisms, particularly with regard to high-resolution input data.

Overview

AFNO operates in the Fourier domain, leveraging the concept of operator learning to facilitate token mixing as a form of global convolution. The AFNO framework builds on Fourier Neural Operators (FNO) by introducing important modifications, such as block-diagonal structures on channel mixing weights and adaptive weight sharing across tokens, along with sparsification of frequency modes via soft-thresholding. These modifications aim to optimize computational and memory efficiency while maintaining or enhancing model expressivity and generalization.

Key Contributions

1. Token Mixing as Operator Learning: The authors cast token mixing within the framework of operator learning, drawing on global convolution ideas to improve how different tokens interact within a vision transformer's layers. This approach circumvents dependency on input resolution, offering benefits such as zero-shot super-resolution by enabling the model to adapt seamlessly across different image resolutions.
2. Architectural Innovations: The paper introduces enhancements specific to the vision domain, such as the imposition of block-diagonal structures for weight matrices. This aims to maintain expressivity while reducing computational overhead. Adaptive weight sharing and mode sparsification are employed to improve model performance efficiency further.
3. Experimental Validation: AFNO demonstrates superior performance over other mixers like self-attention, GFNs, and LS transformations, particularly for few-shot semantic segmentation tasks. Performance metrics include efficiency gains—illustrated by a 30% improvement in FLOPs—and accuracy over state-of-the-art methods for tasks such as Cityscapes segmentation.

Implications and Future Directions

Practically, AFNO's architectural benefits might enable more resource-efficient models suitable for deployment in environments where computational resources and power are limited without compromising on accuracy. Theoretically, the adaptation of operator learning for vision transformers underscores a shift toward a more geometry-aware understanding of image data, which could influence a broader range of neural network design principles.

As a future direction, exploration of alternative transform bases, such as wavelets, could further enhance locality and reduce computational complexity without sacrificing coverage in the frequency domain. Moreover, AFNO's capacity to handle high-resolution inputs suggests additional utility in applications beyond traditional image classification, like super-resolution and generative modeling, where resolution scale can vary significantly across data sets.

In conclusion, by addressing token mixing with a Fourier neural operator framework, the paper sets a compelling precedent for efficient and powerful transformer models. This work has the potential to inspire further exploration into leveraging continuous mathematical operators for discrete data processing, a direction that aligns with ongoing advances at the cutting edge of machine learning and data science.