SOFT: Softmax-free Transformer with Linear Complexity
Authors: Jiachen Lu
Jinghan Yao
Junge Zhang
Xiatian Zhu
Hang Xu
Weiguo Gao
Chunjing Xu
Tao Xiang
Li Zhang
Venue: NeurIPS 2021 Spotlight (2021)
Keywords: Transformer

Abstract: Vision transformers (ViTs) have pushed the state-of-the-art for various visual recognition tasks by patch-wise image tokenization followed by self-attention. However, the employment of self-attention modules results in a quadratic complexity in both computation and memory usage. Various attempts on approximating the self-attention computation with linear complexity have been made in Natural Language Processing. However, an in-depth analysis in this work shows that they are either theoretically flawed or empirically ineffective for visual recognition. We further identify that their limitations are rooted in keeping the softmax self-attention during approximations.  Specifically, conventional self-attention is computed by normalizing the scaled dot-product between token feature vectors. Keeping this softmax operation challenges any subsequent linearization efforts. Based on this insight, for the first time, a softmax-free transformer or  SOFT is proposed. To remove softmax in self-attention,  Gaussian kernel function is used to replace the dot-product similarity without further normalization. This enables a full self-attention matrix to be approximated via a low-rank  matrix decomposition. The robustness of the approximation is achieved by calculating its Moore-Penrose inverse using  a  Newton-Raphson method. Extensive experiments on ImageNet show that our SOFT significantly improves the computational efficiency of existing ViT variants. Crucially, with a linear complexity, much longer token sequences are permitted in SOFT, resulting in superior trade-off between accuracy and complexity.

