Efficient Token Mixing for Transformers via Adaptive Fourier Neural Operators

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 . Vision transformers have recently shown promise in producing rich contextual representations for recognition and generation tasks . However , a major challenge is posed by long sequences from high resolution images and videos . Here , long-range and multiway dependencies are crucial to understand the compositionality and relationships among the objects in a scene . A key component for the effectiveness of transformers is attributed to proper mixing of tokens . Finding a good mixer is however challenging as it needs to scale with the sequence size , and systematically generalize to downstream tasks . Recently , there has been extensive research to find good token mixers ; see e.g. , Tay et al . ( 2020b ) and references therein . The original self-attention imposes graph structures , and uses the similarity among the tokens to capture the long-range dependencies Vaswani et al . ( 2017 ) ; Dosovitskiy et al . ( 2020 ) . It is parameter efficient and adaptive , but suffers from a quadratic complexity in the sequence size . To achieve efficient mixing with linear complexity , several approximations have been introduced for self-attention ; see Section 2 . These approximations typically compromise accuracy for the sake of efficiency . For instance , long-short ( LS ) transformer aggregates a long-range attention with dynamic projection to model distant correlations and a short-term attention to capture local correlations Zhu et al . ( 2021 ) . Long range dependencies are modeled in low dimensions , which can limit expressiveness . More recently , alternatives have been introduced for self-attention that relax the graph assumption for efficient mixing . Instead , they leverage the geometric structures using Fourier transform Rao et al . ( 2021 ) ; Lee-Thorp et al . ( 2021 ) . For instance , the Global Filter Networks ( GFN ) proposes depthwise global convolution for token mixing that enjoys an efficient implementation in the Fourier domain Rao et al . ( 2021 ) . GFN mainly involves three steps : ( i ) spatial token mixing via fast Fourier transform ( FFT ) ; ( ii ) frequency gating ; and ( iii ) inverse FFT for token demixing . GFN however lacks adaptivity and expressiveness at high resolutions since the parameter count grows with the sequence size , and no channel mixing is involved in ( ii ) . Our Approach . To address these shortcomings , we frame token mixing as operator learning that learns mappings between continuous functions in infinite dimensional spaces . We treat tokens as continuous elements in the function space , and model token mixing as continuous global convolution , which captures global relationships in the geometric space . One way to solve global convolution efficiently is through FFT . More generally , we compose such global convolution operations with nonlinearity such as ReLU to learn any general non-linear operator . This forms the basis for designing Fourier Neural operators ( FNOs ) which has shown promise in solving PDEs Li et al . ( 2020a ) . We thus adopt FNO as a starting point for designing efficient token mixing . Designing AFNO . Adapting FNO from PDEs to vision needs several design modifications . Images have high-resolution content with discontinuities due to edges and other structures . The channel mixing in standard FNO incurs a quadratic complexity in the channel size . To control this complexity , we impose a block-diagonal structure on the channel mixing weights . Also , to enhance generalization , inspired by sparse regression , we sparsify the frequencies via soft-thresholding Tibshirani ( 1996 ) . Also , for parameter efficiency , our MLP layer shares weights across tokens ( see Table 1 ) . We term the resulting model as adaptive FNO ( AFNO ) . We perform extensive experiments with pretraining vision transformers for upstream classification and inpainting that are then finetuned for downstream segmentation . Compared with the state-of-theart , our AFNO using the ViT-B backbone outperforms existing GFN , LS , and self-attention for fewshot segmentation in terms of both efficiency and accuracy , e.g , compared with self-attention , AFNO achieves slightly better accuracy while being 30 % more efficient . For Cityscapes segmentation with the Segformer-B3 backbone , AFNO achieves state-of-the-art and beats previous methods , e.g . AFNO achieves more than 2 % better mIoU compared with efficient self-attention Xie et al . ( 2021 ) , and is also competitive with GFN and LS . Key Contributions . Our main contributions are summarized as follows : • We establish a link between operator learning and high-resolution token mixing and adapt FNO from PDEs as an efficient mixer with a quasi-linear complexity in the sequence length . • We design AFNO in a principled way to improve its expressiveness and generalization by imposing block-diagonal structure , adaptive weight-sharing , and sparsity . • We conduct experiments for pretraining and finetuning . AFNO outperforms existing mixers for few-shot segmentation . For Cityscapes segmentation with the Segformer-B3 backbone , AFNO ( sequence : 65k ) achieves state-of-the-art , e.g. , with 2 % gain over the efficient self-attention . 2 RELATED WORKS . Our work is at the intersection of operator learning and efficient transformers . Since the inception of transformers , there have been several works to improve the efficiency of self-attention . We divide them into three lines of work based on the structural constraints . Graph-Based Mixers primarily focus on finding efficient surrogates to approximate self-attention . Those include : ( i ) sparse attentions that promote predefined sparse patterns ; see e.g. , sparse transformer Child et al . ( 2019 ) , image transformer Parmar et al . ( 2018 ) , axial transformer Ho et al . ( 2019 ) , and longformer Beltagy et al . ( 2020 ) ; ( ii ) low-rank attention that use linear sketching such as linformers Wang et al . ( 2020 ) , long-short transformers Lian et al . ( 2021 ) , Nyströmformer Xiong et al . ( 2021 ) ; ( iii ) kernel methods that approximate attention with ensemble of kernels such as performer Choromanski et al . ( 2020 ) , linear transformer Katharopoulos et al . ( 2020 ) , and random feature attention Peng et al . ( 2021 ) ; and ( iv ) clustering-based methods such as reformer Kitaev et al . ( 2020 ) , routing transformer Roy et al . ( 2021 ) , and Sinkhorn transformer Tay et al . ( 2020a ) . These surrogates however compromise accuracy for efficiencyy . MLP-Based Mixers relax the graph similarity constraints of the self-attention and spatially mix tokens using MLP projections . The original MLP-mixer Tolstikhin et al . ( 2021 ) achieves similar accuracy as self-attention . It is further accelerated by ResMLP Touvron et al . ( 2021 ) that replaces the layer norm with the affine transforms . gMLP Liu et al . ( 2021a ) also uses an additional gating to weight tokens before mixing . This class of methods however lack scalability due to quadratic complexity of MLP projection , and their parameter inefficiency for high resolution images . Fourier-Based Mixers apply the Fourier transform to spatially mix tokens . FNet Lee-Thorp et al . ( 2021 ) resembles the MLP-mixer with token mixer simply being pre-fixed DFT . No filtering is done to adapt the data distribution . Global filter networks ( GFNs ) Rao et al . ( 2021 ) however learn Fourier filters to perform depthwise global convolution , where no channel mixing is involved . Also , GFN filters lack adaptivity that could negatively impact generalization . In contrast , our proposed AFNO performs global convolution with dynamic filtering and channel mixing that leads to better expressivity and generalization . Operator Learning deals with mapping from functions to functions and commonly used for PDEs . Operator learning can be deployed in computer vision as images are RGB-valued functions on a 2D plane . This continuous generalization allows us to permeate benefits from operators . Recent advances in operator learning include DeepONet Lu et al . ( 2019 ) that learns the coefficients and basis of the operators , and neural operators Kovachki et al . ( 2021 ) that are parameterized by integral operators . In this work , we adopt Fourier neural operators Li et al . ( 2020a ) that implement global convolution via FFT which has been very successful for solving nonlinear and chaotic PDEs . 3 PRELIMINARIES AND PROBLEM STATEMENT . Consider a 2D image that is divided into a h × w grid of small and non-overlapping patches . Each patch is represented as a d-dimensional token , and the image can be represented as a token tensor X ∈ Rh×w×d . Treating image as a token sequence , transformers then aim to learn a contextual embedding that transfers well to downstream tasks . To end up with a rich representation , the tokens need to be effectively mixed over the layers . Self-attention is an effective mixing that learns the graph similarity among tokens . It however scales quadratically with the sequence size , which impedes training high resolution images . Our goal is then to find an alternative mixing strategy that achieves favorable scaling trade-offs in terms of computational complexity , memory , and downstream transfer accuracy . 3.1 KERNEL INTEGRATION . The self-attention mechanism can be written as a kernel integration ( Tsai et al. , 2019 ; Cao , 2021 ; Kovachki et al. , 2021 ) . For the input tensor X we denote the ( n , m ) -th token as xn , m ∈ Rd . For notation convenience , we index the token sequence as X [ s ] : = X [ ns , ms ] for some s , t ∈ [ hw ] . Define also N : = hw as the sequence length . The self-attention mixing is then defined as follows : Definition 1 ( Self Attention ) . Att : RN×d → RN×d Att ( X ) : = softmax ( XWq ( XWk ) > √ d ) XWv ( 1 ) where Wq , Wk , Wv ∈ Rd×d are the query , key , and value matrices , respectively . Define K : = softmax ( 〈XWq , XWk〉/ √ d ) as the N × N score array with 〈· , ·〉 being inner product in Rd . We then treat self-attention as an asymmetric matrix-valued kernel κ : [ N ] × [ N ] → Rd×d parameterized as κ [ s , t ] = KWv , which can be viewed as a kernel summation . Att ( X ) [ s ] : = N∑ t=1 κ [ s , t ] X [ t ] ∀s ∈ [ N ] . ( 2 ) This kernel summation can be extended to continuous kernel integrals . The input tensor X is no longer a finite-dimensional vector in the Euclidean space X ∈ RN×d , but rather a spatial function in the function space X ∈ ( D , Rd ) defined on domain D ⊂ R2 which is the physical space of the images . In this continuum formulation , the neural network becomes an operator that acts on the input functions . This brings us efficient characterization originating from operator learning . Definition 2 ( Kernel Integral ) . We define the kernel integral operator K : ( D , Rd ) → ( D , Rd ) as K ( X ) ( s ) = ∫ D κ ( s , t ) X ( t ) dt ∀s ∈ D. ( 3 ) with a continuous kernel function κ : D ×D → Rd×d Li et al . ( 2020b ) . For the special case of the Green ’ s kernel κ ( s , t ) = κ ( s− t ) , the integral leads to global convolution defined below . Definition 3 ( Global Convolution ) . Assuming κ ( s , t ) = κ ( s− t ) , the kernel operator admits K ( X ) ( s ) = ∫ D κ ( s− t ) X ( t ) dt ∀s ∈ D. ( 4 ) The convolution is a smaller complexity class of operation compared to integration . The Green ’ s kernel has beneficial regularization effect but it is also expressive enough to capture global interactions . Furthermore , the global convolution can be efficiently implemented by the FFT .

Summary: Vision transformers scale quadratically with the number of pixels. To cope with this challenge, this paper proposes Adaptive Fourier Neural Operator (AFNO) as an efficient token mixer that learns to mix in the Fourier domain. This is achieved by modifying FNO, including 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 has a quasi-linear complexity and linear memory in the sequence size.