Year: 2021

Abstract:
Transformers achieve remarkable performance in various domains, including NLP, CV, audio processing, and graph analysis. However, they do not scale well on long sequence tasks due to their quadratic complexity w.r.t. the input’s length. Linear Transformers were proposed to address this limitation. However, these models have shown weaker performance on the long sequence tasks comparing to the original one. In this paper, we explore Linear Transformer models, rethinking their two core components. Firstly, we improved Linear Transformer with Shift-Invariant Kernel Function SIKF, which achieve higher accuracy without loss in speed. Secondly, we introduce FastRPB which stands for Fast Relative Positional Bias, which efficiently adds positional information to self-attention using Fast Fourier Transformation. FastRPB is independent of the self-attention mechanism and can be combined with an original self-attention and all its efficient variants. FastRPB has O(N logN) computational complexity, requiring O(N) memory w.r.t. input sequence length N . We compared introduced modifications with recent Linear Transformers in different settings: text classification, document retrieval, and image classification. Extensive experiments with FastRPB and SIKF demonstrate that our model significantly outperforms another efficient positional encodings method in accuracy, having up to x1.5 times higher speed and requiring up to x10 times less memory than the original Transformer.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RECENT WORKS

heading: 2.1 ATTENTION MECHANISM
text: 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.

heading: 2.2 EFFICIENT ATTENTION MECHANISM
text: 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 .

heading: 2.3 KERNEL FUNCTION VARIANTS
text: 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.

heading: 2.4 POSITIONAL INFORMATION
text: 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.

heading: 3 APPROACH

heading: 3.1 SHIFT-INVARIANT KERNEL FUNCTION (SIKF)
text: We hypothesize that the shift-invariance property of softmax function (i.e., the fact that softmaxi(x + c) = softmaxi(x), where x is some vector, c is a constant, which is added to every component of x) is an important property which makes original Transformer perform better than Linear Transformer with an arbitrary kernel. Based on this assumption, we propose SIKF as φ(x) = exp (x), which satisfies the property of shift-invariance. If we substitute this function in the linear attention from Equation 3, then for every real-valued constants c and d we will get:
φ(qm + c) T ∑ n φ(kn + d)vn
φ(qm + c)T ∑ m φ(kn + d) = ecφ(qm)
T ∑ n e
dφ(kn)vn ecφ(qm)T ∑ n e dφ(kn) = φ(qm)
T ∑ n φ(kn)vn
φ(qm)T ∑ n φ(kn)
(9)
Thus, attention in Linear Transformer with exp(·) kernel function holds the same shift-invariance property as plain softmax.
Based on our experiments, we conclude that SIKF is faster than Performer and DPFP, simultaneously has comparable accuracy, and does not provide an extra memory footprint, which is essential for scaling Linear Transformer on extremely long sequences.

heading: 3.2 FAST RELATIVE POSITIONAL BIAS (FASTRPB)
text: Although adding positional information in the attention mechanism is beneficial for model accuracy, current approaches are relatively inefficient in long sequences in terms of speed and memory footprint. In this context, there is a desire to design an approach that will add relative positional information to attention efficiently and simultaneously be compatible with various efficient attention modifications. To achieve this goal, we propose FastRPB5 as a separate term for attention.
More formally, output matrix Y of attention layer with FastRPB is defined as:
Y = AttentionVariant(Q,K,V ) +WV (10)
here matrix W ∈ RM×N consists of learnable weights Wm,n representing relative distances between m and n embedding vectors from matrix V . Note that Equation 10 is invariant of choosing a specific attention mechanism and could be used with both vanilla attention and its linear variants (Equations 1 and 3 respectively).
One can think of the matrix W as a bias term to the usual attention matrix A from equation 1, correcting the attention weights according to the relative distance between corresponding tokens. However, adding positional bias term in the Equation 10 still requires O(NMD) computations due to the matrix product and O(NM) memory to store the bias matrix W 6. In this regard, in the following two subsections, we will construct FastRPB positional bias terms matrices W1d and W2d for different types of sequences, that can be efficiently multiplied by V . W1d utilized in the case of 1D sequences (e.g., natural language texts), and its coefficients will correspond to distances between words in 1D sequences. For 2D sequences we will utilize W2d, which coefficients will represent distances between elements of 2D sequences (i.e., pixels). We will show that these specific matrices W1d and W2d could be multiplied with V using only O(DN logN) computations, and requiring only O(N) memory. Further, we work with self-attention — a variant of attention mechanism, where input and output sequences lengths are the same, i.e. N = M . In the general case of attention, when we have an input sequence of length N and an output sequence of length M , we can pad the longer one to make the input and output lengths match.

