Rethinking Attention with Performers

1 INTRODUCTION AND RELATED WORK . Transformers ( Vaswani et al. , 2017 ; Dehghani et al. , 2019 ) are powerful neural network architectures that have become SOTA in several areas of machine learning including natural language processing ( NLP ) ( e.g . speech recognition ( Luo et al. , 2020 ) ) , neural machine translation ( NMT ) ( Chen et al. , 2018 ) , document generation/summarization , time series prediction , generative modeling ( e.g . image generation ( Parmar et al. , 2018 ) ) , music generation ( Huang et al. , 2019 ) , and bioinformatics ( Rives et al. , 2019 ; Madani et al. , 2020 ; Ingraham et al. , 2019 ; Elnaggar et al. , 2019 ; Du et al. , 2020 ) . Transformers rely on a trainable attention mechanism that identifies complex dependencies between the elements of each input sequence . Unfortunately , the regular Transformer scales quadratically with the number of tokens L in the input sequence , which is prohibitively expensive for large L and precludes its usage in settings with limited computational resources even for moderate values of L. Several solutions have been proposed to address this issue ( Beltagy et al. , 2020 ; Gulati et al. , 2020 ; Chan et al. , 2020 ; Child et al. , 2019 ; Bello et al. , 2019 ) . Most approaches restrict the attention mechanism to attend to local neighborhoods ( Parmar et al. , 2018 ) or incorporate structural priors on attention such as sparsity ( Child et al. , 2019 ) , pooling-based compression ( Rae et al. , 2020 ) clustering/binning/convolution techniques ( e.g . ( Roy et al. , 2020 ) which applies k-means clustering to learn dynamic sparse attention regions , or ( Kitaev et al. , 2020 ) , where locality sensitive hashing is used to group together tokens of similar embeddings ) , sliding windows ( Beltagy et al. , 2020 ) , or truncated targeting ( Chelba et al. , 2020 ) . There is also a long line of research on using dense attention matrices , but defined by low-rank kernels substituting softmax ( Katharopoulos et al. , 2020 ; Shen et al. , 2018 ) . Those methods critically rely on kernels admitting explicit representations as dot-products of finite positive-feature vectors . The approaches above do not aim to approximate regular attention , but rather propose simpler and more tractable attention mechanisms , often by incorporating additional constraints ( e.g . identical query and key sets as in ( Kitaev et al. , 2020 ) ) , or by trading regular with sparse attention using more ∗Equal contribution . Correspondence to { kchoro , lcolwell } @ google.com . Code for Transformer models on protein data can be found in github.com/google-research/ google-research/tree/master/protein_lm and Performer code can be found in github.com/ google-research/google-research/tree/master/performer . Google AI Blog : https : // ai.googleblog.com/2020/10/rethinking-attention-with-performers.html layers ( Child et al. , 2019 ) . Unfortunately , there is a lack of rigorous guarantees for the representation power produced by such methods , and sometimes the validity of sparsity patterns can only be verified empirically through trial and error by constructing special GPU operations ( e.g . either writing C++ CUDA kernels ( Child et al. , 2019 ) or using TVMs ( Beltagy et al. , 2020 ) ) . Other techniques which aim to reduce Transformers ’ space complexity include reversible residual layers allowing one-time activation storage in training ( Kitaev et al. , 2020 ) and shared attention weights ( Xiao et al. , 2019 ) . These constraints may impede application to long-sequence problems , where approximations of the attention mechanism are not sufficient . Approximations based on truncated back-propagation ( Dai et al. , 2019 ) are also unable to capture long-distance correlations since the gradients are only propagated inside a localized window . Other methods propose biased estimation of regular attention but only in the non-causal setting and with large mean squared error ( Wang et al. , 2020 ) . In response , we introduce the first Transformer architectures , Performers , capable of provably accurate and practical estimation of regular ( softmax ) full-rank attention , but of only linear space and time complexity and not relying on any priors such as sparsity or low-rankness . Performers use the Fast Attention Via positive Orthogonal Random features ( FAVOR+ ) mechanism , leveraging new methods for approximating softmax and Gaussian kernels , which we propose . We believe these methods are of independent interest , contributing to the theory of scalable kernel methods . Consequently , Performers are the first linear architectures fully compatible ( via small amounts of fine-tuning ) with regular Transformers , providing strong theoretical guarantees : unbiased or nearly-unbiased estimation of the attention matrix , uniform convergence and lower variance of the approximation . FAVOR+ can be also applied to efficiently model other kernelizable attention mechanisms beyond softmax . This representational power is crucial to accurately compare softmax with other kernels for the first time on large-scale tasks , that are beyond the reach of regular Transformers , and find for them optimal attention-kernels . FAVOR+ can also be applied beyond the Transformer scope as a more scalable replacement for regular attention , which itself has a wide variety of uses in computer vision ( Fu et al. , 2019 ) , reinforcement learning ( Zambaldi et al. , 2019 ) , training with softmax cross entropy loss , and even combinatorial optimization ( Vinyals et al. , 2015 ) . We test Performers on a rich set of tasks ranging from pixel-prediction through text models to protein sequence modeling . We demonstrate competitive results with other examined efficient sparse and dense attention methods , showcasing the effectiveness of the novel attention-learning paradigm leveraged by Performers . We emphasize that in principle , FAVOR+ can also be combined with other techniques , such as reversible layers ( Kitaev et al. , 2020 ) or cluster-based attention ( Roy et al. , 2020 ) . 2 FAVOR+ MECHANISM & POSITIVE ORTHOGONAL RANDOM FEATURES . Below we describe in detail the FAVOR+ mechanism - the backbone of the Performer′s architecture . We introduce a new method for estimating softmax ( and Gaussian ) kernels with positive orthogonal random features which FAVOR+ leverages for the robust and unbiased estimation of regular ( softmax ) attention and show how FAVOR+ can be applied for other attention-kernels . 2.1 PRELIMINARIES - REGULAR ATTENTION MECHANISM . Let L be the size of an input sequence of tokens . Then regular dot-product attention ( Vaswani et al. , 2017 ) is a mapping which accepts matrices Q , K , V ∈ RL×d as input where d is the hidden dimension ( dimension of the latent representation ) . Matrices Q , K , V are intermediate representations of the input and their rows can be interpreted as queries , keys and values of the continuous dictionary data structure respectively . Bidirectional ( or non-directional ( Devlin et al. , 2018 ) ) dot-product attention has the following form , where A ∈ RL×L is the so-called attention matrix : Att↔ ( Q , K , V ) = D −1AV , A = exp ( QK > / √ d ) , D = diag ( A1L ) . ( 1 ) Here exp ( · ) is applied elementwise , 1L is the all-ones vector of length L , and diag ( · ) is a diagonal matrix with the input vector as the diagonal . Time and space complexity of computing ( 1 ) are O ( L2d ) and O ( L2 + Ld ) respectively , because A has to be stored explicitly . Hence , in principle , dot-product attention of type ( 1 ) is incompatible with end-to-end processing of long sequences . Bidirectional attention is applied in encoder self-attention and encoder-decoder attention in Seq2Seq architectures . Another important type of attention is unidirectional dot-product attention which has the form : Att→ ( Q , K , V ) = D̃ −1ÃV , Ã = tril ( A ) , D̃ = diag ( Ã1L ) , ( 2 ) where tril ( · ) returns the lower-triangular part of the argument matrix including the diagonal . As discussed in ( Vaswani et al. , 2017 ) , unidirectional attention is used for autoregressive generative modelling , e.g . as self-attention in generative Transformers as well as the decoder part of Seq2Seq Transformers . We will show that attention matrix A can be approximated up to any precision in time O ( Ld2 log ( d ) ) . For comparison , popular methods leveraging sparsity via Locality-Sensitive Hashing ( LSH ) techniques ( Kitaev et al. , 2020 ) have O ( Ld2 logL ) time complexity . In the main body of the paper we will describe FAVOR+ for bidirectional attention . Completely analogous results can be obtained for the unidirectional variant via the mechanism of prefix-sums ( all details in the Appendix B.1 ) . 2.2 GENERALIZED KERNELIZABLE ATTENTION . FAVOR+ works for attention blocks using matrices A ∈ RL×L of the form A ( i , j ) = K ( q > i , k > j ) , with qi/kj standing for the ith/jth query/key row-vector in Q/K and kernel K : Rd × Rd → R+ defined for the ( usually randomized ) mapping : φ : Rd → Rr+ ( for some r > 0 ) as : K ( x , y ) = E [ φ ( x ) > φ ( y ) ] . ( 3 ) We call φ ( u ) a random feature map for u ∈ Rd . For Q′ , K′ ∈ RL×r with rows given as φ ( q > i ) > and φ ( k > i ) > respectively , Equation 3 leads directly to the efficient attention mechanism of the form : Âtt↔ ( Q , K , V ) = D̂ −1 ( Q′ ( ( K′ ) > V ) ) , D̂ = diag ( Q′ ( ( K′ ) > 1L ) ) . ( 4 ) Here Âtt↔ stands for the approximate attention and brackets indicate the order of computations . It is easy to see that such a mechanism is characterized by space complexity O ( Lr + Ld+ rd ) and time complexity O ( Lrd ) as opposed to O ( L2 + Ld ) and O ( L2d ) of the regular attention ( see also Fig . 1 ) . The above scheme constitutes the FA-part of the FAVOR+ mechanism . The remaining OR+ part answers the following questions : ( 1 ) How expressive is the attention model defined in Equation 3 , and in particular , can we use it in principle to approximate regular softmax attention ? ( 2 ) How do we implement it robustly in practice , and in particular , can we choose r L for L d to obtain desired space and time complexity gains ? We answer these questions in the next sections . 2.3 HOW TO AND HOW NOT TO APPROXIMATE SOFTMAX-KERNELS FOR ATTENTION . It turns out that by taking φ of the following form for functions f1 , ... , fl : R → R , function g : Rd → R and deterministic vectors ωi or ω1 , ... , ωm iid∼ D for some distribution D ∈ P ( Rd ) : φ ( x ) = h ( x ) √ m ( f1 ( ω > 1 x ) , ... , f1 ( ω > mx ) , ... , fl ( ω > 1 x ) , ... , fl ( ω > mx ) ) , ( 5 ) we can model most kernels used in practice . Furthermore , in most cases D is isotropic ( i.e . with pdf function constant on a sphere ) , usually Gaussian . For example , by taking h ( x ) = 1 , l = 1 and D = N ( 0 , Id ) we obtain estimators of the so-called PNG-kernels ( Choromanski et al. , 2017 ) ( e.g . f1 = sgn corresponds to the angular kernel ) . Configurations : h ( x ) = 1 , l = 2 , f1 = sin , f2 = cos correspond to shift-invariant kernels , in particular D = N ( 0 , Id ) leads to the Gaussian kernel Kgauss ( Rahimi & Recht , 2007 ) . The softmax-kernel which defines regular attention matrix A is given as : SM ( x , y ) def = exp ( x > y ) . ( 6 ) In the above , without loss of generality , we omit √ d-renormalization since we can equivalently renormalize input keys and queries . Since : SM ( x , y ) = exp ( ‖x‖ 2 2 ) Kgauss ( x , y ) exp ( ‖y‖2 2 ) , based on what we have said , we obtain random feature map unbiased approximation of SM ( x , y ) using trigonometric functions with : h ( x ) = exp ( ‖x‖ 2 2 ) , l = 2 , f1 = sin , f2 = cos. We call it ŜM trig m ( x , y ) . There is however a caveat there . The attention module from ( 1 ) constructs for each token , a convex combination of value-vectors with coefficients given as corresponding renormalized kernel scores . That is why kernels producing non-negative scores are used . Applying random feature maps with potentially negative dimension-values ( sin / cos ) leads to unstable behaviours , especially when kernel scores close to 0 ( which is the case for many entries of A corresponding to low relevance tokens ) are approximated by estimators with large variance in such regions . This results in abnormal behaviours , e.g . negative-diagonal-values renormalizers D−1 , and consequently either completely prevents training or leads to sub-optimal models . We demonstrate empirically that this is what happens for ŜM trig m and provide detailed theoretical explanations showing that the variance of ŜM trig m is large as approximated values tend to 0 ( see : Section 3 ) . This is one of the main reasons why the robust random feature map mechanism for approximating regular softmax attention was never proposed . We propose a robust mechanism in this paper . Furthermore , the variance of our new unbiased positive random feature map estimator tends to 0 as approximated values tend to 0 ( see : Section 3 ) . Lemma 1 ( Positive Random Features ( PRFs ) for Softmax ) . For x , y ∈ Rd , z = x + y we have : SM ( x , y ) = Eω∼N ( 0 , Id ) [ exp ( ω > x−‖x‖ 2 2 ) exp ( ω > y−‖y‖ 2 2 ) ] = ΛEω∼N ( 0 , Id ) cosh ( ω > z ) , ( 7 ) where Λ = exp ( −‖x‖ 2+‖y‖2 2 ) and cosh is hyperbolic cosine . Consequently , softmax-kernel admits a positive random feature map unbiased approximation with h ( x ) = exp ( −‖x‖ 2 2 ) , l = 1 , f1 = exp and D = N ( 0 , Id ) or : h ( x ) = 1√2 exp ( − ‖x‖2 2 ) , l = 2 , f1 ( u ) = exp ( u ) , f2 ( u ) = exp ( −u ) and the same D ( the latter for further variance reduction ) . We call related estimators : ŜM + m and ŜM hyp+ m . In Fig . 2 we visualize the advantages of positive versus standard trigonometric random features . In critical regions , where kernel values are small and need careful approximation , our method outperforms its counterpart . In Section 4 we further confirm our method ’ s advantages empirically , using positive features to efficiently train softmax-based linear Transformers . If we replace in ( 7 ) ω with√ d ω‖ω‖ , we obtain the so-called regularized softmax-kernel SMREG which we can approximate in a similar manner , simply changing D = N ( 0 , Id ) to D = Unif ( √ dSd−1 ) , a distribution corresponding to Haar measure on the sphere of radius √ d in Rd , obtaining estimator ̂SMREG + m. As we show in Section 3 , such random features can also be used to accurately approximate regular softmax-kernel .

Summary: The paper proposed a theoretically grounded O(N) approximation of the softmax attention. The key idea is to interpret attention as a kernel function and construct the random feature projection that can reproduce this kernel. It is highly non-trivial to derive a feature mapping that can accurately approximate the softmax kernel. To better approximate the softmax kernel, the author proposed some important design choices, all of which are supported by theoretical and empirical evidences. The author showed that 1) adopting non-negative random features is very essential to the approximation and the proposed Positive Random Features (PRF) can effectively reduce the variance when the attention values are small, 2) drawing orthogonal random matrices can further reduce the variance of the approximation, 3) the final proposed Performer model runs faster, takes less memory, and has better performance than other O(N) and O(N logN) attention methods.