Efficient Token Mixing for Transformers via Adaptive Fourier Neural Operators
Authors: John Guibas
Morteza Mardani
Zongyi Li
Andrew Tao
Anima Anandkumar
Bryan Catanzaro
Venue: ICLR 2022 Poster (2022)
Keywords: self attention
linear complexity
high-resolution inputs
operator learning
Fourier transform

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.

Published as a conference paper at ICLR 2022
ADAPTIVE FOURIER NEURAL OPERATORS: EFFICIENT TOKEN MIXERS FOR TRANSFORMERS
John Guibas Morteza Mortesa Mardani', Zongyi Andrew Tao',
Anima Aanandkumar": Bryan Catanzaro'
NVIDIA', California Institute of Technology?. Stanford University@
Stguibas@stanforddo edu, (mmardani, zongyil, atac, beat anzaro,
com
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)
efficient token mixer that learns mix in the Fourier domain. AFNO is based on a principled foundation of operator learning which allows to frame token mixing as continuous global convolution with- out any dependence on the input resolution. This principle was previously used design FNO, which solves global convolution efficiently in the Fourier domain and has shown promise in learning challenging PDEs. To handle challenges in vi- sual representation learning such as discontinuities 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 the channel mixing weights, adaptively sharing weights across to- kens, and sparsifying the frequency modes via soft-thresholding and shrinkage. The resulting model is highly parallel with quasi-linear complexity and has lin- ear 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-Bi backbone, AFNO can handle a sequence size 65k and outperforms other self-attention mechanisms. Code is available'
INTRODUCTION
Vision transformers have recently shown promise in producing rich contextual represen- tations 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 depen- dencies are crucial to understand the composi- tionality and relationships among the objects in a scene. A key component for the effective- ness of transformers is attributed to proper mix- ing of tokens. Finding a good mixer is how- ever challenging it needs to scale with the sequence size, and systematically generalize to downstream tasks.
80
79-
-9- ESA
A-10
GEN
76
Swin
MobileNer22
Recently, there has been extensive research to find good token mixers; see e.g., Tay (020b) and references therein. The origi- self-attention imposes graph structures, and uses the similarity among the tokens to capture
/5
MobiloVatV2
ResNet
2x10 3x xx10
10t
Poramater Count (Milloon)
Figure 1: Parameter count and mloU for Segformer, Swin, and other models different scales. AFNO con- sistently outperforms other mixers (see Section 5.7).
Joint first authors, contributed equally. The first author has done this work during internship NVIDIA, and the second author was leading the project." Code: httthh/wwwiiggaaaasa/ppieoeooeee

---