heading: 3.2.1 1D SEQUENCE CASE
text: Suppose we have a 1D sequence with N tokens. In such a sequence, there are exactly 2N − 1 relative distances between tokens7. Let’s assign a learnable parameter wi ∈ R for each relative distance i ∈ {−N + 1, ...,−1, 0, 1, ..., N − 1}. We then will obtain a set of parameters:
{w−N+1, ..., w−1, w0, w1, ..., wN−1} (11)
Next we will construct matrix W1d using parameters {wi}N−1i=−N+1. The basic intuition is to make (n,m)-th element of matrix W1d to be assigned to the relative distance between from m-th token to
5There was a desire to name FastRPB as FastRPE to represent that it is like a faster RPE. However, we change one letter to emphasize that FastRPB is orthogonal to the selection of an attention algorithm and could be seen as a separate bias term to the attention map.
6O(N2D) and O(N2) respectively in the case of self-attention 7Relative distance from m-th token to n-th token is m − n, which can have both positive and negative
values, in this regard we have exactly 2N − 1 learnable parameters
n-th token, i.e., wm−n. Therefore matrix W1d will have the following structure:
W1d =  w0 w1 w2 · · · wN−1 w−1 w0 w1 · · · wN−2
... ...
... . . .
w−N+1 w−N+2 w−N+3 · · · w0  (12) By definition, W1d is a Toeplitz matrix (Gray, 2001). A naive way to calculate product W1d · V
requires O(N2D) computations in case of self-attention8. It turns out that it can be efficiently multiplied by a matrix V according to the following proposition:
Proposition 3.1 For every Toeplitz matrix W1d ∈ RN×N and for every matrix V ∈ RN×D, matrix product W1d · V requires O(DN logN) operations and O(N) memory. Here N is length of the input sequence.
Using proposition 3.1 we can claim that FastRPB for 1D sequence will requireO(DN logN) computational operations. Moreover, we only needO(N) memory for storing parameters {wi}N−1i=−N+1, generating Toeplitz matrix W1d. For the proof and a more detailed explanation of proposed properties see Appendix B.2.

heading: 3.2.2 2D SEQUENCE CASE
text: In the case of 2D sequences (e.g., images), a similar matrix to W1d could be defined. We will call this matrix W2d, and it will consist of learnable weights assigned to pairwise distances from each pixel of the image to the rest of the pixels. Here we will consider only the case of square images of size N ×N 9. The natural way to process images of the sizeN×N in the Transformer model is to flatten them into a vector of size N2. In this regard, matrix of pairwise distances W2d needs to be of size N2 ×N2. For simplicity we will assume images to be presented as aN×N matrix, and W2d will be expressed as a tensor W2d of size (N×N)× (N×N), which (n,m, l, k) component represents distance from pixel (l, k) to pixel (n,m).
We will assume the distance between two pixels to be a sum of two distances: the vertical and horizontal10. In this regard, tensor W2d can be decomposed on vertical and horizontal tensor terms X and Y, respectively, as W2d = X + Y. Similar to the 1D case, we will assign shared learnable parameters {wi}N−1i=−N+1 for horizontal and vertical distances. Then to compute a matrix product of tensor W2d of size N2×N2 with matrix V of size N2×D we will simply flatten tensors X and Y, and obtain matrices Xflat and Yflat of shape N2 ×N2, and compute W2dV as XflatV + YflatV . It turns out that the structure of matrices Xflat and Yflat is very similar to Toeplitz matrices from the Section 3.2.1. In this regard, W2d can be multiplied by V efficiently, according to the following proposition:
Proposition 3.2 Product of matrix W2d ∈ RN 2×N2 with matrix V ∈ RN2×D using matrices Xflat and Yflat requires O(DN logN) and O(N) memory.
Using the above proposition 3.2 we conclude that FastRPB 2D will require O(DN logN) computational operations. Moreover, it is not essential to store whole tensors X and Y to compute the product, we only need O(N) of memory for parameters {wi}N−1i=−N+1 generating this tensors. (See Appendix B.3 for the proof).

