Skyformer: Remodel Self-Attention with Gaussian Kernel and Nystr\"om Method
Authors: Yifan Chen
Qi Zeng
Heng Ji
Yun Yang
Venue: NeurIPS 2021 Poster (2021)
Keywords: efficient transformers
attention
NLP applications

Abstract: Transformers are expensive to train due to the quadratic time and space complexity in the self-attention mechanism. On the other hand, although kernel machines suffer from the same computation bottleneck in pairwise dot products, several approximation schemes have been successfully incorporated to considerably reduce their computational cost without sacrificing too much accuracy. In this work, we leverage the computation methods for kernel machines to alleviate the high computational cost and introduce Skyformer, which replaces the softmax structure with a Gaussian kernel to stabilize the model training and adapts the Nyström method to a non-positive semidefinite matrix to accelerate the computation. We further conduct theoretical analysis by showing that the matrix approximation error of our proposed method is small in the spectral norm. Experiments on Long Range Arena benchmark show that the proposed method is sufficient in getting comparable or even better performance than the full self-attention while requiring fewer computation resources.


Skyformer: Remodel Self-Attention with Gaussian Kernel and Nyström Method
Yifan Chen; Qi Zeng, Heng Ji, Yun Yang
University of Illinois trbana-Champaign (yifanc10, qizeng2, hengji, yy84)eillinois.co edu
Abstract
Transformers are expensive to train due to the quadratic time and space complexity in the self-attention mechanism. On the other hand, although kernel machines suffer from the same computation bottleneck pairwise dot products, several approximation schemes have been successfully incorporated to considerably reduce their computational cost without sacrificing too much accuracy. In this work, we leverage the computation methods for kernel machines to alleviate the high computational cost and introduce Skyformer, which replaces the softmax structure with Gaussian kernel to stabilize the model training and adapts the Nyström method to a non-pon non-positive semidefinite matrix to accelerate the computation. We further conduct theoretical analysis by showing that the matrix approximation error our proposed method is small the spectral norm. Experiments on Long Range Arena benchmark show that the proposed method is sufficient in getting comparable or even better performance than the full self-attention while requiring fewer computation resources.
1 Introduction
The cost language model training increases exponentially, Among different models, Transformer- based language models (Vaswani et al. 2017 Devlin et al. 2019 Liu al. 2019 Lewis al.) 2020 are shown to enjoy state-of-the-art (SOTA) performances on many Natural Language Processing (NLP) tasks despite their enormous training cost. One of the computation bottlenecks lies in the self-attention mechanism, which is known to be resource-intensive with quadratic time and space complexity (O(n) where the input sequence length). Consequently, Transformers cannot support long sequence processing and large batch size with limited resources.
The challenge of improving computational efficiency of Transformers has motivated several recent studies on attention acceleration, using either sparse attention pattern [Qiu et al. 2020 Child et al. 2019 Zaheer al. 2020 Beltagy et al. [0200] Kitaev et al.) 2020) or low-rank approximation Choromanski et al. 2020 Wang et al. [020]. However, there is usually a lack of theoretical analysis
on the approximation error of these methods due the complex softmax structure, which makes the
theoretical comparison between the efficiency each method infeasible. It is also unclear in theory how to set the hyper-parrmetters of those methods attain desired level of approximation accuracy.
Another issue of Transformers is the training instability that small perturbations in parameter updates tend to be amplified, resulting in significant disturbances in the model output [Liu et al. 2020a]. Transformers on some NLP tasks have shown to be sensitive to hyper-parameters, learning schedulers, even random seeds, which usually demands time-costly grid search for the best configuration in real-world applications. It has also been observed in our experiments that a slight change in the learning rate may cause the failure of convergence for some models. We conjecture that the instability
Equal contribution.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

in Transformer training comes from the softmax structure, as the un-normalized attention score matrices before softmax tend have extremely large condition numbers due its fast singular value decay.
To alleviate the instability issue, an extra factor of 1/ 'D in the softmax kernel SM is suggested by Vaswani et al. [2017) to restrain the scale variation; Liu et al. [2020a] proposes a new scheme to control the magnitude of output change and stabilize the training in early stages. In practice, we also need to consider the lower numerical precision of GPU implementation in model training, which further deteriorates the stability.
Kernel methods may be the answer to both challenges. As pointed out by Choromanski et al. (000)]. the softmax structure is closely related to Gaussian kernels up to diagonal matrix multiplicttions, as the pairwise dot products naturally appear when expanding the squared & distance. We further notice some important connections between self-attention and Gaussian kernels. First, the un-normalized attention score matrix can be formed via basic matrix operations on an empirical Gaussian kernel matrix. Moreover, the form of Gaussian kernels has the natural interpretation of assigning "attention" to different tokens. Compared to the softmax function, Gaussian kernels automatically perform the normalization as softmax does (c.f. Section 4.1. These observations motivate us replace the softmax structure with Gaussian kernels. As we demonstrated in this paper, the new attention model, Kernelized Attention, empirically stabilizes the model training while being comparable to self-attentin attention in model accuracy.
To further improve the efficiency, propose Skyformer (Symmetrization of Kernelized attention for NYström method) to accelerate kernelized attention. Skyformer adapts the Nyström method Willliams and Seeger 2001 Drineas et al. 2005) to the non-PSD empirical Gaussian kernel matrix (as query matrices in general do not equal to key matrices), by instead lifting the kernelized attention score matrix into a large PSD matrix that contains the un-normalized attention score matrix as the off-diagonal block. We further conduct theoretical analysis by showing that Skyformer has a small matrix approximation error on kernelized attention in the spectral norm. Our experiments on the LRA benchmark show that Skyformer consistently uses less space and time while achieving better accuracy than other baseline methods.
In summary, our main contributions are:
(1) We revisit the intrinsic connection between self-attention and kernel methods, and explore new kernel-based structure, kernelized attention, to stabilize the training of Transformers.
(2) We propose Skyformer, which approximates the kernelized attention via low dimensional ran- domized sketches by adapting the Nyström method to non-PSD matrix. We provide the theoretical guarantee that the matrix multiplication error is small in term of spectral norm.
(3) Extensive experiments show that Skyformer achieves comparable performance to the original self-attention with fewer computational costs
2
Related Work
Among all the transformer acceleration methods, including attention layer simplification by pruning redundant attention heads [Voita al. 2019 Michel al. 2019] and model size reduction with knowledge distillation [Jiao et al. 2020 Tang et al. 2019 Liu et al. 2020b], we focus on attention approximation models, which are closely related to kernel methods.
To reduce the time and space complexity by avoiding exhaustive computation over the attention metric, recent studies propose to apply sparse attention patterns to limit the numbers of elements
participating in matrix multiplications [Qiu et al. [020 Child et al.) | Zaheer et al.) 2020 Beltagy
et al. 2020 Beyond limiting the attention to fixed patterns, some approaches learn the patterns by
determining token assignments to relevant groups (Kitaev et al., 2020 Roy et al., 2021 Those models utilize local and global information in the attention score matrix to perform approximation, which coincides with the attempt to accelerate the computation in Gaussian processes [Snelson and Ghahramani] 2007].
"Our code is released at http:://ggtttttccooo///////ggg//
com

---