Published as a conference paper at ICLR 2022
L layers
FFT-based Global Conv.
GFNet
Channel Mixing
X
W
IFFT
(b, w. d)
(h, w, d)
Spatial Mixing
FNO
Multiply Frequency
X
Wi
Patch and Position Embedding
(h, W.
(h. w.
AFNO
FFT
X
$,O
d) k*MLP(d/k.d/k)
Figure 2: The multi-layer transformer network with FNO, GFN, and AFNO mixers. GFNet performs element- wise matrix multiplication with separate weights across channels (k). FNO performs full matrix multiplica- tion that mixes all the channels. AFNO performs block wise channel mixing using MLP along with soft- thresholding. The symbols h, d, and ls refer to the height, width, channel size, and block count, respectively.
Models
Complexity (FLOPs)
Parameter Count
Interpretation
Self-Attention
N 3Nd Nd + Nd N
3d
Graph Global Conv.
GFN
Nd
Depthwise Global Conv.
FNO
Nd + Ndlog N
Nd²
Global Conv.
AFNO (ours)
Nd /k No
(1+ 4/k)d 2 + 4d
Adaptive Global Conv.
Table 1:
Complexity, parameter count, and interpretation for FNO, AFNO, GFN, and Self-Attention. N hw, d, and refer the sequence size, channel size, and block count, respectively.
the long-range dependencies Vaswani et al. (2017); Dosovitskiy et al. (2020) It is parameter effi- cient and adaptive, but suffers from a quadratic complexity 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-shooe (LS) transformer aggregates long-range attention with dynamic projection to model distant correlations and a short-term attention to capture local correlations Zhu et al. (22T Long range dependencies 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 ct al. (2021); Lee-Thhrp et al. (202T). For instance, the Global Filter Networks (GFN) proposes depthwise global convolution for token mixing that enjoys efficient implementation in the Fourier domain Rao et al. (2021). GFN mainly involves three steps: (i) spatial token mixing via fast Fourier transform (FFT); (zt) 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 convolu- tion, which captures global relationships in the geometric space. One way to solve global convolu- tion efficiently is through FFT. More generally, nonlinearity such as ReLU to learn any general non-linear operator. This forms the basis for design- ing Fourier Neural operators (FNOs) which has shown promise in solving PDEs Li al. (2020a). We thus adopt FNO as a starting point for designing efficient token mixing.
compose such global convolution operations with
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 quadratic complexity in the channel size. To control this com-
2

---

Published as a conference paper at ICLR 2022
plexity, we impose block-diagonal structure the channel mixing weights. Also, to enhance generalization, inspired by sparse regression, we sparsify the frequencies via toft-thresholding Tib- shirani (1996). Also, for parameter efficiency, MLP layer shares weights across tokens (see Table 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-the- art, our AFNO using the ViT-B backbone outperforms existing GFN, LS, and self-attention for few- shot 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-B B3 backbone, AFNO achieves state-of-the-arr and beats previous methods, e.g. AFNO achieves more than 2% better mloU compared with efficient self-attention Xie et (2021) and is also competitive with GFN and LS.
Key Contributions. Our main contributions are summarized as follows:
We establish link between operator learning and high-resolution token mixing and adapt FNO from PDEs as an efficient mixer with quasi-linear complexity in the sequence length.
We design AFNO in a principled way to improve its expressiveness and generalization by impos- ing block-diagonal structure, adaptive -haring, 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 tatt--rr, 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 trans- formerCChid et al. (2019), image transformer] PareformerrPPrmarr et al. (2018), axial transformer et al. (2019), and longformer Beltagy al. (2020): (ti) low-rank attention that use linear sketching such lin- formers Wang et al. (2020). Long-shor transformers Lian et al. (2021), Nyströmformer Xiong et al. 2022 | (1xv) 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 at- tention Peng et al. (2021); and (žv) clustering- based methods such as reformer Kitaev et al. (2020),
routing transformer Roy al. (2021), and Sinkhorn transformer Tay al. (020a). These surrogates however compromise accuracy for efficiency.
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. (202T) achieves similar accuracy as self-attention. It is further accelerated by ResMLP Touvron al. (202T) that replaces the layer norm with the affine transforms. gMLP Liu et al. (2021a) also uses an additional gating 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 spatially mix tokens. FNet Lee-Thorp et al. (02T) 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 al. (02TT) 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 PDBs. Operator learning can be deployed in computer vision as images are RGB- RGB-valued functions on a 2D plane. This continuous generalization allows to permeate benefits from operators. Recent advances operator learning include DeepONet Lu et al. (2019) that learns the coefficients and basis of the operators, and neural operators Kovachki al. (202)) that are parameterized by integral
3

---

Published as a conference paper at ICLR 2022
operators. In this work, 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 2D image that is divided into ht x grid of small and non-overlapping patches. Each patch represented as a d-dimensional token, and the image can be represented as a token tensor X E pxxwxd Treating image token sequence, transformers then aim to learn contextual embedding that transfers well to downstream tasks. To end up with a rich representation, the tokens need be effectively mixed over the layers.
Self-attention is an effective mixing that learns the graph similarity among tokens. 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 kernel integration (Tsai al. 2019 Cao 2021
Kovachki et al., 2021). For the input tensor X we denote the (n, m) m) token as Zn.m I For notation convenience, we index the token sequence as X[s] == X [m m,) for some s, t E [hw]].
Define also N : hw as the sequence length. The self-attention mixing is then defined as follows:
Definition (Self Attention). Att: Nxd - Rxd
A == softmax (XW()
XW.
Vd
where Wy We, We Raxd are the query, key, and value matrices, respectively. Define K softmax (XW, XW))/vd) as N N score array with (:--) being inner product in Rd. We then treat self-attention as an asymmetrie matrix-valued kernel x: [N] [N] + Raxd parameterized
as k[s, K [s, f] W. (where [[s, is scalar valued and *." is scalar-matrix multiplication). Then the self-attention can be viewed as a kernel summation.
N
Att(x)[s) = Σ [[]s[sst]]
Vs [N].
Where [[t)x[[s t] [st] tWW, = ( This kernel summation can be extended to continuous kernel integrals. The input tensor X is no longer finitedimennsionl vector in the Euclidean space X E R xd but rather a spatial function in the function space E (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 (Kernel Integral). We define the kernel integral operator 1C: (D, Rd) - (D, RRd) as
(x))(s)= = -f, n((s di
E
with continuous kernel function D x D Raxd Liet al. Green's kernel x(s, () = *(s 1), the integral leads to global convolution defined below.
2020b). For the special case of the
Definition 3 (Global Convolution). Assuming k(s,t) x(s t), the kernel operator admits
(xX)(s) f0000=[[(0 1((s-1) (0000000n((001000000 de
Vs E D.
The convolution is smaller complexity class of operation compared integration. The Green's kernel has beneficial regularization effect but is also expressive enough to capture global interac- tions. Furthermore, the global convolution can be efficiently implemented by the FFT.

---

Published as a conference paper at ICLR 2022
3.2 FOURIER NEURAL OPERATOR AS TOKEN MIXER
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 the convolution theorem Soliman & Srinath (1990), global convolution in the spatial domain amounts multiplication in the eigen transform domain. popular example of such eigen functions the Fourier transform. Accordingly, one can define the Fourier neural operator (FNO) Li et al. (2020a).
Definition 4 (Fourier Neural Operator). For the continuous input X D and kernel k, the kernel integral at token s is found as
(x))(s)= = FF(( ((X))
Ws E D,
where denotes matrix multiplication, and F, F.F-1 denote the continous Fourier transform and its inverse, respectively.
Discrete FNO. Inspired by FNO, for images with finite dimension on discrete grid, our idea is to mix tokens using the discrete Fourier transform (DFT). For the input token tensor X E Rxxwxd define the complex-valued weight tensor W == DFT(s) E Chxwxdxd to parameterize the kernel.
FNO mixing then entails the following operations per token (m, n) E [h] x [w]
step (1). token mixing
Em,n [DFT(X)
step (2). channel mixing
8m,n W m. 2m,
step (3). token demixing
Ym.. = [IDFT(2) (2)
Local Features. DFT assumes a global convolution applied on periodic images, which is not typ- ically true for real-world images. To compensate local features and non-periodic boundaries, can add a residual term Fm.n (can also be parameterized as a simple local convolution) to the token demixing step 3 in FNO; see also Wen et al. (2021).
Resolution Invariance. The FNO model is invariant to the discretization h, w. It parameterizes the tokens function via Fourier bases which invariant the underlying resolution. Thus, after train- ing on one resolution it can be directly evaluated at another resolution (zero-shot super-resolution) Further, the FNO model encodes the higher- feeerr-rr information in the channel dimension. Thus, even after truncating the higher frequency modes 2m,n, FNO can still output the full spectrum.
is also important to recognize step (2) of FNO, where d weight matrix Wm. mixes the channels. This implies mixed-channel global convolution. Note that the concurrent GFN work Rao ct al. 202I) is a special case FNO, when Wm.n is diagonal and the channels are separable.
The FNO incurs O(N log( dd²)) complexity, and thus quasi-linear in the sequence size. The pa- rameter count is however ) as each token has its own channel mixing weights, which poorly scales with the image resolution. In addition, the weights m. are static, which can negatively impact the generalization. The next section enhances FNO to cope with these shortcomings.x
4 ADAPTIVE FOURIER NEURAL OPERATORS FOR TRANSFORMERS
This section fixes the shortcomings of FNO for images to improve scalability and robustness.
Block-Diagonal Structure on W. FNO involves d x weight matrices for each token. That results in d²) parameter count that could be prohibitive. To reduce the paramater count we impose block diagonal structure W, where it is divided into weight blocks of size d/k d/k. The kernel then operates independently on each block as follows
(()) Wl() (e) m.n mmn
t=
(5)
m,
The block diagonal weights are both interpretable and computationally parallelizable. In essence, each block can be interpreted a head as in multi-head self attenttennion, which projects into a subspace of the data. The number of blocks should be chosen properly so each subspace has sufficiently 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.
5

---

Published as a conference paper at ICLR 2022
def (x)
Tenaor [b,
bias
W_1, W_2
ComplexTent [ky d/k, 4/k]
RF'T22
b_2 Comp exmpeeeoesss (k, d/k]
X. reshape (b, h, w//211, 3, 6/3) BlockMLP (x)
x. reshhpp h, w/ w//21, d)
daf BlackKLP (x)
SortShrink (x)
MatMol (X, RL1) b_1
IRFF12 (x)
RelU
returr. bias
return Ma:Mul (x, R_2) + b_2
Figure 3: Pseudocode for AFNO with adaptive weight sharing and adaptive masking.
Weight Sharing. Another caveat with FNO is that the weights are static and once learned they will be adaptively changed for the new samples. Inspired by self-attention we want the tokens 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 two-layer perceptron that is supposed to approximate any function for a sufficiently large hidden layer. For (n, m-)tth token, it admits
mn ML.((mm.n) W2 (2 2m,.)) +
(6)
Note, that the weights W1, W2, b shared for all tokens, and thus the parameter count can be significantly reduced.
Soet-Thresholding and Shrinkage. Images 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 representing the important tokens. To sparsify the tokens, instead of linear combination as (5)
the nonlinear LASSO Tibshirani (1996) channel mixing as follow
min |mm W Wm, + Al|
This can be solved via soft- noft-hheesholding and shrinkage operation
Em, (W.mn Rm.
(8)
that defined as (x) sign max [x| - 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 modifications, the overall AFNO mixer module is shown in Fig along with the pseudo code Fig Also, for the sake of comparison, the AFNO compared against FNO, GFN, and self-attention in Table terms of interpretation, memory, and complexity.
5 EXPERIMENTS
We conduct extensive experiments to demonstrate the merits of our proposed AFNO transformer. Namely, 1) we evaluate the srreiinncy-acccurae trade-off between AFNO and alternative mixing mechanisms inpainting and classification 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 wide-range of datasets, including ImageNet-1k, CelebA-Facees LSUN-Cats, ADE-Cars, and Cityscapes as in
Deng al. (2009); Liu et al. (2015) Yu et al. (2015) Krause et al. (2013) Cordts al. (2016).
5.1 IMAGENET-IK K INPAINTING
We conduct image inpainting experiments which compare AFNO to other competitive mixing mech- anisms. The image inpainting task defined 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 Sigaal-toNoois 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.

---

Published as a conference paper at ICLR 2022
Backbone
Mixer
Params
GFLOPs
Latency(sec)
SSIM
PSN
VIT-B/4
Self- SeffAAtentionn
87M
357.2
1.2
0.931
27.06
VIT-B/4
LS
87M
274.2
1.4
0.920
26.18
VIT-B/4
GEN
87M
177.8
0.7
0.928
26.76
VIT-B/4
AFNO (ours)
87M
257.2
0.8
0.931
27.05
Table Inpainting PSNR SSIM for ImageNet- Ik validation data. AFNO matches the performance of Self-Attention despite using significantly less FLOPs.
Inpainting Results. PSNR and SSIM are reported in Table 2 for AFNO versus alternative mix- ers. It appears that AFNO is competitive with self-attention. However, AFNO uses significantly less GFLOPs than Self-Attention. Compared to both LS and GFN, AFNO acheives significantly better PSNR and SSIM. More importantly, AFNO achieves favorable downstream transfer, which is elaborated in the next section for few-shot segmentation.
5.2
FEW SHOT SEGMENTATION
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 valida- tion images from CelebA-Faces, ADE-Cars, and LSUN-Cats as in Zhang et al. (2021b). The model trained using cross-entropy loss. We measure mloU the validation set. More details about the experiments are deferred to the Appendix.
Backbone
Mixer
Params
GFLOPs
LSUN-Cats
ADÉ-Cars
CelebA-Faces
ViT-B/4
Self-Attention
87M
357.2
35.57
49.26
56.91
VIT-B/4
LS
87M
274.2
20.29
29.66
41.36
ViT-B/4
GFN
87M
177.8
34.52
47.84
55.21
VIT-B/4
AFNO (ours)
87M
257.2
35.73
49.60
55.75
Table 3: Few-shot segmentation mloU for AFNO versus alternative mixers. AFNO surpasses Self-Attention for datasets while using less flops.
Few-Shot Segmentation Results, Results are reported in Table 3 is evident that AFNO performs on par with self-attention. Furthermore, for out-of-domain datasets such as ADE-Cars or LSUN- Cats it slightly outperforms self-attention, which partly attributed to the sparsity regularization endowed in AFNO.
5.3 CITYSCAPES SEGMENTATION
To test the scalability of AFNO for high resolution images with respect to alternative mixers, we evaluate high-resolution (1024 x 1024) semantic segmentation for the Cityscapes dataset. We use the SegFormer-Bi BB backbone which is hicarchical vision transformer Xie et al. (202I). We train the model using the cross-entropy loss and measure performance via reporting mloU over the validation set. More details about the experiments and the model are available in the Appendix.
Backbone
Mixer SA
Params
Total GFLOPs
Mixer GFLOPs
mIoU
Segformer-B33 B3/4
45M
N/A
825.7
N/A
Segformer-B35 B3/4
Efficient SA
45M
380.7
129,9
79.7
Segformer B3/4
LS
45M
409.1
85.0
80.5
Segformer B3/4
GFN
45M
363.4
2.6
80.4
Segformer B3/4
AFNO-100% (ours)
45M
440.0
23.7
80.9
Segformer-B3/4
AFNO-25% (ours)
45M
429.0
12.4
80.4
Table 4: mloU and FLOPs for Cityscapes segmentation 1024 x 1024 resolution. Note, both the mixer and
total FLOPs included. For GFN and AFNO, the MLP layers the bottleneck for the complexity. Also, AFNO-25% only keeps 25% low frequency modes. while AFNO- 100% keeps all the modes. Results for self-attention cannot be obtained due to the long sequence length the first few layers.
7

---

Published as a conference paper at ICLR 2022
Cityscapes Segmentation Results. We report the final numbers for Cityscapes semantic segmenta- tion Table 4 AFNO-100% outperforms all other methods in terms of mloU. Furthermore, we find that the AFNO-25% model which truncates 75% of high frequeney modes during finetuning only loses 0.05 mloU and competitive with the other mixers. important to note that the majority of computations is spent for the MLP layers after the attention module.
5.4 IMAGENET-IK K CLASSIFICATION
We run image classification experiments with the AFNO mixer module using the VIT backbone on ImageNet- IK dataset containing 1.28M training images and 50K validation images from 1,000 classes at 224 x 224 resolution. We measure performance via reporting top-1 and top-5 valida- tion accuracy along with theoretical FLOPs of the model. More details about the experiments are provided in the appendix.
Backbone
Mixer
Params
GFLOPs
Top-1 Accuracy
Top-5 Accuracy
VIT-S/4
LS
16M
15.8
80.87
95.31
ViT-S/4
GFN
16M
6.1
78.77
94.4
VIT-S/4
AFNO (ours)
16M
15.3
80.89
95.39
Table 5: ImageNet-IK IK classification efficiencyy-accuraay trade-off when the input resolution is 224 x 224.
Classification Results. The classification accuracy for different token mixers are listed in Table 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 classification.
5.5 ABLATION STUDIES
We also conduct experiments to investigate how different components of AFNO contribute to per- formance.
25.65
25.60
PSNR
25.50
26.45
25.35
Sparsity Theshold (A)
Number of Blocks
Figure 4: Ablations for the sparsity thresholds and block count measured by inpainting validation PSNR. The results suggest that soft thresholding and blocks are effective
Sparsity Threshold. We vary the sparsity threshold À from 0 to 10. For each pretrain the network first, and then finetune for few-shot segmentation on the CelebA-Faces dataset. We report both the inpainting PSNR from pretraining and the segmentation mloU. The results are shown in Figure 3. X = 0 corresponds to no sparsity. It is evident that the PSNR/mIoU peaks at À 0.01, indicating that the sparsity 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 finetuning.
Number of Blocks. We vary the number of blocks used when impose structure W. To make the comparison fair, we simultancously adjust the hidden size so the overall parameter count of the model equal. We vary the number of blocks from to 64 and measure the resulting inpainting PSNR on ImageNet- K. is seen that 8 blocks achieves the best PSNR. This shows that blocking is effective.
weights
Impact of Adaptive Weights. We evaluate how removing adaptive weights and instead using static weights affects the performance of AFNO for ImageNet-I 1K inpainting and few-shot segmentation.
8

---

Published as a conference paper at ICLR 2022
The results presented in Table 6. The results suggest that adaptive weights are crucial to AFNO's performance.
Backbone
Mixer
Parameter Count
PSNR
CelebA-Facce mloU
ViT-XS/4
FNO
16M
24.8
39.27
ViT-XS/4
AFNO INon-Adaptive Weights]
16M
25.1
44.04
ViT-XS/4
AFNO [Hard Thresholding 35%]
16M
23.58
34.17
ViT-XS/4
AFNO
16M
25.69
49.49
Ablations for AFNO versus FNO, AFNO without adaptive weights, and hard thresholding. Results are on inpainting pretraining with 10% ImageNet along with few-show segmentation mlaU on CelebA- Faces. Hard thresholding only keeps 35% of low frequency modes. AFNO demonstrates superior performance for the same parameter count in both tasks.
Table 6:
Comparison to FNO. To show that AFNO's modifications fix the shorteomings of FNO for images, we directly compare AFNO and FNO on ImageNet- K inpainting pretraining and few-shot segmen- tation on CelebA-Faces. The results are also presented in Table 6. The results suggest that AFNO's modifications are crucial to performance in both tasks.
5.6 COMPARISON WITH DIFFERENT TRUNKS AT DIFFERENT SCALES
In order to provide more extensive comparison with the state-of-the-art efficient 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 segmen- tation the benchmark that is challenging task due to the high 1024 x 2048 resolution of images. For the trunks we adopt: i) the Segformer Xie et al. (2021) backbones B0, B1, B2, B3, under three different mixers namely AFNO, GFN and efficient self-attention (ESA); ii) Swin backbones Liu iet al. (2021b) (T, S. B), and iii) ResNet He et al. (2016) and MobileNetV2 Sandler et al. (2018). Results for LS and self-attention are not reported due instability issues with half-precision train- ing and quadratic memory usage with sequence length respectively. Note that efficient self-attention self-attention but with a sequence reduction technique introduced in Wang et al. (202T). It is meant to be a cheap approximation to self-attention. Numbers for ResNet and MobileNetV2 are directly adopted from Xie et al. 202T). For training Segformer we use the same recipe as discussed in Section A.3 which consists of pretraining on ImageNet-1K KK classification for 300 epochs. For training Swin, we use pretrained classification checkpoints available from the original authors and then combine Swin with the Segformer head Xie et al. (2021). We use the same training recipe as the Segformer models for Swin.
The mloU scores listed in Fig. versus the parameter size. It is first 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.
6 CONCLUSIONS
We leverage the geometric structure of images in order to build an efficient token mixer in compar- ison to self-attention. Inspired by global convolution, we borrow Fourier Neural Operators (FNO) from PDEs for mixing tokens and propose principled architectural modifications to adapt FNO for images. Specifically, we impose a block diagonal structure on the weights, adaptive weight shar- ing, and sparsify the frequency with soft-thres ooft-thresholding and shrinkage. We call the proposed mixer Adaptive Fourier Neural Operator (AFNO) and incurs quasi-linear linear complexity in sequence length. Our experiments indicate favorable accuracy-efficient competitive high-resolution segmentation compared with state-of-the-arr, 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. (202I).
trade-off for few-shot segmentation, and

---

Published as a conference paper at ICLR 2022
REFERENCES
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint 200.x::0004.05/500 2020.
Shuhao Cao. Choose a transformer: Fourier or galerkin. arXiv preprint arXiv:2 14995, 2021.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint 1.XIv:1904.100000 2019.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 14794, 2020.
Marius Cordts, Mohamed Omran, Sebastian Ramos, Timo Rehfeld, Markus Enzweiler, Rodrigo Benenson, Uwe Franke, Stefan Roth, and Bernt Schiele. The cityscapes dataset for semantic urban scene understanding. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3213-3223, 2016.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hi- erarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pp. 248-255 Icee, 2009.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929, 2020.
Gaurav Gupta. Xiongye Xiao, and Paul Bogdan. Multiwavelet- based operator learning for differen- tial equations. arXiv preprint 10t.:21099/34550 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp- 770-778, 2016.
Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans. Axial attention in multidi- mensional transformers. arXiv preprint 192i: 2019.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In International Conference on Ma- chine Learning, pp. 5156-5165. PMLR, 2020.
Nikita Kitaev, tukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint wrXiv:2001.0445/ 2020.
Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadeneshelli Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces. arXiv preprint arXhv:2108.08481, 2021.
Jonathan Krause, Michael Stark, Jia Deng, and Fei-Fei. 3d object representations for fine-grained categorization. In 4th International IEFE Workshop 3D Representation and Recognition (3dRR-13), Sydncy, Australia, 2013.
James Lee- Lee-Thorp, Joshua Ainslie, Ilya Eckstein, and Santiago Ontanon. Fnet: Mixing tokens with fourier transforms. arXiv preprint arXiv:2105.088 033844, 2021.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint arXiv:2010.08899 2020a.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Neural operator: Graph kernel network for partial differ- ential equations. arXiv preprint arXiv:2003.0344 2020b.
10

---

Published as a conference paper at ICLR 2022
Dongze Lian, Zehao Yu, Xing Sun, and Shenghua Gao. As-mlp: An axial shifted mlp architecture for vision. arXiv preprint arXiv:2107.0839) 2021.
Hanxiao Liu, Zihang Dai, David R So, and Quoe V Le. Pay attention to mlps. arXiv preprint arXiv:2105.08050 2021a.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. arXiv preprint arXiv:2103. 14030, 2021b.
Ziwei Liu, Ping Luo. Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
Lu Lu, Pengzhan Jin, and George Em Karniadakis. Deeponet: Learning nonlinear operators for iden- tifying differential equations based on the universal approximation theorem operators. arXiv preprint
2019.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International Conference on Machine Learning, pp- 4055- 4064. PMLR, 2018.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. arXiv preprint arXiv:2103.02143. 2021.
Yongming Rao, Wenliang Zhao, Zheng Zhu, Jiwen Lu, and Jie Zhou. Global filter networks for image classification. arXiv preprint Sxkiv:2177.0064: 2021.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient bonsedt-bbsed sparse attention with routing transformers. Transactions of the Association for Computational Linguis- tics, 9:53-68, 2021.
Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chieh Chen. Mo- bilenetv2: Inverted residuals and linear bottlenecks. Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 45 4510-4520, 2018.
Samir S Soliman and Mandyam D Srinath. Continuous and discrete signals and systems. Englewood Cliff',, 1990.
Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da-Cheng Juan. Sparse sinkhorn attention. In International Conference on Machine Learning, pp. 9438-9447, PMIR, 2020a.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey, arXiv preprint arXiv:2009.067. 2020b.
Robert Tibshirani. Regression shrinkage and selection via the lasso. Journal of the Royal Statistical 58(1):267-288) 58()))2
Sociery: Series B (Methodological),
1996.
Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Un- terthiner, Jessica Yung. Daniel Keysers, Jakob Uszkoreit, Mario Lucie, et al. Mlp-mixer: An all-mlp architecture for vision. arXiv preprint arXiv:2105.00660 2021.
Hugo Touvron, Piotr Bojanowski, Mathilde Caron, Matthieu Cord, Alaaeldin El-Nouby, Edouard Grave, Armand Joulin, Gabriel Synnaeve, Jakob Verbeek, and Hervé Jégou. Resmlp: Feedforward networks for image classification with data-efficient training. arXiv preprint arXiv:2105.03404, 2021.
Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philippe Morency, and Ruslan Salakhutdinov. Transformer dissection: unified understanding of transformer's attention via the lens of kernel. arXiv preprint arXiv:1908, 11775, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 1998-6008. 2017.

---

Published as a conference paper at ICLR 2022
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768, 2020.
Wenhai Wang, Enze Xie, Xiang Li, Deng-Ping Fan, Kaitao Song, Ding Liang, Tong Lu, Ping Luo, and Ling Shao. Pyramid vision transformer: A versatile backbone for dense prediction without convolutions. arXiv preprint arXiv:2102.12122 2021.
Gege Wen, Zongyi Li, Kamyar Azizzadenesheli, Anima Anandkumar, and Sally M Benson. U- fno-an an enhanced fourier neural operator based-deep learning model for multiphase flow. arXiv preprint arXiv:2109.03660 2021.
Enze Xie, Wenhai Wang. Zhiding Yu, Anima Anandkumar, Jose M Alvarez, and Ping Luo. Seg- former: Simple and efficient design for semantic segmentation with transformers. arXiv preprint trXiv:2/05.15203. 2021.
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nystr\," omformer: A nystr\ om-based algorithm for approximating self-attention. arXiv preprint arXiv:2/02.03902, 2021.
Fisher Yu, Yinda Zhang, Shuran Song, Ari Seff, and Jianxiong Xiao. Lsun: Construction of large-scale image dataset using deep learning with humans in the loop. arXiv preprint arXiv:1506.03365. 2015.
Yuxuan Zhang, Huan I Ling, Jun Gao, Kangxue Yin, Jean-Francois Lafleche, Adela Barriuso, Antonio Torralba, and Sanja Fidler. Datasetgan: Efficient labeled data factory with minimal human effort. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10145-10155, 2021a.
Yuxuan Zhang, Huan Ling, Jun Gao, Kangxue Yin, Lean-Francois Lafleche, Adela Barriuso, Antonio Torralba, and Sanja Fidler. Datasetgan: Efficient labeled data factory with minimal human effort. In CVPR, 2021b.
Chen Zhu, Wei Ping, Chaowei Xiao, Mohammad Shoeybi, Tom Goldstein, Anima Anandkumar, and Bryan Catanzaro. Long-short transformer: Efficient transformers for language and vision. arXiv preprint 1rXiv:2107.02192 2021.
12

---

Published as a conference paper at ICLR 2022
A
APPENDIX
This section includes visualizations of AFNO as well as the details of the experiments.
A.1 VISUALIZATION AFNO
To gain insight into how AFO works, we produce visualizations of AFNO's weights and represen- tations. In particular, we visualize the spectral clustering the tokens and sparsity masks from soft thresholding.
AFNO clustering versus other mixers. We show the clustering of tokens after each transformer layer. For 10-layer transformer pretrained with 10% of ImageNet-Ik, lk, apply spectral clustering on the intermediate features when we use k-NN kernel with k 10 and the number of clusters also set to 4. It appears that AFNO clusters are good 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
2 that show the superior performance of AFNO over the alternatives.
SA
GFN
LS
AFNO
Figure 5: Spectral clustering of tokens for different token mixers. From top to bottom, it shows the input and the layers 2, 4, 6, 8, 10 for the inpainting pretrained model.
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 Figure 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.
Figure 6: Log magnitude of tokens (56x56) after wef-thresholding and shrinkage (A 0.1) and the sparsity mask averaged over channels for inpainting pretrained network. Left to right shows layers to 5, respectively.
A.2 INPAINTING
We use the V/T-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 Super(Dosovitski al. (222)). 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 hidden size of 768.
13

---

Published as a conference paper at ICLR 2022
Long-Short Short transformer Zhu et al. (2021) (LS) uses window-size of 4 and dynamic pro- jection rank with a hidden size 768.
Long-
Global Filter Network Rao et al.
2021 (GFN) uses a hidden size 768.
Adaptive Fourier Neural Operator (AFNO) uses block, hidden dimension of 750, and sparsity threshold of 0.1, and a convolution layer as the bias.
The training procedure can be summarized as follows. Given an image I X from ImageNet-IK, we randomly mask pixel intensities to zero by initially sampling uniformly from (( and then apply the transition function T(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.
UniformSample( (( 1,j), (i +
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 minimum learning rate of 10 We use gradient clipping threshold of 1.0 and weight-decay of 0.01.
A.3 CITYSCAPES SEGMENTATION
We use the SegForme--B3 backbone for the Cityscapes segmentation experiments. The SegForme-- B3 backbone 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. (202I). 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].
Efficient 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] order to match the parameter count of the other networks. Because GFN is not resolution invaraint, we use bilinear interpolation the forward pass to make the filters match the input resolution of the tokens.
Long-Short uses 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-BB B3 backbone ImageNet- IK classification 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, batch size of 1024, gradient clipping threshold of 1 1.0, and learning rate warmup for 6250 iterations. We then finetune these models on Cityscapes for 450 epochs using learning rate of 1.2 10 We train on random 1024x crops and also evaluate 1024x1024.
A.4 FEW-SHOT SEGMENTATION
The models used for few-shot segmentation described in B.I. We the inpainting pretrained models and finetune them on few-shot segmentation on CelebA-Faces, ADE-Cars, and LSUN-Cats 224x224 resolution. The samples for the few-shot dataset selected as done in DatasetGAN in Zhang et al. (2021a). To train the network, we use the per- per-pixel cross entropy loss.
We finetune the models the few-shot datasets for 2000 epochs with learning rate of 10-4 for self-attention and 1 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.
14

---

Published as a conference paper at ICLR 2022
A.5 CLASSIFICATION
The models for classification are based on the Global Filter Network GFN-XS models but with 4x4 patch size. In particular, we utilize 12 transformer layers and adjust the hidden size and attention- specific 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 the last layer to produce output softmax probabilities for the 1, 000 classes in ImageNet- lk. More details about each of the models provided below.
Self-Attention Vaswani et al. 2017) uses 12 attention heads and a hidden size 324.
Long-Short transformer Zhu et (2021) (LS) uses window-size of 4 and dynamic pro- jection rank of 8 with 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, 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-entrope loss using the learning rate of (BatchSize/512) 10 for the models. We also five epochs of linear learning-rate warmup. and after cosine-decay schedule to the minimum value 10-5. Along with this, the gradient norm clipped not to exceed 1.0 and weight-decay regularization is 0.05.
A.6 ABLATION
For the ablation studies, have layers and have attention-specife hyperparameters adjusted to reach parameter count of 16M. Details of these models are described below.
backbone we denote ViT-XS which refers to models which only
For FNO, we use hidden size of 64 and five layers to make the parameter count 16M.
For AFNO with Static Weights, we use hidden size 124, four blocks, and 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 hidden size of 124, four blocks and sparsity theshold.
For AFNO, we use a hidden size of 584, four blocks, and sparsity theshold of 0.01.
These models are inpaint pretrained on a randomly chosen subset of only 10% of ImageNet-1 K and trained for 100 epochs. For finetuning, we use the same setup the few-shot segmentation experiments described in A.4. We only evaluate on the CelebA-Faces dataset.
15


Reviews:
[{"review_id": "ogTrM22nngO", "reviewer": "Reviewer_UFjk", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper proposes an alternative mechanism for self-attention in the context of vision transformers. The proposed method reduces the mixing complexity from quadratic to N log N, which could be important for long token sequences, e.g. high-resolution images.\n\nThe paper uses a well-established theory of Fourier transform motivating the proposal, which is neat.\n\nIt is hard from the paper presentation to compare with other efficient attention methods - the comparison is done using two metrics, accuracy and flops. Therefore it is not possible to judge if the paper achieved the state-of-the-art. \n\nThe novelty is somewhat limited. The differences with GFN, i.e. adding MLP and channel mixing, do not seem particularly insightful or significant.\n\nOne weakness is that, at least on the chosen applications, the quadratic complexity of attention is not a bottleneck of the networks. At least it is not clear from the paper. For example, table 4 suggests that the mixer takes at most 35% of total compute, with the rest consumed by the point-wise MLP. From this angle, the significance of the work is not convincingly demonstrated. \n\nI find it weird that the paper on vision transformers does not cite the original Dosovitsky et al. ", "summary_of_the_paper": "The paper proposes a new Adaptive Fourier Neural Operator (AFNO) for mixing tokens in visual transformers. The idea is based on Fourier neural operators (FNO) that transform feature flow in the Fourier space. The difference w.r.t. existing FNO is in two modifications. First, the weight matrix is block-diagonal (analog of multi-head) and, second, use of an MLP instead of just linear weighting. The experiments show that the proposed method is competitive and often achieves results as good as with original self-attention (with fewer flops).", "summary_of_the_review": "Despite limited novelty and significance, I found the work interesting and recommend acceptance.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "1: The main claims of the paper are incorrect or not at all supported by theory or empirical results.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "i5lz9WUlBuh", "reviewer": "Reviewer_PwZa", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The idea of this paper is novel and interesting. My main concern is its experiments: I think that the experiments are insufficient to support the conclusion and demonstrate the effectiveness of the proposed method. This paper should provide extensive experiments and comparisons for image classification, object detection, instance segmentation, and semantic segmentation. The ablation studies should also be conducted on the ImageNet dataset for image classification. These tasks are the most fundamental ones, which can thus reflect the ability of the represent learning of networks. \n\nHowever, this paper provides results for some uncommon tasks, like image inpainting and few-shot segmentation. Hence, I have reasons to guess that the proposed method is only effective for these carefully chosen tasks. Although this paper provides some results for semantic segmentation and image classification, the experiments for semantic segmentation are only conducted using Segformer-B3 and the experiments for image classification are only conducted using ViT-S. This seems that this paper carefully chooses one model for each task, and the comparison to state-of-the-art transformer models is also missing. Therefore, I think that current experiments are insufficient to demonstrate the effectiveness of the proposed method.\n", "summary_of_the_paper": "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.", "summary_of_the_review": "My main concern is about the insufficient experiments. Please see the above comments!", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "38LKkJHINMj", "reviewer": "Reviewer_H7fs", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "It's an interesting idea and a good attempt that aims to propose a solution to reduce the computing burden while maintaining the performance of visual transformer models. The motivation of introducing a Fourier Neural operator is reasonable to improve its expressiveness and generalization by imposing block-diagonal structure, adaptive weight-sharing, and sparsity. However, from my view, the main weakness of this study is lacking comparison experiments on general computer vision tasks with state-of-the-art visual transformer models such as image classification, object detection (these tasks have been included in another visual transformer model \"Liu, Ze, et al. \"Swin transformer: Hierarchical vision transformer using shifted windows.\" arXiv preprint arXiv:2103.14030 (2021).\"). For example, performance of swin transformer on imagenet 1k and 22k are both superior to this paper. According to the advantages shown in this study, it's better to be titled \"an improved transformer model for few-shot segmentation\" rather than the existing one which seems much more ambitious. When it comes to more details of this paper, the proposed methods perform very similar with or even worse than some existing ones while needing more network parameters (Table 4).\n", "summary_of_the_paper": "As is known, vision transformer has been becoming a more and more popular topic in the area of computer vision which is inspired by the success of this deep neural network fashion in other fields. However, the computing burden is a common disadvantage across the majority of transformer models compared with other deep neural networks. From a view of adaptive network operation in transformer, authors proposed to improve its expressiveness and generalization by imposing block-diagonal structure, adaptive weight-sharing, and sparsity.", "summary_of_the_review": "It's an interesting study that aims to propose a solution to a common challenge in the visual transformer model. However, I can't support its acceptance considering the presented results. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: Overall, this paper receives positive reviews. The reviewers find the technical novelty and contributions are significant enough for acceptance at this conference. The authors' rebuttal helps address some issues. The area chair agrees with the reviewers and recommend it be accepted at this conference.
Decision: Accept (Poster)