heading: 4 EXPERIMENTS
text: Long Range Arena. We evaluate proposed methods in the Long Range Arena (Tay et al., 2020), a benchmark for efficient Transformers with several text and image long sequence tasks. The
8Matrix V has size N ×D, where D is a hidden size 9If we work with non-square images of size N ×M , we can simply pad them with zeros to make it square.
10If we consider two pixels p1 = (3, 2) and p2 = (0, 1) of image of size 4 × 4, the horizontal relative distance from pixel p2 to pixel p1 then will be 2− 1, and the vertical will be 3− 0
main challenge of these tasks is dictated by the large sequence lengths, which average varies from 1K to 16K token11. In our experiments, we used the following tasks from this benchmark: (1) ListOps, which test if a model is capable of parsing hierarchical expressions (Nangia & Bowman, 2018); (2) TC, which consist of movie review sentiment analysis on the IMDB corpus (Maas et al., 2011); (3) All About NLP (AAN), which evaluates the model performance in matching and retrieval tasks (Radev et al., 2013); and (4) CIFAR10, image classification dataset (Krizhevsky, 2009).
We compared Vanilla Transformer, Linear Transformer with all kernels observed in section 2.3 with SIKF, combined with different positional encodings, namely: sineSPE, convSPE, FastRPB. We also report results of experiments w/o adding any relative positional information. All models also used trainable Absolute Positional Encodings.
All experiments and hyperparameters were conducted following instructions for the LRA dataset. We also used LRA tasks to measure memory during the evaluation and computational footprints during the training.
MNIST. Due to the fact that in LRA CIFAR10 experiment only single layer transformer is used, we conducted another image recognition experiment with larger networks. We evaluated all the above models w/ and w/o FastRPB on classical image classification dataset MNIST (Lecun et al., 1998). In this experiment, we did not compare FastPRE with other positional encoding methods since, as we observed in LRA, they require dozens of times more training time in experiments with multi-layer transformers with large size of hidden states.
For all experiments, we used a model with 8 layers, 8 attention heads, a hidden size equal to 256, and batch equal to 160. We trained models using AdamW optimizer and made 20 runs of Bayesian hyperparameter search to find the optimal learning rate and than trained all models for 25 epochs. Parameters are presented in Appendix A. We linearly decayed the learning rate to 0 during the training. Final results are averaged over 10 runs with different values of random seed.
11We do not include another synthetic image classification task, Pathfinder, since we were unable to reproduce the results, obtained in the original paper (Tay et al., 2020).