SOFT: Softmax-free Transformer with Linear Complexity
Jiachen Lu
Jinghan Yao
Junge Zhang'
Xiatian Zhu?
Hang Xu
Weiguo Gao
Chunjing Xu
Tao Xiang"
Li Zhang'
Fudan University
Universitty of Surrey
Huawei Noah's Ark Lab
https: http///uaan//////ee/ea zvg github..b......................b 10/SOFT
Abstract
Vision transformers (ViTs) have pushed the statc-of-the-art art for various visual recog- nition tasks by patch-wise image tokenization followed by self-attention. However, the employment of self-attention modules results in quadratic complexity in both computation and memory usage. Various attempts on approximating the self- attention computation with linear complexity have been made in Natural Language Processing. However, an in-depth analysis in this work shows that they are either theoretically flawed or empirically ineffective for visual recognition. We further identify that their limitations are rooted keeping the softmax self-ultention during approximations. Specifically, conventional self-attention is computed by normaliz- ing the scaled dot-product between token feature vectors. Keeping this softmax operation challenges any subsequent linearization efforts. Based on this insight, for the first time, a softmax-free transformer or SOFT is proposed. To remove soft- max in self-attention, Gaussian kernel function is used to replace the dot-product similarity without further normalization. This enables a full self-attention ma- trix to be approximated via low-rank matrix decomposition. The robustness of the approximation is achieved by calculating its Moore-Penrose inverse using Newton- Newton-Raphson method. Extensive experiments ImageNet show that our SOFT significantly improves the computational efficiency of existing ViT variants. Crucially, with linear complexity, much longer token sequences are permitted SOFT, resulting in superior trade-off between accuracy and complexity.
Introduction
Recently the step change brought by Transformers [33] in natural language processing (NLP) [10, 4] 411,41, 41, 47, 46]. Indeed, with less inductive bias in its architecture design than Convolution neural networks (CNNs), pure Vision Transformer (ViT) [11] and its variants have shown to be able to outperform CNNs on various vision tasks [8, 15]. However, there is a bottleneck in any Transformer based model, namely its quadratic complexity in both computation and memory usage. This intrinsic to the self-attention mechanism: given a sequence of tokens (e.g., words or image patches) as input, the self-attention module iteratively learns the feature representations by relating one token to all other tokens. This results in quadratic complexity O(n²) with the token sequence length n in both computation (time) and memory (space) since an sized attention matrix needs to be computed and saved during inference. This problem is particularly acute in vision: 2D image after tokenization will produce far longer sequence than those in NLP even with moderate spatial resolution. This quadratic complexity thus prevents a ViT model from modeling images high spatial resolutions, which are often crucial for visual recognition tasks.
seems to have arrived in vision [11,
"Li Zhang (lizhangfd @ fudan. .al.aanndducn is the corresponding author with School of Data Science, Fudan University.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Imagenet Top-1 Accuracy (56
84
SOFT-Smal)
ovr Bogberr
Medium
LoFg
14000
Linformer
Transfomer
Nystromformer
PVT-Wedium
Del 5
Performer
SOFT
SorT-Tlny
Memory (MBB
78
ResMet-50
SAX18
26
8000
.VVT- kry
6000
4000
70
2000
50
784x 784 784 4 784 784 x
Parameter (M)
Token sequence length
(a)
(b)
Figure I: Topl-Accuracy on ImageNet [9] validation set with respect to parameters and the memory
usage corresponding to the token sequence length in practice compared to other methods. (a) Com- parison with CNN models: RegNet [26], ResNet [13] and Transformer models: PVT [35], DeiT [31], VIT [11]. T2T- VIT [41], Twins-SV [6] and SANIO [45]; (b) Comparison with Transformer [33], Linformer [34], Nyströformer [39] and Performer [5]. The memory usage is measured with a batch
size of on 16GB Tesla V100.
A natural solution is reduce the complexity of self-attention computation via approximation. Indeed, there have been a number of attempts NLP [34, 5, 18, 39]. For example, [34] takes a naive approach by shortening the length of Key and Value via learnable projections. Such coarse approximation would inevitably cause performance degradation. In contrast, [5, 17] both leverage the kemel mechanism approximate softmax normalization to linearize the computation in self-attention. [18] instead adopts a hashing strategy to selectively compute the most similar pairs. Recently, [39] uses Nyström matrix decomposition reconstruct the full attention matrix with polynomial iteration for approximating the pseudo-inverse of the landmark matrix. Nonetheless, softmax normalization simply duplicated across the matrix decomposition process, which is theoretically unsound. We empirically found that none of these methods are effective when applied to vision (see Sec. 4.2).
In this work, identify that the limitations of existing efficient Transformers caused by the of softmax self-attention, and for the first time propose a softmax-free Transformer. More specifically, in all existing Transformers (with or without linearization), softmax normalization is needed top of scaled dot-product between token feature vectors [33]. Keeping this softmax operation challenges any subsequent linearization efforts. To overcome this obstacle, introduce novel softmax-free self-attention mechanism, named as SOFT, with linear complexity O(n) in both space and time. Specifically, SOFT uses Gaussian kernel define the similarity (self-attention) function without the need for subsequent softmax normalization. With this softmax-fre attention matrix, we further introduce novel low-rank matrix decomposition algorithm for approximation. The robustness of the approximation is theoretically guaranteed by employing a Newton-Raphson Raphson method for reliably computing the
inverse of the matrix.
We make the following contributions. (I) We introduce a novel softmax-free Transformer with linear space and time complexity. (II) Our attention matrix approximation is achieved through novel matrix decomposition algorithm with theoretical guarantee. (III) To evaluate our method for visual recognition tasks, design family of generic backbone architectures with varying capacities using SOFT as the core self-attention component. Extensive experiments show that with a linear complexity (Figure lb), SOFT models can take in input much longer image token sequences. As result, with the same model size, our SOFT outperforms the state-of-the-aa CNNs and ViT variants on ImageNet [9] classification in the accuracy/complexity trade-off (Figure la).
Related work
Vision Transformers There is a surge of research interests recently in exploiting Transformers for visual recognition tasks [36, 35, 41, 31, 44], inspired by their remarkable success in NLP [33, 10, 4].

---

Core to these NLP and vision transformers is the same self attention mechanism [33] that computes a selffattention matrix by exhaustively comparing token pairs. This means a quadratic complexity with the sequence length in both space and time, which thus limits the scalability of Transformers in dealing with long sequences. This limitation is more serious in vision than NI.: To process an image with at least thousands of pixels, patch-wise tokenization is must for Transformers to control the computational cost. Given higher resolution images, the patch size also needs to be enlarged proportionally sacrificing the spatial resolution. This limits the capability of Transformers, e.g., learning fine-grained feature representation as required in many visual recognition tasks.
Linear Transformers Recently, there have been a number of linear/efficient variants [5, 34, 17, 18, 30, 24, 16] of Transformers in NLP. For example, [34] learns to shrink the length of Key and Value based on low-rank assumption. [18] adopts a hashing strategy to selective the most similar pairs and only compute attention among them. [5, 17] utilize different kernel functions for approximating softmax-based self-attention matrix. [24] applies random feature mapping on the sequences to approach the original softmax function. [16] decreases the time and memory consumption of the attention matrix by replacing the softmax function with its linear-complexity recurrent alternative. When applied to visual recognition tasks, however, show that these models have considerable performance degradation compared to the standard Transformers [33] (see Sec. 4.2).
The most related work to SOFT is [39] which uses the Nyström matrix decomposition to avoid computing the full attention matrix. However, this method suffers from several theoretical defects: (1) As the standard self-attention needs to apply row-wise softmax normalization on the full attention matrix, a direct application of matrix decomposition infeasible. As a workaround, softmax is simply applied to all the ingredient matrices in [39]. Such approximation ís not guaranteed theoretically. (2) With a polynomial iteration method, it is not guaranteed that the generalized attention matrix inverse can be computed when the matrix is nearly singular one in practice. In contrast to all the above methods, in this paper we propose a softmax- self-attention mechanism that facilitates matrix decomposition for complexity minimization with theoretical guarantees.
3 Method
3.1 Softmax-free self-attention formulation
A schematic illustration of our model is given in Figure 2. Let's first look at our attention module design. Given sequence of n tokens X Rnxd with each token represented by a d-dimensional feature vector, self-attention [33] aims to discover the correlations of all token pairs exhaustively.
Formally, X is first linearly projected into three dimdimenso spaces (query, key, and values) as:
Q = == E R" R'xd, K X E Raxd,
V= XW, E Rxdd,
(I)
where WG, WR, W. E Rdxd, are learnable matrices. Self-attention can expressed in a generic formulation as:
Vi. Σ a(i.,, Kj+) o V3.,
where is the Hadamard product, and i,j {1..... n. index the tokens. The key self-attention function Rd x Rd -> IR is composed of nonlinear function 3 IR. I and relation function 7 Rdd x Rdd > R. A dominant instantiation of a is the scaled dot-product based softmax self-attention [33], defined as
1
8(-) softmax()), ((Qi,,
QT,K,
(3)
VdT
Whilst this softmax self-attention has been the de facto choice and seldomly questioned, as discussed earlier it is not necessarily suited for linearization. To facilitate the design of linear self-artention, we introduce a softmax-free self-attention function with the dot- dot-product replaced by a Gaussian kernel as:
g'(-) = exp(.), 2'i. Kz.) = TT 11Q - K2
(4)
3

---

Multi-head concat.
Image Embedding
Gaussian Kernel
P.E.
Sampling
Gaussian Kernel
Nowton-Raphson
(c)
HoM kd
Iow dh
IXMI Ioow
haw x IhX
IRW d
Gaussian Kernel
Feed
HXW
Inxe HaX
Forward
h
(b)
(8)
(d)
Figure 2: Schematic illustration of the proposed softmax- free self-attention (SOFT) method. P. E.: Position embedding. Dash lines: linear projection. dh: the hidden dim of each attention head. denotes the matrix dot product.
To preserve the symmetric property of attention matrix as in Eq (3), we set the project matrices W and We Eq (1) identical (i.e., Q K). Our self aelf-attention matrix then written
Suj exp (-
2 IQQ..
K,,:11
(5)
For notation simplicity, we define the matrix formulation as: S exp k).
Remarks Our self-attention matrix has three important properties: (1) is symmetric; (2)
the elements lie in unit range of [0, 1]; (3) All diagonal elements hold the largest value 1 (self- reinforced), with the bottom ones (corresponding
most dissimilar token pairs) being close As Gaussian kernel is positive definite kernel [12], S is deemed Gram matrix. However, we find that when using our kernel-based self-attention matrix S without lincarization, the training of transformer fails to converge. This might explain why softmax dot-product based self-attention [33] so popular in vanilla transformers.
3.2 Low-rank regularization via matrix decomposition with linear complexity
To solve the convergence and quadratic complexity problems, we leverage matrix decomposition as a unified solution with low-rank regularization. In particular, we consider Nyström [38], which is originally low-rank matrix approximation algorithm. This enables our model's complexity to be reduced significantly without computing the full self-attention matrix S.
We make this choice because our S is positive semi-definite (i.e., a Gram matrix) without follow-up normalization which all necessary conditions for Nyström. contrast, [39] totally ignores these requirements, leading to theoretical flaw in its approximation.
To define the Nyström method formally, let express S exp (Q K) block matrix:
A
B
s
E Ruxn
(6)
C
where A Rmxm B E (m C ((n-m)x(n-m) with < Through Nyström
decomposition (see derivative details in Supplementary A.1), an approximation can be represented as:
A B
A P where P= B].
and At the Moore-PPnroos (a generalized) inverse of A.
Sampling In the standard Nyström formulation, A and B are sub-matrices of S obtained by randomly sampled tokens, denoted Q. We call the sampled Q bottleneck tokens. However,

---

Algorithm 1: SOFT: Softmax-free attention
Algorithm 2: NR: Newton-Raphsoe iteration
Input: Q E Raxds sampling function fs Sampling Q t f
Input: A E Rmxm, and T E Z a 2/|| I for k ffrrm 1 to T do
Initialize Ao aA;
exp e Q). P t exp( Q); + (NNR(A)P:
A 2AR-1 AA-1AA
end Output: Ay
Output: $
find empirically that random sampling is considerably sensitive the choice of m. We hence explore two additional options by leveraging the structural prior of visual data: (1) Using one convolutional layer with kernel size k and stride k to learn Q, and (2) Using average pooling with kernel size & and stride l to generate Q. For both, need to reshape Q to the form of I xd, Each slide of convolution or pooling produces a token. We set & according to the length of Q such that tokens be obtained. Our experiments show that convolution layer performs better in accuracy. We therefore use a convolution layer by default.
= Q. Given these m tokens, we then compute A and P as: As K identical to Q, we have
A = exp Ô e K),
P= exexp( @ K).
(8)
We finally obtain the regularized self-attention matrix $ of SOFT as:
$ = exp ( (Q (exp (ee exp (ēok).
(9)
leading Algorithm The low-rank regularization is conducted follows. For computing the attention score between any two tokens, we first correlate each of them with sampled tokens using our self-attention function (Eq (5)); With this correlation representation we then compute their similarity under the modulation of the generalized inverse of Q's 's correlation matrix. Similar as standard Nyström, our design associates the input tokens w.r.t. small space spanned by sampled tokens, giving a proper estimation of the original attention relationships subject to low-rank constraint. The correctness of this method is proved in Supplementary A.1.
Moore-Penrose inverse inverse is to use Singular Value Decomposition (SVD). Given A E Rmxm and its SVD form UVVT where U, V are m X m unitary matrices and Σ is a m X m diagonal matrix, the Moore-Penrose inverse of A is At VEtU*. Nevertheless, SVD is not friendly to the training process on GPU hence harming the model training efficiency. To solve this issue, we adopt the Newton-Raphson method. It is an iterative algorithm with the (k + 1) 1-tth iteration formulated given the previous iteration as:
accurate and commonly used way to calculate the Penorse
AAAAA and AD aA. ARt1 2AA
(10)
We now prove that A finally converges to Moore-Penrose inverse of Amxms if a is sufficiently small [3]. Theorem When a is sufficiently small, A+t1 2A, - AAAA, Aş converges to At.
Though a A/| which ensures good convergence behavior in Algorithm 2 (see more details in Supplementary A.2.1), in practice, we find that using an alternative form gives more stable training and faster convergence. Specifically, in || A n
M || S 1 where / equals to 0.5, we find the 29"h
smallest n; that holds this inequality. Then, we initialize a as a T
The following proposition comes with the proof of Theorem 1: Proposition 1 |A4 A All and || A decreases to 0 monotonously, if a is sufficiently small.
The detail of proposition is shown in Supplementary A.2.2. This ensures that our estimated inverse is sufficiently accurate for matrix decomposition, subject to that SOFT attention is regularized.
5

---

Methods
Complexity
Memory
Params
FLOPs
Throughpul (img/s)
Top-1 %
Transformer [33]
((n²)
19.0GBt
13M
3.9G
1073 3240
79.1
Linformer [34]
(n))
11.7GB
13M
1.9G
2767 / 3779
78.2
Performer [5]
(nn)
15.0GB
13M
2.2G
2037 / 3657
76.1
Nyströmformer [39]
((n)
17.2GB
13M
2.0G
1891 / 3518
78.6
SOFT
(
15.8GB
13M
1.9G
1730 / 3436
79.3
Table Comparison different linear/efficient transformer variants on ImageNet [9], based on our multi-stage Tiny configuration (see Table 2). The memory usage is measured with the batch size of 1024 which our standard training setting. Transformer is tested at batch size 256, which is the maximal number possible with the GPU resource at our disposal. Throughput is in format as Train throughput inference throughput.
Complexity We summarize the complexity of SOFT in space and time. For time complexity, involves: (1) Sampling: ((dd)). (2) Calculating three decomposed matrices: O - mnde - m²de) (2mmdd + m² ): (3) Moore-Penross inverse: O(T m²) O(Tm²), where T is the iteration steps. (4) All matrix multiplication: Onnm² + mnd, + mndd) = O + 2mmdd).
The total time complexity is O(dd + 4mde + m² ) Tm + d.m²). The space complexity decided by four decomposed matrices with O m.) O(m x m) + O(m x n) + O(n x de)
((2mm + de) ) + m²). As we keep m (m < n) fixed constant in our model, both time and space complexity are (nz), making SOFT a linear self-attention.
3.3 Instantiations
Figure 2 shows how our proposed softmax-free self-attention block (SOFT block) can be implemented in a neural network. We replace the self atff-attention block with our SOFT block the traditional Transformer, that is, stack SOFT block with feed forward residual block [11] to form softmax-free Transformer layer (SOFT layer).
Focusing on the general image recognition tasks, we integrate our SOFT layer into the recent pyrami- dal Transformer architecture [35] to form our final model SOFT. Further, several improvements are introduced in patch embedding (i.e., tokenization). Specifically, unlike [35] that uses a combination of non-overlapping convolution and layer normalization [1], we adopt a stack of overlapping convolu- tions, batch normalization [14] and ReLU non- non-lineeritty. Concretely, the STEM is implemented by units of 3x3 Conv->BN- to each of three following down-samplin operations with stride of 2 in the multi-stage architecture.
with the stride of 2. 1, respectively. Then, one such unit is applied
The architecture syper-parameters of SOFT are: d: the input channel dimension of SOFT layer. de: the embedding dimension of tokens in SOFT block. In practice, we set de d. h: the head number of SOFT block. d, the channel dimension of each head and d, = de/h. /h. n: the input token sequence length of SOFT block. the bottleneck token sequence length of SOFT block. sp. the sampling ratio token sequence length sampling, which the ratio between input token sequence length and the bottleneck token sequence length. e: the expansion ratio of the 2-layer feed forward block. In SOFT, for all the stages we set d/, = 32, e = 4 and m = 49, sp varies in each stage according to the input token sequence length. Table details the family of our SOFT configurations with varying capacities (depth and width).
4 Experiments
4.1
Setup
We evaluate the proposed SOFT on the ILSVRC-2012 ImageNet-1K IK dataset [9] with 1.28M training images and 50K validation images from 1,000 classes. Following the common practice, we train model on the training set and evaluate on the validation set. Metrics: For model performance, top-1 accuracy on a single crop is reported. To assess the cost-effectiveness, we also report the model size and floating point operations (i.e., FLOPs). Implementation details: We use the code base [37] with the default setting to train and test the models. Specifically, we use weight decay of 0.05 and 10 epochs of linear warm-up. We conduct 300 epochs training with an AdamW optimizer
Dataset:

---

Tiny
Small
Medium
Large
Huge
C33-BN-ReLU, 64-d
Stage
sp. 8x8, 64-d, 2-h
sp. 8x8,
sp. 8x8,
sp. 8x8,
sp. 8x8,
64-d, 2-h
64-d,
64-d, 2-h
64-d. 2-h
C31 131-BN-B ReeU,, 128-d
Stage 2
sp. 4x4, 128-d, 4-h
sp. 4x4,
sp. 4x4,
sp. 4x4, 128-d, 4-h
sp. 4x4,
x
x 3
x
x
x
128-d. d, 4-h
128-d, 4-h
128-d, 4-h
320-d or 288-d
Stage
sp. 2x2, 320-d, 10-h
sp. 2x2, 320-d, 10-h
sp. 2x2,
sp. 2x2,
sp. 2x2,
x
x
x 29
x 40
x 49
288-d, 9-h
320-d, 10-h
352-d, 11-h
C31-BN-ReI.U 512-d
Stage 4 w. cls token
sp. 1xl,
sp. 1xl, x
sp. 1xl,
sp. 1x1,
sp. 1xl, 512-d. 16-h
x
5
x
x
512-d, 16-h
512-d. 16-h
512-d. 16-h
512-d. 16-h
Table 2: Architecture specifications of SOFT variants. sp.: sampling ratio. -d: the hidden dimension. -h: the number of heads in the self-attention block. C33-BN-ReL.U: three 3x3 Conv-BN-ReL.U with the stride of 2, 1, respectively. C31-BN-ReLU: one 3x3 Conv-BN-ReLU with stride of
and decreasing learning rate with the cosine annealing schedule. During training, random flipping. mixup [43] and cutmix [42] are adopted for data augmentation. Label smoothing [28] is used for loss calculation. All variants are trained with batch size of 1024 on 32G NVIDIA V100 GPUs. We also implement our method using the Mindspore [22].
4.2 Comparison with existing linear Transformers
compare our method with three existing linear Transformer models: Linformer [34], Performer [5], Nyströmformer [39] in terms of model complexity and accuracy.
Two experimental settings are adopted. Under the first setting, for all methods we use the same Tiny (Table 2) architecture for fair comparison. That is, we replace the core self-attention block in SOFT with each baseline's 's own attention block with the rest of the architecture unchanged. Note that the spatial reduction module of [35] is special case of Linformer [34]. We set the reduction ratio to be identical to ours. With the same uniform sampling idea, we replace the 1D window averaging of Nyströmformer [39] (for NLP tasks) with 2D average pooling (for images). The downsampling ratio remains identical to ours. It is also worth mentioning that there no official code released for Reformer [18] and the local Sensitive Hash (LSH) module has strict requirements the length of input tokens. We thus do not include this method our comparison.
From Table we can make the following observations: (i) Linear Transformer methods substantially reduce the memory and FLOPs while maintain similar parameter size comparing to the Transformer on the Tiny architecture; (ii) Our approach SOFT achieves the best classification accuracy among all the linearization methods. (iii) Our inference speed on-par with other compared linear Transformers and our training speed is slightly slower than Nystromformer and both are slower than Performer and Linformer. Note that the slow training speed of our model mostly due the Newton-Raphen iteration which can only be applied sequentially for ensuring the accuracy of Moore-Penrose inverse. In summary, due to the on-par inference speed we consider the training cost increase price worth paying for our superior accuracy.
Under the second setting, we focus on the memory efficiency of SOFT against the baselines. Here we follow the ViT [11] network structure, stacking 12 attention layers with hidden dimension d = 384, heads h 12, bottleneck token sequence length m 49. Different attention blocks from the three linearized Transformer variants, Linformer [34], Performer [5], and Nyströmformer [39] are studied. For each Transformer variant, we adjust token sequence length n in linear increment. Specifically, we use token sequence length of 784 X where p = 1, 2, 3, 4, 5, 6, 7,8 and set batch size 1 to verify whether the memory consumption increases "quadratically" or "linearly". Figure 1b shows all compared transformer variants including our SOFT indeed have linear memory usage complexity. This is in contrast with the standard Transformer which cannot cope with long token sequences with quadratic complexity.

---

4.3
Comparison with state-of-the-art CNNs and ViTs
Model
Style
Resolution
M.S. Out.?
Params
FLOPs
Top-1 %.
ResNct-18 [13]
ConvNets
224
11M
1.9G
69.8
PVT-Tiny [35]
Transformers
224
13M
1.9G
75.1
Coat-Lite Mini [40]
Transformers
224
11M
2.0G
78.9
LambdaNets-50
Transformers
224
V
16M
78.9
SOFT-Tin
SOFT
224
V
13M
1.9G
79.3
ResNet-50 [13]
Convolution
224
25M
4.1G
78.5
PVT-Small [35]
Transformer
224
24M
4,0GF
79.8
Deit-Small [31]
Transformer
224
22M
4.6G
79.9
T2T-ViT;-14 [41]
Transformer
224
X
21M
5.2G
80.7
CPVT-Small [7]
Transformer
224
22M
79.9
Twins-SVT-S [6]
Hybrid
224
24M
3.7G
81.7
SOFT-Small
SOFT
224
24M
3.3G
82.2
ResNet-101 [13]
Convolution
224
44M
7.9G
79.8
PVT-Medium [35]
Transformer
224
44M
7.0G
81.2
ViT-Small/16 [11]
Transformer
224
X
48M
9.9G
80.8
SOFT-Medium
SOFT
224
45M
7.2G
82.9
ResNet-152 [13]
Convolution
224
60M
11.6G
80.8
PVT-Large [35]
Transformer
224
/
61M
10.1G
81.7
T2T-ViT;-24 [41]
Transformer
224
x
64M
13.2G
82.2
Bo-Nee--1-1110271
Hybrid
224
55M
82.8
SOFT-Large
SOFT
2212
64M
11.0G
83.1
CaiT-S36[32]
Transformer
224
88M
13.9G
83.3
Swin-B[22
Transformer
224
88M
15.4G
83.3
Twins-SVT-L.
Hybrid
2212
99M
14.8G
83.3
SOFT-Huge
SOFT
221
87M
16.3G
83.3
Table 3: Evaluation results ILSVRC-2012 ImageNet-1K [9] validation set. We report the results using the input size of 224x224 pixels center cropped from resized images with 256x256 pixels. M.S.Out. stands for whether the model is designed for multi-scale output. t: Corrected FLOPs by taking into account the cost of attention matrix multiplication overlooked in the origin paper.
We compare with state-of-the-an alternatives and report the top-1 accuracy on the ImageNet-Ie 1K validation set. FLOPs are calculated at batch size 1024. From Figure 1a and Table 3, the following observations are made: (i) Overall, VIT and its variants yield better classification accuracy over CNNs. (ii) We achieve the best performance among the recent pure vision Transformer based methods including ViT [1]] and DeiT [31], J. as well as the state-of-the-art CNN RegNet [26]. (iii) Our SOFT outperforms the most similar (in architecture configuration) Transformer counterparts PVT [35] at all variants. Since the attention module is the main difference, this validates directly the effectiveness of model. (iv) We also beat the latest VIT variants Twins [6] which is designed to address the efficiency limitation of ViT. We have done so with less parameters and fewer float point computation.
To gain some insights into how attention is learned using our SOFT and the alternatives, Figure shows the attention masks various compared models. For each model, we show the output from the first two attention heads. It is evident that SOFT exhibits robustness and versatility in capturing local and long distance relations among pixels. It is interesting note that, although SOFT trained on an object categorization dataset in ImageNet [9], seems to be able to learn both semantic concepts shared across instances the same category and instance specific features. For instance, in the bottom-right example of bird class, one attention head focuses on the black bird only, while the other attend to both birds in the image. More examples are shown in Supplementary A.4.
4.4 Ablation studies
Pyramidal architecture: Unlike the earlier non-pyramidal vision Transformers (e.g., VIT [11],, most recent pyramidal (multi-scale) Transformers (e.g., PVT [35]) use convolution layers to reduce

---

100
162
341
141
(4)
Figure 3: Comparing the attention heatmaps of a query patch (marked by the cross "+") against all the patches of an image, produced by (a) Transformer [33]. (b) Performer [5], (c) Nystromformer [39] and (d) Our SOFT. See Supplementary A.4 for more examples.
the spatial resolution (i.e., token sequence length) between stages. In this study, we ablate SOFT with pyramidal architecture (our default SOFT-Small), SOFT w/o pyramidal architecture and DeiT-S [31] (no pyramidal architecture either). We replace the Transformer layer with a SOFT layer to get SOFT w/o pyramidal architecture. Note all three variants have similar parameters and FLOPs. Table 5a shows that the conv-based pyramidal architecture clearly superior non-pyramidal design, and our non-pyramidal counterpart is even slightly better than DeiT-S [31] whilst enjoying linear complexity.
Bottleneck token sequence length: In this study, examine how the bottleneck token quence length m, sampled from n tokens, influ- ences the model's performance. We change the bottleneck token sequence length in all stages to 36, 49, 64, 81. Table shows that longer bottle- neck token would increase the memory cost and the computational overhead. m 49 seems to
0.8
|AА, - A|b|/|A||
give the best trade-off between the performance and computational overhead. The memory us- age is measured with the batch size of 128.
0.4
0.2
Token sampling: The sampling function in SOFT can assume different forms. Convolu- tion: The sequence Q Raxde is first reshaped to feature map H xW xd. r x convolution kernel with stride of is applied for downsam- pling, where r = /sp. The output channel size also kept and no bias used. At last, the feature map is reshaped back the sequence. Average pooling: using a randomly picked from tokens. Biased sampling: We pick m tokens with biased policy. Here, the first m tokens are picked. Table 4b shows that average pooling yields the best performance while with less computational overhead comparing to convolution. Biased sampling can miss the most salient samples, and there is no guarantee that random sampling can keep the uniformity of the chosen samples. This result thus justifies the choice of using average pooling in SOFT.
20
30
40
Iterations
Figure 4: Convergence analysis for the approxima- tion of the Moore-Penrose inverse on SOFT-Tiny.
kernel and r stride, where r /sp. Random sampling: m tokens are
Bottleneck
Memory
FLOPs
Top-1 %
Sampling methods
Params
FLOPs
Top- %
36
15.1GB
1.9G
79.0
Convolution
13.07M
2.0G
79.3
49
15.8GB
1.9G
79.3
Random sampling
12.96M
1.9G
79.3
64
16.9GB
2.0G
79.3
Biased sampling
12.96M
1.9G
79.0
81
18.5GB
2.1G
78.9
Average pooling
12.96M
1.9G
79.3
Table (a) Ablations on bottleneck token sequence length. (b) Ablations on sampling methods.

---

Overlapped convolution: We ablate SOFT with overlapped convolution (our default choice, same as many recent works) and SOFT with non-overlapped convolution our Tiny configuration. Table 5b shows that SOFT with overlapped convolution performs better than SOFT without overlapped convolution. Our non-overlapped convolution variant still outperforms the PVT [35] which also has the same non-overlapped convolution by clear margin.
Methods
Pyramidal?
Top-1 %
Methods
Overlapped?
Top- %
DeiT-S [31]
X
79.8
PVT [35]
X
75.1
SOFT
X
80.1
SOFT
X
77.4
SOFT
82.2
SOFT
79.3
Table 5: (a) Ablations on pyramidal architecture. (b) Ablations on overlapped convolution.
Newton- Raphson's convergence: We study how many iterations the Newron-Raphson method needs to converge when computing the Moore-Penrose inverse AT, We use ||AA& A Allp/ with (see Proposition 1) as the convergence metric to quantify the difference between A₁ and A'. Figure 4 shows that our approximation converges within 20 iterations across all stages.
4.5 Additional experiments NLP tasks
In this section, we evaluate our method against other linear counterparts on four tasks of the Long Range Arena (LRA) [29] benchmark e,g., Listops [23], byte-level IMDb reviews text classifica- tion [21], byte-level document retrieval [25]. and image classification on sequences of pixels [19].
Implementations. We use the Pytorch version of LRA [29] benchmark, provided by [39]. For the evaluation protocol, we strictly follow [29,39], We omit the Pathfinder(1K K) task as we cannot replicate the result of Nyströmformer [39]. For our SOFT,
39].
simply the average pooling with window size 4, stride 4 to sample the bottlenecks. We follow the configurations of [39], with 2 layers, 64 and 128 hidden dimension respectively, and 2 attention heads. The results in Table 6 shows that our SOFT outperforms both the standard and alternative efficient Transformers on three out of four tasks, well the average result.
Methods
Listops(2K)
Text(4K)
Retrieval(4K)
Image(
Avg.
Transformer [33]
37.10
65.02
79.35
38.20
54.92
Reformer [18]
19.05
64.88
78.64
43.29
51.47
Linformer [34]
37.25
55.91
79.37
37.84
52.59
Performer [5]
18.80
63.81
78.62
37.07
49.58
Nyströmformer [39]
37.15
65.52
79.56
41.58
55.95
SOFT
37.40
63.49
81.77
46.91
57.39
Table 6: Comparison of different lincar/efficient Transformer variants on Long Range Arena [29],
based on its official configuration. Our SOFT surpasses previous efficient methods on three tasks.
5 Conclusions
We have introduced a novel softmax-free self-attention (SOFT) mechanism for linearizing Trans- former's complexity in space and time. Unlike existing linear Transformers that aim approximate the conventional softmax based self-attention, SOFT employs Gaussian kernel based attention which eliminates the need for softmax normalization. This design enables full self-attention matrix to be approximated via low-rank matrix decomposition. The robustness of the approximation is achieved by calculating its Moore-Penrose inverse using Newton- Raphsso method. Extensive experiments show that SOFT yields superior trade-off in accuracy and complexity.
Acknowledgment This work was funded in part by Shanghai Municipal Science and Technology Major Projects (No.201 8SHZDZX01 and No. 1o.20215HH2DZX00 Foundation of China under Grant No. 1690013, 71991471 and the scientifi-ttecnnlooical innovation plan program of Universities guided by the Ministry of Education, China.
Mindspore, National Science
10

---

References
[1] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint. 2016. 6
[2] Irwan Bello. Lambdanetworks: Modeling long-range interactions without attention. preprint, 2021.
arXiv
[3] Adi Ben-Israel and Dan Cohen. On iterative computation of generalized inverses and associated projections. SIAM Journal on Numerical Analysis, 1966.
[4] Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. In NeurIPS, 2020. 1,2
[5] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreca Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. ICLR, 2021. 2, 3, 6, 7,9, 10
[6] Xiangxiang Chu, Zhi Tian, Yuqing Wang, Bo Zhang, Haibing Ren, Xiaolin Wei, Huaxia Xia, and Chunhua Shen. Twins: Revisiting the design of spatial attention in vision transformers.
arXiv preprint, 2021. 2, 8
[7] Xiangxiang Chu, Zhi Tian, Bo Zhang. Xinlong Wang, Xiaolin Wei, Huaxia Xia, and Chunhua Shen. Conditional positional encodings for vision transformers. arXiv preprint, 2021. 8
[8] Stéphane d' Ascoli, Hugo Touvron, Matthew Leavitt, Ari Morcos, Giulio Biroli, and Levent Sagun. Convit: Improving vision transformers with soft convolutional inductive biases. In ICML, 2021. 1
[9] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In CVPR, 2009. 2, 6, 8
[10] Jacob Devlin, Ming-Wös Wing Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In ACL, 2018. 1.2
[11] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR, 2021.
1, 2, 6, 8
[12] Gregory E Fasshauer. Positive definite kernels: past, present and future. Dolomites Reseurch Notes on Approximation, 2011.
[13] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In CVPR, 2016. 2,
[14] Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In ICML, 2015. 6
[15] Andrew Jaegle, Felix Gimeno, Andrew Brock, Andrew Zisserman, Oriol Vinyals, and Joao Carreira. Perceiver: General perception with iterative attention. arXiv preprint, 2021.
[16] Jungo Kasai, Hao Peng, Yizhe Zhang, Dani Yogatama, Gabriel Ilharco, Nikolaos Pappas, Yi Mao, Weizhu Chen, and Noah Smith. Finetuning pretrained transformers into mns. arXiv preprint arXiv:2103.13076 2021.
[17] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In /CML, 2020. 2,
[18] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In ICLR, 2020. 3, 7, 10
[19] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Citeseer, 2009. 10
11

---

[20] 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. 2021. 8
[21] Andrew Maas, Raymond E Daly, Peter T Pham, Dan Huang, Andrew Y Ng, and Christopher Pous. Learning word vectors for sentiment analysis. In Proceedings of the 49th annual meeting of the association for computational linguistics: Human language technologies, 2011. 10
[22] Mindspore. https: /www mindsporene...
7
[23] Nikita Nangia and Samuel R Bowman. Listops: A diagnostic dataset for latent tree learning. arXiv preprint, 2018. 10
[24] Hao Peng. Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. arXiv preprint arXiv:2103.02143, 2021. 3
[25] Dragomir R Radev, Pradeep Muthukrishnan, Vahed Qazvinian, and Amjad Abu-Jbara. The acl anthology network corpus. Language Resources and Evaluation, 2013. 10
[26] Iljaa Radosavovic, Raj Prateek Kosaraju, Ross Girshick, Kaiming He, and Piotr Dollár. Design- ing network design spaces. In CVPR, 2020. 8
[27] Aravind Srinivas, Tsung-Yi Yi Lin, Niki Parmar, Jonathon Shlens, Pieter Abbeel, and Ashish Vaswani. Bottleneck transformers for visual recognition. arXiv preprint. 2021. 8
[28] Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. Re- thinking the inception architecture for computer vision. In CVPR, 2016. 7
[29] Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. arXiv preprint, 2020. 10
[30] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. arXiv preprint, 2020. 3
[31] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. arXiv preprint, 2020. 2, 8, 9, 10
[32] Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou. Going deeper with image transformers. arXiv preprint, 2021.
[33] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NeurlPS, 2017. 1, 2, 4,
9, 10
[34] Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint, 2020. 2, 3, 6, 10
[35] 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, 2021. 2, 6, 7, 10
[36] Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He. Non-local neural networks.
In CVPR. 2018. 2
[37] Ross Wightman. Pytorch image models. pytorch-LLage-modd 1s, 2019. 6
https: httpss://qithub..com/vvvv/
[38] Christopher Williams and Matthias Seeger. Using the nyström method to speed up kernel machines. In NeurlPS, 2001. 4
[39] Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A nyström-based algorithm for approximating self-attention. In AAAI, 2021. 2, 4. 6, 7, 10
12

---

[40] Weijian Xu, Yifan Xu, Tyler Chang, and Zhuowen Tu. Co-scale conv-attentional image transformers. arXiv preprint, 2021. 8
[41] Li Yuan, Yunpeng Chen, Tao Wang, Weihao Yu, Yujun Shi, Zihang Jiang, Francis EH Tay, Jiashi Feng, and Shuicheng Yan. Tokens-to-token vit: Training vision transformers from scratch on imagenet. arXiv preprint, 2021. 1,2,
[42] Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. Cutmix: Regularization strategy train strong classifiers with localizable features. In ICCV, 2019.
[43] Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. arXiv preprint, 2017.
[44] Li Zhang, Dan Xu, Anurag Arnab, and Philip HS Torr. Dynamic graph message passing networks. In CVPR, 2020.
[45] Hengshuang Zhao, Jiaya Jia, and Vladlen Koltun. Exploring self-attention for image recognition. In CVPR, 2020.
[46] Sixiao Zheng, Jiachen Lu, Hengshuang Zhao, Xiatian Zhu, Zekun Luo, Yabiao Wang, Yanwei Fu, Jianfeng Feng, Tao Xiang, Philip HS Torr, and Li Zhang. Rethinking semantic segmentation from a sequence-fo-sequence perspective with transformers. In CVPR, 2021.
[47] Xizhou Zhu, Weijie Su, Lewei Lu, Bin Li, Xiaogang Wang, and Jifeng Dai. Deformable detr: Deformable transformers for end-to-ent object detection. In ICLR, 2021.
13


Reviews:
[{"review_id": "T7c3miaThhW", "reviewer": "Reviewer_M8ep", "summary": "This paper proposes a novel method named SOFT to calculate the token similarity for self-attention without softmax operation. The authors state that the computational complexity could be reduced to O(n) with SOFT. A family of backbones are designed and evaluated on ImageNet, which achieve SOTA performance. ", "questions": "", "limitations": "", "rating": 7, "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": "This paper follows the work \"Nyströmformer: A Nyström-Based Algorithm for Approximating Self-Attention\" and modifies its computation of attention score to fit in vision task. But this paper is of novelty. \n\nI have some concerns. \n\n- In 3.2 the authors state that the time complexity is O(m^2n + m^3) and the space complexity is O((2m+d_e)m + m^2). However, as m = n/k, which is linear w.r.t. n, the complexity could not be linear. The complexity could be reduced by k^2 or k according to the order of m. The paper also states that they could keep m as a fixed constant, but with image size increased, based on the method of average pooling or stride convolution, m would increase as well. \n\n- There are some writing problems. The complexity in line 167 lacks m as a multiplier. In table 1 there is no introduction to C33 and C31. \n\n- The convolution layers between stages are not included in many other transformers, and there is no ablation on these convolutional layers. \n\n- There are other works on ViT like botnet, swin, cvt e.t.c. and CNNs like efficientnet that could be add in table 3 for comparison. \n\n- As the complexity decreases with SOFT, there could be some experiments to show that the training and inference speed may increases. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "My concerns on the limitations of this paper have been presented in \"Main Review\". The authors do not discuss the potential negative societal impact of their work. Adding more related discussions will be helpful.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8 hours"}, "extra_text": {}}, {"review_id": "QGnb1ILmaxv", "reviewer": "Reviewer_aC3r", "summary": "This paper proposes a novel linear complexity self-attention module that\n(1)\tfirst replaces the softmax attention with the Gaussian kernel\n(2)\tthen apply Nystrom approximation to calculate the gram matrix with O(n) complexity.\n(3)\twhose bottleneck samples are calculated using average pooling or convolutional layer.\n\nThe proposed method demonstrates good top-1 accuracy on ImageNet classification with low space and time computational complexity.", "questions": "", "limitations": "", "rating": 7, "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": "Strength\n\nThe paper proposes a novel linear complexity self-attention architecture that applies kernel approximation technique to the attention weight.\n-\tConsiders the attention weight calculated by Gaussian kernel.\n-\tApproximate the gram matrix using Nystrom method.\n-\tPropose a bottleneck sampling scheme using average pooling and convolutional layer.\n-\tPropose to calculate the inverse matrix of the gram matrix w.r.t. bottleneck samples using Newon-Raphson method, which is suitable for GPU calculation.\n\nThe paper experimentally compares the proposed method with several self-attention with linear complexity. The proposed method demonstrates both good accuracy and low computation complexity.\n\nThe proposed method demonstrates comparable accuracy to the state-of-the-art methods on ImageNet dataset.\n\nThe author conducts several ablation study with respect to the size of bottleneck samples, sampling methods, and the number of iteration for Newton-Raphson method.\n\n\nWeakness\n\nIn the remarks, the author says that the model with Gaussian kernel weight without does not work well. Since the proposed method is derived as the approximation about the Gaussian kernel, I would like to see the discussion why the true weight does not work well while approximated weight shows good performance.\nRegarding this, I would like to see the approximation performance of the gram matrix with the proposed approximation.\n\nIt looks like Nystromformer is the method that tries to approximate the true softmax attention. Since softmax weight works well while non-approximated Gaussian weight does not work well, I would like to know the reason of this reversal in the performance after the approximation.\n\n\nAnother Comments\n\nIn Table (4) Random Sampling works the same performance as Average Pooling. Though the author prefers average pooling, I personally think random sampling is better since random sampling does not violate the permutation equivariance property of the whole the layer while average pooling depends on the position information.\n\nThe proposed Gaussian weighted self-attention may be regarded as the variant of kernel density estimation.\n\n\nEvaluation\n\nConsiders the theoretical and empirical performance of the proposed method, I recommend acceptance.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitation of the gaussian weight is briefly described in the remarks paragraph in L111~117.\nThere are no discussion aout the potential negative social impact. However, since it is the algorithmic paper, the possibility that the paper causes the negative impact may small.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "1meFn-wI9rx", "reviewer": "Reviewer_T2zb", "summary": "The paper substitutes softmax self-attention with a Gaussian kernel function. This allows to approximate the self-attention matrix with a low-rank approximation. The robustness of the approximation is supported by theoretical guarantees, while previous works that don't provide these guarantees. The proposed method is evaluated on ImageNet, arguably the most popular benchmark for evaluating computer vision models, and shows impressive results when compared against other Linear Transformer alternatives and even when compared against the standard (full) self-attention approach.", "questions": "", "limitations": "", "rating": 7, "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": "=== UPDATE AFTER DISCUSSION PERIOD ===\n\nIn addition to the strengths below, the auhors have clarified my questions and addressed satisfactorily most of my concerns. Thus, I'm advocating for accepting the paper (I'm increasing my score from 6 to 7).\n\n=== Strengths ===\n\n1) The paper addresses an important topic, i.e. efficient self-attention mechanisms, with numerous applications in computer vision, and natural language processing. \n\n2) The strong results shown on ImageNet accuracy, compared to other popular efficient attention alternatives, makes it a good candidate to be useful to practitioners and influential to future research.\n\n3) The method presented in the paper is built from strong theoretical foundations (assuming the proof of Theorem 1 is correct, see clarifications below).\n\n4) The paper carefully analyses the problems in many existing \"efficient\" (Linear) transformer alternatives. Mainly focusing on Nyströmformer which is the most similar existing approach.\n\n5) The experimental details are well detailed in section 4.1, which makes easier for other to build on and compare with the results of the paper.\n\n=== Weaknesses ===\n\n1) I find the definition of $n$ a little bit confusing in some parts of the text. $n$ is first defined as the \"sequence length\" (line 94), but then is referred to as \"token length\" (Figure 1b and sections 3.3 and 4.2). \n\n2) Assuming that in all cases $n$ is referring to the same concept (i.e. sequence length, or number of tokens per image), then some of the experiments don't make much sense. In particular the experiments where SOFT is used with the VIT architecture (i.e. Figure 1b). Lengths of 784 x p (with p=1...8) are evaluated, (in all cases using m=49?) but these lengths are unreasonable since the VIT paper uses patches of 32x32 or 16x16, which results in inputs of 49 or 196 tokens in total to the attention layer.\n\n3) In general, I think that a plot/table showing the optimal m (w.r.t. accuracy/time or accuracy/flops) for different sequence lengths should be reported. This is important because the value of m chosen is 49, which is already \"large\" for typical non-pyramidal transformers. Thus, in this regime it might be that the proposed method does not provide any significant speed-up. An alternative experiment would be to use the PVT architecture with different input resolutions (but same downscaling factors in each stage): e.g. 150,224,280,380,448,...\n\n4) Figure 1a shows accuracy vs. # of paramaters vs. accuracy. This plot seems of little relevance for the goal of this work. A plot showing accuracy vs. training/inference flops; or even better accuracy vs. training/inference time would be preferable. Accuracy and flops are reported in Table 3, but given that the proposed method uses an iterative algorithm to approximate the pseudo-inverse, and iterative algorithms are not ideal for modern accelerators (GPUs/TPUs), I would strongly recommend the authors to report runtime.\n\n5) One of the limitations of the presented of method, as described in lines 115-117, is that \"we find that when using our kernel-based self-attention matrix $S$ without linearization, the training of a transformer fails to converge\". Do you have an hypothesis that explains why this happens?\n\n6) All the experiments are limited to ImageNet recognition. It's unclear whether they run a single run for each presented result or the reported numbers are the average of multiple independent runs. In any case, I would recommend the authors to perform additional experiments (if time permits) on a different benchmark. For instance, the benchmark presented in the paper \"Long Range Arena: A Benchmark for Efficient Transformers\" (https://arxiv.org/pdf/2011.04006.pdf), which also highlights that existing \"efficient\" attention mechanisms are far from perfect, seems like a good benchmark to evaluate the proposed method.\n\n=== Clarifications ===\n\nI can't follow the proof of Theorem 1. In particular going from Eq. (18) in the appendix to Eq. (19).  Eq. (18) reads:\n$\n||A^\\dagger - A_{k+1}|| =  ||(A^\\dagger - A_{k})A(A^\\dagger - A_{k})||\n$\n\nLine 36 reads: \"And we left multiply A on the both sides of (17), then norm the equation\". Then Eq. (19) reads:\n$\n||A A^\\dagger - A A_{k+1}|| =  ||A(A^\\dagger - A_{k})A(A^\\dagger - A_{k})||\n$\n\nThis seems to imply that if $A$ is a symmetric semi-definite positve matrix and $||X|| = ||Y||$, then $||A X|| = ||A Y||$. However, this is false in general. See a counter example: X = [[1, 2], [3, 4]], Y = [[1, 3], [2, 4]], A = [[2, 1], [1, 2]] (apologies, it seems that matrices are not rendered well). $A$ is positive definite.\n$||X|| = ||Y|| \\approx 5.4649857$ but \n$||A X|| \\approx 15.42234 \\not\\approx 16.18217 \\approx ||A Y||$.\n\nI tried to use the fact that $||A B|| \\leq ||A||\\cdot||B||$ (in the case of matrix $p$-norms), but didn't get too far. Given that $||A|| > 0$ if $A$ is semi-definite positive:\n\n$\n||A^\\dagger - A_{k+1}|| =  ||(A^\\dagger - A_{k})A(A^\\dagger - A_{k})|| \\Leftrightarrow\n||A|| \\cdot ||A^\\dagger - A_{k+1}|| =  ||A||\\cdot||(A^\\dagger - A_{k})A(A^\\dagger - A_{k})||\n$\n\nThen:\n\n$\n||AA^\\dagger - AA_{k+1}|| \\leq ||A|| \\cdot ||A^\\dagger - A_{k+1}|| \\text{  ???  }\n||A(A^\\dagger - A_{k})A(A^\\dagger - A_{k})||\n\\leq ||A||\\cdot||(A^\\dagger - A_{k})A(A^\\dagger - A_{k})||\n$\n\nTo summarize, I need additional indications from the authors to follow the proof.\n\n=== Other comments ===\n\n Please, carefully review the paper looking for typos. Some that I found:\n- Main text, line 45: \"identity\" -> \"identify\".\n- Main text, line 200: \"AadmW\" -> \"AdamW\".\n- Appendix, equation 19: The first line ends with $|$ to denote the norm, it should end with $||$.\n- Appendix, line 41: \"multiple\" -> \"multiply\"\n\nSuggestion: $\\times$ is typically used to denote the cross-product, I would recommend using $\\cdot$ for denoting the product of a scalar by a matrix (Eq. 3, 4, 5) to avoid confusion, but feel free to ignore this suggestion if you strongly prefer otherwise.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The paper does not discuss any potential negative societal impacts of the work. The paper is more on the theory side of the theory-application spectrum of papers, but it targets (and tests the proposed method on) an image recognition task, i.e. ImageNet. In any case, I do not think that this work raises specific concerns that are not present in a large fraction of NeurIPS submissions, or in virtually all publications applying Machine Learning to Computer Vision.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "9"}, "extra_text": {}}, {"review_id": "1l5CwodlHw", "reviewer": "Reviewer_RD3k", "summary": "The paper proposes a softmax-free self-attention that replaces the previous dot-product between the query matrix and the key matrix with a Gaussian kernel. To facilitate the proposed softmax-free self-attention, this paper further proposes to decompose the self-attention matrix with Singular Value Decomposition. The robustness of the proposed low-rank decomposition can be theoretically proved by computing the Moore-Penrose inverse with the Newton-Raphson method.", "questions": "", "limitations": "", "rating": 7, "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": "Introducing kernels/Gaussian kernels to handle self-attention is not new [1]. However, this paper theoretically proves that it’s possible to use low-rank matrix decomposition to approximate such kernel function, which is quite novel and improves the self-attention mechanism. Since I am not an expert in Newton-Raphson methods, I have to kindly ask the other reviewers to verify the correctness of the equations. As for the other parts, the proposed pyramid architecture has outperformed the competitors by large margins in ImageNet classification task and the attention visualizations seem quite intuitive and reasonable to me.\n\nThe only question here is on the effectiveness of stacking overlapping convolutions and batch normalizations in L178-L179. It would be better to show some ablation studies comparing the tokenization methods using overlapping/non-overlapping convolution.\n\n[1] Peng, H., Pappas, N., Yogatama, D., Schwartz, R., Smith, N. A., & Kong, L. (2021). Random feature attention. ICLR 2021\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3 hours"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers unanimously agree that the presented softmax-free transformer and the kernel approximation techniques are novel and interesting. The paper also demonstrates good performance with reduced computation complexity on ImageNet. After the rebuttal, most of concerns from the reviewers are addressed. The authors are encouraged to further improve the paper by 1) comparing with other ViT works/architectures (https://arxiv.org/abs/2103.14030, https://arxiv.org/abs/2101.11605) 2) having more comprehensive discussion regarding related works with softmax-free attention (https://arxiv.org/abs/2103.02143, https://arxiv.org/abs/2103.13076, etc..) 3) discussing potential negative societal impacts of the work.
Decision: Accept (Spotlight)