The attention score matrix is known to exhibit a very fast rate of singular value decay [Bhojnapaalii
al., 2020 Dong et al. 202T], similar to that of an empirical kernel matrix (Yang et al., 2017].
This near singular property motivates many low-rank attention approximation methods to skillfully
leverage the computation techniques kernel methods. Among them, Linformer (Wang al..
2020] compresses the size of the key and value matrix with random projections based on the Linnson-Lindenstrauss transform, common randomized sketching method in Gaussian processes Yang et al.
2017]: Reformer [Kitaev et al. 2020) applies locality-sensitee hashing (LSH) [Har-Peled to simplify the computation the attention score matrix, which is widely used kernel
ct al. 2012
density estimation [Charikar and Siminelakis 2017 Backurs et al.) 2019]; Performer [Choromanski Rahimi et al. 2007
it al.
2020 projects both query and key matrix through random Fourier features
heavily exploiting Bochner Theorem for stationary kernels.
The most related papers to ours are linear attention (Katharopoulos et al. 2020 Synthesizer Tay et al. 2020a], and Nyströmformer (Xiong et al.) [202]. Linear attention takes the softmax structure self-attention as measure of similarity and replaces it with the dot product separately activated query and key matrices; Synthesizer aims modify the original self-attention by replacing the dot product before softmax with Synthetic Attention, which generates the alignment matrix independent of token-token dependencies. Their attempts indicate that the softmax structure in self-attention is not the only feasible choice, and justify our usage of kernelized attention. Rather than remodeling self-attention, Nyströmformer applies the Nyström method [Williams and Seeger 2001 Drineas et al., 2005], a powerful and effective method for large-scale kernel machines acceleration, to approximate the attention score matrix. However, Nyströmformer applies the Nystrõm method matrix, and thus fails to utilize the full potential of the Nyström method. This issue is resolved in proposed Skyformer by instead lifting the kernelized attention score matrix into large PSI) matrix which contains the target non-PSD matrix as its off-diagonal block. For more details on attention approximation methods, we refer readers a survey paper on efficient transformers [Ta et al. 2020c].
non-PSD
3 Preliminaries and notations
3.1 Revisiting self-attention
For given input sequence X E guxdo of length n and embedding dimension do, The dot-product attention for a single head in Transformer Vaswani al. 2017 is defined as
QKT
Attention( Q, K, V) softmax
V
vp
where Q XWe. K XWK, and XWr, and Wo. Wx and Wv are the query, key, and value weight metrics that linearly project the input X of do dimension to an output tensor of dimensions.
To simplify the future analysis, the left softmax term can be rewritten into D-1A, where A := exp((KK /vD) is the un-normalized attention score matrix; D is diagonal matrix whose diagonal exp(QKT/ /p) (by convention is a size-n vector with all elements being 1). Following the notation in Performer [Choromanski et al. 2020), we define SM(q, k) := exp(q" (//p""k//))) as the softmax kernel function, and represent A by the notation SM(Q, K), which means the element a.j from the i-th row and j-th column in A is equal to SM(q;, ky). Throughout this paper q; (resp. k;) means the i-th (resp. j-th) row in Q (resp. K).
We close this subsection with short lemma to show SM(..)) a positive semidefinite (PSD) kernel function Wainwright 2019 Definition 12.6] by relating it to Gaussian kemelss
Lemma 1. SM(:,.) :) is a PSD kernel function. Equivalently, for all integers n 2 1 and elements { in C RR, the n-by-n matrix C SM(Q, Q) is PSD.
Proof. We first state an important equation to connect the softmax kernel and Gaussian kernels as follows:
q 4;
11
|1g,
12
||,11² 1
SMigi, 43) exp
= exp
exp
exp
/P
2/P
2/P
2.P
3

---