heading: 5 RESULTS
text: Long Range Arena. See Table 1 for the evaluation results. Linear Transformer with SIKF kernel comes out in the top two results for every dataset, except CIFAR10, which we will discuss separately. Memory footprint (see Table 3) of SIKF is indeed very tight to the ReLU, DPFP, and Performer, while DPFP and Performer constantly performed slower (up to 1.4x times).
Linear Transformer equipped with FastRPB showed significantly higher results on ANN and TC both in memory, speed, and accuracy, achieving even better results than Original Transformer. Moreover, architectures with FastRPB, confirming the above propositions 3.1 and 3.2, proved to have memory and computation consumption comparable with the default architectures. For the ListOps dataset, the best performance was obtained by the model without any relative positional encoding. We attribute this result to the fact that relative distances can be confusing in parse hierarchical structure, such as expressions for ListOps or source code (e.g., the distance between IF and ELSE in source code can be pretty large, however, these statements are inwardly connected). We measured memory and computational footprint of models (see Table 3), according to which FastRPB requires up to 30x less time to train than convSPE, and up to 3x less then sinSPE. Simultaneously, FastRPB has 1.5x smaller memory footprint on evaluation then sinSPE and convSPE. Thus we conclude that FastRPB is the fastest and the most accurate method among others.
The learned weights of FastRPB for different text tasks sequence is presented in Figure 1, where x-axis represents relative distance m−n from n-th token to m-th token, and color denotes the value of (W1d)n,m = wm−n. In the AAN task FastRPB forces the model to attend more to the very end of the text. Such observation can be attributed to the fact that AAN mainly consists of scientific texts, which final part usually contains a conclusion. As for ListOps, learned FastRPB weights are usually relatively small, which supports the hypothesis that relative positional encodings in such tasks should be designed using the text hierarchy information. In the TC task, learned weights mainly draw the model’s focus forward and backward to enable the model to capture long-range dependencies.
In experiments with CIFAR, usage of FastRPB was shown to decrease the model’s performance, and the best result was obtained with convSPE, which outperformed others by a significant margin. We attributed this to the experiment setup, where a single-layer network is utilized. In this regard, we conduct the experiments with a more extensive network on the MNIST dataset.
MNIST. In this task, each of the above models equipped with FastRPB showed superior performance (see Table 2), requiring quite a small amount of additional memory and computational time (see Figures 3a and 3b respectively). As can be observed from the plots, of Linear Transformer with FastRPB is up to 10x times less than the original ones. Moreover, in terms of speed Linear Transformer with FastRPB is 5x time faster in evaluation time compared to original Transformer.
The detailed view of trained FastRPB could be found in Figure 2, where we learned slice (W2d):,:,12,10, which represents pairwise distances from pixel (12, 10) to every other pixel of 28×28 MNIST image. We observed that FastRPB enforced the model to look at more distinct pixels rather than to close ones.

heading: 6 CONCLUSION
text: We presented two novel approaches aimed to increase the accuracy of the Linear Transformer model without additional memory footprint and significant loss in speed. Contribution is two-fold: we first make linear attention shift-invariance and then add a bias term to attention scores, representing pairwise distances between tokens of the sequence. We computed this bias term efficiently and achieved O(N logN) complexity and O(N) memory w.r.t. sequence length. We demonstrate the superiority of our approach among others using four long sequence tasks from the Long Range Arena benchmark and on the MNIST dataset. Our model performs significantly better than previous approaches, obtaining the best accuracy on several tasks while being almost as efficient in terms of speed and memory as plain Linear Transformer.
We believe that the principles presented in this work can serve as a basis for future research on the role of positional information encoding in transformer architectures. To this end, we open-source all the code and trained models.

heading: A MNIST HYPERPARAMTERTS SEARCH

heading: B PROPOSITION PROOFS

heading: B.1 CIRCULANT MATRICES
text: To design a more efficient positional encoding method, we leveraged circulant matrices – a subclass of matrices with some special properties due to their relation to the Fast Fourier Transform (FFT) and circular convolution Bamieh (2020). Here we will only focus on the property that allows performing matrix-vector product fast and efficient in terms of speed and memory. Given an vector c = (c0, c1, ..., cn−1) we will define the associated n × n circulant matrix C = circ(c) which first column is exactly c, and each subsequent column is obtained by a circular shift of the previous column:
C =  c0 cn−1 cn−2 · · · c1 c1 c0 cn−1 c2 c2 c1 c0 c3 ... . . . . . . ...
cn−1 cn−2 cn−3 · · ·  (13) It can be shown that for every vector x of size nmatrix-vector product Cx requires onlyO(n log n) computation Rosowski (2021). Moreover, to compute the above product, it is not necessary to store the whole matrix C in memory; it is enough to only keep in memory O(n) parameters of vector c. Further, we will prove that FastRPB and FastRPB 2D can be expressed through circular matrices, and hence relative positional information can be embedded efficiently in the self-attention mechanism.

