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: There are two major contributions of this paper. First, it proposes a method to incorporate relative positional embeddings (RPE) into kernelized attention, which is based on Fast Fourier Transformation (FFT) and Toeplitz matrix multiplication and leads to a $O(n \log n)$ time complexity compared to $O(n^2)$ using the conventional self-attention with RPE. Second, it proposes to normalize the queries and keys by their `$\ell_2$ norms, which makes the optimization more stable.