E FFICIENT T OKEN M IXING FOR T RANSFORMERS VIA A DAPTIVE F OURIER N EURAL O PERATORS
Anonymous authors
Paper under double-blind review
A BSTRACT
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 efﬁcient 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 efﬁciently 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 modiﬁcations to FNO which results in memory and computational efﬁciency.
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 efﬁciency and accuracy.
For Cityscapes segmentation with the Segformer-B3 backbone, AFNO can handle a sequence size of 65k and outperforms other efﬁcient self-attention mechanisms.
I NTRODUCTION
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 ﬁnd 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. ; Dosovitskiy et al.
It is parameter efﬁcient and adaptive, but suffers from a quadratic complexity in the sequence size.
To achieve efﬁcient mixing with linear complexity, several approximations have been introduced for self-attention; see Section 2.
These approximations typically compromise accuracy for the sake of efﬁciency.
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.
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 efﬁcient mixing.
Instead, they leverage the geometric structures using Fourier transform Rao et al.
(2021); Lee-Thorp et al.
For instance, the Global Filter Networks (GFN) proposes depthwise global convolution for token mixing that enjoys an efﬁcient implementation in the Fourier domain Rao et al.
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
Models
[Equation]
Graph Global Conv.
[Equation]
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 inﬁnite 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 efﬁciently 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 efﬁcient token mixing.
Designing AFNO .
Adapting FNO from PDEs to vision needs several design modiﬁcations.
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 efﬁciency, 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 classiﬁcation and inpainting that are then ﬁnetuned 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 efﬁciency and accuracy, e.g, compared with self-attention, AFNO achieves slightly better accuracy while being 30% more efﬁcient.
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 efﬁcient 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 efﬁcient 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 ﬁnetuning.
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 efﬁcient self-attention.
R ELATED WORKS
Our work is at the intersection of operator learning and efﬁcient transformers.
Since the inception of transformers, there have been several works to improve the efﬁciency of self-attention.
We divide them into three lines of work based on the structural constraints.
Graph-Based Mixers primarily focus on ﬁnding efﬁcient surrogates to approximate self-attention.
Those include: ( i ) sparse attentions that promote predeﬁned sparse patterns; see e.g., sparse transformer Child et al.
, 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.
, long-short transformers Lian et al.
(2021), Nystr ¨ omformer 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 efﬁciencyy.
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 afﬁne 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 inefﬁciency 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-ﬁxed DFT.
No ﬁltering is done to adapt the data distribution.
Global ﬁlter networks (GFNs)
Rao et al.
(2021) however learn
Fourier ﬁlters to perform depthwise global convolution, where no channel mixing is involved.
Also,
GFN ﬁlters lack adaptivity that could negatively impact generalization.
In contrast, our proposed
AFNO performs global convolution with dynamic ﬁltering 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 beneﬁts from operators.
Recent advances in operator learning include DeepONet Lu et al.
(2019) that learns the coefﬁcients 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.
P RELIMINARIES AND P ROBLEM S TATEMENT
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 ∈ R h × 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 ﬁnd an alternative mixing strategy that achieves favorable scaling trade-offs in terms of computational complexity, memory, and downstream transfer accuracy.
3. K ERNEL I NTEGRATION
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 x n,m ∈ R d .
For notation convenience, we index the token sequence as X [ s ] :=
X [ n s , m s ] for some s, t ∈ [ hw ] .
Deﬁne also N := hw as the sequence length.
The self-attention mixing is then deﬁned as follows:
[Equation]
where W q , W k , W v ∈ R
d × d are the query, key, and value matrices, respectively.
Deﬁne K := softmax( (cid:104) XW q , XW k (cid:105) / √ d ) as the N × N score array with (cid:104)· , ·(cid:105) being inner product in R d .
We then treat self-attention as an asymmetric matrix-valued kernel κ : [ N ] × [ N ] → R d × d parameterized as κ [ s, t ] =
KW v , which can be viewed as a kernel summation.
[Equation]
This kernel summation can be extended to continuous kernel integrals.
The input tensor X is no longer a ﬁnite-dimensional vector in the Euclidean space X ∈ R N × d , but rather a spatial function in the function space X ∈ ( D, R d ) deﬁned on domain D ⊂ R 2 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 efﬁcient characterization originating from operator learning.
Deﬁnition 2 (Kernel Integral) .
We deﬁne the kernel integral operator K : ( D, R d ) →
( D, R d ) as
[Equation]
with a continuous kernel function κ : D × D → R d × d Li et al.
(2020b).
For the special case of the
Green’s kernel κ ( s, t ) =
κ ( s − t ) , the integral leads to global convolution deﬁned below.
Deﬁnition 3 (Global Convolution) .
Assuming κ ( s, t )
= κ ( s − t ) , the kernel operator admits
[Equation]
The convolution is a smaller complexity class of operation compared to integration.
The Green’s kernel has beneﬁcial regularization effect but it is also expressive enough to capture global interactions.
Furthermore, the global convolution can be efﬁciently implemented by the FFT.
3. F OURIER N EURAL O PERATOR AS T OKEN M IXER
The class of shift-equivariant kernels has a desirable property that they can be decomposed as a linear combination of eigen functions.
Eigen transforms have a magical property where according to the convolution theorem Soliman & Srinath (1990), global convolution in the spatial domain amounts to multiplication in the eigen transform domain.
A popular example of such eigen functions is the
Fourier transform.
Accordingly, one can deﬁne the Fourier neural operator (FNO) Li et al.
(2020a).
Deﬁnition 4 (Fourier Neural Operator) .
For the continuous input X ∈ D and kernel κ , the kernel integral at token s is found as
[Equation]
where · denotes matrix multiplication, and F , F − 1 denote the continous Fourier transform and its inverse, respectively.
Discrete FNO .
Inspired by FNO, for images with ﬁnite dimension on a discrete grid, our idea is to mix tokens using the discrete Fourier transform (DFT).
For the input token tensor X ∈ R h × w × d , deﬁne the complex-valued weight tensor W :
= DFT( κ ) ∈
C h × w ×
d to parameterize the kernel.
FNO mixing then entails the following operations per token ( m, n ) ∈
[ h ] × [ w ]
[Equation]
Local Features .
DFT assumes a global convolution applied on periodic images, which is not typically true for real-world images.
To compensate local features and non-periodic boundaries, we can add a residual term x m,n (can also be parameterized as a simple local convolution) to the token demixing step 3 in FNO; see also Wen et al.
Resolution Invariance .
The FNO model is invariant to the discretization h, w .
It parameterizes the tokens function via Fourier bases which are invariant to the underlying resolution.
Thus, after training on one resolution it can be directly evaluated at another resolution (zero-shot super-resolution).
Further, the FNO model encodes the higher-frequency information in the channel dimension.
Thus, even after truncating the higher frequency modes z m,n , FNO can still output the full spectrum.
It is also important to recognize step (2) of FNO, where d × d weight matrix W m,n mixes the channels.
This implies mixed-channel global convolution.
Note that the concurrent GFN work Rao et al.
(2021) is a special case of FNO, when W m,n is diagonal and the channels are separable.
The FNO incurs O ( N log( N ) d 2 ) complexity, and thus quasi-linear in the sequence size.
The parameter count is however O ( Nd 2 ) as each token has its own channel mixing weights, which poorly scales with the image resolution.
In addition, the weights W m,n are static, which can negatively impact the generalization.
The next section enhances FNO to cope with these shortcomings.
A DAPTIVE F OURIER N EURAL O PERATORS FOR T RANSFORMERS
This section ﬁxes the shortcomings of FNO for images to improve scalability and robustness.
Block-Diagonal Structure on W .
FNO involves d × d weight matrices for each token.
That results in O ( Nd 2 ) parameter count that could be prohibitive.
To reduce the paramater count we impose a block diagonal structure on W , where it is divided into k weight blocks of size d/k × d/k .
The kernel then operates independently on each block as follows
[Equation]
The block diagonal weights are both interpretable and computationally parallelizable.
In essence, each block can be interpreted a head as in multi-head self-attention, which projects into a subspace of the data.
The number of blocks should be chosen properly so each subspace has a sufﬁciently large dimension.
In the special case, that the block size is one, FNO coincides with the GFN kernels.
Moreover, the multiplications in (5) are performed independently, which is quite parallelizable.
Weight Sharing .
Another caveat with FNO is that the weights are static and once learned they will not be adaptively changed for the new samples.
Inspired by self-attention we want the tokens to be adaptive.
In addition, static weights are independent across tokens, but we want the tokens interact and decide about passing certain low and high frequency modes.
To this end, we adopt a two-layer perceptron that is supposed to approximate any function for a sufﬁciently large hidden layer.
For ( n, m ) -th token, it admits
[Equation]
Note, that the weights W 1 , W 2 , b are shared for all tokens, and thus the parameter count can be signiﬁcantly reduced.
Soft-Thresholding and Shrinkage .
Images are inherently sparse in the Fourier domain, and most of the energy is concentrated around low frequency modes.
Thus, one can adaptively mask the tokens according to their importance towards the end task.
This can use the expressivity towards [Figure] representing the important tokens.
To sparsify the tokens, instead of linear combination as in (5), we use the nonlinear LASSO Tibshirani (1996) channel mixing as follow
[Equation]
This can be solved via soft-thresholding and shrinkage operation
[Equation]
that is deﬁned as S λ ( x ) = sign( x ) max {| x | − λ, 0 } , where λ is a tuning parameter that controls the sparsity.
It is also worth noting that the promoted sparsity can also regularize the network and improve the robustness.
With the aforementioned modiﬁcations, the overall AFNO mixer module is shown in Fig 1 along with the pseudo code in Fig 2.
Also, for the sake of comparison, the AFNO is compared against
FNO, GFN, and self-attention in Table 1 in terms of interpretation, memory, and complexity.
E XPERIMENTS
We conduct extensive experiments to demonstrate the merits of our proposed AFNO transformer.
Namely, 1) we evaluate the efﬁciencyy-accuracy trade-off between AFNO and alternative mixing mechanisms on inpainting and classiﬁcation pretraining tasks; and then 2) measure performance on few-shot semantic segmentation with inpainting pretraining; and 3) evaluate the performance of
AFNO in high resolution settings with semantic segmentation.
Our experiments cover a wide-range of datasets, including ImageNet-1k, CelebA-Faces, LSUN-Cats, ADE-Cars, and Cityscapes as in
Deng et al.
; Liu et al.
Yu et al.
; Krause et al.
; Cordts et al.
5. I MAGE N ET -1K I NPAINTING
We conduct image inpainting experiments which compare AFNO to other competitive mixing mechanisms.
The image inpainting task is deﬁned as follows: given an input image X of size [ h, w, d ] , where h, w, d denote height, width, and channels respectively, we randomly mask pixel intensities to zero based on a uniformly random walk.
The loss function used to train the model is mean squared error between the original image and the reconstruction.
We measure performance via the Peak
Signal-to-Noise Ratio (PSNR) and structural similarity index measure (SSIM) between the ground truth and the reconstruction.
More details about the experiments are provided in the appendix.
[Table]
Inpainting Results.
PSNR and SSIM are reported in Table 2 for AFNO versus alternative mixers.
It appears that AFNO is competitive with self-attention.
However, AFNO uses signiﬁcantly less GFLOPs than Self-Attention.
Compared to both LS and GFN, AFNO acheives signiﬁcantly better PSNR and SSIM.
More importantly, AFNO achieves favorable downstream transfer, which is elaborated in the next section for few-shot segmentation.
5. F EW S HOT S EGMENTATION
After pretraining on image inpainting, we evaluate the few-shot sematic segmentation performance of the models.
We construct three few-shot segmentation datasets by selecting training and validation images from CelebA-Faces, ADE-Cars, and LSUN-Cats as in Zhang et al. (2021b). The model is trained using cross-entropy loss. We measure mIoU over the validation set.
More details about the experiments are deferred to the Appendix.
[Table]
Few-Shot Segmentation Results.
Results are reported in Table 3.
It is evident that AFNO performs on par with self-attention.
For out-of-domain datasets such as ADE-Cars or LSUN-Cats it slightly outperforms self-attention, which is attributed to the sparsity regularization in AFNO.
5. C ITYSCAPES S EGMENTATION
To test the scalability of AFNO for high resolution images with respect to alternative mixers, we evaluate high-resolution ( 1024 × 1024 ) semantic segmentation for the Cityscapes dataset.
We use the SegFormer-B3 backbone which is a hiearchical vision transformer Xie et al.
We train the model using the cross-entropy loss and measure performance via reporting mIoU over the validation set.
More details about the experiments and the model are available in the Appendix.
[Table]
Cityscapes Segmentation Results.
We report the ﬁnal numbers for Cityscapes semantic segmentation in Table 4.
AFNO-100% outperforms all other methods in terms of mIoU.
Furthermore, we ﬁnd that the AFNO-25% model which truncates 75% of high frequency modes during ﬁnetuning only loses 0.05 mIoU and is competitive with the other mixers.
It is important to note that the majority of computations is spent for the MLP layers after the attention module [Table]
5. I MAGE N ET -1K C LASSIFICATION
We run image classiﬁcation experiments with the AFNO mixer module using the ViT backbone on ImageNet-1K dataset containing 1 .
28 M training images and 50 K validation images from 1 ,classes at 224 × 224 resolution.
We measure performance via reporting top-1 and top-5 validation accuracy along with theoretical FLOPs of the model.
More details about the experiments are provided in the appendix.
Classiﬁcation Results.
The classiﬁcation accuracy for different token mixers are listed in Table 5.
It can be observed that AFNO outperforms GFN by more than 2% top-1 accuracy thanks to the adaptive weight sharing which allows for a larger channel size.
Furthermore, our experiments demonstrate that AFNO is competitive with LS for classiﬁcation. [Table]
5. A BLATION S TUDIES
We also conduct experiments to investigate how different components of AFNO contribute to performance.
Sparsity Threshold.
We vary the sparsity threshold λ from 0 to 10.
For each λ we pretrain the network ﬁrst, and then ﬁnetune for few-shot segmentation on the CelebA-Faces dataset.
We report both the inpainting PSNR from pretraining and the segmentation mIoU.
The results are shown in
Figure 3.
λ = 0 corresponds to no sparsity.
It is evident that the PSNR/mIoU peaks at λ = 0 .
01 , indicating that the sparsity is effective.
We also compare to hard thresholding (always removing higher frequencies as in FNO) in Table 6.
We truncate 65% of the higher frequencies for both inpainting pretraining and few-shot segmentation ﬁnetuning.
Number of Blocks.
We vary the number of blocks used when we impose structure on our weights
W .
To make the comparison fair, we simultaneously adjust the hidden size so the overall parameter count of the model are equal.
We vary the number of blocks from 1 to 64 and measure the resulting inpainting PSNR on ImageNet-1K. It is seen that 8 blocks achieves the best PSNR.
This shows that blocking is effective.
Impact of Adaptive Weights.
We evaluate how removing adaptive weights and instead using static weights affects the performance of AFNO for ImageNet-1K inpainting and few-shot segmentation.
The results are presented in Table 6.
Adaptive weights are crucial to AFNO’s performance.
[Table]
Comparison to FNO.
To show that AFNO’s modiﬁcations ﬁx the shortcomings of FNO for images, we directly compare AFNO and FNO on ImageNet-1K inpainting pretraining and few-shot segmentation on CelebA-Faces.
The results are also presented in Table 6.
The results suggest that AFNO’s modiﬁcations are crucial to performance in both tasks.
5. V ISUALIZATION OF AFNO
To gain insight into how AFO works, we produce visualizations of AFNO’s weights and representations.
We visualize the spectral clustering of the tokens and sparsity masks from soft thresholding.
AFNO clustering versus other mixers .
We show the clustering of tokens after each transformer layer.
For a 10-layer transformer pretrained with 10% of ImageNet-1k, we apply spectral clustering on the intermediate features when we use k-NN kernel with k = 10 and the number of clusters is also set to 4 .
It appears that AFNO clusters are as good as self-attention.
AFNO clusters seem to be more aligned with the image objects than GFN ones.
Also, long-short transformer seems not to preserve the objectness in the last layers.
This observation is consistent with the the few-shot segmentation results in section 5.2 that show the superior performance of AFNO over the alternatives.
[Figure]
Sparsity Masks .
We also explore how the sparsity mask affects the magnitude of the values in tokens.
We calculate the fraction of values over the channel dimension and blocks that have been masked to zero by the Softshrink function.
As one can see in Figure 5, it is clear that the input images are very sparse in the Fourier domain.
Furthermore, this sparsity suggests that we can aggressively truncate higher frequencies and maintain performance.
C ONCLUSIONS
We leverage the geometric structure of images in order to build an efﬁcient token mixer in comparison to self-attention.
Inspired by global convolution, we borrow Fourier Neural Operators (FNO) from PDEs for mixing tokens and propose principled architectural modiﬁcations to adapt FNO for images.
Speciﬁcally, we impose a block diagonal structure on the weights, adaptive weight sharing, and sparsify the frequency with soft-thresholding and shrinkage.
We call the proposed mixer Adaptive Fourier Neural Operator (AFNO) and it incurs quasi-linear complexity in sequence length.
Our experiments indicate favorable accuracy-efﬁciency trade-off for few-shot segmentation, and competitive high-resolution segmentation compared with state-of-the-art.
There are still important avenues to explore for the future work such as exploring alterantives for the DFT such as the Wavelet transform to better capture locality as in Gupta et al.
A A PPENDIX
This section includes the details of the experiments.
A. I NPAINTING
We use the ViT-B/4 backbone for the ImageNet-1K inpainting experiments.
The ViT-B backbone has 12 layers and is described in more detail in the original paper Dosovitskiy et al.
Importantly, we use a 4x4 patch size to model the long sequence size setting.
• Self-Attention Vaswani et al.
(2017) uses 16 attention heads and a hidden size of 768 .
• Long-Short transformer Zhu et al.
(2021) (LS) uses a window-size of 4 and dynamic projection rank of 8 with a hidden size of 768 . • Global Filter Network Rao et al.
(2021) (GFN) uses a hidden size of 768 . • Adaptive Fourier Neural Operator (AFNO) uses 1 block, a hidden dimension of 750 , and a sparsity threshold of 0 . 1 , and a 1D convolution layer as the bias.
The training procedure can be summarized as follows.
Given an image x ∈ R 3 × 224 × 224 from
ImageNet-1K, we randomly mask pixel intensities to zero by initially sampling uniformly from { ( i, j ) } h,wi,j =1 and then apply the transition function T (( i, j ))
=
UniformSample ( { ( i − 1 , j ) , ( i + 1 , j ) , ( i, j − 1) , ( i, j + 1) } ) for 3136 steps.
We use a linear projection layer at the end to convert tokens into a reconstructed image.
Our loss function computes the mean squared error between the ground truth and reconstruction of the masked pixels.
We measure performance via the Peak
Signal-to-Noise Ratio (PSNR) and structural similarity index measure (SSIM) between the ground truth and the reconstruction.
We train for 100 epochs using the Adam optimizer with a learning rate of 10 − 4 for self-attention and 10 − 3 for all the other mechanisms using the cosine-decay schedule to a minimum learning rate of 10 − 5 .
We use gradient clipping threshold of 1.0 and weight-decay of
A. C ITYSCAPES S EGMENTATION
We use the SegFormer-B3 backbone for the Cityscapes segmentation experiments.
The SegFormerB3 backbone is a four-stage architecture which reduces the sequence size and increase the hidden size as you progress through the network.
The model is described in more detail in Xie et al.
More details about how the mixing mechanisms are combined with SegFormer-B3 is described below.
All models do not modify the number layers in each stage which is [3, 4, 18, 3].
• Efﬁcient Self-Attention uses a hidden size of [64, 128, 320, 512] and [1, 2, 5, 8] for the four stages • Global Filter Network uses a hidden size of [128, 256, 440, 512] in order to match the parameter count of the other networks.
Because GFN is not resolution invaraint, we use bilinear interpolation in the forward pass to make the ﬁlters match the input resolution of the tokens.
• Long-Short uses a hidden size of [128, 256, 360, 512] to match the parameter count of the other networks and [1, 2, 5, 8] attention heads.
• Adaptive Fourier Neural Operator uses [208, 288, 440, 512] to match the parameter count of the other networks.
It uses [1, 2, 5, 8] blocks in the four stages.
We pretrain the SegFormer-B3 backbone on ImageNet-1K classiﬁcation for 300 epochs.
Our setup consists of using the Adam optimizer, a learning rate of 10 − 3 with cosine decay to 10 − 5 , weight regularization of 0 .
05 , a batch size of 1024 ,
gradient clipping threshold of 1.0, and learning rate warmup for 6250 iterations.
We then ﬁnetune these models on Cityscapes for 450 epochs using a learning rate of 1 . 2 · 10 − 4 .
We train on random 1024x1024 crops and also evaluate at 1024x1024.
A. F EW -S HOT S EGMENTATION
The models used for few-shot segmentation are described in B.1.
We use the inpainting pretrained models and ﬁnetune them on few-shot segmentation on CelebA-Faces, ADE-Cars, and LSUN-Cats at 224x224 resolution.
The samples for the few-shot dataset are selected as done in DatasetGAN in
Zhang et al.
(2021a).
To train the network, we use the per-pixel cross entropy loss.
We ﬁnetune the models on the few-shot datasets for 2000 epochs with a learning rate of 10 − 4 for self-attention and 10 − 3 for other mixers.
We use no gradient clipping or weight decay.
We measure validation performance every 100 epochs and report the maximum across the entire training run.
A. C LASSIFICATION
The models for classiﬁcation are based on the Global Filter Network GFN-XS models but with 4x patch size.
In particular, we utilize 12 transformer layers and adjust the hidden size and attentionspeciﬁc hyperparameters to reach a parameter count of 16M.
Due to some attention mechanisms not being able to support class tokens, we use global average pooling at the last layer to produce output softmax probabilities for the 1 , 000 classes in ImageNet-1k.
More details about each of the models is provided below.
• Self-Attention Vaswani et al.
(2017) uses 12 attention heads and a hidden size of 324 .
• Long-Short transformer Zhu et al.
(2021) (LS) uses a window-size of 4 and dynamic projection rank of 8 with a hidden size of 312 . • Global Filter Network Rao et al.
(2021) (GFN) uses a hidden size of 245 .
The hidden size is smaller due to the need to make all the models have the same parameter count.
• Adaptive Fourier Neural Operator (AFNO) uses 16 blocks, a hidden dimension of 384 , sparsity threshold of 0 . 1 , and a 1D convolution layer as the bias.
We trained for 300 epochs with Adam optimizer and cross-entropy loss using the learning rate of ( BatchSize / 512) × 5 × 10 − 4 for the models.
We also use ﬁve epochs of linear learning-rate warmup, and after a cosine-decay schedule to the minimum value 10 − 5 .
Along with this, the gradient norm is clipped not to exceed 1 . 0 and weight-decay regularization is set to 0 .
A. A BLATION
For the ablation studies, we use a backbone we denote ViT-XS which refers to models which only have 5 layers and have attention-specifc hyperparameters adjusted to reach a parameter count of 16M. Details of these models are described below.
• For FNO, we use a hidden size of 64 and ﬁve layers to make the parameter count 16M. • For AFNO with Static Weights, we use a hidden size of 124, four blocks, and a sparsity theshold of 0.01.
• For AFNO-35%, we hard threshold and only keep the bottom 35% frequencies.
In practice, this means we keep 32/56 frequncies of the tokens along each spatial dimension.
We use a hidden size of 124, four blocks and no sparsity theshold.
• For AFNO, we use a hidden size of 584, four blocks, and a sparsity theshold of 0.01.
These models are inpaint pretrained on a randomly chosen subset of only 10% of ImageNet-1K and trained for 100 epochs.
For ﬁnetuning, we use the same setup as the few-shot segmentation experiments described in A.3.
We only evaluate on the CelebA-Faces dataset.
A. C OMPARISON WITH DIFFERENT TRUNKS AT DIFFERENT SCALES
In order to provide more extensive comparison with the state-of-the-art efﬁcient transformers we have included experiments for different trunks at different scales.
Since the primary motivation of this work is to deal with high resolution vision, we focus on the task of Cityscapes semantic segmentation a the benchmark that is a challenging task due to the high 1024 × 2048 resolution of images.
For the trunks we adopt: i) the Segformer Xie et al.
backbones B0, B1, B2, B3, under three different mixers namely AFNO, GFN and efﬁcient self-attention; ii)
Swin backbones Liu et al.
(2021b)
(T, S, B), and iii)
ResNet He et al.
and MobileNetV2 Sandler et al.
Numbers for ResNet and MobileNetV2 are directly adopted from Xie et al.
For training Segformer we use the same recipe as discussed in Section A.2 which consists of pretraining on ImageNet1K classiﬁcation for 300 epochs.
For training Swin, we use pretrained classiﬁcation checkpoints available from the original authors and then combine Swin with the Segformer head Xie et al.
We use the same training recipe as the Segformer models for Swin.
The mIoU scores are listed in Fig.
6 versus the parameter size.
It is ﬁrst observed that AFNO outperforms other mixers when using the same Segformer backbone under the same parameter size.
Also, when using AFNO with the hierarchical segformer backbone, it consistently outperforms the
Swin backbone for semantic segmentation.