heading: B.2 PROPOSITION 1
text: In 3.2.1 we introduced a Toeplitz matrix W1d of shape N ×N :
W1d =  w0 w1 w2 · · · wN−1 w−1 w0 w1 · · · wN−2 w−2 w−1 w0 wN−3
... . . . . . . ... w−N+1 w−N+2 w−N+3 · · · w0
 (14)
AAN
Relative distance between tokens, i
−0.02
0.00
0.02
0.04
0.06
A ss
ig ne
d w
ei gh
t, w i
(a) Learned in AAN task weights wi assigned to pairwise distances between tokens i in FastRPB 1D.
ListOps
Relative distance between tokens, i
−0.04
−0.02
0.00
A ss
ig ne
d w
ei gh
t, w i
(b) Learned in ListOps task weights wi assigned to pairwise distances between tokens i in FastRPB 1D.
TC
Relative distance between tokens, i
−0.01
0.00
0.01
A ss
ig ne
d w
ei gh
t, w i
(c) Learned in TC task weights wi assigned to pairwise distances between tokens i in FastRPB 1D.
Our goal is to efficiently multiply W1d by arbitrary matrix V of shape N × D. As was stated in section Circulant matrices, a special class of matrices, namely circulant matrices, can be multiplied by a vector efficiently in O(N logN) operations and requires O(N) memory. We will extend matrix W1d with additional rows and columns and thus obtain circulant matrix W ext1d . Then we will introduce V ext, a modified version of matrix V , which W ext1d will be multiplied by. Finally, we will select a slice from the product W ext1d · V ext, which will be exactly W1d · V . The first step is to define W ext1d :
W ext1d =  w−N+1 w−N+2 · · · w0 w1 w2 · · · wN−1 wN−1 w−N+1 · · · w−1 w0 w1 · · · wN−2 wN−2 wN−1 · · · w−2 w−1 w0 wN−3 ... . . . . . . . . . ... w1 w2 · · · w−N+1 w−N+2 · · · w−1 w0 ... . . . . . . . . . . . . ...
w−N+2 w−N+3 · · · · · · · · · · · · wN−1 w−N+1
 (15)
As can be seen, constructed matrix W ext1d is indeed circulant. Moreover, the right upper corner of W ext1d is essentially W1d. Hence, W1d can be expressed as a slice of W ext 1d following numpy
notation: W1d = ( W ext1d ) N : , :N+1 .
Now we want to calculate matrix product W1d ·V using matrix W ext1d of size (2N − 1)× (2N − 1). Due to this fact we will need to multiply W ext1d with appropriate matrix V
ext of size 2N − 1 × D. Since, as we seen, W1d is a slice of W ext1d , we only will need first N rows of the resulting product W ext1d ·V ext. More formally we need to find such matrix V ext that: W1d ·V = ( W ext1d ·V ext ) :N+1 , :
. To achieve with we will basically pad V with N − 1 additional rows filled with zeros as follows:
V ext =  0 0 · · · 0 0 0 · · · 0 ... ... . . . ... 0 0 · · · 0 v0,0 v0,1 · · · v0,D−1 v1,0 v1,1 · · · v1,D−1
... ...
. . . ...
vN−1,0 vN−1,1 · · · vN−1,D−1
 (16)
To complete the proof lets explicitly show that W1d · V = ( W ext1d · V ext ) :N+1 , :
. Lets assume D = 1, generalization for bigger dimensions can be done using similar operations:
( W ext1d · V ext ) :N+1 , : =  w−N+1 · · · w0 w1 · · · wN−1 wN−1 · · · w−1 w0 · · · wN−2 wN−2 · · · w−2 w−1 wN−3
... . . . . . . ... w1 · · · w−N+1 w−N+2 · · · w0
  0 0 ... 0 v0 v1 ...