The middle part exp et al. (020) have more discussion on the findings.)
2 [|q 9:11":
exactly Gaussian kernel with bandwidth p
Choromanski
Through this equation, can rewrite C as
Q Q p!/4
C = D2
D²²
(1)
11/4
where DQ is a diagonal matrix with elements (DQ) exp
| vp
Vi E [no], and k(g., ==
exp (-11g: qz /21) is the standard Gaussian kernel function.
Q p1/4 pi/a
We prove the lemma by using the fact that is PSD kernel and
is PSD matrix.
3.2 Nyström method
Due to the intrinsic low-rankness of an empirical kernel matrix B, the so-called Nyström method that replaces B with its low-rank approximation B. has been applied to accelerate kernel methods (Gittens and Mahoney 2016 Kumar et al. 2009 Williams and Seeger 2001]. Specifically, the Nyström approximation of B is the matrix B BS(s" BS) B, where (-) denotes the Moore-Penrose pseudoinverse of a matrix, and S E Rnd is zero-one sub-sampling matrix whose columns are subset of the columns in I, indicating which d observations have been selected. The formal definition of the uniform sub-sampling matrix is given as follows:
Definition (Uniform sub- sub-sampling matrix). For random matrix S E Raxd if has i.i.d. columns and the j-th column ( can randomly be v / with probability to2 where e, the i-th column of the n-by-n identity matrix In, then S is called uniform sub-sampling.
We close this subsection with a remark that it is not appropriate directly extend the Nyström method from kernel method to self-attention due to core requirement that B should be PSD with consideration of approximation performance improvement. We will show the next section how to address this challenge and properly adapt Nyström method to non-PSD matrices.
3.3 Approximation evaluation
Beyond the time and space complexity, attention acceleration methods have been mostly evaluated with empirical experiment results, such as the perplexity of pretrained language models and the fine- tuned performance downstream natural language understanding tasks. Specifically, Long Range Arena benchmark [Tay et al.] (020b) has been proposed to systematically evaluate the performance of efficient transformers with ten NLP tasks in long-context scenarios. However, such empirical results are indirect for theoretical analysis. Therefore, we introduce a common criterion used in matrix approximation, spectral norm, ease the future discussion on performance.
Definition 2 (Spectral norm guarantee for matrix approximation (MA)). Given a matrix M E
I Xm2, two constants 0,8 8 < 1, we say that its approximation matrix M E Raixma satisfies (e, 6)-M property for M, if
PP PIIM - Ml > =| < <8
(2)
In previous works, the direct analysis of the approximation error to the entire output D AV in the (c, -M manner is usually spared due to the difficulty caused by the complex softmax structure. In this paper, with the new kernelized attention, we are allowed to perform the analysis through the existing theoretical results in kernel methods. Consequently, in Section 4.5 we are able to give relatively precise error analysis on the approximation of Skyformer to the entire kernelized attention, which eases the future comparison with other methods approximating kernelized attention.

---

Method
4.1 Kernelized Attention
Kernelized Attention replaces the softmax structure in vanilla self-attention with a Gaussian kerel,, and the new attention model is stated as:
Q K V, Kernelized-Attentionnn (Q -) = p1/4" pz/4
(
(3)
where we define the n-by-n matrix C' as the kernelized attention score matrix K xQ/p1/4 (Q/p"////////)
The justification for using the kernelized attention model is as follows. A significant advantage of softmax attention is that tokens are allowed to attend to limited number of other important tokens in the sequence. We observe that Gaussian kernel function can play similar role. The expression of a Gaussian kernel is k(g,,k; k;) == exp (-1lg: k, k2)l Via this expression, for token in the query, Gaussian kernel assigns a large attention score the token j when k, is close to qi- The distance-based weight assignment is indeed considered as major reason why kernel methods are powerful. The form of kernelized attention also leads to an automatic normalization. Based on Equation ( the new attention model can be rewritten in terms of the un-normalized attention score matrix A
Kernelized-AAtteniion K. = (Da AA ADK
1/22 V,
where Do (resp. DR) diagonal matrix with elements (Dq):: exp
1q: vP
(resp. (Dx)a
[ Vi [n]. We remark the kernelized attention model can thus be formally taken as
exp
vp
variant of the original self-attention, which instead normalizes the matrix A in form of D-1 A. The intrinsic normalization allows kernelized attention to have more reasonable condition number than self-attention, which benefits the stability of model training. To demonstrate the improvement in stability, we additionally provide toy experiment in Appendix Fi which shows the "condition number" of kernelized attention smaller than self-attention. Moreover, empirical evaluation in Section 5 supports our claim that the new attention model can attain comparable performance the original attention model.
4.2 Skyformer: a modified Nyström method
Before jumping into details of Skyformer, we first propose a method to apply Nyström method to approximate an asymmetric (and thus non-PSD) empirical kernel matrix B constructed with any PSD kernel -(---). Specifically, with two different n-by-p design matrices Q and K, its element by from the i-th row and j-th column in B equal to of qi: k;), where qi (resp. k;) is the i-th (resp. j-th) row in Q (resp. K). We remark this type of empirical kernel matrices involves the un-normalized attention score matrix := SM(Q. K), and the empirical Gaussian kernel matrix C == x<(//p/44 K| /A). Therefore this method leads to a low-rank approximation to the output of either self-attention D AV Kernelized Attention CV. (D in self-attention can be obtained by computing 1, and thus a low-rank approximation to A also implies an approximation to D.)
Computational details are stated as follows. To tackle the challenge of approximating non-PSD matrix B, our first step is complete the matrix into PSD matrix B:
Q K
))
Q K
B
Then we approximate B with B through
B BS BS) B,
(5)
where S is a bn- uniform sub-sampling matrix as defined in Definition I The final approximation will be given as
B = (I,0)
(6)
5

---

The original matrix B can be well well-approximated by B due to the following inequality
B B| (B (), < B Bl|,
and thus we show our task of approximating the non-PSD matrix B boils down to well approximating the PSD matrix B.
Remark. The reason why B can well approximated by a low-rank B that as an empirical kernel matrix the cigenvalues in B usually decay fast, and thus there are many small cigenvalues in the long tail. In this case, theoretically a low-rank matrix (e.g. truncated singular value decomposition (SVD) of B) has enough potential to well approximate the original matrix B (and B accordingly) in terms of spectral norm.
With the derivation above, we officially introduce our proposed Skyformer as an approximation to Kemeliied Attention, which applies the modified Nyström method the kernelized attention score matrix C. The next two subsections will continue our discussion on it, and respectively state the theoretical analysis its approximation error and some details of its implementation in practice.
4.3 Error analysis of Skyformer
As mentioned, an implicit advantage of using Kernelized Attention is that we can leverage the existing conclusions for kernel methods analyze the theoretical properties of the model. In this subsection, we aim to provide some theoretical analysis of its approximation error.
We state high probability bound the size d of the sub- sampling matrix used in Skyformer to attain (e, - property for the kernelized attention score matrix C' by the following theorem. We refer the readers to the proof in Appendix to take a closer look at our claim that the matrix to be approximated should be PSD is a key to the theoretical guarantee of Nyström method.
Theorem 2 (Adapted from Lemma and Theorem 3 [Mssco and Musco 2017)) Consider the query. key, and value matrix Q, K, V E RRPP and two positive constants & < 1, < 2- For the empirical Gaussian kernel matrix C x(Q/p'/4, /4, ",A,KK/ K..KPpp (+)+) defined above, let X == =||C|| < ||C|
where C is the completion of C (similar to B, constructed as substituting the Gaussian kernel with bandwidth 22/1 for the arbitrary kernel function áb in Equation (4). We comment X serves as the regularization coefficient as well as the approximation error bound. To ease the analysis, we specifically define the i-th diagonal element of CCC + AI2n) as leverage score E,, Vi L.... 2n, and define their sum Tr (C(C + x))- as the statistical dimension datas, which increases with 1/e as X xx E. Suppose S is uniform sub sumpling matrix, and assume there exists constant 3 E (0, 1] such that 3 S huss Vi l..... 2n, For the approximation matrix C constructed with C, S as in
Int.
Equation 5 there exists a constant C such that if
d 2 C dotat
n log
8
then ė C&C Ō + with probability 1- 8. Here * denotes the Loewner ordering: B A means A- - B is positive semidefiitte. Farthermore, for our approximation Ö in Equation 6 the kernelized atlention score C, we have
IC-CIss>=ell
This theorem implies the time and space complexity of our proposed approximation depends on the statistical dimension dstat- If we directly use the conclusion from Gaussian kernels, dstat should be ((1) (complexity modulo poly-log term) [Yang et al. 2017 due to the exponential eigenvalue decay rate of Gaussian kernels, which is comparable to the complexity of most other efficient transformers. However, different than the case in the classical kernel methods, the distribution of the query and key matrix Q and K changes during the training procedure, which may invalidate the conclusion about dytat- We leave the exact non-asymptotic analysis of the computational complexity for future work.
4.4 Workaround in implementation
potential limitation with the implementation of the proposed method lies in the tricky fact that the matrix inversion on GPU is much slower and numerically less stable than the same operation on CPU

---

due to the different back-end libraries in the two platforms. We attempt to circumvent the problem by adapting the strategy Nystrômforner [Xiong et al.) 202T] to our setting. Specifically, we use the matrix-produme product--ased iterative method [Razavi et al.) for finding approximate inverses, instead of some revision-bbsed methods (such as the conjugate gradient method) which induces some instability in model training.
To apply the iterative method and inverse matrix M = CS, we need to satisfy its assumption RRazavi et al. Theorem 2] that | - M < 1. In practice, we instead pass the matrix D 1/1/2(1 (M M 1/2 as input the iterative method, where 7 > is small constant and the diagonal matrix DM is defined as diag ((M + +))).. We give the following lemma to justify our practical usage of the method. The proof is deferred Appendix E
D
Lemma 3. Given a constant y > O, if matrices M is constructed as STCS, and DM are defined as above, then all the singular values of DM 1/2 (M + +YI)DM" 1/2
are within (0, 1), which implies
that D 1/2 (M + +71) 11/2 <.
We further comment that numerically an implicit risk of the Schulz-type iterative method we use is the unintended consequence of "zero fill-in". . If we some sparse kernels (e.g. test functions with bounded support) other than Gaussian kernels, the empirical kernel matrices are sparse while the approximate inverse will converge to dense matrix, which increases the computational cost.
4.5 Empirical approximation evaluation
Spectral norm, the maximum singular value of matrix, is computation-light indicator of matrix approximation performance. In this work, we compare the spectral norm of the difference between the outputs from attention functions and the output from vanilla self-attention with the same input.
We use the initialized and pretrained bert-base-casee models from Huggingface's implementa- tion (Wolf et al. [019] The input vector X is embedded from the tokenized raw text in Wikitext-2 dataset [Merity et al. (017). The query; key and value weight matrices initialized or pretrained models transform input X into Q, K. of different distributions. We compare the results with differ- ent sequence lengths and different numbers of features used attention approximation methods. We set the number of features in the range of 24 to 28, More features usually require more computation resources.
Figure shows the performance of the modified Nyström method on approximation error with regards the number of features. We conclude that for Skyformer the approximation significantly better with the increased number of features, while for other methods the gain is not obvious. The good performance of the modified Nyström method also validates our previous claim that the Nyström method is currently one of the most powerful methods in large-scale kernel machines acceleration.
Remark. Although in single step the modified Nyström method in Section 4.2 can give low approximation error, we do not recommend directly applying it to the original self-attention. With some exploratory experiments
classification tasks, we find the variant suffers more severe gradient explosion issue than usual transformers. We speculate that it is because the matrix S AS (in the middle of Equation (6) inherits the high condition number of the original attention score matrices A, while the derivative of matrix inverse ((-1)) - -AAAAAAA further amplifies the condition number during backpropagation.
5 Experimental Results
Tasks and Datasets. We evaluate the proposed methods on five classification tasks on LRA bench- mark [Tay et al.] 2020b), which focuses on model quality under long-context scenarios: ListOps (Nangia and Bowman] (2088) Text Classification on IMDb review dataset [Maas al. 2011] Doc- ument Retrieval on AAN dataset [Radev et al., 2013]. Pathfinder [Linsley et al. 2018 and Image Classification on CIFAR-10 10 (Krizhevsky et al. 2009). The LRA benchmark covers diverse long- sequence tasks in sequence length, task difficulty, and inspected model abilities. For example, ListOps and Pathfinder evaluate the abilities to capture the long-range hierarchical dependency and spatial dependency, respectively, which poses challenges for sparse attention pattern based methods. We
7

---

Syfomer
Nattonnmmm
.
Hommr
intiialized n12
intialized n=1024
intiiaized n048
intiialized
1.5
1.5
2.25
3.5
3.75
3.0
3.0
2.5
2.0
0.0
0.4
0.30
50
100
150
200
250
50
100
150
200
252
30
100
150
100
232
102
150
200
250
n-1024
n-2048
pretrolned n=0096
2.22
2.00
1.75
2.00
1.5
Soccool Norm
2.0
1.25
2.5
1.0
100 150 200
255 50
100
150
200
150
200
255
100
150
200
250
Figure 1: Spectral norm results with different sequence lengths under different Wo. WK. Wy settings, either from initialized or pretrained BERT models. All methods are approximating the original self-attention output. Y axis: Lower spectral loss means better approximation. X axis: Higher d (number of features) means visiting more elements in the original matrix and bringing more computation costs. The label "Skyformer" here means that we use the algorithm behind Skyformer, mainly Eq. (5), to approximate the raw attention score matrix A in self attentionno In this experiment, "Skyformer" also needs to first approximate A, and then approximate D, as Performer does.
64524
partfnder
Irage
0.45
(.43
0.2
0.50
2.22
WAD T101 1100112011
2000
2300
2008
1900 10:0
Figure 2: Validation accuracy changes with respect to training time for 50k steps. X axis: training time (s). Y axis: classification accuracy.
Table Classification accuracy (%) on LRA benchmark.
Model
Text
ListOps
Retrieval
Pathfinder
Image
AVG.
Self-Attention
61.95
38.37
80.69
65.26
40.57
57.37
Kernelized Attention
60.22
38.78
81.77
70.73
41.29
58.56
Nystromformer
64.83
38.51
80.52
69.48
41.30
58.93
Linformer
58.93
37.45
78.19
60.93
37.96
54.69
Informer
62.64
32.53
77.57
57.83
38.10
53.73
Performer
64.19
38.02
80.04
66.30
41.43
58.00
Reformer
62.93
37.68
78.99
66.49
48.87
58.99
BigBird
63.86
39,25
80.28
68.72
43.16
59.05
Skyformer
64.70
38.69
82.06
70.73
40.77
59.39
8

---

Table 2: Running time (hour) and peak memory usage (GB). TC: Text Classification. LO: ListOps. RE: Retrieval. PF: Pathfinder. IC: Image Classification. KA Kernelized Attention.
Time (h) RE
Memory (GB)
Model
TC
LO
PF
IC
TC
LO
RE
PF
IC
Self-Attention
4.30
2.24
8.33
2.57
4.22
10.37
5.37
10.77
5.74
11.47
KA
3.91
1.99
7.46
2.42
4.05
5.73
5.94
10.46
6.38
6.38
Nystromformer
0.71
0.71
1.29
1.49
2.70
1.21
1.37
2.39
3.35
6.71
Linformer
0.65
0.60
1.13
1.09
2.19
0.99
0.99
1.89
1.97
3,94
Informer
1.60
1.19
2.91
2.39
3.90
5.12
4.85
5.77
4.75
9.51
Performer
0.77
0.73
1.41
1.40
2.55
1.09
1.09
2.16
2.20
4.39
Reformer
0.94
0.85
1.73
1.70
3.08
1.61
1.61
2.98
3.21
6.42
BigBird
2.00
1.88
3.81
3.39
6.53
2.83
2.71
4.97
4.97
9.95
Skyformer
1.02
1.29
1.86
2.03
3.40
1.59
1.75
3.15
4.13
8.26
report the classification accuracy on the test set, training time, and peak memory usage during training for each task.
Baselines. Aside from the vanilla quadratic self-attention, we compare with Big Bird [aheer et al. 2020], Performer [Choromanski et al.) 2020 Linformer [Wang et al. 2020). Nyströmformer [Xiong et al. 202I], Informer [Zhou et al. 2020]. and Reformer [Kitaev et al. 2020J. Most methods are approximating the vanilla full attention for efficiency and thus are not expected to have better performance. As it is not realistic to exhaustively fine-tune all models and search for the best performance under limited computation resources, we instead only replace the atffeattention module with the various attention methods and keep other experimental settings the same for fair comparisons.
Tesla V100 SXM2 16GB. We use the LRA evaluation benchmark reimplemented in Py'Torch by Xiong et al. (02T). We use a 2-layer transformer model with 64 embedding dimension, 128 hidden dimension. 2 attention heads, and mean pooling for classification. Batch size is selected conditioned on the memory requirements of the standard self-attention method, which leads to 16 for Text Classification, 32 for ListOps, 16 for Document Retrieval, 128 for Pathfinder, and 256 for Image Classification. Learning rate is set to le 4 for Text Classification, ListOps, and Image Classification, and 2e for Retrieval and Pathfinder. Each model on each task is trained for 50k steps, during which the best checkpoint with the highest accuracy on the development set will be saved for evaluation. For comparable computation complexity, we control the number of features to be 128 used in all methods (except Big Bird), under which setting the models will visit 128 n elements in the attention matrix. For numerical consistency, all experiment results are averaged across three runs with different random seeds.
Implementation Details. We conduct each experiment
We do not follow all settings in [Xiong et al. 2021 due the hardware limitation. The compromises, such as approximation dimension and gradient accumulations steps, might bring performance dif- ferences comparing to results reported in [Xiong et al. 2021]. The training instability problem also helps explain the performance gap.
Results. The training process of the standard softmax-based method is unstable observed in Figure it takes more steps to reach the stationary distribution of its long-time limit, and it is more easily getting stuck in local minimum. Runs with different random seeds may bring divergent performances, and probably leads to lower averaged scores. We have also tried directly approximating the self-attention method with the Nyström method and observed numerical instability during training.
Replacing the softmax structure with Gaussian kernel somehow alleviates this instability problem with boosted performance as shown in Table T However, the time and space requirement of Kemeliizdd Attention is not significantly improved compared to the original version, which serves as the motivation to approximate Kernelized Self-Attention with Nyström method.
Though not necessarily the fastest, our proposed Skyformer can efficiently converge to the long- time limit with comparable general performance in classification accuracy (Table and resource consumption (T). The advantages over the standard self-attention are significant with consis- tently less training time and generally better performance. For example, Skyformer brings nearly

---

times speed-up up on text classification and document retrieval while with 2.75% and 1.37% accuracy improvement over the standard self-attention.
Limitations. The applications of Skyformer might be limited to long sequence tasks because for small sequence length the statistical dimension da might be close to
To make the claim above clear, we first reiterate that the efficiency of Skyformer is related to d, tat. As implied by Theorem 2 the intrinsic difficulty of approximating raw attention score matrix concluded as dstat. which corresponds to the effective rank of matrix C. The complexity of Skyformer depends on the sub-sample size d (the size of the sub-sampling matrix S). A large dstat leads to a large d, and an inefficient application of the Nyström method.
The classical theory for statistical dimension only guarantees that dstat is small (compared to n) when n is large enough, and is possible the statistical dimension associated with a short sequence might be even close to the sequence Therefore a large serves as a condition to make the method work. Figure empirically shows that our method performs better with larger n's.
Conclusions and future work
Motivated by the connection between kernel methods and self-attention, introduce Kernelized Attention, which replaces the softmax structure in self-attention with a Gaussian kernel. We also propose Skyformer, which adapts the Nyström method Kernelized Attention improve its efficiency. We expect the new model can enjoy more stable training while inheriting the strong performance from self-attention. Extensive experiments verify our intuitions and show that both Kernelized Attention and its Nyström approximation variant have comparable accuracy to the original Transformer the LRA benchmark.
Direct development of this work is the incorporation of further computation tricks kernel methods, such the local and global approximation for gram matrix (Snelson and Ghahramani, [007]] and the importance sampling in Nyström methods [Musco and Musco 2017 Chen and Yang [2011ba] a]. Other related questions include the choice of the kernel other than the Gaussian kernel our kernelized attention model. It is expected that for different tasks there will be specific kernels more proper than the original self-attention. The results in this work also shed new light the design of the attention mechanism, which may benefit board downstream NLP tasks.
Acknowledgments and Disclosure of Funding
This research is based upon work in part supported by the Office of the Director of National Intelligence (ODNI), Intelligence Advanced Research Projects Activity (IARPA), via contract No. FA8650-17-C-9116 and U.S. DARPA KAIROS Program No. 189750-1922-106 This work also in part supported by NSF grant DMS- 1 10831. The views and conclusions contained herein are those of the authors and should not be interpreted as necessarily representing the official policies, either expressed or implied, of DARPA, ODNI, IARPA, or the U.S. Government. The U.S. Government authorized to reproduce and distribute reprints for governmental purposes notwithstanding any copyright annotation therein.
10

---

References
Arturs Backurs, Piotr Indyk, and Tal Wagner. Space and time efficient kernel density es- timation in high dimensions. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelz- imer, Florence d' d'Alché-Buc, Emily B. Fox, and Roman Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 15773-15782, 2019. URL https://proceedding http:///weeennoeeeeereeeeeeeee/p wttp://epp....pppp//////////////
Neural Information
htt.:/7066652280aaaaa66662223660 Abstraatt.. html
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. CoRR, abs/2004.051555 2020. URL htttpe://arxiv.org///a 00150
Srinadh Bhojanapalli, Chulhee Yun, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Low- rank bottleneck in multi-head attention models. In International Conference on Machine Learning. pages 864-873. PMLR, 2020.
Moses Charikar and Paris Siminelakis. Hashing-based--stingg based-estimators for kernel density in high di- mensions. In 2017 IEEE 58th Annual Symposium on Foundations of Computer Science (FOCS). pages 1032-1043. IEEE, 2017.
Yifan Chen and Yun Yang. Accumulations of projections- a unified framework for random sketches in kernel ridge regression. In International Conference on Artificial Intelligence and Statistics, pages 2953-2961. PMLR, 2021a.
Yifan Chen and Yun Yang. Fast statistical leverage score approximation in kernel ridge regression. In International Conference on Artificial Intelligence and Statistics, pages 9935-2943. PMLR, 2021b.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. CoRR, abs/1904. 10509, 2019. URL httpp://aarri..rrr//b//19944.
Krzysztof Choromanski. Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamás Sarlós, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy Colwell, and Adrian Weller, Rethinking attention with performers. CoRR, abs/2009. 14794, 2020. URL httts::/aaxii wwwrexiv..oggaaae www.abg.//a./o....o..oo....... aba/2009 14794
Jacob Devlin, Ming-Wer Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pages 4171-4186 Minneapolis, Minnesota, URL https/
June 2019, Association for Computational Linguistics. doi: 10.11853/v1/N19-111 //www. aclweb. org/ wholooogtoh1199 1423
Yihe Dong, Jean- Baptistee Cordonnier, and Andreas Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. arXiv preprint 1X3v:2203303444, 2021.
Petros Drineas, Michael W Mahoney, and Nello Cristianini. On the nyström method for approximating gram matrix for improved kernel-based learning. journal of machine learning research, 6( (12), 2005.
Alex Gittens and Michael W Mahoney. Revisiting the nyström method for improved large-scale machine learning. The Journal of Machine Learning Research, 17(1):39774404) 2016.
Sariel Heled,, Piotr Indyk, and Rajeev Motwani. Approximate nearest neighbor: Towards removing the curse of dimensionality. Theory of computing, 8(1):321-350, 2012.
Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. Tinybert: Distilling BERT for natural language understanding. In Trevor Cohn, Yulan He, and Yang Liu, editors, Proceedings of the 2020 Conference on Empirical Methods Natural Language Processing: Findings. EMNLP 2020, Online Event, 16-20 November 2020, pages 4163-4177 Association for Computational Linguisties, 2020. doi: 18.18653/v 1/2020 1ta.5//////2.....innnnemmmmmmmm findings-emnlp.372.
URL stttps://oo.. org/10. 18653/11/20022 windinggnn-iii emnlp.33 372.
11

---

Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnss: Fast autoregressive transformers with linear attention. International Conference on Machine Learning, pages 5156-5165. PMLR, 2020.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.ner, 2020. IiRLlittpaa:/oppnnevieees wtt.:///ffoommmtrrrrgghhhtttta
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Sanjiv Kumar, Mehryar Mohri, and Ameet Talwalkar, Sampling techniques for the nystrom method. In Artificial Intelligence and Statistics, pages 304-311, 2009.
Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Veselin Stoyanov, and Luke Zettlemoyer. BART: denoising sequence-to-sequeece pre-training for natural language generation, translation, and comprehension. In Dan Jurafsky, Joyce Chai, Natalie Schluter, and Joel R. Tetreault, editors, Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, ACL 2020, Online, July 5-10, 2020, pages 7881-7880. Association for Computational Linguistics, 2020. doi: 10.18653/v 1/2.5533112222.ccc acl- madd-mmi..7... URL acl-main. 703
https://doi. 10./10. 1855//11222222
Drew Linsley, Junkyung Kim, Vijay Veerabadran, Charles Windolf, and Thomas Serre. Learn- ing long-range spatial dependencies with horizontal gated recurrent units. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Gar- nett, editors, Advances Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pages 152-164, 2018 URL Abstract. html
019956677899787111177cccct
Liyun Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the difficulty of training transformers. In Bonnie Webber, Trevor Cohn, Yulan He, and Yang Liu, editors, Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, EMNLP 2020, Online, November 16-20, 2020, pages 5747-5763, Association for Computational Linguistics, 2020a. doi: 10. 10.18663/1122220emmmlpp m63n.463. URL https 7/doi //doi.org//0. 1w1111666112222 http://222..emmlpmmmmemmmme mmmlp-main.c.
Weijie Liu, Peng Zhou, Zhiruo Wang. Zhe Zhao, Haotang Deng, and Qi Ju. Fastbert: a self-distiling BERT with adaptive inference time. In Dan Jurafsky, Joyce Chai, Natalie Schluter, and Joel R. Tetreault, editors, Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, ACL 2020, Online, July 501 2020, pages 6035-6044, Association for Computational Linguisties, 2020b. doi: (0.18653/v1/2022..c aaaamm...... URL https: http:///oooooo/g/// org/10......./ 18653/v1/ 2020. 200.aacl-00 22200acc-main.55 macccmmain.53) 537
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zeutlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized BERT pretraining approach. CoRR, abs/1907.1 1692, 2019. http:tttt////rr..ooo wwtp:///////b/////// 10t...9..........
Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Dekang Lin, Yuji Matsumoto, and Rada Mihalcea, editors, The 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, Proceedings of the Conference, 19-24 June, 2011, Portland, Oregon. USA, pages 152-150. The Association for Computer Linguistics, 2011. URL https://www aclweb. orgg/anthology/I11---11
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mix- ture models. In 5th International Conference on Learning Representations, ICL.R 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. URL https //openrevvew.neteeeee nww htt.eee...ee//ommmmm//./y//aa.a.
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Bue, Emily B. Fox, and Roman Gamett, editors, Advances in Neural Information Processing Systems 32: Annual
12

---

Conference on Neural Information Processing Systems 2019, NewrIPS 2019, December 8-14. 2019, Vancouver, BC, Canada, pages 44014-14024, 2019. URL http://papers nips cc/paeer/ 9551-are-
Setty--beteer-tate than-one
Cameron Musco and Christopher Musco. Recursive sampling for the nystrom method. In Isabelle Guyon, Ulikee von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vish- wanathan, and Roman Gamett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9. 2017, Long Beach, CA, USA, pages 3833-3845, 2017. URL http::/papers. mips. cc/paper/
1973-recursiver samp ing waauunng-----aee for-ttee-nsseeeeeeeeeeeeeeeeee-e paystroo-mettinn
Nikita Nangia and Samuel R. Bowman. Listops: A diagnostic dataset for latent tree learning. In Silvio Ricardo Cordeiro, Shereen Oraby, Umashanthi Pavalanathan, and Kyeongmin Rim, editors, Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics, NAACL-HLT 2018, New Orleans, Louisiana, USA, June 2-4, 2018, Student Research Workshop. pages 92.99. Association for Computational Linguistics, 2018. doi: 10.18653/v1/n 1/n18-4013, URL https://doi. org/10. 18663/v1/118-400000
Jiezhong Qiu, Hao Ma, Omer Levy, Wen-tau Yih, Sinong Wang, and Jie Tang. Blockwise self- attention for long document understanding. In Trevor Cohn, Yulan He, and Yang Liu, editors, Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: Findings, EMNLP 2020, Online Event, 16-20 November 2020, pages 2555-2565. Associa- tion for Computational Linguistics, 2020. doi: 10. 18665//12002..fnmding-emm/ findingss-eeeeeeee emnlp. 232
URL
Dragomir R. Radev, Pradeep Muthukrishnan, Vahed Qazvinian, and Amjad Abu-Jbara. The ACL anthology network corpus. Lang. Resour. Evaluation, 47:919-94 9-494, 2013. doi: 10.1007/ s10579-012-9211-2. /utpss:::///oii wtt.://o.//..11//100151111111110 9211-2
Ali Rahimi, Benjamin Recht, et al. Random features for large-scale kemel machines. In NIPS, volume 3, page 5, Citeseer, 2007.
M Kafaci Razavi, Asghar Kerayechian, Mortaza Gachpazan, and Stanford Shateyi. A new iterative method for finding approximate inverses of complex matrices. In Abstract and Applied Analysis, volume 2014. Hindawi.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Trans. Assoc. Comput. Linguistics, 9:53-68, 2021. URL https://ttrasaacl.org http://aaacc....g//////ddeeeee https//inne...pppttaaaeeeeeeeeee phtp:/tacl/aticle//tteeeeeeeeeee
Edward Snelson and Zoubin Ghahramani. Local and global sparse gaussian process approximations.
In Artificial Intelligence and Statistics, pages 524-531. PMLR, 2007.
Raphael Tang, Yao Lu, Linqing Liu, Lili Mou, Olga Vechtomova, and Jimmy Lin. Distilling task- specific knowledge from BERT into simple neural networks. CoRR, abs/1903.12211 12136, 2019. URL httt:///rxiv..rg/ www.///////a/.....o.oo........e/ abt//193 12136
Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng. Synthesizer: Rethinking self-attention in transformer models. CoRR. 000703,, 2020a. URL https
//arxiv.org 200500 00743
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. CoRR, abs/2011.00400 2020b. URL http:p/://aa....or//aba///////// 2000000006
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. CoRR. abs/2009.06732. 2020c. URL HtLLLhtpss:/// http::///rxxvvvv/// wttpe//wwwrg//////eeee///// 007222009000602222
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 5998 00086600,60000 2017. URL BRLLhttp::/// htttp:/ppapers. nips http/appr//11111. att.eiionnttnntntnneeieeo is all-alllouo yollyyuu-ae you--ned
13

---

Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. In Anna Korho- nen, David R. Traum, and Lluís Mårquez, editors, Proceedings of the 57th Conference of the Association for Computational Linguisties, ACI. 2019, Florence, Italy, July 28- August 2, 2019, Volume 1: Long Papers, pages 5797-5808. Association for Computational Linguistics, 2019. doi: 10.18653/v1/p19-115558 1580. URL httpt::/t//// org/10, /ttppp////..o..oo000//1/115///1/ 1580
Martin J Wainwright. High-dimensional statistics: A non-asymptotic viewpoint, volume 48. Cam bridge University Press, 2019.
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear linear complexity. CoRR, abs/2006.04768, 2020. httpttts: http///www..rgaabaa/a/aaaa .04768
Christopher Williams and Matthias Seeger. Using the nyström method to speed up kernel machines. In Proceedings of the 14th annual conference on neural information processing systems, number CONF, pages 682-6822666 688, 2001.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, and Jamie Brew. Huggingface's transformers: State-of-the-art natural language processing. CoRR, abs/ Abs/1990,,3771, 2019. URL
http://arxiv.org// wttp://///..aaa0/00ga//// 1t1/119911100..71
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. CoRR, abs/2102.03902, 2021. URL https://arxiv, org/abs/ /a///bbb/211000 https//212..0000 03902
Yun Yang. Mert Pilanci, Martin J Wainwright, et al. Randomized sketches for kernels: Fast and optimal nonparametric regression. The Annals of Statisties, 45(3)9911-1023, 2017.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, San- tiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang, I.i Yang, and Amr Ahmed. Big Transformers for longer sequences. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tice Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020. URL pttpss:/ppocceeding: wwwphhgg.....riss.ocoo hptp://220haaarr/aaaaa/222//5522 Abt.laabbbbbbbttat.....mm html
bird:
Haoyi Zhou, Shanghang Zhang, Jieqi Peng, Shuai Zhang, Jianxin Li, Hui Xiong, and Wancai Zhang. Informer: Beyond efficient transformer for long sequence time-series forecasting. CoRR,
abs/2012.07436, 2020. URL https: //rpxvaa./vvv//////////////a/a org/abs///12/
14


Reviews:
[{"review_id": "kYTwuHkY4un", "reviewer": "Reviewer_5adD", "summary": "The paper proposes a method (Skyformer) to apply Nystrom approximation to the attention matrix. It embeds the attention matrix inside a larger PSD matrix (unlike Nystromformer), allowing Nystrom method to work well. Theoretical and empirical validation shows that the method is competitive with other forms of efficient attention (e.g. BigBird, Performer, Reformer).", "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": "Strengths:\n1. The paper is well-written and easy to understand. Section 4 provides valuable intuition, and I enjoyed reading it.\n2. The method can significantly speed up attention computation, which is necessary for scaling to longer sequences.\n3. The method is novel to the best of my knowledge.\n\nWeaknesses:\n1. Limitation of application: it's not clear to me how widely the proposed method can be apply, as there's no discussion of its limitations. I'm not sure which conditions will allow it to work well. \n2. The empirical validation is not sufficient. As the paper proposes a new variant of attention (based on Gaussian kernel), one has to demonstrate this works on different tasks (e.g. language modeling, translation, etc.), not just classification as in the long-range arena benchmark.\n\nQuestions:\n1. Section 4.2, instead of the block matrix [[phi(Q, Q), phi(Q, K)], [phi(K, Q), phi(K, K)]], would it also work to embed phi(Q, K) in the block matrix [[0, phi(Q, K)], [phi(K, Q), 0]]? I think this matrix is also PSD, and it has a lot of zeros so it might save some computation.\n2. Figure 2: (a) Isn't Skyformer approximating a different matrix (the matrix C instead of the attention matrix)? Why would Skyformer have low approximation error wrt the attention matrix here?\n(b) Maybe Performer is having problem because it also tries to approximate the normalization constant D^{-1}, which Skyformer does not do. Would that make this an unfair comparison, because Skyformer is approximating a different matrix and Performer is doing more work to approximate the denominator?\n\nComments:\n- Line 129, (D_Q)_ii = ||q_i||^2 / sqrt(p) should be (D_Q)_ii = exp(-||q_i||^2 / sqrt(p)).\n- Equation after line 177 equation, this definition is for the matrix C, not for Kernelized-Attention(Q, K, V). I assume kernelized-attention is the product of C and V.\n- Line 178, (D_Q)_ii = ||q_i||^2 / sqrt(p) should be (D_Q)_ii = exp(-||q_i||^2 / sqrt(p)).\n\n============== Post-rebuttal\n\nThank you for the response to my questions. That has helped clarified some of my misunderstanding of the method. \n\nI would still encourage the authors to validate the method on a wider variety of applications to make a stronger case.\n\nI've increased my rating.", "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"}, "extra_text": {}}, {"review_id": "ZDzmbAOnrzv", "reviewer": "Reviewer_4vy6", "summary": "The authors reformulate soft max structure of the self-attention mechanism in Transformers in terms of a Gaussian kernel evaluation. They claim that such a model is empirically stable during training, while allowing the use of various Kernel approximation strategies to improve the quadratic computational complexity. \n\nIn this work, approximation strategy used by the authors is the popular Nyström approximation. However, since the kernel matrix formed between the query and the key matrices is not symmetric (hence not PSD), the Nyström approximation is not directly applicable. The authors propose to form a PSD kernel matrix by first evaluating the Gaussian kernel on the row concatenated  query and key matrices which is PSD, performing the Nyström approximation and then choosing the appropriate block in the resulting matrix.\n\nBy adapting an existing theorem [Musco, Musco 17],  the authors claim that such a low-rank approximation of the kernelized attention score is bounded in terms of the relative spectral norm. This is claimed to be due to the observed fast decay of the eigenvalues in an empirical kernel matrix. \n\nThe matrix inversion required for the Nyström approximation, however is a potential issue in terms of numerical stability on a GPU. To circumvent the problem, instead of a standard conjugate gradient method, they use a Schulz type iteration which only uses matrix products, without any floating point divisions.  \n\nEmpirical evaluation shows marginal improvement on most tasks on the LRA benchmark, as compared to the naive Nyström approximation [Xiong et al 2021]. However, this is at the cost of relatively larger time and memory usage. ", "questions": "", "limitations": "", "rating": 7, "confidence": 5, "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": "Originality: \nThe main contribution of this work seems to be the proposed fix to the theoretical issue of Nyström approximation in [Xiong et al 2021]. The rest of the methods used here are priorly known techniques, though presented in a cogent manner. They cite most important of the previous works, however at times their contributions are not clearly delineated from existing methods. For example, the Kernelized attention model (line 54). Reference to [Chormanski et al 2020] for Lemma 1 is also missing. \n\nQuality: \nThe submission is technically sound and correctly applies the Nyström approximation to the kernelized attention mechanism. The claim of good approximation in spectral norm is supported by a stated theorem as well as empirical evaluation. However, the use of Schulz type iteration for inverse may have the unintended consequence of \"zero fill in\", i.e. the approximate inverse converges to a dense matrix even when the kernel matrix was sparse. \n\nClarity: \nThe paper is well written and clearly organised. Claims of speed-up though (line 313) might be considered misleading as they direct competitor is the method in [Xiong et al 2021]. The experimental setup is fairly well described with the code being adapted from [Xiong et al 2021]. \n\nSignificance: \nThe paper is good contribution towards the development of theoretically sound approximations and reformulations of the self-attention mechanism. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors did not point out or address any potential negative societal impact of their work. Being essentially a theoretical improvement over previous work, there is limited potential for excess social impact as compared to previous work in this field. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "WohPGo87byn", "reviewer": "Reviewer_Dpwz", "summary": "The author proposes a modified Nyström method named Skyformer to approximate the Kernelized Attention. They also provide a theoretical guarantee that the approximation error is small in terms of the spectral norm. Extensive experiments are also conducted to show that Skyformer achieves comparable performance to vanilla self-attention with fewer computational costs.", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "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 is well written and easy to follow. In this paper, A new variant of the Nyström method, Skyformer, is proposed to estimate the Kernelized Attention. Specifically, they expand the original Kernelized Attention to a larger PSD Kernelized Attention and then use the Nyström method to approximate the larger PSD matrix. The theoretical result shows that the complexity of Skyformer is comparable to the complexity of most other efficient transformers. Experiment results also facilitate the author's claim. Although I admire the theoretical contribution of this work, I am worried about the practice of this method since it enlarges the original dimension n to 2n and needs to calculate the inverse of a matrix.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "- In Theorem 2, it seems that the sampling size d is independent of the approximate accuracy \\epsilon, which is unreasonable.\n\n- For Figure 3, it is not evident that Skyformer is more stable than the other competitors. Does the author have more solid evidence to show that Skyformer is more stable than softmax?\n\n- For Figure 3, It is fairer if the author uses the real-time as X-axis, instead of steps. Also, in Table 2, it seems that Skyformer is not the fastest algorithm. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "RarYeaX-ZTk", "reviewer": "Reviewer_ugqB", "summary": "This paper proposes Skyformer a new xformer variant that uses Gaussian kernels and the Nystorm method. ", "questions": "", "limitations": "", "rating": 5, "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 does sound interesting and novel, the experiments on long range arena look satisfactory. \n\nResults look borderline but overall better than some of the existing baselines while being faster/less memory hungry.\n\nThe improvements of Nystromformer is that: \"Nyströmformer blindly applies the Nyström method to a non102 PSD matrix, and thus fails to ultize the full potential of Nyström method.\" the authors solve this by \"instead lifting the kernelized attention score matrix into a large PSD\nmatrix which contains the target non-PSD matrix as its off-diagonal block.\"\n\nI have some minor concerns.\n\n1) I find that the intro is an over nag. The paper starts off by telling us how expensive GPT-3 training is (or how long it takes). I don't think the proposed method solves this problem. I would advise the authors to keep the exposition short. Most reviewers in the space know the problems with Transformers and the quadratic complexity problem.\n2) What happened to the image CIFAR task on LRA? Why was that omitted? \n3) A nit but why is the model called Skyformer? What does SKY mean here?\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1"}, "extra_text": {}}, {"review_id": "J27Mhb1iY11", "reviewer": "Reviewer_7p5f", "summary": "The paper proposes a new form of self-attention for use in transformer models using Gaussian Kernels. They show that the new attention method is comparable in accuracy to traditional softmax self-attention. They then go on to show that kernel based attention can be approximated by taking a Nystrom sketch of the kernel matrix. They give rigorous bounds on the spectral norm of the approximation (assuming that the kernel matrix has a high condition number). ", "questions": "", "limitations": "", "rating": 6, "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": "Pros\n* Proposal of a new attention method (kernel based instead of soft-max)\n* Rigorous proof and a proper sketch of the scenarios where one can use Nystrom method. Mathematically founded approximation of the kernel matrix unlike earlier approach (Nystromformer)\n* Mostly clear and easy to follow\n\nCons\n* Work seems incremental (Nystromformer is very similar). But the nice thing is the authors do a good job of giving proofs motivating their algorithm\n\nNit\n* Page 1: Introduction: O(n^2) not O(n)", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Don't see any potential negative societal impact from reducing resources required for training and inferencing", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "CeRwWfWuYI3", "reviewer": "Reviewer_oeCF", "summary": "The paper presents an alternative approach to attention in the Transformer. There are two modeling contributions:\n\n1. Kernelized attention that omits the typical normalization term used for softmax attention.\n2. A Nystrom method approach that allows efficiently approximating the output of the attention mechanism\n\nThe contributions also include theoretical results regarding the accuracy of the Nystrom approximation, and experimental results showing better convergence, higher accuracy, and medium efficiency gains from the approach.\n\nBelow I summarize my understanding/takeaways from the two key methods; please correct me in the author response if these are inaccurate:\n1. As I understand it, \"kernelized attention\" mean that target keys don't compete with each other, the only thing that matters is their distance to the query. In standard softmax attention, suppose that all key vectors are orthogonal to a query vector -- in that case, attention would weigh all values uniformly. In kernelized attention, on the other hand, if key vectors are all far away from the query vector, then all value vectors would be attended to with near-zero weight.\n2. The Nystrom approximation involves sampling some queries and keys, and then computing kernels between all pairs. But whereas Nystromformer computes similarity for query+key pairs only, the proposed approach also computes kernels for query+query and key+key pairs as well.", "questions": "", "limitations": "", "rating": 9, "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": "Overall this is a very solid paper. The methods presented are novel in the context of efficient attention, and they are justified from both a theoretical and experimental standpoint. I particularly appreciate that the error caused by the Nystrom approximation rapidly shrinks as the number of features increases (as shown in Figure 2), and I believe that it should approach zero as the number of features grows to match the length of the sequence. The paper is clear and well-structured, with some minor exceptions noted below. The overall research program of constructing efficient attention variants has gained a lot of interest in the community, and this paper appears to represent significant progress on this research direction.\n\nThe paper offers two contributions that are in principle orthogonal: kernelized attention, and the Nystrom approximation. Indeed, kernelized attention (without any approximations) is presented as a standalone ablation. Why about evaluating the Nystrom approximation in isolation? I believe this would just involve taking the proposed Skyformer and re-introducing a $D^{-1}$ normalization term. Does this make sense conceptually, and have the authors tried it in practice?\n\nIn terms of writing, the only point where I was confused was Theorem 2, which I still don't understand. Line 222 simultaneously introduces a new scalar lambda, uses a matrix C never previously mentioned in the theorem, and *also* introduces a new concept C-bar. There is just too much happening all at once. It would help to introduce all terms one by one, specify in words what lambda is, and otherwise clarify.\n\nI have some questions regarding the LRA evaluations. First, is there a reason the Image task is omitted? Without an explanation in the text it's possible to interpret this as cherry-picking. Next, Retrieval scores are a full 20 points higher than the LRA paper. Nystromformer has similar results, so I trust that the results presented are correct, but it would be helpful to have a comment in the paper regarding this situation. Finally, Pathfinder scores are actually *lower* than both Tay et al (LRA) and Xiong et al. (Nystromformer). Do you know why this might be?\n\nOverall LRA benchmarking is becoming a bit of a mess, because results from two different papers can't be compared against each other. This situation is not limited to only the present paper, but it's still disappointing because it makes it possible for individual papers to use cherry-picked baselines or poor hyperparameter choices for competing methods, without no easy way to flag it using cross-paper comparisons. I understand that it's hard for any one paper to address systemic benchmarking issues, but the results would certainly be stronger if the authors had a way to mitigate these concerns. If anything a little more discussion to justify the evaluation would help.\n\nMinor\n- Why the name Skyformer? I'm just curious, because the naming choice is never discussed in the paper.\n- line 27: I would not characterize Reformer as a low-rank approximation, since it can learn full-rank sparsity patterns\n- line 77: \"By reduce\" is ungrammatical\n- Equation on line 177: there is no V term on the right side of the equality, so I think this equation isn't quite right\n- Figure 1: $K_S^T$ should have shape $d_S$ x $n$", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations and impact are adequately addressed.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "7"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers agree that this is a solid paper re-introducing the Nystrom method for Transformers to address the quadratic space and time complexity of the regular attention module. Thus, it leverages the fruitful area of research on using kernel approaches to improve regular Transformers. The novelty given that Nystrom method was previously used in that context is limited, but the empirical results are strong and furthermore the authors explain in detail approximation guarantees of their method. What is missing is equally detailed analysis of the computational time of the presented method. Nystrom algorithm is in general expensive, but the authors manage to bypass some of its complexities, as showed in the experimental section. More detailed discussion regarding Nystrom method in theory and practice would strengthen the paper.
Decision: Accept (Poster)