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 paper presents two contributions 1. It expands on Linear Transformers by introducing a feature map \phi = exp(x) function whose dot product approximates the softmax kernel (dubbed Shift Invariant Kernel Function SIKF); whose motivation comes from the shift invariant property of the softmax function. They show that it performs comparably to other proposed approximations in literature while improving training efficiency and reducing the memory footprint 2. It presents a way to use relative position bias with Linear Transformers. Dubbed Fast Relative Position Bias (FastRPB), the method proposes to learn a scalar for pairwise difference of positions, and draw a connection between the distance matrix and Circulant Matrices. They use this to efficiently compute the relative position bias in O(ND log N ) operations with O(N) memory. They show that FastRBP outperforms other relative positional bias methods while not substantially increasing compute.