vN−1
 = (17)
 w0 w1 · · · wN−1 w−1 w0 · · · wN−2 w−2 w−1 wN−3
... . . . . . . ... w−N+1 w−N+2 · · · w0
  v0 v1 ...
vN−1  = W1d · V (18) The last thing we have to do is to calculate the complexity of matrix product of circulant matrix W ext1d with V ext. Since matrix W ext1d is circulant of size (2N − 1) × (2N − 1), according to the
section Circulant matrices, it requires O(N) memory to it, and O(N logN) operations to perform a matrix-vector product with vector if size 2N − 1. In this regard to compute matrix product with matrix V of size N × D, we will need to perform D times operations more, i.e. O(DN logN) operations.

heading: B.3 PROPOSITION 2
text: In the following sections, we will be considering an example of 3 × 3 image. In section Structure of pairwise distances tensors we will study the general structure of tensors X and Y, which were introduced in section . In section Flattening of the tensors we will reshape these tensors and obtain a new pair of tensors Xflat and Yflat, which will be then efficiently multiplied by V matrix in the final section Efficient matrix product.

heading: B.3.1 STRUCTURE OF PAIRWISE DISTANCES TENSORS
text: First of all, to gain a deeper understanding of structure of tensors X and Y, we will explicitly write down components of this tensors for an image of size 3 × 3. Consider X1,1,:,: and Y1,1,:,: which contains weights assigned to vertical and horizontal relative distances from pixel (1, 1) to all other pixels:
X1,1,:,: = ( w−1 w−1 w−1 w0 w0 w0 w1 w1 w1 ) , Y1,1,:,: = ( w−1 w0 w1 w−1 w0 w1 w−1 w0 w1 ) (19)
One can observe, that Xn,m,:,: and Yn,m,:,: have the property of symmetry, through which it can be proven that:
Proposition B.1 Xn,i,:,: = Xn,j,:,: and Yi,n,:,: = Yj,n,:,: for every n, i, j ∈ {0, ..., N − 1}, and this property holds for image of every size.
Taking advantage of the proposition B.1 we can write out explicit form of tensors X and Y in case of 3× 3 images. We will introduce the following notation:
A = ( w−2 w−2 w−2 w−1 w−1 w−1 w0 w0 w0 ) , B = ( w−1 w−1 w−1 w0 w0 w0 w1 w1 w1 ) , C = ( w0 w0 w0 w1 w1 w1 w2 w2 w2 ) (20)
It can be seen that different slices of tensor X can be expressed using matrices A,B,C: X0,0,:,: = X0,1,:,: = X0,2,:,: = C (21) X1,0,:,: = X1,1,:,: = X1,2,:,: = B (22) X2,0,:,: = X2,1,:,: = X2,2,:,: = A (23)
Moreover, are matrices A,B,C also applicable for tensor Y: Y0,0,:,: = Y0,1,:,: = Y0,2,:,: = C T (24)
Y1,0,:,: = Y1,1,:,: = Y1,2,:,: = B T (25)
Y2,0,:,: = Y2,1,:,: = Y2,2,:,: = A T (26)

