FastRPB: a Scalable Relative Positional Encoding for Long Sequence Tasks

1 INTRODUCTION . Transformer architecture ( Vaswani et al. , 2017 ) originally proposed for machine translation tasks has shown impressive results in a wide range of domains , including natural language processing , image recognition , audio captioning , graph analysis , and bioinformatics ( Lin et al. , 2021 ) . However , in applications that require processing long sequences , the benefits of transformers are often accompanied by high consumption of computational and memory resources . The main bottleneck is the transformer ’ s core component , the self-attention mechanism . Self-attention computes similarity scores for all pairs of tokens in the input sequence , and therefore , it has a quadratic complexity O ( N2 ) in computations and memory relative to the length of the input sequence N 1 . Recently , several approaches have been introduced to reduce the computational complexity and memory footprint of self-attention . Some works utilize the sparsity of the attention map ( Beltagy et al. , 2020 ) , others express self-attention as a linear dot-product of kernel feature maps φ ( · ) ( Katharopoulos et al. , 2020 ) , or utilize random feature vectors ( Choromanski et al. , 2020 ) . Proposed approaches reduce the computational complexity toO ( N ) 2 . One of the promising variants of a transformer is the Linear Transformer ( Katharopoulos et al. , 2020 ) since , along with linear complexity , it requires constant O ( 1 ) memory in auto-regressive language modeling . Experiments with the long sequence benchmark Long Range Arena ( LRA ) ( Tay et al. , 2020 ) 3 have indeed shown that 1The full complexity of self-attention also depends on attention head size D. For the original self-attention , complexity is O ( N2D ) 2In contrast , for Linear Transformer ( Katharopoulos et al. , 2020 ; Choromanski et al. , 2020 ) , the complexity of linear self-attention is O ( ND2 ) . In long sentences , N is assumed to be around thousands of tokens . Therefore , switching to linear self-attention appears beneficial . 3In benchmark sequences ranging from 1K to 16K tokens -3600 -2800 -2000 -1200 -400 400 1200 2000 2800 3600 A A N -1800 -1400 -1000 -600 -200 200 600 1000 1400 1800 Li st O ps -900 -700 -500 -300 -100 100 300 500 700 900 Relative distance between tokens , i TC −0.02 −0.01 0.00 0.01 0.02 A ss ig ne d w ei gh t , w i Figure 1 : Learned weights wi assigned to pairwise distances between tokens i in FastRPB 1D for different text LRA tasks . the Linear Transformer is 5x times faster than the vanilla Transformer in training speed . However , the drawback of this architecture is lower performance compared to the original Transformer . One way to reduce the performance gap between the Linear Transformer and the original one is to select a more suitable kernel function φ ( · ) in linear attention ( Choromanski et al. , 2020 ; Schlag et al. , 2021 ) . The poor performance of efficient transformers on LRA can also be attributed to the model ’ s ability to capture positional information . The original Transformer model utilizes only absolute positional information , which is added through positional embeddings to contextual embeddings of the tokens . Other approaches , which enrich self-attention with additional information about relative distances between tokens , have recently shown visible improvements in performance . Some of them directly add a matrix of relative distances to the attention map ( Shaw et al. , 2018 ) , others compute separate attention scores between positional embeddings ( He et al. , 2020 ) . We hypothesize that adding relative positional information could improve efficient transformers . However , most of the current implementations possess quadratic computational complexity , which neutralizes all efficiency of the Linear Transformer . To deal with this problem , a linear complexity stochastic positional encoding ( SPE ) was proposed ( Liutkus et al. , 2021 ) . Despite linear asymptotic , SPE remains relatively inefficient in training time due to its stochastic nature , while the improvement in accuracy it brings is relatively small on several LRA tasks . The contribution of this paper is two-fold . At first , we propose the Shift-Invariant Kernel Function ( SIKF ) . It could be used as a kernel for the Linear Transformer model and holds the shift-invariance property of softmax in the original attention . Second , we propose Fast Relative Positional Bias ( FastRPB ) — Fast Fourier Transform-based bias for self-attention that represents relative positional information within sequences , hasO ( N logN ) complexity and requires onlyO ( N ) memory . FastRPB is orthogonal to the self-attention mechanism and can be combined with both efficient and original implementations . We observed that SIKF is comparable to more complex kernels ( Choromanski et al. , 2020 ; Schlag et al. , 2021 ) while being as fast as the original one ( Katharopoulos et al. , 2020 ) . We also evaluated FastRPB under different long-context scenarios , such as image classification and Long Range Arena tasks . Through a comprehensive study , we showed that the proposed technique outperforms the prior fast positional encoding method ( Liutkus et al. , 2021 ) by a significant margin without adding a substantial computational footprint . 2 RECENT WORKS . 2.1 ATTENTION MECHANISM . The core component of Transformer ( Vaswani et al. , 2017 ) is the attention layer , which computes attention weights Am , n , measuring how important the role of n-th key word is in shaping the meaning of m-th output word . Using Am , n we can construct attention matrix A ∈ RM×N , and rewrite the equation using matrix notation . The output of the attention layer Y is defined based on three matrices Q ∈ RM×D , K ∈ RN×D and V ∈ RN×D ( Queries , Keys , and Values ) as follows : Y = AV = softmax ( A ) V = softmax ( QKT / √ D ) V ( 1 ) In vanilla Transformer , the attention matrix A is computed explicitly , which leads to a O ( MND ) complexity , and O ( MN ) memory to store the matrix4 . 2.2 EFFICIENT ATTENTION MECHANISM . Linear Transformer variants ( Katharopoulos et al. , 2020 ; Choromanski et al. , 2020 ) are a way to reduce the complexity of attention from quadratic to linear using the associative property of matrix products and kernel reformulation of attention . By substituting the softmax function in Equation 1 , we obtain m-th row ym of matrix Y : ym = ∑ n exp ( q T mkn/ √ D ) vn∑ n exp ( q T mkn/ √ D ) = ∑ n sim ( qm , kn ) vn∑ n sim ( qm , kn ) ( 2 ) where exp ( qTmkn/ √ D ) is generalized by any arbitrary defined similarity function sim ( qm , kn ) . The core idea of Linear Transformer is to replace sim ( qm , kn ) with a dot-product using kernel function φ ( · ) and then use an associative property of matrix products as follows : ym = ∑ n φ ( qm ) Tφ ( kn ) vn∑ n φ ( qm ) Tφ ( kn ) = φ ( qm ) T ∑ n φ ( kn ) vn φ ( qm ) T ∑ n φ ( kn ) ( 3 ) Original attention has O ( N2D ) time complexity , where N represents the sequence length , and O ( N2 ) for the memory footprint , while linear attention has time and memory complexityO ( ND2 ) , which scales linearly with sequence length N . 2.3 KERNEL FUNCTION VARIANTS . An open question is selecting an appropriate kernel function for Linear Transformer since different kernel functions dramatically affect trained model accuracy and speed . ELU + 1 . Originaly proposed kernel is an element-wise ELU ( · ) + 1 ( Katharopoulos et al. , 2020 ) : φ ( x ) = ELU ( x ) + 1 = { x+ 1 , x > 0 exp ( x ) , x ≤ 0 ( 4 ) The choice of ELU ( · ) +1 over ReLU ( · ) was prompted by its non-zero gradients for negative values . 4In case of self-attention , M equals to N , and thus the complexity is O ( N2D ) and memory requirements is O ( N2 ) . Performer . The core idea is to approximate the softmax on average using random features ( Choromanski et al. , 2020 ) . The kernel function is evaluated as : φ ( x ) = h ( x ) √ m [ exp ( Rx ) exp ( −Rx ) ] , where h ( x ) = 1√ 2 exp ( −1 2 ‖x‖ ) ( 5 ) Here [ exp ( Rx ) exp ( −Rx ) ] states for concatenation of vectors exp ( Rx ) and exp ( −Rx ) along feature dimension , each row r ∈ RD of matrix R ∈ RR×D is sampled from normal distribution N ( 0 , ID ) , and dimension size R is a hyperparameter . The main drawback of Performer is that sampling of matrix R requires extra computations and introduces variance into the model ’ s output . DPFP . Deterministic parameter-free projection is an alternative approach ( Schlag et al. , 2021 ) . The kernel function , designed to facilitate orthogonality in the projected space RDproj , is described as : φi·ν ( x ) = ReLU ( [ x −x ] ) i ReLU ( [ x −x ] ) i+ν , where φ : RD → RDproj ( 6 ) here i ·ν indicates the index of vector φ ( x ) , i ∈ { 1 , 2 , ... , 2D } is an index and ν ∈ { 1 , 2 , ... , 2D−1 } is a hyperparameter , controlling the capacity of kernel function φ ( · ) . Linear Transformer with DPFP model outperforms model with default kernel and Performer , even if Dproj is relatively small . Also , DPFP showed to be faster than models utilizing random features , but still slightly slower than ELU + 1 . 2.4 POSITIONAL INFORMATION . Attention is permutation-invariant , which means that the attention layer does not make use of the order of the sequence . There exist different ways to encode positional information in the attention : Absolute Positional Encoding ( APE ) proposed in original Transformer architecture uses realvalued vector pi ∈ RD assigned to each positions i . Some approaches , such as vanilla Transformer ( Vaswani et al. , 2017 ) , use predefined vectors , while others employ learnable vectors , e.g. , in BERT ( Devlin et al. , 2018 ) . Relative Positional Encoding ( RPE ) is complement to the absolute positional encoding , which explicitly adds relative positional information between vectors ( Shaw et al. , 2018 ) to the model . Raffel et al . ( 2019 ) proposed to directly embed positional information to the matrix A ( see the Equation 1 ) . This approach was then improved by separating semantic correlation of words and their positions correlation by Ke et al . ( 2020 ) . The componentAm , n of matrix A then calculated as : Am , n = 1√ D qTmkn + 1√ D ( UQpm ) T ( UKpn ) ( 7 ) where pn and pm are embeddings of corresponding positions n and m , and UQ , UK ∈ RD×D are learnable projection matrices for the positional embedding . By design , these approaches have quadratic computational complexity . Thus their usage with Linear Transformer is challenging since the naive application will neutralize all effectiveness of linear computation time . Stochastic Positional Encoding ( SPE ) proposed by Liutkus et al . ( 2021 ) is , to the best of our knowledge , currently the only positional encoding method compatible with Linear Transformer variants due to its linear complexity . The key idea for SPE is to represent the attention relative distances matrix as a covariance . Following the notation from equation 1 , we can express Am , n as : Am , n = D∑ d=1 Qm , d · Pd ( m , n ) ·Kn , d/ √ D , where Pd ( m , n ) = E [ qd ( m ) · kd ( n ) ] ( 8 ) here Qm , d and Kn , d are components of matrices Q and K respectively . qd ( m ) and kd ( n ) are two real and zero-mean random variables such that their covariance function matches Pd . Varying the structure of matrices Pd authors designed two variants of SPE : sinSPE and convSPE . The first one yields periodic covariance functions , which showed to be beneficial in such tasks as music generation . The second utilizes vanishing covariance functions , a promising concept introduced in Wang et al . ( 2020 ) , which yields notably smaller validation losses in some SPE experiments . Although SPE was beneficial in some music generation tasks , it still requires many computations due to its stochastic nature . In practice , it could be dozens of times slower than the original Transformer , as we will show feather .

Summary: The proposed method is mainly based on Linear Transformer. The authors propose Shift-Invariant Kernel Function (SIKF) by replacing the kernel function in Linear Transformer with exp (x), which satisfies the property of shift-invariance. Besides, the authors also proposed FastRPB which adds Relative Positional Bias to self-attention weights. The experiments show that both methods are important to improve Linear Transformer. Moreover, the proposed method is fast and memory efficient.