Skyformer: Remodel Self-Attention with Gaussian Kernel and Nystr\"om Method

1 Introduction . 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 can not 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 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 ] . 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 ] . 3 Preliminaries and notations . 3.1 Revisiting self-attention . 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 . ♦

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