heading: B.3.2 FLATTENING OF THE TENSORS
text: In transformer architecture before processing the image of size N × N , it is usually flattened into 1-dimensional vector of size N2. To define flatten operation, consider an arbitrary matrix M of shape 3× 3, than its flattened version mflat will have the following structure:
M = ( m0,0 m0,1 m0,2 m1,0 m1,1 m1,2 m2,0 m2,1 m2,2 ) flattening−−−−−→ mflat =  m0,0 m0,1 m0,2 m1,0 m1,1 m1,2 m2,0 m2,1 m2,2  (27)
Due to the flatten of images in transformer, matrix V in will have shape N2 × D and hence it is essential to reshape tensors X and Y from size (N × N) × (N × N) to size N2 × N2. We will denoted reshaped versions of tensors X and Y as Xflat and Yflat respectively. Reshaping of the above tensors can be decomposed into two consecutive flatten operations firstly applied to last two dims of tensors X and Y and then to the first two ones. Flattening of the last dimensions is equivalent to flattening of each of the matrix A,B,C, after this operation in case of 3× 3 images we will obtain the following three vectors:
aflat =  w0 w0 w0 w1 w1 w1 w2 w2 w2  , bflat =  w−1 w−1 w−1 w0 w0 w0 w1 w1 w1  , cflat =  w−2 w−2 w−2 w−1 w−1 w−1 w0 w0 w0  (28)
Than after next flattening operation we will get:
Xflat = X.reshape(N2, N2) =  w0 w0 w0 w1 w1 w1 w2 w2 w2 w0 w0 w0 w1 w1 w1 w2 w2 w2 w0 w0 w0 w1 w1 w1 w2 w2 w2 w−1 w−1 w−1 w0 w0 w0 w1 w1 w1 w−1 w−1 w−1 w0 w0 w0 w1 w1 w1 w−1 w−1 w−1 w0 w0 w0 w1 w1 w1 w−2 w−2 w−2 w−1 w−1 w−1 w0 w0 w0 w−2 w−2 w−2 w−1 w−1 w−1 w0 w0 w0 w−2 w−2 w−2 w−1 w−1 w−1 w0 w0 w0  (29)
Yflat = Y.reshape(N2, N2) =  w0 w1 w2 w0 w1 w2 w0 w1 w2 w−1 w0 w1 w−1 w0 w1 w−1 w0 w1 w−2 w−1 w0 w−2 w−1 w0 w−2 w−1 w0 w0 w1 w2 w0 w1 w2 w0 w1 w2 w−1 w0 w1 w−1 w0 w1 w−1 w0 w1 w−2 w−1 w0 w−2 w−1 w0 w−2 w−1 w0 w0 w1 w2 w0 w1 w2 w0 w1 w2 w−1 w0 w1 w−1 w0 w1 w−1 w0 w1 w−2 w−1 w0 w−2 w−1 w0 w−2 w−1 w0  (30)
Note that rows of the Xflat and Yflat are aflat, bflat and cflat.
Now we have flattened representation Xflat and Yflat and the last step we have to take is to compute matrix product Xflat · V and Yflat · V .

heading: B.3.3 EFFICIENT MATRIX PRODUCT

heading: Calculation of Yflat · V
text: In case of 3 × 3 images it is easy to see that the matrix Yflat, presented in formula 30 consists of 9 identical blocks, we will denote this blocks as R:
Yflat = ( R R R R R R R R R ) , where R = ( w0 w1 w2 w−1 w0 w1 w−2 w−1 w0 ) (31)
Note that matrix R is just a block of Y , not its element. Also a very important fact is that can matrix R is a Toeplitz, which means that it can be efficiently multiplied by a vector, according to Proposition 1. Our goal is to efficiently multiply Yflat with a matrix V of shape N2×D. Since V is basically D times stacked vectors of size N2 it will be enough to consider matrix-vector product
of Yflat with a vector v of shape N2. The product can be expressed as:
Yflat · V = ( R R R R R R R R R )  v0 v1 v2 v3 v4 v5 v6 v7 v8  =  R · ( v0 v1 v2 ) +R · ( v3 v4 v5 ) +R · ( v6 v7 v8 ) ... ...  = (32)
 R · ( v0 + v3 + v6 v1 + v4 + v7 v2 + v5 + v8 ) ... ...  (33) Lets introduce an additional notation, we will denote as Vm a tensor of size N ×N ×D, which was obtained as a reshape of matrix V of size N2 ×D. Vm will have the following structure:
