Stable, Fast and Accurate: Kernelized Attention with Relative Positional Encoding
Authors: Shengjie Luo
Shanda Li
Tianle Cai
Di He
Dinglan Peng
Shuxin Zheng
Guolin Ke
Liwei Wang
Tie-Yan Liu
Venue: NeurIPS 2021 Poster (2021)
Keywords: Attention
Transformer
Relative Positional Encoding
Fast Fourier Transform
Language Pre-training
Language Modelling
Machine Translation
Image Classification

Abstract: The attention module, which is a crucial component in Transformer, cannot scale efficiently to long sequences due to its quadratic complexity. Many works focus on approximating the dot-then-exponentiate softmax function in the original attention, leading to sub-quadratic or even linear-complexity Transformer architectures. However, we show that these methods cannot be applied to more powerful attention modules that go beyond the dot-then-exponentiate style, e.g., Transformers with relative positional encoding (RPE). Since in many state-of-the-art models, relative positional encoding is used as default, designing efficient Transformers that can incorporate RPE is appealing. In this paper, we propose a novel way to accelerate attention calculation for Transformers with RPE on top of the kernelized attention. Based upon the observation that relative positional encoding forms a Toeplitz matrix, we mathematically show that kernelized attention with RPE can be calculated efficiently using Fast Fourier Transform (FFT). With FFT, our method achieves $\mathcal{O}(n\log n)$ time complexity. Interestingly, we further demonstrate that properly using relative positional encoding can mitigate the training instability problem of vanilla kernelized attention. On a wide range of tasks, we empirically show that our models can be trained from scratch without any optimization issues. The learned model performs better than many efficient Transformer variants and is faster than standard Transformer in the long-sequence regime.

Stable, Fast and Accurate: Kernelized Attention with Relative Positional Encoding
Shengjie Luo'r, Shanda Li Tianle Cai', Di He
Dinglan Peng", Shuxin ? Guolin Ke Liwei Wang Yin Liu Center for Data Science, Peking University
2 Key Laboratory of Machine Perception, MOE, School of EECS, Peking University
"Institute for Artificial Intelligence, Peking University *Princeton University University of Science and Technology of China "Microsoft Research juosj0stu, pku. edu. cn, lishanda@pku. edu. cn, tianle. wownle..a@@rrinettoon edu, pengdinglan an@mail edd.. en, {dihe, shuz, guoke, tyliu))microsoff com, vanglušpku.ede edu.
Abstract
The attention module, which is crucial component in Transformer, cannot scale efficiently to long sequences due to its quadratic complexity. Many works focus on approximating the thennexxponeniatee softmax function the original atten- tion, leading to sub-quadratic or even shear-complexity Transformer architectures. However, we show that these methods cannot be applied more powerful attention modules that go beyond the relative positional encoding (RPE). Since in many state-of-the-art models, relative positional encoding is used as default, designing efficient Transformers that can incorporate RPE is appealing. In this paper, we propose a novel way to accelerate attention calculation for Transformers with RPE on top of the kernelized attention. Based upon the observation that relative positional encoding forms Toeplitz ma- trix, we mathematically show that kernelized attention with RPE can be calculated efficiently using Fast Fourier Transform (FFT). With FFT, our method achieves further demonstrate that properly
style, e.g., Transformers with
On log n.) time complexity. Interestingly, using relative positional encoding can mitigate the training instability problem vanilla kernelized attention. On wide range of tasks, empirically show that our models can be trained from scratch without any optimization issues. The learned model performs better than many efficient Transformer variants and is faster than standard Transformer in the long-sequence regime.
Introduction
Transformer (45) has achieved tremendous success variety of tasks in natural language processing [ computer vision 22 10 42) and speech [13]. Despite the great success, Transformers are usually inefficient in modeling long-son 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
To reduce the cost, many works focus on approximating the dot-then tithen--xpoenttitee function in the attention module of the original Transformer [41. Most works can be roughly categorized into
three classes: sparse attention 3
33
15
46 7], low-rank approximation-ba attention
49 48
and kernelized attention
17 Sparse attention variants design either pre-defined
"Equal contribution.
1 "Correspondence Di He wdihetmiicoosot.. com> and Shuxin Zheng wwhuz@miccoooff..com
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

