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