Skyformer: Remodel Self-Attention with Gaussian Kernel and Nyström Method
Year: 2021

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.

heading: 1 Introduction
text: The cost of language model training increases exponentially. Among different models, Transformerbased language models [Vaswani et al., 2017, Devlin et al., 2019, Liu et al., 2019, Lewis et 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 n is 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 et al., 2020, Beltagy et al., 2020, Kitaev et al., 2020] or low-rank approximation [Choromanski et al., 2020, Wang et al., 2020]. However, there is usually a lack of theoretical analysis on the approximation error of these methods due to the complex softmax structure, which makes the theoretical comparison between the efficiency of each method infeasible. It is also unclear in theory how to set the hyper-parameters of those methods to attain a 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, or even random seeds, which usually demands a 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 (NeurIPS 2021).
in Transformer training comes from the softmax structure, as the un-normalized attention score matrices before softmax tend to have extremely large condition numbers due to its fast singular value decay. To alleviate the instability issue, an extra factor of 1/ √ p 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. [2020], the softmax structure is closely related to Gaussian kernels up to diagonal matrix multiplications, as the pairwise dot products naturally appear when expanding the squared `2 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 to 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-attention in model accuracy.
To further improve the efficiency, we propose Skyformer (Symmetrization of Kernelized attention for NYström method) to accelerate kernelized attention. Skyformer adapts the Nyström method [Williams 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 a new kernel-based structure, kernelized attention, to stabilize the training of Transformers.
(2) We propose Skyformer, which approximates the kernelized attention via low dimensional randomized sketches by adapting the Nyström method to a 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

heading: 2 Related Work
text: Among all the transformer acceleration methods, including attention layer simplification by pruning redundant attention heads [Voita et al., 2019, Michel et 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., 2020, Child et al., 2019, 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].
2Our code is released at https://github.com/pkuzengqi/Skyformer
The attention score matrix is known to exhibit a very fast rate of singular value decay [Bhojanapalli et al., 2020, Dong et al., 2021], 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 in kernel methods. Among them, Linformer [Wang et al., 2020] compresses the size of the key and value matrix with random projections based on the Johnson–Lindenstrauss transform, a common randomized sketching method in Gaussian processes [Yang et al., 2017]; Reformer [Kitaev et al., 2020] applies locality-sensitive hashing (LSH) [Har-Peled et al., 2012] to simplify the computation of the attention score matrix, which is widely used in kernel density estimation [Charikar and Siminelakis, 2017, Backurs et al., 2019]; Performer [Choromanski et al., 2020] projects both query and key matrix through random Fourier features [Rahimi et al., 2007], 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., 2021]. Linear attention takes the softmax structure in self-attention as a measure of similarity and replaces it with the dot product of separately activated query and key matrices; Synthesizer aims to 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 to a non-PSD matrix, and thus fails to utilize the full potential of the Nyström method. This issue is resolved in our proposed Skyformer by instead lifting the kernelized attention score matrix into a large PSD matrix which contains the target non-PSD matrix as its off-diagonal block. For more details on attention approximation methods, we refer readers to a survey paper on efficient transformers [Tay et al., 2020c].

heading: 3 Preliminaries and notations

heading: 3.1 Revisiting self-attention
text: For a given input sequence X ∈ Rn×d0 of length n and embedding dimension d0, The dot-product attention for a single head in Transformer Vaswani et al. [2017] is defined as
Attention(Q,K,V ) = softmax ( QKT √ p ) V
where Q = XWQ, K = XWK , and V = XWV , and WQ, WK and WV are the query, key, and value weight metrics that linearly project the input X of d0 dimension to an output tensor of p dimensions.
To simplify the future analysis, the left softmax term can be rewritten into D−1A, where A := exp(QKT / √ p) is the un-normalized attention score matrix; D is a diagonal matrix whose diagonal is exp(QKT / √ p) · 1 (by convention 1 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(qTk/ √ p) as the softmax kernel function, and represent A by the notation SM(Q,K), which means the element aij from the i-th row and j-th column in A is equal to SM(qi,kj). Throughout this paper qi (resp. kj) means the i-th (resp. j-th) row in Q (resp. K).
We close this subsection with a short lemma to show SM(·, ·) is a positive semidefinite (PSD) kernel function [Wainwright, 2019, Definition 12.6] by relating it to Gaussian kernels. Lemma 1. SM(·, ·) is a PSD kernel function. Equivalently, for all integers n ≥ 1 and elements {qi}ni=1 ⊆ Rp, 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:
SM(qi, qj) = exp ( qTi qj√ p ) = exp ( ‖qi‖2 2 √ p ) exp ( −‖qi − qj‖ 2 2 √ p ) exp ( ‖qj‖2 2 √ p ) .
The middle part exp ( ‖qi−qj‖2
2 √ p
) is exactly a Gaussian kernel with bandwidth p 1 4 . (Choromanski
et al. [2020] have more discussion on the findings.)
Through this equation, we can rewrite C as
C = D 1/2 Q · κ
( Q
p1/4 ,
Q
p1/4
) ·D1/2Q , (1)
where DQ is a diagonal matrix with elements (DQ)ii = exp ( ‖qi‖2√
p
) ,∀i ∈ [n], and κ(qi, qj) :=
exp ( −‖qi − qj‖2/2 ) is the standard Gaussian kernel function.
We prove the lemma by using the fact that κ is a PSD kernel and κ (
Q p1/4 , Q p1/4
) is a PSD matrix. ♦

heading: 3.2 Nyström method
text: 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(STBS)†STB, where (·)† denotes the Moore-Penrose pseudoinverse of a matrix, and S ∈ Rn×d is a zero-one sub-sampling matrix whose columns are a 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 1 (Uniform sub-sampling matrix). For a random matrix S ∈ Rn×d, if S has i.i.d. columns and the j-th column S(j) can randomly be √ 1 dei with probability 1 n , where ei is the i-th column of the n-by-n identity matrix In, then S is called a uniform sub-sampling.
We close this subsection with a remark that it is not appropriate to directly extend the Nyström method from kernel method to self-attention due to a core requirement that B should be PSD with consideration of approximation performance improvement. We will show in the next section how to address this challenge and properly adapt Nyström method to non-PSD matrices.

heading: 3.3 Approximation evaluation
text: 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 finetuned performance on downstream natural language understanding tasks. Specifically, Long Range Arena benchmark [Tay et al., 2020b] 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, to ease the future discussion on performance.
Definition 2 (Spectral norm guarantee for matrix approximation (MA)). Given a matrix M ∈ Rn1×n2 , two constants ε > 0, δ < 12 , we say that its approximation matrix M̃ ∈ R
n1×n2 satisfies (ε, δ)-MA property for M , if
P { ‖M − M̃‖ > ε‖M‖ } < δ. (2)
In previous works, the direct analysis of the approximation error to the entire output D−1AV in the (ε, δ)-MA 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 a 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.

heading: 4 Method

heading: 4.1 Kernelized Attention
text: Kernelized Attention replaces the softmax structure in vanilla self-attention with a Gaussian kernel, and the new attention model is stated as:
Kernelized-Attention(Q,K,V ) = CV := κ ( Q
p1/4 , K p1/4
) V , (3)
where we define the n-by-n matrix C as the kernelized attention score matrix κ(Q/p1/4,K/p1/4).
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 a limited number of other important tokens in the sequence. We observe that Gaussian kernel function can play a similar role. The expression of a Gaussian kernel is κ(qi,kj) := exp ( −‖qi − kj‖2/2 ) . Via this expression, for token i in the query, Gaussian kernel assigns a large attention score to the token j when kj is close to qi. The distance-based weight assignment is indeed considered as a major reason why kernel methods are powerful. The form of kernelized attention also leads to an automatic normalization. Based on Equation (1), the new attention model can be rewritten in terms of the un-normalized attention score matrix A as
Kernelized-Attention(Q,K,V ) = ( D −1/2 Q ·A ·D −1/2 K ) V ,
where DQ (resp. DK) is a diagonal matrix with elements (DQ)ii = exp ( ‖qi‖2√
p
) (resp. (DK)ii =
exp ( ‖ki‖2√
p
) ), ∀i ∈ [n]. We remark the kernelized attention model can thus be formally taken as
a variant of the original self-attention, which instead normalizes the matrix A in a form of D−1A. The intrinsic normalization allows kernelized attention to have a more reasonable condition number than self-attention, which benefits the stability of model training. To demonstrate the improvement in stability, we additionally provide a toy experiment in Appendix F, which shows the “condition number" of kernelized attention is smaller than self-attention. Moreover, empirical evaluation in Section 5 supports our claim that the new attention model can attain a comparable performance to the original attention model.

heading: 4.2 Skyformer: a modified Nyström method
text: 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 bij from the i-th row and j-th column in B is equal to φ(qi,kj), where qi (resp. kj) 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 A := SM(Q,K), and the empirical Gaussian kernel matrix C := κ(Q/p1/4,K/p1/4). Therefore this method leads to a low-rank approximation to the output of either self-attention D−1AV or Kernelized Attention CV . (D in self-attention can be obtained by computing A · 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 a non-PSD matrix B, our first step is to complete the matrix into a PSD matrix B̄:
B̄ := φ (( Q K ) , ( Q K )) . (4)
Then we approximate B̄ with ˜̄B through
˜̄B = B̄S(ST B̄S)†ST B̄, (5)
where S is a 2n-by-d uniform sub-sampling matrix as defined in Definition 1. The final approximation will be given as
B̃ := (I,0) ˜̄B(0, I)T . (6)
The original matrix B can be well-approximated by B̃ due to the following inequality
‖B − B̃‖ = ‖(I,0)(B̄ − ˜̄B)(0, I)T ‖ ≤ ‖B̄ − ˜̄B‖, 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 be well approximated by a low-rank B̃ is that as an empirical kernel matrix the eigenvalues in B̄ usually decay fast, and thus there are many small eigenvalues 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 Kernelized Attention, which applies the modified Nyström method to the kernelized attention score matrix C. The next two subsections will continue our discussion on it, and respectively state the theoretical analysis of its approximation error and some details of its implementation in practice.

heading: 4.3 Error analysis of Skyformer
text: As mentioned, an implicit advantage of using Kernelized Attention is that we can leverage the existing conclusions for kernel methods to analyze the theoretical properties of the model. In this subsection, we aim to provide some theoretical analysis of its approximation error.
We state a high probability bound on the size d of the sub-sampling matrix used in Skyformer to attain (ε, δ)-MA property for the kernelized attention score matrix C by the following theorem. We refer the readers to the proof in Appendix D 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 9 and Theorem 3 [Musco and Musco, 2017]). Consider the query, key, and value matrix Q,K,V ∈ Rn×p and two positive constants ε < 1, δ < 12 . For the empirical Gaussian kernel matrix C := κ(Q/p1/4,K/p1/4) defined above, we let λ := ε‖C‖ < ‖C̄‖, where C̄ is the completion of C (similar to B̄, constructed as substituting the Gaussian kernel with bandwidth p1/4 for the arbitrary kernel function φ in Equation (4)). We comment λ 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 C̄(C̄ + λI2n)−1 as leverage score `i,∀i = 1, . . . , 2n, and define their sum Tr ( C̄(C̄ + λI)−1 ) as the statistical dimension dstat, which increases with 1/ε as λ ∝ ε. Suppose S is a uniform sub-sampling matrix, and assume there exists a constant β ∈ (0, 1] such that β ≤ dstat2n`i ,∀i = 1, . . . , 2n. For the approximation matrix
˜̄C constructed with C̄,S as in Equation (5), there exists a constant C such that if
d ≥ C dstat β log n δ
then ˜̄C 4 C̄ 4 ˜̄C + λI with probability 1 − δ. Here 4 denotes the Loewner ordering: B 4 A means A−B is positive semidefinite. Furthermore, for our approximation C̃ in Equation (6) to the kernelized attention score C, we have
‖C̃ −C‖ ≤ λ = ε‖C‖.
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 Õ(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 dstat. We leave the exact non-asymptotic analysis of the computational complexity for future work.

heading: 4.4 Workaround in implementation
text: A 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 in Nyströmformer [Xiong et al., 2021] to our setting. Specifically, we use the matrix-product-based iterative method [Razavi et al.] for finding approximate inverses, instead of some division-based methods (such as the conjugate gradient method) which induces some instability in model training.
To apply the iterative method and inverse matrix M = ST C̄S, we need to satisfy its assumption [Razavi et al., Theorem 2] that ‖I −M‖ < 1. In practice, we instead pass the matrix D−1/2M (M + γI)D
−1/2 M as an input to the iterative method, where γ > 0 is a small constant and the diagonal
matrix DM is defined as diag ((M + γI)1). We give the following lemma to justify our practical usage of the method. The proof is deferred to Appendix E.
Lemma 3. Given a constant γ > 0, if matrices M is constructed as ST C̄S, and DM are defined as above, then all the singular values of D−1/2M (M + γI)D −1/2 M are within (0, 1), which implies that ‖I −D−1/2M (M + γI)D −1/2 M ‖ < 1.
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 use 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 a dense matrix, which increases the computational cost.

heading: 4.5 Empirical approximation evaluation
text: Spectral norm, the maximum singular value of a matrix, is a 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-cased models from Huggingface’s implementation [Wolf et al., 2019] . The input vector X is embedded from the tokenized raw text in Wikitext-2 dataset [Merity et al., 2017]. The query, key and value weight matrices in initialized or pretrained models transform input X into Q,K, V of different distributions. We compare the results with different sequence lengths and different numbers of features used in attention approximation methods. We set the number of features in the range of 24 to 28. More features usually require more computation resources.
Figure 1 shows the performance of the modified Nyström method on approximation error with regards to the number of features. We conclude that for Skyformer the approximation is 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 a 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 on classification tasks, we find the variant suffers a more severe gradient explosion issue than usual transformers. We speculate that it is because the matrix ST ĀS (in the middle of Equation (6)) inherits the high condition number of the original attention score matrices A, while the derivative of matrix inverse ( ( A−1 )′ = −A−1A′A−1) further amplifies the condition number during backpropagation.

heading: 5 Experimental Results
text: Tasks and Datasets. We evaluate the proposed methods on five classification tasks on LRA benchmark [Tay et al., 2020b], which focuses on model quality under long-context scenarios: ListOps [Nangia and Bowman, 2018], Text Classification on IMDb review dataset [Maas et al., 2011], Document Retrieval on AAN dataset [Radev et al., 2013], Pathfinder [Linsley et al., 2018], and Image Classification on CIFAR-10 [Krizhevsky et al., 2009]. The LRA benchmark covers diverse longsequence 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
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 [Zaheer et al., 2020], Performer [Choromanski et al., 2020], Linformer [Wang et al., 2020], Nyströmformer [Xiong et al., 2021], Informer [Zhou et al., 2020], and Reformer [Kitaev et al., 2020]. 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 self-attention module with the various attention methods and keep other experimental settings the same for fair comparisons.
Implementation Details. We conduct each experiment on one Tesla V100 SXM2 16GB. We use the LRA evaluation benchmark reimplemented in PyTorch by Xiong et al. [2021]. 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 1e − 4 for Text Classification, ListOps, and Image Classification, and 2e − 4 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.
We do not follow all settings in [Xiong et al., 2021] due to the hardware limitation. The compromises, such as approximation dimension and gradient accumulations steps, might bring performance differences 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 as observed in Figure 2: it takes more steps to reach the stationary distribution of its long-time limit, and it is more easily getting stuck in a 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 1. However, the time and space requirement of Kernelized 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 longtime limit with comparable general performance in classification accuracy (Table 1) and resource consumption (Table 2). The advantages over the standard self-attention are significant with consistently less training time and generally better performance. For example, Skyformer brings nearly 4
times speed-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 n the statistical dimension dstat might be close to n.
To make the claim above clear, we first reiterate that the efficiency of Skyformer is related to dstat. As implied by Theorem 2, the intrinsic difficulty of approximating a raw attention score matrix is 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 it is possible the statistical dimension associated with a short sequence might be even close to the sequence n. Therefore a large n serves as a condition to make the method work. Figure 1 empirically shows that our method performs better with larger n’s.

heading: 6 Conclusions and future work
text: Motivated by the connection between kernel methods and self-attention, we 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 to Kernelized Attention to 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 on the LRA benchmark.
Direct development of this work is the incorporation of further computation tricks in kernel methods, such as the local and global approximation for gram matrix [Snelson and Ghahramani, 2007] and the importance sampling in Nyström methods [Musco and Musco, 2017, Chen and Yang, 2021b,a]. Other related questions include the choice of the kernel other than the Gaussian kernel in 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 on 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. FA8750-19-2-1004. This work is also in part supported by NSF grant DMS-1810831. 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 is authorized to reproduce and distribute reprints for governmental purposes notwithstanding any copyright annotation therein.

Review summaries:
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).

The paper presents an alternative approach to attention in the Transformer. There are two modeling contributions:  1. Kernelized attention that omits the typical normalization term used for softmax attention. 2. A Nystrom method approach that allows efficiently approximating the output of the attention mechanism  The 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.  Below I summarize my understanding/takeaways from the two key methods; please correct me in the author response if these are inaccurate: 1. 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. 2. 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.

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.

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). 

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.   In 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.  By 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.   The 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.    Empirical 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.