or learnable patterns 22 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 which introduces kernel view on the dot-then-exponentiae
B 30 15
44 17], 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 com- plicated than dot-then-eNt exponentiiiiinnn For example, the relative positional encoding (RPE) [38 35]
has become standard in state-of fatte-of-tie of-the-art models [35] [8] [4] Besides using the dot-products between queries and keys, RPE incorporates 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 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-then- exponentiation, i.e., for some attention with RPE, there exists no dot-then- thhen-exponentitee function that can represent it. Therefore, previous approaches cannot be directly applied to approximate Transformers with RPE and obtain non- non-trivial speed-up. In this paper, we develop a novel attention computation that achieves an almost- 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 not-then-expooentiate 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 X n. RPE position-correeation matrix with kernelized feature matrices, which leads to 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 Toeplitz matrix. In numerical linear algebra, it is well-known that multiplication between a Toeplitz matrix and row vectors be accelerated using Fast Fourier Transform (FFT). And by using FFT, we can obtain a fast attention computation with On log n) time cost.
Beyond the efficiency, we also find that our method that incorporates RPE into kernllized 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 we have to train standard Transformer first and convert to kernelized attention with finetuning. We first shed light upon the reason for this phenomenon by proving that these kernelized attentions have unacceptable variance when the lo 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 not constrained, the whole model can still have enough capacity to achieve good performance.
We conduct experiments 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 sequene regime.
wide range of tasks spanning language pre-training, machine translation
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 usu- ally formulated querying dictionary with key-value pairs, e.g., Attention( Q, K, V) softmax (KW VWV, where d is the dimensionality of the hidden representations. Q vi (Query), K (Key), V (Value) are specified as the hidden representations of the previous layer and Wº, W and WV are projection matrices for any specific head.

---

To be more concrete, we denote [21, 22 .nn) as the input to the attention module, where
the length of the sequence and I; E IRd is a row vector denoting the contextual representation of the token at position Denote (21, 22 .. , Zn) as the oulput. Then the attention module can be written as
exp(oa5)
f==1 exp(a (z, W where dij
(a,
As see, the calculation of attention for each position requires (n)) time and memory cost. Thus, the time and memory cost for the whole sequence is O(n²), which impedes the attention mechanism scaling to long sequences. Recently, [7] introduce a kernel view the attention mechanism and derive a general kernelized attention framework as follows:
k( 2,, j== =f== 2 WK) (z,WW),
2,
where x(.,-): : Rd x Rd + R is any positive-definite kernel measuring the pairwise similarity. Based on this kernel view, the attention mechanism can be alternatively modeled linear dot-product of kernelized features:
(1, 2, 2- Wo) (2,33""" (Z, 6(I,W9) E oa, (2, ( WK)
(3)
where (6.) Rd + Rm is the feature map and m is the dimension the feature space. Since can be reused for each position, the time and both 2 0(I, (x,)W) and - 6(r,
memory complexity of the kemnelized attention is reduced to (nn). 117 simply uses elu(.) + as o(.). Followed-up works 32 4 further propose two (randomized) feature maps with which the kernels 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 Trigonometric Random Feature (TRF) and Positive Random Feature (PRF) respectively:
(T) exp( |==1
sin (1 T), .... sin(um J), cos(wiæT) ),,.., cos(WmmT))
(4)
||= 2
exp(
(P)
Texp(w (1p x. =T),..., ex...eeppwuu... 2m
(5)
V
778
where each vector W/ independently sampled from (00, Ia). Our work aims to design an efficient Transformer with relative positional encoding 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 ineffective in capturing relative word orders. They propose an embedding matrix in which the values of the elements
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 24]. In this model (35), the relative positional encoding is served as a bias term added to the dot- dot-product between queries and keys:
rel 1
(Z, + bj-i.
(6)
For each j - i, by-s is learnable scalar and shared across all layers.
3

---

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 negative result: We prove that RPE expands the expressiveness of the original attention module. So those efficient approaches that can only 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 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
teteen-xxponentiatee functions cammor approximate
3.1 Attention with RPE Goes Beyond the Dol-then-exponentiate Function Class
For simplicity, we study the single-head attention demonstration where W rets W &dxd All the conclusions remain the same for the multi-head attention setting. The following proposition shows that Transformers with RPE are strictly more expressive than those without RPE.
Proposition 1. If > d + 1, there exist weight matrices W rel: W E TRdxd and RPE weights bz- as in Equation 6 such that there does not exist weight matrices W WK "onsill"" vanilla and input vectors
F1 Fn E _d that satisfy
exp(a Ek exp(ate))
exp(ara panilla) Ek=1 expla paniita Vi,j.
where
Oij
rel
(z,
W2 rel (x, WRR) + by-
(8)
Ri
vanilla
(x, vanstta (x, W vanstta)
(9)
This theoretical finding consistent with all empirical observations that Transformers equipped with RPE are more powerful and usually perform better than the original Transformer. At the same time, such a result suggests that Transformers with RPE cannot well-approximated by techniques developed in the previous works as those methods are under the simple dot-then seeene--xeenniattt attention setting. accelerate Transformers with RPE, new techniques need be explored.
3.2 Fast Attention Calculation using Fast Fourier Transform
We propose an efficient method that accelerates the computation of attention modules with RPE. The approach is built on top of the kernelized attentions, i.c.. Equation 3 When relative positional encoding is further implemented, the kernelized attention will be formulated as:
((x;WQ) Ej e 6(I, W (T (z,W%) 2;
(10)
(lx:WQ) Ef ((,,
From Equation we can see that to calculate the output (z) (z) we need to compute Dj and D2 Σ - ((,, :=1 To better illustrate
E> o0 W [= the calculation step, we reshape D1 D2 into the form where elements are vectorized as row vectors and stacked into matrices:
vec Σn e ( kK) (x,,)())))))
vec (E e (b-tb(a,W/K]]
D1 =
.D =
voc
- (6
ṭ² -
vec En a - (
(1I)

---

80
zemmxx aslandoo
RRT-R m-2
app approxíma oprrxi/mation attion error
2.0
R-1
60
RRT-RPI m-4
1.5
R-2
forward time
RRT-RPEIPD. AaPiIPD.mm12 m=11
RT-.RPDDD..ooo m-@
R=4
40
PRT-RP mm12
1.0
Rm8
0.5
0
0.0
0
5
10
15
20
25
30
35
40
16
64
256
1024
sequence length (x10*)
feature map dimension
(a)
(b)
Figure 1: (a). Model forward speed of vanilla Transformer and our model with different feature map dimensions (m). (b). Approximation error of PRF with different query/key norms (R)..
where vec (.) denotes vectorization operation. Let C₁ e.t. Di, can be factorized as the multiplication of two known matrices:
vec J
40
C1
C2
Gn-1
vec (0(x2W (2awr))
C-1
Co
C1
C--2
D1
(12)
e (-(n-1)
(-(n-2)
(-(n-3)
co
vee (K(z (2n (2)
Cp
C1
C2
Cn-1
vec
(Ø(1s WK)
vec (0(x2WK)]
C-1
Co
Cn-2
D2
(13)
C-(n-1)
C-(n-2)
C- (n(n33)
co
vec (6(zn TT)
Denote C {cg-s) the positional-ccorelation matrix. It can be easily seen that
naive matrix multiplication to obtain D1, D2. But it leads to quadratic time cost with respect to the sequence length ni and makes advantage of using kernelized attention disappear. Interestingly, we notice that C' is a diagonal-constan matrix in which each descending diagonal from left to right is constant, i.e., a Toeplitz matrix. numerical linear algebra [T,, left multiplying a Toeplitz matrix be done in On log n) time using Fast Fourier Transform (FFT). This fact leads to efficient algorithm to compute kernelized attention with RPE in O log n. time.
We empirically evaluate the forward cost of the kernelized Transformer with RPE using FFT and the standard Transformer. For simplicity, we set the number of heads to 1 and the size of embedding to 64. We use synthetic input sequences with lengths ranging from 1k to 40k. For our model, we experiment with different feature map dimensions. All the experiments are run on Tesla V100 GPU with 16GB memory for ten trials. The averaged result is shown in Figure Ia It can be seen that while forward cost of standard attention becomes much significant on longer sequences, our model still remains efficient (almost linear) in the long-sequence regime.
3.3 RPE Enables Stable Training of Kernelized Attention
Recent works 32 show that PRF and TRF attention can approximate the dot-thhen-exponentiaee attention with theoretical guarantees and achieves much better performance than other methods. However, the training of these models is challenging. As demonstrated by several works [16 A] kernelized attention can hardly be optimized from scratch for large-scale tasks and may even diverge during training. This significantly limits the practical impact of the model a standard Transformer first and convert it to the kernelized version. [16] further observed that the model conversion usually leads to considerable performance drop which suggests that the
has to train
One limitation of the proposed method is that it cannot accelerate sequence generation models during the inference step.
approach, Dj and Da need be pre-computed using FFT for all the queries as whole. However, the inference step the sequence generation, we cannot pre-compute Dj and D2 as the tokens are generated on the fly and the FFT operation needs to be used each time step. But our method can still be used to accelerate the training of the generation model which uses teacher forcing by setting bi. -oo (i.e., Ci-j 0) when > j. We leave how to accelerate the inference step using method future work.
5

---

approximation seems to be inaccurate. To understand the problem, we start with Lemma 2 below from 4. which analyzes the variance of the
estimation with the PRF feature map.
Lemma 2 (Variance of the PRF (Lemma 2 in 44). For Φrnr defined in Equation 5
we have
Varl)()))(())) 1 (exp( (or + y1l)) - -1) exp(ay)))
(14)
It can be seen from the above lemma that the variance of the estimation is related to the scale of the queries and keys. If the scale of the queries and keys is large, the variance will be significant. Our following theorem precisely characterizes that one may need an exponential number of feature map dimensions to obtain a well-approximated PRF kernelized attention. Similar results also hold for the TRF kernel
Theorem 3 (Sample Complexity of Approximation), For query 9 and n keys k,, k.m let A, Á be the attention score obtained by the standard softmax and the PRF approximation respectively. Assume that En-norms of the query and keys are upper-bounded by R. Then for any 0, if feature
map dimension m = O n expp (ex((1R)) we have Pr (114 - A|l 2 2 < 8.
828
We also perform numerical simulations to verify the theoretical findings. We set 64 as the query/key dimension and independently sample a query vector and 1024 key vectors from the uniform distribution over a d dimensional unit hypersphere. Therefore, the f--norm of the query and the keys we sampled equals one. Denote R 2 1 as a scaling factor. For each value of R. we rescale the query and the keys with R and calculate the softmax score A and the approximated score Å using PRF with different feature map dimensions. We use I|A - Áll to quantify the approximation error. The results are shown in Figure lb It can be seen that when R is large, the approximation error is huge and can only be slightly improved as the feature map dimension grows from 4 to 1024. But when R = 1, the approximation error is small and drops significantly when we increase the feature map dimension. This numerical result is consistent with our theory, which implics that the scale of the queries and keys significantly matters the approximation quality.
Based on the results, we can understand why the model conversion (well-trainel standard attention + kernelized attention) usually leads to performance drop as observed [66] For a Transformer, some attention heads generate sharp attention distributions (e.g., attending to the next token (6) To generate a sharp attention distribution, queries and keys with a large norm are essential. If we convert those attention heads to their kernelized version, the output can be significantly different due to the large variance, and the performance will be poor. Furthermore, queries and keys with large norm may also be the root cause of the training instability 4. As discussed in (45), at initialization, the parameters are sampled from Gaussian distribution and unbounded. If the norms of the queries or keys happen to be larger than one, the kernelized attention (forward pass) will have large variance, leading to a large-variance gradient (backward pass) and making the training unstable.
Normalized Kernelized Attention with RPE. A straightforward solution to fix the problem is normalizing queries and keys by their & norm so the variance of the kernelized attention will be well controlled. However, this would bound the dot-products between queries and keys into range [-1, 1]. Then the attention mechanism can no longer represent sharp distributions, which will hurt the model performance. Fortunately, when equipped with RPE, the model is endowed with the ability to model sharp attention distributions even if the queries and keys are normalized, since relative positional encoding bi-; remains unbounded. For ease of reference, we call our model Normalized PRF Transformer with RPE and use NPRn-Transfrreer with RPE for short. A pseudocode implementation is provided in Algorithm 1. In the next section, we will show that
'In 4 the authors provide sample complexity bound of the approximation error between the TRF kernel and the dot-then-expooenniiae function without considering the denominator term the softmax. In this paper, we directly provide the sample complexity bound on the approximation error between the PRF-induced attention distribution and the softmax attention distribution.
'The TRF paper [32] also proposes to normalize queries and keys. In 32 the normalized queries and keys are multiplied with Gaussian vector w which is rescaled using a learnable vector a E Rdd This is equivalent to rescaling the normalized q and & using a and then feeding them into Φ(-). This doesn't 't control the norms of queries and keys since the scale of a can be large, which doesn't 't address the optimization issue.

---

Algorithm Efficient Normalized Kernelized Attention with RPE using FFT
function RPE_NKA Q, K, V, Wº, W WV (b-(n-1))"" bn-1), Ø):
# X; denotes the i-th element of X Q.W9
Q [1@: i=1 K t exp(b), A1 (T(K))VV)))) A2 (
K.W |1K
V + (V.W))
i=
# FFTMatrixMul () is FFT-based algori thm for multiplying Toeplitz matrix as in Eq. (12), (33) We omit the reshaping steps and assume the second input variable and the output have the same shape.
D1 < FFTMatrixMul(, Ai), D2 + FFTMatrixMul(,,
X
d(Q.)..1
22Q,)
return X.
end
Table I: GLUE scores dev set.
"** indicates the best performance.
Method
Complexity
CoLA
RTE
MRPC
STS-B
MNLI
QNLI
QQP
SST-2
AVG
Finetune from RoBERTa-base
SMYRE
logn)
68.6
87.7
89.7
85.0
91.1
89.7
93.2
FCA 46
O(n)
59.8
49.8
43.6
78.9
79.4
74.6
89.4
94.4
71.2
Improved FCA 46
(n)
60.1
70.4
87.3
90.0*
88.0*
93.0
91.5
94.7
84.4
Pre-train from scratch
Linformer 48
(n)
/
/
91.2
90.8
93.1
Nyströmformer
O(n)
80.9
88.7
86.3
Oin log
64.7%
75.1
86.4
94.8*
85,2
Tran-Tranffrrer with RPE can achieve comparable performances with vanilla Transformers and can be trained from scratch for all tasks due to using the low-variance kernels.
Experiments
In this section, we conduct experiments to verify the effectiveness of our approach on several benchmark datasets covering language pre-training, language modeling, machine translation, and image classification. Then we provide the ablation study and discussions the design choices. Due to space limitations. detailed description of the experiment settings are presented in Appendix A
4.1 Language Pre-training
We use the BERT-base architecture 9 in our experiments, which consists of 12 Transformer layers. For each layer, the hidden size is set to 768, and the number of attention heads is set 12. We follow 23 to construct the data corpus and train the model using masked language modeling task. We use the GLUE (General Language Understanding Evaluation) dataset (47 the downstream tasks evaluate the performance the pre-trained models. All codes are implemented based on fairseq (8 64 NVIDIA Tesla V100 GPUs with mixed-precisioe 26)
in PyTorch 31 All models
To show the effectiveness of our proposed Transfrresformmr with RPE, we choose several com- petitive baselines in literature, including Linformer [48), Nyströmformer [49], SMYRF [7]. and Fast Clustered Attention (FCA) [46]. All the baselines use the same architecture as ours. The details
'The expressive power of NPRI-Transformer with RPE be weaker than vanilla Transformer can only produce context-agnoste sharp distribution (from RPE) may not produce context-awar aware sharp patterns. However, this does not conflict with the observation that TraasTrrasformer with RPE outperforms Transforrrrmmr in practice. Although PRF Transformer seems to have more flexibility, its training instability issue hinders the model from learning well.

---

Table Language model perplexity scores on WikiText- 103 validation set. We use indicate the best performance. All the results of the baseline methods are reported in [32]
Model
Perplexity
Vanilla Transformer
33.0
Linear Transformer
38.4
Tra-fraasformer
33.6
T'RF-Traanformer-GGATTE
31.3
NPRF-Transformerr w/ RPE (Ours)
30.6*
about these baselines can be found in Appendix A. For baselines, we report the number their original papers. We also tried to train the PRF Transformer model from scratch or continue to train a PRF Transformer from the released RoBERTa checkpoint [23]. But all the efforts failed due to the optimization instability issue. For model, we directly modify the self-attention layer using normalized kernelized attention with RPE and train it from scratch.
The overall comparison results are presented in Table D ("/" indicates the GLUE score of the corresponding downstream task is not reported). It can be easily seen that our NPRF-Transformer with RPE outperforms the previous works in terms of the average GLUE score, with a light computational also worth noting that our model can be stably
overhead. Beyond the superior performance. pre-trained from scratch, which is particularly important in practice.
4.2 Language Modeling
We conduct experiments on WikiText- language modeling task to demonstrate the effectiveness
of our proposed method. We compare our model with the following baselines: 1. Transformer with softmax attention. 2. Linear Transformer proposed in [17], which uses kernelized attention with clu(-) + as its feature map. 3. TRF-Transformer using TRF kernel (called RFA in (32) 4. TRF-Tramfformer-GAATE
32 which uses gating mechanism to further boost performance.
Following (32), the sequence length is 512 during both training and evaluation. All models are trained without access to the context from previous mini-batches for a fair comparison. The model architecture consists of 6 decoder layers. The number of attention head is set 8. The hidden dimension is set to 512. The dimension of feed-forward layer is set to 2048. The dropout ratio and the weight decay are set to 0.1 and 0.01, respectively. The batch size is set 64. The feature map dimension is set to 64. We use Adam [19 as the optimizer, and set its hyperparameter e to le 6 and (51, 52) to (0.9, 0.98). The peak learning rate set to 2c The model is trained for 150k steps with a 6k-step warm-up stage followed by an inverse square-root learning rate scheduler. All models are trained on 8 NVIDIA Tesla V100 GPUs.
The results are shown in Table 2 Our SPRF-Transfommen with RPE outperforms all the baselines by large margin, e.g., it achieves 30.6 perplexity score which is 3.0 lower than the baseline TRF- Transfororrr and 0.7 lower than TRF-Trans TRF-Transformer-G.. GATE. It is worth noting that we mainly make architectural changes and should compare our model with TRRF-Tansformet In [32]. set of training tricks (such as gating) are introduced for improving the performance. Our method can be combined with those tricks for further improvements.
4.3 Machine Translation
In machine translation, we evaluate our method on the widely used public dataset IWSLT14 German<-English and French+-Englise For each language pair, we use vocabulary of 10K tokens based on a joint source and target byte pair encoding (BPE) (36). All the experiments use Transformer mecoder-decoder architecture, with 6-layer encoder and 6-layer decoder. We set the embedding dimension to 512 and the number of heads 4. For a fair comparison, we set the feature map dimension to 16 in the encoder and 24 in the decoder whenever kernelized attention is used. BLEU (29] is used as the evaluation measure of the model performance. In this small-scale task, we find the optimization for all the model variants is stable, and all models are trained from scratch.

---

Top-5 acc Table 3:
Test BLEU scores on machine translation tasks. We use " to indicate the best performance
obtained by the standard model / efficient model. "Enc/dec" denotes "encoder/decoder" for short.
Complextity encoder/ decoder
Encoder softmax PRF
Decoder softmax PRF
BLEU
Transformer model type
de-cn
en-dc
fi-cn
en-fr
Standard enc-doc
34.7
29.2
40.0
40.2
36.0
Standard enc + PRF
(OO
34.6
29.3"
40.3"
40.6
36.2"
PRF enc-dec
((n)
32.1
27.6
38.0
38.2
34.0
NPRF enc-dec w/ RPE (Ours)
OK log (t
34.3"
29.4
39.9"
40.4"
36.0"
Table Test accuracy on ImageNet. We use */-* to indicate the best performance obtained by the standard model efficient model.
Query/key
Attention
ImageNet
Model
Complexity
standard normalized
softmax
PRF
Top-1
ViT-hase
O(n")
77.9
/
ViT-large
O(n²)
76.5
DeiT-base
O(n²)
81.2*
95.0'
PRF DeiT-base (fine-tuned)
O(n)
79.5
94.1
NPRF DeiT-base
O(n)
77.7
93.3
NPRF DeiT-base RPE (Ours)
On log
80.9
94.8"*
As the translation task uses an encoder-decoder framework, it is natural to check whether the attention modification (i.e., kernelized attention) be applied to both encoder and decoder. In [32], the authors only show that using TRF attention in the decoder will not hurt the model performance. But long source sequence in translation,
argue that long target sequence usually corresponds and modifying the attentions in both encoder and decoder should be equally important.
We give systematic study the model performance of using PRF in the encoder/decoder and summarize the results shown in T We first observe that models using PRF kernel in the decoder comparable with or even slightly better than standard ones (the 2nd line v.s. the 1st line), which consistent with [32]. The performance improvement may be attributed to regularization effects should still be O(n?) as the
brought by the randomness in PRF However, the overall time standard attention is used in the encoder. When applying the kernelized attention to both encoder and decoder, we observe significant performance drop (the 3rd line the 1st line). This may indicate that when using more layers of kernelized attention, the model becomes harder to train due to accumulated approximation error through layers. On the contrary, our method can incorporate kernelized attention in both encoder and decoder, achieving the same average BLEU score as the standard Transformer.
4.4 Image Classification on ImageNet
Beyond the natural language tasks, we also extend our study to image classification. We use the DeiT-base [42] our backbone architecture, which consists of 12 Transformer layers. For each layer, the hidden size is set to 768, and the number of attention heads is set to 12. We set the feature map dimension to 32 whenever kernelized attention is used. Following 10 42 we benchmark our method on ImageNet- 1K (8, which contains 1.28M training images and 50K validation images from 1,000 classes. For evaluation, the top-1/top-5 accuracy on a single crop is reported. For our proposed method, we use the 2-dimensional relative positional encoding developed in [24).
We compare our proposed model with several baseline models using standard attentions, including WUT-base/ViT-lann proposed in [10] and the DeiT-base (42). We also compare our model with finetuned PRF attention model from DeiT, and an NPRF attention model without RPE. The result is shown in Table 4 First, we can see that our model achieves competitive performance compared to standard attention models (81.2 v.s. 80.9 in terms of t accuracy, 95.0 v.s. 94.8 in terms of top-5 accuracy). Second, we can see that compared with other efficient variants, our model achieves the best performance, and the use of normalized queries/keys and RPE are both helpful.
4.5 More Analyses
We conducted more studies the IWSLT14 Horman-to-Englise (De-En) task to ablate our designs.

---

Standard attention
Normalized attention
Standard attention w/ RPE
Narmalized attention RPE
valication BLEU
validettion OLEU
voliibetion BLEU)
15
15
25
vollidation
25
88
20
20
10
srade
15
32
14
120 255
512
ordde 30
32
04
320 254 512
orec 16 32 04 128 256 512
10
crac
15
32
04 120 256 512
reature map dimension
feature map dimension
feature map dimension
Testure map dimens on
Figure 2: Performance of kernelized attention models converted from well-trained Transformers.
"Oracle" represents the BLEU score of the well-trained model for reference.
velidation BLEU
velidation BUEU S4
map dimension
lype map
(a)
Figure 3: (a). Performance of the
with RPE using different feature map dimen- sions. (b). Performance of the models using different feature map functions.
Normalized attention easy to be converted to its kernelized version We conduct experiments to study whether using normalized attention practically mitigate the performance drop issue observed in the previous works. We train four model variants: Transformer with standard/normalized attention and with/without RPE. After the training finishes, we replace the attention softmax functions by Equation 5 without further finetuning. For each setting, we run with 5 random seeds. Figure 2 summarizes the experiment results terms of the validation BLEU scores with confidence intervals.
We can see that Transformer with standard attention suffers a considerable performance degrada- tion after the conversion. In contrast, Transformer with normalized attention only suffers small performance drop. This clearly shows that normalization reduces the approximate error significantly. Interestingly, observe that RPE universally improves the performance of the converted model for both standard and normalized attention. For our NPRF-Transformer with RPE, the feature map dimension is not sensitive any longer, e.g., the converted model only has 0.5 performance drop when the feature map dimension is only 32.
Affects of feature map dimensions and other choices of feature map The dimension feature map é is tlyee--aramett that balances the accuracy and efficiency. For our NPRF-Transfonin Transformer with RPE, we train models from scratch with different feature map dimensions and check the validation BLEU, We can see from Figure 3a that the feature map dimension is not sensitive regarding the final performance. Setting the dimensionality to 16 even a slightly better choice.
Our approach is general which can be applied to any feature map function. The main body of the paper showcases models using the PRF feature map, and here we provide some empirical comparisons with other feature maps. Specifically, 2) the Sphere-PRF feature map, which designed in 4 as an alternative choice of PRF, where all w; in Equation 5 are sampled independently from Unif(vds** also designed in 4 where ... W in Equation 5 are forced to be orthogonal to each other. For all the choices, we train models with normalized kernelized attention and RPE under the same setting
experiment with 1) the TRF feature map defined by Equation 4
3) the ORF feature map, which is
in Section 4.3 and report the best validation BLEU. The results shown in Figure 3b) We can
see that when we normalize the attention and use RPE, all models obtain similar performance.
5 Conclusion
In this paper, we propose novel way accelerate attention calculation for Transformers with relative positional encoding and achieve O(n log n) time complexity. Our method is built top of the kernelized attention. Using the fact that relative positional encoding forms Toeplitz matrix, we mathematically show that kernelized attention with RPE can be calculated efficiently using Fast Fourier Transform (FFT). We further demonstrate the additional benefit of using relative positional encoding from the optimization perspective. As mentioned in the main paper, we will extend the current method the decoder model in generation tasks and explore more strategies to accelerate attention in the long- sequence regime.
10

---

Acknowledgements
We thank all the anonymous reviewers for their constructive comments. This work was supported by National Key R&D Program of China (2018YFB (4FB140))00) Key-Area Research and Development Program of Guangdong Province (No.2019B12124421 BJNSF (L172037) and Beijing Academy of Artificial Intelligence. Project 2020BD006 supported by PKU-Baidu Fund.
References
[1] Hangbo Bao, Li Dong, Furu Wei, Wenhui Wang, Nan Yang. Xiaodong Liu, Yu Wang, Jianfeng Gao, Songhao Piao, Ming Zhou, al. Unilmv2: Pseudo-masked language models for unified language model pre-training. In International Conference on Machine Learning, pages 642-652.
PMLR, 2020.
[2] Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European Conference on Computer Vision, pages 213-229. Springer, 2020.
[3] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019,
[4] Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song. Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy Colwell, and Adrian Weller. Rethinking attention with performers. In International Conference on Learning Representations, 2021.
[5] Patryk Chrabaszez, Ily Loshchilov, and Frank Hutter. A downsampled variant of imagenet as an alternative to the cifar datasets. arXiv preprini 1707:0/707.0880 2017.
[6] Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D Manning. What does bert look at? an analysis of bert's s attention. ACL 2019, page 276, 2019.
[7] Giannis Daras, Nikita Kitaev, Augustus Odena, and Alexandros G Dimakis. Smyrf-efficient attention using asymmetric clustering. Advances in Neural Information Processing Systems, 33, 2020.
[8] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image datahase. In 2009 IEEE conference on computer vision and pattern recognition, pages 248-255. leee, 2009.
[9] Jacob Devlin, Ming- Wei Chang. Kenton Lec, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pages 4171-4186, 2019.
[10] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
[11] Robert M Gray. Toeplitz and circulant matrices: A review. now publishers inc, 2006.
[12] Matej Greić, Ivan Grubišić, and Siniša Šegvić. Densely connected normalizing flows. arXiv preprint arXiv:2106. 04627, 2021.
[13] Anmol Gulati, James Qin. Chung-Cheng Chiu, Niki Parmar, Yu Zhang, Jiahui Yu, Wei Han, Shibo Wang, Zhengdong Zhang, Yonghui Wu, et al. Conformer: Convolution-auggenntte transformer for speech recognition. Proc. Interspeech 2020, pages 5036-5040, 2020.
[14] Pengcheng He, Xiaodong Liu, Jianfeng Gao, and Weizhu Chen. Deberta: Decoding enhanced recoding-enhanced bert with disentangled attention. In International Conference on Learning Representations, 2021.
[15] Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans. Axial attention in multidimensional transformers. 2019.
[16] Jungo Kasai, Hao Peng, Yizhe Zhang, Dani Yogatama, Gabriel Ilharco, Nikolaos Pappas, Yi Mao, Weizhu Chen, and Noah A. Smith. Finetuning pretrained transformers into rnns, 2021.
11

---

[17] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mnss: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pages 5156-5165. PMIR, 2020.
(18) Guolin Ke, Di He, and TieYaan Liu. Rethinking positional encoding in language pre-training. In International Conference on Learning Representations, 2021.
[19] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR (Poster), 2015.
[20] Diederik P Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. arXiv preprint arXiv:2 107.00070000 10000000, 2021.
[21] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference Learning Representations, 2019.
[22] Taku Kudo and John Richardson. Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 66-71,
2018.
[23] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 1692, 2019.
[24] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. arXiv preprint arXiv:2103.1400 2021.
[25] Ilya Loshchilov and Frank Hutter. Fixing weight decay regularization in adam, 2018
[26] Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, and Hao Wu. Mixed precision training, 2018.
[27] Sharan Narang, Hyung Won Chung, Yi Tay, William Fedus, Thibault Fevry, Michael Matena, Karishma Malkan, Noah Fiedel, Noam Shazeer, Zhenzhong Lan, et al. Do transformer modifi- cations transfer across implementations and applications? arXiv preprint arXiv:2 2021.
[28] Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. In Proceedings of NAACL-HLLT 2019: Demonstrations, 2019.
[29] Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. Bleu: method for automatic evaluation of machine translation. In Proceedings of the 40th annual meeting on association for computational linguistics. pages 31 311-318. Association for Computational Linguisties. 2002.
[30] Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International Conference on Machine Learning, pages 4055-4064 PMLR, 2018
[31] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in Neural Information Processing Systems, 32:8026 8037, 2019.
[32] Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah Smith, and Lingpeng Kong. Random feature attention. International Conference on Learning Representations, 2021.
[33] Jiezhong Qiu, Hao Ma, Omer Levy, Wen-tau Yih, Sinong Wang, and Jie Tang. Blockwise self-attention for long document understanding. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: Findings, pages 2555-2565, 2020.
[34] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
[35] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer, 2020.
12

---

[36]
Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Pupers), pages 1715-1725, 2016.
[37] Rico Senarich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. In ACI., 2016.
[38]
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position repre- sentations. arXiv preprint 1rXiv:1803.02155. 2018.
[39] Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. arXiv e-prints, pages arXiv-2101, 2021.
[40] Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. Re- thinking the inception architecture for computer vision. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 28 2818-2826, 2016.
[41] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey,
2020.
[42] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers distillation through attention, 2021.
[43] Trieu H Trinh and Quoc V Le. A simple method for commonsense reasoning. arXiv preprint arXiv:1806 180..02877 2018.
[44] Yao- Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis Phill-p Morency, and Ruslan Salakhutdinov. Transformer dissection: An unified understanding for transformer's attention via the lens of kernel. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pages 1335-43442 2019.
[45] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez. Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pages 5998-6008, 2017.
[46] Apoorv Vyas, Angelos Katharopoulos, and François Fleuret. Fast transformers with clustered attention. Advances in Neural Information Processing Systems, 33, 2020.
[47] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. 2019. In the Proceedings of ICLR.
[48] Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity, 2020.
[49] Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A nyström- based algorithm for approximating self-attention. In Proceedings of the AAAI Conference Artificial Intelligence, 2021.
[50] Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding. Advances in Neural Information Processing Systems, 32:5753-5766 2019.
[51] Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In Proceedings of the IEEE international conference on computer vision, pages 19-27, 2015.
13


Reviews:
[{"review_id": "ppiMwYa-WMW", "reviewer": "Reviewer_LQ92", "summary": "The paper introduces a novel way to add relative positional encoding to efficient transformers with kernelized attention. The authors make use of the efficient matrix multiplication of Toeplitz matrices using the fast Fourier transform.  The authors perform several experiments on masked language modeling, machine translation and even vision transformers.", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths\n-------------\n\n- Using the fast multiplication of Toeplitz matrices based on the Fourier transform to implement relative positional encodings is brilliant and very interesting\n- The experimental evaluation shows that kernelized transformers with the proposed RPE can actually train and often achieve better results than the same model without RPE\n- The theoretical analysis regarding the norm of the queries and keys as well as the ablations and analyses in section 4.4 are great and of independent interest. In particular, I find Figure 2 very informative that concisely showcases both the fact that reducing the norm of the queries and keys helps the random Fourier feature approximations as well as that RPE accounts for a large part of the attention distribution (when used) and helps even more when approximating.\n- The authors provide clean code with their submission\n\nWeaknesses\n-----------------\n\n### Asymptotic complexity vs real world speed\n\nThe main weakness of the paper stems from the misrepresentation of the complexity of the attention computation using the proposed algorithm. The algorithm in its provided form, requires storing in memory a matrix of size $O(ND^2)$ (namely $\\phi(K_i) V_i^T$ for all $i$) per sample per head. This is the same issue as showcased in [15] for autoregressive kernelized attention. The result is a theoretically fast algorithm but a slow one in practice.\n\nThis can be seen also from Figure 1a. Even at 32k sequence lengths and using only the forward pass, the proposed method needs to down-project the queries and keys to half the dimensions in order to be faster than the **full softmax** attention. In comparison, the simplest kernelized attention by [15] is **10x times** faster in the same experiment. Down-projecting using PRF could make it even faster.\n\nFinally, evaluation of the computational cost is not provided even for a single one of the real world experiments. For instance what is the computational cost of an epoch in GPU wall clock time for the models in tables 1, 2 and 3. Since the proposed model is an approximation of the softmax transformer it should at least be verified that it is faster than that.\n\n### Minor comments\n\n1. The wording and notation in section 3.2 is unnecessarily complicated. The $vec(\\cdot)$ notation can be omitted and if needed intermediate matrices can be properly defined together with their dimensions.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors do not discuss the societal impact of their work.\n\nThe limitations are not adequately discussed as mentioned in the weaknesses section of the main review.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "mkydVbxMwTZ", "reviewer": "Reviewer_npFH", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "First, I summarize the pros and cons of this paper.\n\nPros:\n1. The proposed method is simple and novel. \n2. Adding the proposed RPE improves the performance over the linear kernelized attentions.\n3. The authors conduct experiments on a variety of tasks: language pre-training, image classification, and machine translation and show that the proposed method is more accurate than the kernelized attention baselines.\n4. The authors demonstrate that the proposed normalized kernelized attention is not more more accurate and more stable.\n\n\nCons:\n1. No inference time provided on real datasets (See Limitations And Societal Impact for details)\n2. Not showing the performance-efficiency tradeoff (See Limitations And Societal Impact for details)\n3. The proposed method operates at only $O(n \\log n)$ time, but the authors do not compare with simple linear baselines which has relative positional information like just using convolutions [1] or a combination of convolutions and kernelized attention. This makes me wonder why we should prefer the proposed $O(n \\log n)$ method rather than a $O(n)$ baseline.\n4. Figure 3 (b) is very hard to read. It would be better to show y-axis between 30 and 35 or 33 and 35.\n\nOverall, the paper is written in high quality and I recommend accept it; however, there are a couple of critical concerns which hinder me from rating this paper higher.\n\nReference\n[1] Tay, Yi, et al. \"Are Pre-trained Convolutions Better than Pre-trained Transformers?.\" ACL 2021.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "While I believe reducing the time complexity is impactful, the impact of the paper is sort of limited by the experimental results it presents for the following two reasons:\n1. The authors only show the forward speed in Fig. 1 (a). In all tables in Sec. 5, the authors only show the time complexity instead of the real inference time on the real datasets. This worries me that this paper may be yet another paper with great time complexity but relatively slow inference time on real use cases.\n2. While the paper is proposing an efficient alternative to self-attention, The authors do not show the performance-efficiency trade-off. I wish I could see that the method is at the Pareto frontier. For example, showing with a similar inference speed, this method outperforms linear attention and softmax attention. Linear attentions have $O(n)$ time complexity, so it fairer to compare with a larger model with linear attention and the same inference time as the proposed PRF + RLE method. Similarly, comparing the proposed method to a smaller softmax attention model with the same inference time is also expected.\n\nI am willing to raise my score to 8 if both of these concerns are addressed.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "LMPAZR5QYte", "reviewer": "Reviewer_mmJ4", "summary": "This paper contributes to the literature on efficient [transformer][]\narchitectures. Specifically, papers that focus on reducing the quadratic\ncomputational complexity of the self-attention mechanism between tokens.\n\nThe problem this paper focuses on is the incompatibility of [relative\npositional encodings (RPE)][rpe] and sub-quadratic efficient attention\nmechanisms. It adapts the kernelized attention used in the [Performer][]\narchitecture to allow the use of RPE. In the process, the authors find that\nthis adaptation allows use of the FFT algorithm to perform the approximate\nattention mechanism, resulting in an $O(N \\log N)$ algorithm, where $N$ is\nthe number of tokens.\n\nBenchmarks are provided on the GLUE dataset according to various metrics\nand on the ImageNet dataset using Image Transformer variants. Performance\nto competing methods chosen is better on most of the metrics presented.\n\n[transformer]: https://papers.nips.cc/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf\n[rpe]: https://arxiv.org/abs/1803.02155\n[performer]: https://arxiv.org/abs/2009.14794", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Originality\n\n> Are the tasks or methods new? Is the work a novel\n> combination of well-known techniques? (This can be valuable!) Is it clear\n> how this work differs from previous contributions? Is related work\n> adequately cited?\n\nThis work is a novel combination of two existing methods: [relative\npositional encodings][rpe] and [kernelized efficient\nself-attention][performer]. In the process they make an independent\ncontribution, discovering that combining the two methods enables use of the\nFFT algorithm to compute the efficient self-attention mechanism. This fills\nits own niche among the many existing efficient self-attention mechanisms.\n\nRelated work cited is a thorough review of efficient self-attention\nmethods, with each correctly sorted into their respective types. The review\nmakes it clear where this method fits into this set.\n\nIn addition the authors make various observations that will be useful to\nmachine learning practitioners:\n\n1. Ablation comparisons of [Performer][] and [Random Feature\nAttention][rfa]\n2. Analysis of the function class of transformers using [RPE][]\n3. Analysis of the approximation error of kernelized attention\n    1. Fourth paragraph of Section 3.3 makes a useful observation about the\ntraining dynamics of transformers using kernelized attention\n    2. This is further put into context in the literature by the footnotes\nin the following paragraph on \"Normalized Kernel Attention with RPE\"\n4. Address the limitations of existing kernelized attention mechanisms,\nsuch as [Performer][], that are unstable when trained from scratch by using\na normalization step\n\n[rpe]: https://arxiv.org/abs/1803.02155\n[performer]: https://arxiv.org/abs/2009.14794\n[rfa]: https://openreview.net/forum?id=QtTKTdVrFBB\n\n## Quality\n\n> Is the submission technically sound? Are claims well supported\n> (e.g., by theoretical analysis or experimental results)? Are the methods\n> used appropriate? Is this a complete piece of work or work in progress? Are\n> the authors careful and honest about evaluating both the strengths and\n> weaknesses of their work?\n\nAll claims in the paper are supported by theoretical inquiry and\nexperiment. The experimental results presented demonstrate the claims\nmade in the text. On the benchmarks chosen (GLUE, machine translation,\nImageNet) the network performs well in comparison to the methods chosen for\ncomparison.\n\nThe technical analysis is in depth and considers a limitation of the work.\nIn a footnote of Section 3.2, the authors mention that the method cannot be\nused to improve the efficiency of the inference step. In practice the\nmemory explosion of the kernel matrix during training is a larger problem\nwhen training transformers than the inference step memory usage. But, it\nmay affect this method's ability to scale to large sequence lengths.\n\nLemma 2 in particular is an important observation about the variance of\nkernelized attention and is backed up in the paper by theoretical analysis\nand precise experimental exploration illustrated in Figure 1. Both of these\ntogether are very effective in convincing the reader that this Lemma is\ncorrect and worthwhile.\n\n[performer]: https://arxiv.org/abs/2009.14794\n\n## Clarity\n\n> Is the submission clearly written? Is it well organized? (If not, please\n> make constructive suggestions for improving its clarity.) Does it\n> adequately inform the reader? (Note that a superbly written paper provides\n> enough information for an expert reader to reproduce its results.)\n\nThe submission is clearly written and well organised. Reading this paper in\nsequence each topic is introduced in a way that is easy to understand and\neach idea follows from the previous. I am confident that any researcher\nfamiliar with transformers will be able to understand this paper.\n\nSection 2 deserves praise for providing a concise and useful description of\nkernelized attention. It helps to make the contribution of the paper clear\nto the reader by drawing together various prior work in kernelized\nattention under a common framework. I would use it as a reference if I were\nwriting about kernelized attention in future.\n\nSome minor suggestions I would make:\n\n1. The related work is discussed in the Introduction, which may not draw\nthe eye for those searching for it.  However, putting it in it's own\nsection would ruin the flow and probably push the paper over the page\nlimit. I would suggest putting a `\\paragraph` bold title indicating that\nparagraph 2 of the Introduction is the Related Work section but if the\nauthor's find it worse then I would not argue.\n2. Section 3 mentions the authors will \"develop new techniques\" but does\nnot indicate what those techniques are. It would be better to be precise\nabout what the new techniques are.\n3. Sentence 3 of Section 3.3 has a grammar mistake that impedes\nreadability, \"those kernelized attention\" should probably be \"kernelized\nattention\"\n\n## Significance\n\n> Are the results important? Are others (researchers or practitioners) likely\n> to use the ideas or build on them? Does the submission address a difficult\n> task in a better way than previous work?  Does it advance the state of the\n> art in a demonstrable way? Does it provide unique data, unique conclusions\n> about existing data, or a unique theoretical or experimental approach?\n\nThe specific niche this work fits into in the literature on efficient\nself-attention mechanisms is useful. This seems like a natural question\nwhen considering [RPE][] and [kernelized attention][performer]. \n\nIn addition, the demonstration of how the FFT algorithm may be used to\nperform a self-attention mechanism appears valuable. It seems like there\nmay be ways to build upon it, such as incorporating it with other relative\npositional encodings, such as [rotary encodings][rotary] and the result may\ncontinue to improve performance. There are likely other ways this work\ncould be adopted and improved into Transformer design.\n\n[rpe]: https://arxiv.org/abs/1803.02155\n[performer]: https://arxiv.org/abs/2009.14794\n[rfa]: https://openreview.net/forum?id=QtTKTdVrFBB\n[rotary]: https://arxiv.org/abs/2104.09864", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Efficient self-attention methods with sub-quadratic computational\ncomplexity are typically applied to problems that have a large number of\ntokens. An example of this can be seen in the experiment in Section 4.5 of\nthe [Performer][] paper, on the ImageNet64 task that cannot be addressed by\ntraditional transformers. This paper does not include experiments with\nnumbers of tokens that would be infeasible for a normal transformer.\n\nThe authors discuss how normalizing the kernelized attention reduces the\ncapacity of the model to express sharp attention distributions in the final\nparagraph of Section 3.3. As the goal of this work is not to lose the extra\ncapacity offered by [RPE][] when using kernelized attention it is a\nlimitation that the authors do not present an experiment investigating how\nthe model performs without normalization and with RPE (ie PRF w/ RPE in the\nnotation of the paper's experiments). It may be necessary to address the\nstability using the same tricks employed by [Performer][] to run this\nexperiment.\n\nThe FFT algorithm proposed to perform the Toeplitz matrix multiplication\nreceives little discussion, save that it is possible. I think this is\nappropriate but there are some concerns that could be addressed:\n\n1. FFT algorithms are not well supported in all deep learning frameworks\nand FFT Toeplitz matrix multiplications are not supported directly\n2. FFT algorithms, unless carefully implemented, require the inputs to be a\npower of 2 number of dimensions\n3. Implementing this using FFT operations in an autograd system will have a\nrelatively large number of intermediate variables that will be stored for\nthe backward pass and could increase memory usage. This could be addressed\nwith custom operations that define the appropriate backward pass directly.\n\n[performer]: https://arxiv.org/abs/2009.14794\n[rpe]: https://arxiv.org/abs/1803.02155", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "4dO2nWjsbRz", "reviewer": "Reviewer_AuHp", "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.\n\nThe 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.\n\nExperiments 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\n\n\n\n \n", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Originality\nThe paper brings three well known techniques together: kernelized attention, RPE, and the Convolution Theorem, to enable kernelized RPE attention. It's a nice idea that leads to an elegant algorithm. To my knowledge it is the first time that this has been proposed. All prior work has been cited appropriately.\n\n## Quality \nThe overall quality of the paper is very good. The derivations are sound and the experiments support the claims made by the authos.\n\n\nThere are two things missing that I think should be discussed: \n\n* How does the model behave with causal attention. I think it can be implemented due to the RPE by setting $c_i=0$ for $i\\ge 0$ in (11) and (12). Has this been done in the experiments and found to work? If so, then I think footnote 1 is not a limitation unique to this algorithm. Sequence generation at inference time always gives another factor $n$, and so does it here. The convolution does not improve on that, bat at least it does not make things worth (at least compared to non-kernelized attention) \n\n* For the *normalized kernelized attention with RPE* ( l 203 - 209) it should be noted that this is not the same as *unnormalized kernelized attention*. You can produce sharp distributions but only at specific distances for which for which $b_{i-j}$ happens to be large. So NPRF w/ RPE has lower expressivity than PRF, still showing better performance (Tab. 3). It would add to the strength of the paper if the authors could elaborate on that.\n\n## Clarity\nThe paper is very well written. The mathematical derivation is comprehensive an easy to follow. All important information is in the main paper\n\n## Significance\nRPE have been shown to have a big impact on the performance of attention based models on various problems. On the other hand, kernelized attention can save a big deal of computational cost and make transformer models scalable to real world large scale problems. To link both approaches in an efficient way clearly advances the field and has a big impact for practitioners working with transformer models.  \n\n**I vote for accepting the paper.**\n\nMinor issues\n* l 79: indicate the axis of the softmax operation\n* Sec. 2.1 ff:  all $x_i$ are row vectors which is unusual and can lead to confusion, in particular when writing $x=(x_1,\\ldots,x_n)$. You should either revise the notation or state explicitly that $x_i\\in\\mathbb R^d$ is a row vector \n* l 88 - 94: wherever you have running variables $j'$ in summations, make sure to also use $j'$ in the summation terms.\n* Fig. 1a would gain from one more sequence length step ($n=48k$ or $n=64k$) showing the (almost) linear relation better\n \n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors don't explicitly address limitations of their model, other than additional speedup for sequence generation is not provided.\nAnother potential limitation, that should be mentioned, could come from the fact that NPRF even with RPE cannot generate sparse attention to the same extend as PRF or plain dot-then-exponentiate attention  (see above)\n\nAlthough not explicitly addressed in the paper,  I don't see any negative societal impact of the work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This is an interesting paper extending linear attention mechanisms proposed in Performers to work with an *arbitrary* relative position encoding (RPE) mask. Previous results showed that it is possible for special instantiations of RPE. This paper shows that as long as a mask has Toeplitz structure (this can be actually further extended to a low displacement rank structure), efficient computation of the attention module can be achieved by combining random feature map tricks with Fast Fourier Transformer mechanism (every RPE is by definition a Toeplitz matrix). The authors furthermore claim, that by incorporating RPE into linear attention modules, they can substantially improve their overall accuracy. According to them, this is possible since by adding RPEs, one can constrain queries and keys to be L2-normalized (resulting in the substantial reduction of the variance of the random feature map estimator that grows exponentially in the lengths of keys/queries). Normally such a normalization would affect the expressiveness of the model, but according to the authors, the RPE (which in principle is not bounded) provides extra expressiveness that would be lost otherwise. 
This claim is controversial and additional detailed experiments showing that RPE can indeed carry out the otherwise-lost expressiveness would further strengthen the paper (in particular an extensive ablation studies over L2-normalized variants with and without RPE would provide extra clarity).  The paper would also benefit from providing more theoretical understanding of this phenomenon.
The other important point that should be discussed in more detail is an efficient implementation on TPUs of Fast Fourier Transform. Unless such an implementation is provided, practical usage of this algorithm will be limited. 
Nevertheless, it is an important result in the field and an elegant approach to incorporating structured masks into kernelizable attention modules. This algorithmic observation is the main contribution of the paper.
Decision: Accept (Poster)