Vm = ( v0 v1 v2 v3 v4 v5 v6 v7 v8 ) = V .reshape(N, N, D) (34)
Then:
V Tm .sum(1) = ( v0 + v3 + v6 v1 + v4 + v7 v2 + v5 + v8 ) (35)
Finally using notation 34 and result of product 33, we can conclude that:
Yflat · V =  R · ( V Tm .sum(1) ) ... ...  (36) Note that the summation by dim = 1, not by dim = −1, since in general case, when D is not 1, the latter dimension refers to hidden size.
R is a Topleitz matrix we worked with in the previous paragraph B.2, that is, the product of matrix R with V Tm .sum(1) can be efficiently computed using the properties of Toeplitz matrix. Finally we can conclude that for N × N image computation of product Yflat · V require O(DN logN) computations and O(N) memory to store a vector, generating matrix R.

heading: Calculation of Xflat · V
text: Here we will apply the following trick: by rearranging the columns of matrix Xflat we will construct matrix X ′flat, and accordingly construct matrix U by rearranging the rows of matrix V in such a way that the equality holds Xflat · V = X ′flat ·U . We still will work with the D = 1 case since, as was noted previously, it can be easily generalized for higher dimensions.
Xflat · V =  w0 w0 w0 w1 w1 w1 w2 w2 w2 w0 w0 w0 w1 w1 w1 w2 w2 w2 w0 w0 w0 w1 w1 w1 w2 w2 w2 w−1 w−1 w−1 w0 w0 w0 w1 w1 w1 w−1 w−1 w−1 w0 w0 w0 w1 w1 w1 w−1 w−1 w−1 w0 w0 w0 w1 w1 w1 w−2 w−2 w−2 w−1 w−1 w−1 w0 w0 w0 w−2 w−2 w−2 w−1 w−1 w−1 w0 w0 w0 w−2 w−2 w−2 w−1 w−1 w−1 w0 w0 w0  ·  v0 v1 v2 v3 v4 v5 v6 v7 v8  = (37)
 w0 w1 w2 w0 w1 w2 w0 w1 w2 w0 w1 w2 w0 w1 w2 w0 w1 w2 w0 w1 w2 w0 w1 w2 w0 w1 w2 w−1 w0 w1 w−1 w0 w1 w−1 w0 w1 w−1 w0 w1 w−1 w0 w1 w−1 w0 w1 w−1 w0 w1 w−1 w0 w1 w−1 w0 w1 w−2 w−1 w0 w−2 w−1 w0 w−2 w−1 w0 w−2 w−1 w0 w−2 w−1 w0 w−2 w−1 w0 w−2 w−1 w0 w−2 w−1 w0 w−2 w−1 w0  ·  v0 v3 v6 v1 v4 v7 v2 v5 v8  = X ′flat ·U =  α α α β β β γ γ γ  (38)
Now lets use already introduced matrix Vm:
R · ( Vm.sum(1) ) = ( w0 w1 w2 w−1 w0 w1 w−2 w−1 w0 )( v0 + v1 + v2 v3 + v4 + v5 v6 + v7 + v8 ) = (39)
R · ( v0 v3 v6 ) +R · ( v1 v4 v7 ) +R · ( v2 v5 v8 ) = ( α β γ ) (40)
And finally:
X ′flat ·U =  α α α β β β γ γ γ  = ( α α α β β β γ γ γ ) .reshape(N2) = (41)
( R · ( Vm.sum(1) ) · · · · · · ) .reshape(N2) (42)

heading: B.3.4 CONCLUSION
text: As we proved in the previous section, products Xflat · V and Yflat · V can be computed using just two products: R · ( Vm.sum(1) ) and R · ( V Tm .sum(1) ) respectively. Moreover, matrix R utilized
in the above products is a Toeplitz, and in this regard this products can be computed efficiently according to Proposition 1. Summarizing it all we get that FastRBP 2D will require be O(DN logN) computations and O(N) memory.

Review summaries:
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.

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.

This paper focused on two existing issues of the Transformer framework: inefficient scaling and weak performance on long sequence tasks. Take the Linear Transformer model as an example, the authors proposed the Shift-Invariant Kernel Function (SIKF) for the weak performance issue and the Fast Relative Positional Bias (FastRPB) for the weak performance issue, respectively. The proposed methods were evaluated on several long sequence tasks.