Stable, Fast and Accurate: Kernelized Attention with Relative Positional Encoding

1 Introduction . Transformer [ 45 ] has achieved tremendous success on a variety of tasks in natural language processing [ 34 , 9 , 50 ] , computer vision [ 2 , 10 , 42 ] and speech [ 13 ] . Despite the great success , Transformers are usually inefficient in modeling long-sequence input since the key component—the attention module— needs to calculate pairwise correlations between all the positions , which comes with quadratic time and memory cost with respect to the sequence length n. To reduce the cost , many works focus on approximating the dot-then-exponentiate function in the attention module of the original Transformer [ 41 ] . Most works can be roughly categorized into three classes : sparse attention [ 3 , 33 , 30 , 15 , 21 , 46 , 7 ] , low-rank approximation-based attention [ 49 , 48 ] and kernelized attention [ 17 , 32 , 4 ] . Sparse attention variants design either pre-defined ∗Equal contribution . †Correspondence to : Di He < dihe @ microsoft.com > and Shuxin Zheng < shuz @ microsoft.com > . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . [ 3 , 33 , 30 , 15 ] or learnable patterns [ 21 , 46 , 7 ] to reduce the amount of the key-value pairs that each query needs to attend to . Low-rank approximation-based attention variants project the key matrix into a length-agnostic low-dimensional space [ 48 ] , or use Nyström approximation [ 49 ] to reduce the quadratic dependency on sequence length . The kernelized attention is first developed in [ 44 , 17 ] , which introduces a kernel view on the dot-then-exponentiate softmax function and comes up with a linear attention mechanism . Followed-up works [ 32 , 4 ] theoretically investigate the choices of different kernel feature maps and propose several random feature maps that achieve much better performance on various applications . However , in recently developed Transformers , the attention mechanism is designed to be more complicated than dot-then-exponentiation . For example , the relative positional encoding ( RPE ) [ 38 , 35 ] has become standard in state-of-the-art models [ 35 , 18 , 14 ] : Besides using the dot-products between queries and keys , RPE incorporates a position-correlation matrix in the softmax exponentiation to encode the distance between any two positions . With RPE , Transformers can effectively capture the relative word orders and achieve superior performance on various tasks . As Transformer with RPE is more powerful , one may wonder whether the model can be still approximated by directly applying previous acceleration approaches . Unfortunately , we mathematically show that such a term expands the expressiveness of dot-thenexponentiation , i.e. , for some attention with RPE , there exists no dot-then-exponentiate function that can represent it . Therefore , previous approaches can not be directly applied to approximate Transformers with RPE and obtain non-trivial speed-up . In this paper , we develop a novel attention computation that achieves an almost-linear speed-up rate for Transformer with RPE . We build our method on top of the kernelized attention in [ 4 , 32 ] . The kernelized attention method replaces the dot-then-exponentiate attention by a simple matrix multiplication between the kernelized features of queries and key-value pairs . When taking RPE into the kernelized attention , one can naively multiply the n× n RPE position-correlation matrix with kernelized feature matrices , which leads to a trivial quadratic time cost . While interestingly , we find that the elements in the RPE position-correlation matrix are determined by the relative distance between positions , and this matrix is precisely a Toeplitz matrix . In numerical linear algebra , it is well-known that multiplication between a Toeplitz matrix and row vectors can be accelerated using Fast Fourier Transform ( FFT ) . And by using FFT , we can obtain a fast attention computation with O ( n log n ) time cost . Beyond the efficiency , we also find that our method that incorporates RPE into kernelized attention can also stabilize the training of such models . As demonstrated by several works [ 16 , 4 ] , although these kernelized attention methods have theoretical guarantees to approximate the softmax attention , they face optimization issues and will diverge when training from scratch for large-scale tasks . This significantly limits the practical impact of the model as we have to train a standard Transformer first and convert it to kernelized attention with finetuning . We first shed light upon the reason for this phenomenon by proving that these kernelized attentions have an unacceptable variance when the ` 2 norms of queries and keys are large . However , queries and keys with large norms usually appear during training , leading to optimization instability or sub-optimal solutions . When incorporating RPE , we can properly normalize queries and keys to obtain a “ low-variance ” attention , making the training stable . Since the value range of the relative positional encodings is not constrained , the whole model can still have enough capacity to achieve good performance . We conduct experiments on a wide range of tasks spanning language pre-training , machine translation and image classification . On all the tasks , we show that our model can be trained from scratch stably , outperforms other competitive efficient models , and is faster than the standard Transformer in the long-sequence regime . 2 Preliminary . 2.1 Attention Module and its Kernel View . The attention module is one of the key components in the Transformer [ 45 , 9 ] . It is usually formulated as querying a dictionary with key-value pairs , e.g. , Attention ( Q , K , V ) = softmax ( QWQ ( KWK ) > √ d ) VWV , where d is the dimensionality of the hidden representations . Q ( Query ) , K ( Key ) , V ( Value ) are specified as the hidden representations of the previous layer and WQ , WK and WV are projection matrices for any specific head . To be more concrete , we denote x = ( x1 , x2 · · · , xn ) as the input to the attention module , where n is the length of the sequence and xi ∈ Rd is a row vector denoting the contextual representation of the token at position i. Denote z = ( z1 , z2 · · · , zn ) as the output . Then the attention module can be written as zi = n∑ j=1 exp ( αij ) ∑n j′=1 exp ( αij′ ) ( xjW V ) , where αij = 1√ d ( xiW Q ) ( xjW K ) > . ( 1 ) As we can see , the calculation of attention for each position requires O ( n ) time and memory cost . Thus , the time and memory cost for the whole sequence is O ( n2 ) , which impedes the attention mechanism scaling to long sequences . Recently , [ 44 , 17 ] introduce a kernel view on the attention mechanism and derive a general kernelized attention framework as follows : zi = n∑ j=1 κ ( xiW Q , xjW K ) ∑n j′=1 κ ( xiW Q , xj′WK ) ( xjW V ) , ( 2 ) where κ ( · , · ) : Rd × Rd → R is any positive-definite kernel measuring the pairwise similarity . Based on this kernel view , the attention mechanism can be alternatively modeled as a linear dot-product of kernelized features : zi = n∑ j=1 φ ( xiW Q ) φ ( xjW K ) > ∑n j′=1 φ ( xiW Q ) φ ( xj′WK ) > ( xjW V ) = φ ( xiW Q ) ∑n j=1 φ ( xjW K ) > ( xjW V ) φ ( xiWQ ) ∑n j′=1 φ ( xj′W K ) > , ( 3 ) where φ ( · ) : Rd → Rm is the feature map and m is the dimension of the feature space . Since both ∑n j=1 φ ( xjW K ) ( xjW V ) and ∑n j′=1 φ ( x ′ jW K ) can be reused for each position , the time and memory complexity of the kernelized attention is reduced to O ( n ) . [ 17 ] simply uses elu ( · ) + 1 as φ ( · ) . Followed-up works [ 32 , 4 ] further propose two ( randomized ) feature maps with which the kernels are unbiased estimators of the dot-then-exponentiate function , therefore achieve similar expressiveness as the standard Transformer . For simplicity , we call the feature maps used in [ 32 ] and [ 4 ] as Trigonometric Random Feature ( TRF ) and Positive Random Feature ( PRF ) respectively : φTRF ( x ) = exp ( ‖x‖ 2 2 ) √ m [ sin ( w1x > ) , ... , sin ( wmx > ) , cos ( w1x > ) , ... , cos ( wmx > ) ] , ( 4 ) φPRF ( x ) = exp ( −‖x‖ 2 2 ) √ m [ exp ( w1x > ) , ... , exp ( wmx > ) ] , ( 5 ) where each vector wi is independently sampled from N ( 0 , Id ) . Our work aims to design an efficient Transformer with relative positional encoding on top of the kernelized attention . Without loss of generality , we use Positive Random Features as the feature map to demonstrate the effectiveness of our approach . Our approach can also be applied to other feature maps such as the Trigonometric Random Features . See Section 4.5 for some ablation studies . 2.2 Relative Positional Encoding . Relative Positional Encoding ( RPE ) is recognized as one of the most successful modifications to the original Transformer model [ 27 ] . RPE is first introduced in [ 38 ] , which suggests that using absolute positional encoding may be ineffective in capturing relative word orders . They propose an embedding matrix in which the values of the elements are determined by the distance between the row index and column index . [ 35 ] further proposes a simple yet more effective relative positional encoding , which is popularly used in many recent state-of-the-art Transformers , including [ 1 , 18 , 24 ] . In this model [ 35 ] , the relative positional encoding is served as a bias term added to the dot-product between queries and keys : αrelij = 1√ d ( xiW Q ) ( xjW K ) > + bj−i . ( 6 ) For each j − i , bj−i is a learnable scalar and shared across all layers . 3 Incorporating RPE into Kernelized Attention . We are curious about whether previous acceleration methods can be directly applied to Transformers with RPE for achieving both speed-up and better accuracy . In this section , we first provide a negative result : We prove that RPE expands the expressiveness of the original attention module . So those efficient approaches that can only approximate dot-then-exponentiate functions can not approximate Transformers with RPE at least in the sense of representation power . To address the problem , we develop new techniques for computing attention with RPE on top of the kernelized attention , and show that we can still obtain non-trivial speed-up using Fast Fourier Transform . Lastly , we study how to stabilize the training of the PRF attention and achieve good performance with the help of RPE . Due to space limitations , we put all the proofs in Appendix B .

Summary: The authors propose a method to compute kernelized attention with relative position encoding (RPE) efficiently in $O(n\log n)$. They show that the RPE family of attention is richer than attention without RPE, and the usual $O(n)$ implementation of kernelized attention is no longer possible.  The derivation gives rise to a convolution, which the authors propose to compute in the Fourier domain, providing $O(n\log n)$ complexity, in contrast to $O(n^2)$  for a naive implementation.  The paper provides a comprehensive derivation of the algorithm, mathematically and in pseudo code. Interestingly, it is found that with the help of RPE training stability of kernelized attention can be improved, thus allowing for the first time to train certain models end-to-end.  Experiments with random toy data where performed to assess the computational load of the algorithm, and experiments on datasets coming from language pre-training, machine translation, and image classification are used to show the effectiveness and the improved training stability of the algorithm on real world problems