{'fraction_non_alphanumeric': 0.03693439215477269, 'fraction_numerical': 0.0139370036774503, 'mean_word_length': 4.61277295842058, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 2, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Skyformer: Remodel Self-Attention with Gaussian Kernel and Nystr\"om Method

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\"om 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.


Introduction
The cost of language model training increases exponentially. Among different models, Transformerbased language models [Vaswani et al., 2017, Devlin 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. 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 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 andSeeger, 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

Related Work
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].
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  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 andSiminelakis, 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 andSeeger, 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].

Revisiting self-attention
For a given input sequence X ∈ R n×d0 of length n and embedding dimension d 0 , The dot-product attention for a single head in Transformer Vaswani et al. [2017] is defined as where Q = XW Q , K = XW K , and V = XW V , and W Q , W K and W V are the query, key, and value weight metrics that linearly project the input X of d 0 dimension to an output tensor of p dimensions.
To simplify the future analysis, the left softmax term can be rewritten into D −1 A, where A := exp(QK T / √ p) is the un-normalized attention score matrix; D is a diagonal matrix whose diagonal is exp(QK T / √ 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(q T k/ √ p) as the softmax kernel function, and represent A by the notation SM(Q, K), which means the element a ij from the i-th row and j-th column in A is equal to SM(q i , k j ). Throughout this paper q i (resp. k j ) 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 Proof. We first state an important equation to connect the softmax kernel and Gaussian kernels as follows: 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 We prove the lemma by using the fact that κ is a PSD kernel and κ Q p 1/4 , Q p 1/4 is a PSD matrix. ♦

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 approximationB, 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 matrixB = BS(S T BS) † S T B, where (·) † denotes the Moore-Penrose pseudoinverse of a matrix, and S ∈ R n×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 ∈ R n×d , if S has i.i.d. columns and the j-th column S (j) can randomly be 1 d e i with probability 1 n , where e i is the i-th column of the n-by-n identity matrix I n , 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.

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 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 ∈ R n1×n2 , two constants ε > 0, δ < 1 2 , we say that its approximation matrix In previous works, the direct analysis of the approximation error to the entire output D −1 AV 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.

Kernelized Attention
Kernelized Attention replaces the softmax structure in vanilla self-attention with a Gaussian kernel, and the new attention model is stated as: where we define the n-by-n matrix C as the kernelized attention score matrix κ(Q/p 1/4 , K/p 1/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 κ(q i , k j ) := exp − q i − k j 2 /2 . Via this expression, for token i in the query, Gaussian kernel assigns a large attention score to the token j when k j is close to q i . 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 . 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 −1 A. 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.

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 b ij from the i-th row and j-th column in B is equal to φ(q i , k j ), where q i (resp. k j ) 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/p 1/4 , K/p 1/4 ). Therefore this method leads to a low-rank approximation to the output of either self-attention D −1 AV 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 matrixB: Then we approximateB withB through where S is a 2n-by-d uniform sub-sampling matrix as defined in Definition 1. The final approximation will be given asB The original matrix B can be well-approximated byB due to the following inequality and thus we show our task of approximating the non-PSD matrix B boils down to well approximating the PSD matrixB.
Remark. The reason why B can be well approximated by a low-rankB is that as an empirical kernel matrix the eigenvalues inB 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) ofB) has enough potential to well approximate the original matrixB (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.

Error analysis of Skyformer
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 ∈ R n×p and two positive constants ε < 1, δ < 1 2 . For the empirical Gaussian kernel matrix C := κ(Q/p 1/4 , K/p 1/4 ) defined above, we let λ := ε C < C , whereC is the completion of C (similar toB, constructed as substituting the Gaussian kernel with bandwidth p 1/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 ofC(C + λI 2n ) −1 as leverage score i , ∀i = 1, . . . , 2n, and define their sum Tr C (C + λI) −1 as the statistical dimension d stat , which increases with 1/ε as λ ∝ ε. Suppose S is a uniform sub-sampling matrix, and assume there exists a constant β ∈ (0, 1] such that β ≤ dstat 2n i , ∀i = 1, . . . , 2n. For the approximation matrixC constructed withC, S as in Equation (5), there exists a constant C such that if Here denotes the Loewner ordering: B A means A − B is positive semidefinite. Furthermore, for our approximationC in Equation (6) to the kernelized attention score C, we have This theorem implies the time and space complexity of our proposed approximation depends on the statistical dimension d stat . If we directly use the conclusion from Gaussian kernels, d stat 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 d stat . We leave the exact non-asymptotic analysis of the computational complexity for future work.

Workaround in implementation
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 = S TC 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/2 M (M + γI)D −1/2 M as an input to the iterative method, where γ > 0 is a small constant and the diagonal matrix D M 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. 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.

Empirical approximation evaluation
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 2 4 to 2 8 . 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 S TĀ 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 −1 A A −1 ) further amplifies the condition number during backpropagation.

Experimental Results
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 Figure 1: Spectral norm results with different sequence lengths under different W Q , W K , W V 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-attention. In this experiment, "Skyformer" also needs to first approximate A, and then approximate D, as Performer does.   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 , 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 d s tat might be close to n.
To make the claim above clear, we first reiterate that the efficiency of Skyformer is related to d s tat.
As implied by Theorem 2, the intrinsic difficulty of approximating a raw attention score matrix is concluded as d stat , which corresponds to the effective rank of matrixC. The complexity of Skyformer depends on the sub-sample size d (the size of the sub-sampling matrix S). A large d stat leads to a large d , and an inefficient application of the Nyström method.
The classical theory for statistical dimension only guarantees that d stat 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.

Conclusions and future work
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 andMusco, 2017, Chen andYang, 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.   The singular values decay fast and thus justify the low-rank approximation, as analyzed by , Dong et al. [2021]. We propose to measure the task difficulty with the singular value decay rate in attention output, as higher intrinsic task difficulty forces the model to output a matrix with more large singular values. Such matrices are considered more informative since they are harder to approximate, requiring more ranks even in the truncated SVD approximation. With the observation in Figure 4, we conclude that the singular values in Document Retrieval and Pathfinder tasks decay slower, and those two tasks are more difficult than Text Classification and ListOps.

C Useful facts
This section introduces some useful facts, which are key in the proof in the next section. To start with, we provide a matrix concentration inequality as follows. Lemma C.1 (Matrix Bernstein Inequality [Tropp, 2012]). Consider a finite sequence {X k } of independent, random, self-adjoint matrices with dimension n. Assume that each random matrix satisfies E X k = 0 and X k ≤ R almost surely. Then, for all t ≥ 0, For a certain n-by-n orthogonal matrix H (HH T is a diagonal matrix) and an n-by-d uniform sub-sampling matrix S (as defined in Definition 1 in the main paper), we denote the sketching matrix Π := √ nS. We aim to show HΠΠ T H T can satisfy ( 1 2 , δ)-MA property for HH T by the following lemma.
Proof. The main idea is to utilize Lemma C.1 by setting t = 1 2 HH T = 1 2 H 2 . Specifically, we denote the matrices We still need two steps to give control of R and σ 2 . For R, we have where {z ki } n i=1 are the indicators of whether the i-th column is chosen. The first inequality of the preceding display holds due to the fact that H is an orthogonal matrix. Using the condition n ≤ H 2 F β H (i) 2 , ∀i = 1, . . . , n, we further have and we thus set R := H 2 F dβ . On the other hand, Again using the condition that n H (i) 2 ≤ H 2 F β , ∀i = 1, . . . , n, we reach and set σ 2 : Finally we plug R and σ 2 into Lemma C.1 and obtain: .
To ensure the right-hand-side is smaller than δ, we just need which validates the lemma. ♦

D Proof of Theorem 2 in the main paper
Proof. The conclusion in the lemma can be divided into two parts, thatC C andC C + λI.
To prove them we first introduce some notations and auxiliary results. SinceC is PSD, there exists a matrix B satisfying BB T =C. We further denote B's SVD decomposition as B = U Σ 1 2 V T (C = U ΣU T ), where both U and V are 2n-by-2n orthonormal matrices. (In this section we slightly abuse the notation that V represents the matrix of right-singular vectors, instead of the value matrix in self-attention.) DefineΣ := Σ + λI, Ψ := U Σ 1 2Σ − 1 2 , which impliesC(C + λI) −1 = ΨΨ T . Also following the notations in the last section, we define the 2n-by-d matrix Π := √ 2nS With those notations,C can be rewritten as BB T Π(Π T BB T Π) † ΠBB T = BP Π B T , where P Π is the orthogonal projection matrix for the column space of B T Π. It is easy to check that C −C = B(I − P Π )B T . Since I − P Π is an orthogonal projection matrix (which is PSD), we haveC −C 0, which proves the first conclusion thatC C .
For the second conclusion, we utilize the following important identity: The proof of the statement above is similar to the proof of the well-known conclusion that graph Laplacian matrix is PSD. For simplicity we denote W := M + γI, and given any vector x ∈ R d we have where the last equation holds due to the fact that (D M ) ii = d j=1 W ij .
Combining the pieces above we can conclude that I − D

F Additional discussions about the stability in model training
For our argument about stability, we mainly refer to the paper [Liu et al., 2020a], which identifies that the amplification of small parameter perturbations in the self-attention module is the root cause of training instability. We take kernelized attention as mitigation since it contains an automatic normalization. We have empirically used Figure 2 and Figure 3 in Appendix A to support our claim.
For further analysis we conduct a toy experiment adapting from Figure 4 in the aforementioned paper [Liu et al., 2020a]. We aim to show that in kernelized attention (and Skyformer) the output changes f (x, W * ) − f (x, W ) for parameter changes W * − W is smaller than in self-attention (and its approximation Nyströmformer). This concept involved is somewhat similar to condition number and below we will formalize it as "instability score".
We show a table of the averaged ratios between the instability scores of kernelized attention (we also add Skyformer and Nyströmformer for reference) and self-attention to conclude our statement about stability. A ratio smaller than 1 means higher stability compared to self-attention. We follow all the settings in Table 1 in the main paper except here we only update the model for 20 steps (we limit the number of steps as suggested by Liu et al. [2020a] to make the results of the same step comparable among different models). In step i for each model we compute the instability score , 20, where f () gives the embedding after two layers, x i is the i-th input sequence batch, W 0 represents the initial parameters, and W i represents the parameters after step i. In each step we compute the ratio of a certain method's τ i to the τ i of self-attention, and finally average the 20 ratios in Table 3 in the appendix.
As we can observe, both kernelized attention and Skyformer consistently have a lower instability score than self-attention, while the instability score of Nyströmformer, an approximation to self-attention, fluctuates around 1 in all the tasks. The results support our claim that the proposed kernelized attention can improve stability.