Random Feature Attention
Authors: Hao Peng
Nikolaos Pappas
Dani Yogatama
Roy Schwartz
Noah Smith
Lingpeng Kong
Venue: ICLR 2021 Spotlight (2021)
Keywords: Attention
transformers
machine translation
language modeling

Abstract: Transformers are state-of-the-art models for a variety of sequence modeling tasks. At their core is an attention function which models pairwise interactions between the inputs at every timestep. While attention is powerful, it does not scale efficiently to long sequences due to its quadratic time and space complexity in the sequence length. We propose RFA, a linear time and space attention that uses random feature methods to approximate the softmax function, and explore its application in transformers. RFA can be used as a drop-in replacement for conventional softmax attention and offers a straightforward way of learning with recency bias through an optional gating mechanism. Experiments on language modeling and machine translation demonstrate that RFA achieves similar or better performance compared to strong transformer baselines. In the machine translation experiment, RFA decodes twice as fast as a vanilla transformer. Compared to existing efficient transformer variants, RFA is competitive in terms of both accuracy and efficiency on three long text classification datasets. Our analysis shows that RFA’s efficiency gains are especially notable on long sequences, suggesting that RFA will be particularly useful in tasks that require working with large inputs, fast decoding speed, or low memory footprints.

Published as a conference paper at ICLR 2021
RANDOM FEATURE ATTENTION
Hao Peng
Nikolaos Pappas* Dani Yogatama* +
Roy Schwartz
Noah A. Smith o
Lingpeng Kong". +
Paul Allen School of Computer Science & Engineering, University of Washington
DeepMind Allen Institute for Artificial Intelligence
School of Computer Science & Engineering, Hebrew University of Jerusalem
Department of Computer Science, The University of Hong Kong
h apeng, apappas, nasmith)ecs washington. edu
dyogat aoaatama@gooooooo @ma@google com, roysêcs
1pk0cs aku.hk
ABSTRACT
Transformers are state-of-the-art models for a variety of sequence modeling tasks. At their core is attention function which models pairwise interactions between the inputs at every timestep. While attention is powerful, it does not scale effi- ciently to long sequences due to its quadratic time and space complexity in the sequence length. We propose RFA, a linear time and space attention that uses random feature methods to approximate the softmax function, and explore its ap- plication in transformers. RFA can be used as a drop-in replacement for conven- tional softmax attention and offers straightforward way of learning with recency bias through an optional gating mechanism. Experiments on language model- ing and machine translation demonstrate that RFA achieves similar or better per- formance compared to strong transformer baselines. In the machine translation experiment, RFA decodes twice as fast as a vanilla transformer. Compared ex- isting efficient transformer variants, RFA is competitive in terms of both accuracy and efficiency on three long text classification datasets. Our analysis shows that RFA's 's efficiency gains are especially notable on long sequences, suggesting that RFA will be particularly useful in tasks that require working with large inputs, fast decoding speed, or low memory footprints.
INTRODUCTION
Transformer architectures (Vaswani et al. 2017) have achieved tremendous success on a variety of sequence modeling tasks (Out al. 2018 Radford al. 2018 Parmar et al. 2018, Devlin al., 2019 Parisotto et al. 2020, inter alia). Under the hood, the key component is attention (Bahdanau al. 2015). which models pairwise interactions of the inputs, regardless of their distances from each other. This comes with quadratic time and memory costs, making the transformers computationally expensive, especially for long sequences. large body of research has been devoted to improving their time and memory efficiency (Tay et al.] 2020c). Although better asymptotic complexity and prominent gains for long sequences have been achieved (Lee et al. 2019 Child et al. 2019) Beltagy ct al. 2020 inter alia), in practice, many existing approaches are less well-suited for moderate- length ones: the additional computation steps required by some approaches can overshadow the
time and memory they save (Kitaev et al. 2020 Wang al. 2020: Roy al. 2020 inter alia).
This work proposes random feature attention (RFA), an efficient attention variant that scales lin- early in sequence length in terms of time and space, and achieves practical gains for both long and moderate length sequences. RFA builds on a kernel perspective of softmax (Rawat et al. 2019).
Using the well-established random feature maps (Rahimi & Recht 2007 Avron et al.] 2016 22
RFA approximates the fol-then-expooentien
function with a kernel trick (Hofmann et al. 2008):
exp(x y) ((x) <bl). Inspired by its connections to gated recurrent neural networks (Hochre-
iter & Schmidhuber 1997 Cho et al., 2014) and fast weights (Schmidhuber 1992). we further
augment RFA with an optional gating mechanism, offering a straightforward way of learning with recency bias when locality is desired.
"The majority of this work done while these authors DeepMind.

---

Published as a conference paper at ICLR 2021
RFA and its gated variant ((3) can be used as drop-in substitute for the canonical softmax attention, and increase the number of parameters by less than 0.1%. We explore its applications in transformers on language modeling, machine translation, and long text classification ($4). Our experiments show that RFA achieves comparable performance to vanilla transformer baselines in all tasks, while out- performing recent related approach (Katharopoulos et al. 2020). The gating mechanism proves particularly useful in language modeling: the gated variant of RFA outperforms the transformer baseline on WikiText-103. RFA shines in decoding, even for shorter sequences. In our head-to-head comparison on machine translation benchmarks, RFA decodes around 2 faster than transformer baseline, without accuracy loss. Comparisons to several recent efficient transformer variants on three long text classification datasets show that RFA is competitive in terms of both accuracy and efficiency. Our analysis (55) shows that more significant time and memory efficiency improvements can be achieved for longer sequences: 12x decoding speedup with less than 10% of the memory for 2,048-Iength outputs.
BACKGROUND
2.1 ATTENTION IN SEQUENCE MODELING
The attention mechanism Babdanau et al. tasks. Its dot-product variant is the key building block for the state-of-the-art transformer archi- tectures (Vaswani et al. 2017). Let {qe ) denote a sequence of N query vectors, that attend to sequences of M key and value vectors. At each timestep, the attention linearly combines the values weighted by the outputs of a softmax:
2013) has been widely used in many sequence modeling
exp (q: k//T) E, exp (q: ky/T)
attn (qe, (k.). (v.)) (k.).(v.))= [vm))=
(1)
is the temperature hyperparameter determining how "flat" the softmax is (Hinton et al. 2015
Calculating attention for a single query takes ((M) time and space. For the full sequence of N queries the space amounts ((MN). When the computation cannot be parallelized across the queries, e.g., in autoregressive decoding, the time complexity is quadratic in the sequence length.
2.2 RANDOM FEATURE METHODS
The theoretical backbone of this work is the unbiased estimation of the Gaussian kernel by Rahimi & Recht (2007). Based on Bochner's theorem (Bochner I955), Rahimi & Recht (2007) proposed random Fourier features approximate desired shift-invariant kernel. The method nonlinearly transforms a pair of vectors x and using a random feature map Φ; the inner product between ((x) and Φvy) approximates the kernel evaluation on x and y. More precisely:
Theorem (Rahimi& & Recht 2007). Let Φ 12d + R2 be a nonlinear transformation:
Φ (x) V1/D sin (w1 x..... sin (WD x), cos (W1 ....... cos (WD xx]
(2)
When d-dimensional random vectors w; are independently sampled from N(,, I2I),
Ew, ( (x) & (y)] = exp (- ||X 10--yll y /20"))
Variance of the estimation is inversely proportional to D (Appendix A.2 Yu et al. 2016).
Random feature methods proved successful in speeding up kernel methods (Oliva et al. 2015 Avron inter alia), and more recently are used to efficiently approximate soft- max Rawat et al. 2019). In $3.1 we use it derive unbiased estimate to exp((: .)) and then an efficient approximation to softmax attention.
ct al. 2017 Sun 2019
MODEL
This section presents RFA 3.I) and its gated variant (3.2). In 33. lay out several design choices and relate REA to prior works. We close by practically analyzing RFA 's complexity (63.4).
M N in self-attention; they may differ, e.g., in the cross attention of sequence-to-sequenee model.
2

---

Published as a conference paper at ICLR 2021
(q.hen
{q: [im
((())
soft
[k: J
O(MNN
(-)
(be
[ka
(
(h: J
[v:)==1
(LM)
{v:
M
O(M))
(a) Softmax attention.
(h) Random feature attention.
Figure I: Computation graphs for softmax attention (left) and random feature attention (right). Here, assume cross attention with source length M and target length N.
3.1 RANDOM FEATURE ATTENTION
RFA builds on an unbiased estimate to exp((- ()) from Theorem
which we begin with:
exp (x y/o²) exp [|x|| /2o² + |lyl /2o² exp (- (-I|x x11 /202)
(4)
exp ||x|] /20² llyl /20² Φ (x) Φ (y).
The last line does not have any nonlinear interaction between (xx) and <by), allowing for linear time/space approximation to attention. For clarity assume the query and keys are unit vectors F
Σ exp (q: k,/o²)
attn (q: (k:), (vi))
E, exp kz/02)
(qr) o v
(5)
Σ, (q)) 6 (ky)
o (q)) Σ, Φ (k:) & REA (q. (ki), (ve)).
q Σ, Φ (k,)
denotes the outer product between vectors, and 6² corresponds to the temperature term 7 in Eq
RFA can be used as a drop-in-replasemmen for softmax-attention.
The input is revealed in full to cross attention and encoder self-attention. Here RFA
calculates attention using Eq. In causal attention RFA attends only the prefix This allows for a recurrent compu- tation. Tuple (S, E 2xd E IR22) is used as the "hidden state" at time step t to keep track of the history, similar those in RNNs. Then RFA(Q (kiises (vilige) (./(((()))))) 2t), where
S + (k,) 3 Vt, 2: = Φ (k:).
(6)
2D denotes the size of ((-). Appendix A.I summarizes the computation procedure of RFA, and [eguparcommpars it against the softmax attention. Appendix derives causal RFA in detail.
Analogously to the softmax attention, RFA has its multiheaded variant (Vaswaní al. 2017). In our experiments we use causal RFA in transformer language model ($4.1). and both cross and causal RFA in the decoder lequeece-to-sequeee machine translation model.
3.2 RFA-GATE: LEARNING WITH RECENCY BIAS
The canonical softmax attention does not have any explicit modeling of distance or locality. In learning problems where such inductive bias is crucial (Ba et al. 2016 Parmar et al.) 2018 Miconi ct al. 2018 Li et al. 2019 inter alia), transformers heavily rely on positional encodings. Answering to this, many approaches have been proposed, e.g., learning the attention spans (Sukhbaatar et al.
Thhis can be achieved by f2-normalizing the query and keys. See $3.3 for a related discussion.
It is also sometimes called "deender self-attention" "autoregressive attention."
3

---

Published as a conference paper at ICLR 2021
2019 Wu al. 2020), and enhancing the attention computation with recurrent (Hao et al. 2019 Chen et al. 2019) or convolutional (Wu et al. 2019 Mohamed et al. 2019) components.
RFA faces the same issue, but its causal attention variant (Eq. 6) offers straightforward way of learning with recency bias. We draw inspiration from its connections to RNNs, and augment RFA with learned gating mechanism (Hochreiter & Schmidhuber 1997
Cho al. 2014 Peng al.)
2018 inter alia):
9: sigmoid( Wa + bg):
S: 9t S + (1 gr) Φ (k,) @ Vr.
gt 2 + (1 9b) Φ (kr)
Wa and by are learned parameters, and is the input representation timestep / By multiplying the learned scalar gates 0 < gt < against the hidden state (S,,z), history exponentially decayed, favoring more recent context.
The gating mechanism shows another benefit of REA: it would be otherwise more difficult to build similar techniques into the softmax attention, where there is clear sense of "recurrence" (Ap- pendix A.5). It proves useful in our language modeling experiments (..
3.3 DISCUSSION
On query and key norms, and learned random feature variance. Eq. 5 assumes both the query and keys are of norm-1. It therefore approximates a softmax attention that normalizes the queries and keys before multiplying them, and then scales the logits by dividing them by o². Empirically, this normalization step scales down the logits /aaswani et al. 2017) and enforces that -1 < k S 1. In consequence, the softmax outputs would be "flattened" if not for σ, which can be set priori as a hyperparameter (Yu et al.) 2016 Avron al.) 2017 f [un 2019 inter alia). Here we instead learn from data with the reparameterization trick (Kingma Welling 2014):
W; ~ (O,
W.
Ia is the d x identity matrix, and denotes elementwise product between vectors. d-dimensional vector or is learned, but random vectors W, are nor
This norm-1 constraint is never mandatory. Rather, we employ it for notation clarity and easier implementation. In preliminary experiments we find it has little impact on the performance when set properly or learned from data. Eq. i2 Appendix presents RFA without imposing it.
Going beyond the Gaussian kernel. More broadly. random feature methods can be applied to a family shift-invariant kernels, with the Gaussian kernel being one of them. In the same family, the order-1 are-cosine kernel (Cho & Saul 2009) can be approximated with feature map: darccos (x) /1/D(ReLU(w, x)....,,Re R..,,,eeUUWw x (Alber et al. 2017) In our exper- iments, the Gaussian and arc-cosine variants achieve similar pertormance. This supplements the exploration of alternatives to softmax in attention (T'sai al., 2019 Gao al., 2019).
Relations to prior work. Katharopoulos et al. (020) inspire the causal attention variant of RFA. They use a feature map based on the exponential linear unit activation (Clevert et al. 2016):
elu(-) + 1. It significantly underperforms both the baseline and RFA in our controlled experiments, showing the importance of properly-chosen feature map. Random feature approximation of atten- tion is also explored by concurrent work (Choromanski et al. (021), with applications in masked language modeling for proteins. They propose positive random features to approximate softmax, aiming for lower variance in critical regions. RFA instead normalizes the queries and keys before random projection to reduce variance. Going beyond both, RFA establishes the benefits of random feature methods as more universal substitute for softmax across all attention variants, facilitating its applications in, e.g., sequence-to-sequence learning.
'In multihcad attention Vaswani et al. 2017 k; and v, are calculated from x, using learned affine trans- formations.
*This departs from Eq. by lifting the isotropic assumption imposed the Gaussian distribution: the difference between the vector ar in Eq. X and the scalar a Eq. 3 We find this improves the performance in practice (34). even though the same result in Theorem may not directly apply. "Apart from replacing the sinusoid functions with ReT.U, it constructs w; in the same way Eq.
4

---

Published as a conference paper at ICLR 2021
There are interesting connections between gated RFA and fast weights (Schmidhuber 1992 1993
Ba et al. 2016 Miconi et al. 2018] inter alia). Emphasizing recent patterns, they learn a temporal memory to store history similarly to E The main difference is that REA additionally normalizes the output using ((q+) z as in Eq. 6 a by-product of approximating softmax's ' partition function. It is intriguing to study the role of this normalization term, which we leave future work.
3.4 COMPLEXITY ANALYSIS
Time. Scaling linearly in the sequence lengths, RFA needs less computation (in terms of number of operations) for long sequences. This implies speedup wherever the quadratic-tima softmax attention cannot be fully aully-parallelized across time steps. More specifically:
Significant speedup can be expected in autoregressive decoding, both conditional (e.g., machine translation) and unconditional (e.g., sampling from language model). For exam- ple, 1.9x speedup is achieved in our machine translation experiments (.4.2) and more for
longer sequences (e.g-, 12x for 2,048-length ones; 5
Some applications (e.g.. language modeling, text classification) reveal inputs to the model in full When there are enough threads to parallelize softmax attention across time steps, hardly any specdup from RFA can be achieved; when there are not, typically for very long sequences (> (,000,,), substantial speed gain is possible. For example, REA does not achieve any speedup when working with 512-length context (44.1) but achieves 5.3x speedup with 4,000-length context (84.2).
Memory. Asymptotically, REA has a better memory efficiency than its softmax counterpart (linear vs. quadratic). To reach a more practical conclusion, we include in our analysis the cost of the feature maps. Ø's memory overhead largely depends on its size D. For example, let's consider the cross attention of a decoder. RFA uses O(LD+2Dd) space to store ((q:). E, (k,)) @ vi, and Σ, (kk,) g In contrast, softmax cross attention stores the encoder outputs with O(Md) memory, with M being the source length. In this case RFA has lower memory overhead when 2D < ... Typically D should be no less than in order for reasonable approximation (Yu et 2016); In transformer model, d is the size of an attention head, which is usually around 64 or 128 (Vaswani et al., 2017 Ott et al. 2018). This suggests that RFA can achieve significant memory saving with longer sequences, which is supported by
(Eq. 5 line 12 of Algo. 2
empirical analysis in (5) Further, using moderate sized feature maps is also desirable, so that its overhead does not overshadow the time and
memory RFA saves. We experiment with D) and 2d; the benefit of using D > 2d is marginal.
Appendix A.6 discusses the time and space complexity in more detail, and Appendix C.2 studies the effect of random feature size on performance.
EXPERIMENTS
We evaluate RFA on language modeling, machine translation, and long text classification.
4.1 LANGUAGE MODELING
Setting. We experiment with Wiki 103 (Merity al. 2017). It is based on English Wikipedia. Table 5 in Appendix B summarizes some of its statistics. We compare the following models:
BaSE is our implementation of the strong transformer-based language model by Baevski
& Auli (2019).
RFA builds on BASE, but replaces the softmax attention with random feature attention. We experiment with both Gaussian and arc-cosine kernel variants.
REA-G GATE additionally learns a sigmoid gate on top of RFA (33.2) also has a Gaussian
kernel variant and a arc-cosine kernel one ?
Helu is a baseline to RFA. Instead of the random feature methods it uses the elu(-) +
feature map, as in Katharopoulos et al. (2020).
A causal masking is usually used to prevent the model from accessing future tokens language models.
& RRF never constructs the M x 2D x d tensor (k.k. & vals, but sequentially processes the sequence. This gating technique specific to RFA variants, in the sense that it is less intuitive to apply it in BASE.
5

---

Published as a conference paper at ICLR 2021
comparable implementations,, tuning, and training procedure. All models use a 512 block size during both training and evaluation, i.c., they read as input a segment of 512 consecutive tokens, without access the context from previous mini-batches. REA variants use 64-dimensional random feature maps. We experiment with two model size settings, small (around 38M parameters) and big (around 242M parameters); they described in Appendix alonng with other implementation details.
To ensure fair comparisons,
Small
Big
Model
Dev.
Test
Dev.
Test
BASE
33.0
34.5
24.5
26.2
Petu (Katharopoulos al. 2020)
38.4
40.1
28.7
30.2
RFA-Gaussian
33.6
35.7
25.8
27.5
RFA-arccos
36.0
37.7
26.4
28.1
RFA-GATE-Gaussian
31.3
32.7
23.2
25.0
RFA- TFA-GATE-arccos
32.8
34.0
24.8
26.3
RFAA-GATTT-Gauusia-SSttllll
29.4
30.5
22.0
23.5
Table I: Language model perplexity (lower is better) on the WikiText-103 development and test sets. Bolded numbers outperform BASE.
Results. compaeess the models' performance in perplexity on Wiki'Text- development and test data. Both kernel variants REA, withou gating, outperform Øelu by more than 2.4 and 2.1 test perplexity for the small and big model respectively, confirming the benefits from using random feature approximation 10 Yet both underperform BASE, with RPA-Gaussian having a smaller gap. Comparing RFA against its gated variants, more than 1.8 perplexity improvement can be attributed the gating mechanism; and the gap is larger for small models. Notably, REA-GATE Gaussiann outperforms BASE under both size settings by at least 1.2 perplexity. In general, RFA models with Gaussian feature maps outperform their arc-cosine counterparts [T From the analysis $3.4 we would not expect speedup by RFA models, do we any in the experiments
Closing this section, we explore a "stateful" variant REA-GATE-GGussian passes the last hidden state (St, 2;) to the next mini-batch during both training and evaluation, a technique commonly used in RNN language models (Merity et al.] 2018). This a consequence of RFA'S RNN-style computation, and is less straightforward to be applicable in the vanilla transformer models! From the last row of Table [ we see that this brings a more than 1.5 test perplexity improvement.
4.2 MACHINE TRANSLATION
Datasets. We experiment with three standard machine translation datasets.
WMT14 EN-DE and EN-FR (Bojar et al. 2014). Our data split and preprocessing follow those of Vaswani et al. (0077) We share the source and target vocabularies within each language pair, with 32,768 byte pair encoding types (BPE; Sennrich et al. 2016).
IWSLT14 4 DE-EN (Cettolo et al. 2014) based on TED talks. The preprocessing follows follows/Edunoy Edunov
et al. 2018). Separate vocabularess of 9K/7K BPE types are used for the source and target.
Table in Appendix Popendix|Bsunnnnn summarizes some statisties of the datasets.
10 models are trained for 150K steps: this could be part of reason behind suboptimal performance of dulu: may need 3 times more gradient updates to reach similar performance to the softmax attention baseline [Katharopoulos et al. 2020).
We observe that RFA Gaussian variants are more stable easier to train than the arc-cosine ones well as Pelu- conjecture that this is because
outputs of the Gaussian feature maps have an fornorm of which can help stabilize training. To see why, sin?(x)) + cos2(z)) = cos(æ xr) 1.
1'nn fact, RFA mains around 15% slower than BASE due the additional overhead from the feature maps.
''Some transformer models use text segment from the previous mini-batch as a prefix (aevski & Auli
2019
Dai et al. 2019). Unlike RFA, this gives the model access to only limited amount of context, and significantly increases the memory overhead.

---

Published as a conference paper at ICLR 2021
Setting. We compare the RFA variants described in $4.1 They build on BASE model that is our implementation of the base-sized transformer (Vaswami et al., 2017). All RFA models apply random feature attention in decoder cross and causal attention, but softmax attention in encoders. This setting yields the greatest decoding time and memory savings ($3.4). We use 128/64 for D in cross/causal attention. RFA-GATE learns sigmoid gates in the decoder causal attention. The Pelu baseline uses the same setting and applies feature map in both decoder cross and causal attention, but the encoders. Further details are described in Appendix B.2
B.2
WMT14
IWSLT14
Model
EN-DE
EN-FR
DE-EN
Speed
BASE
28.1
39.0
34.6
1.0x
Pelu (Katharopoulos et al. 2020)
21.3
34.0
29.9
2.0x
REA UFA-Gaussian
28.0
39.2
34.5
1.8x
RFA-arccos
28.1
38.9
34.4
1.9x
RFA-G GATE-Gaussian
28.1
39.0
34.6
1.8x
RFA-Ga GATE-arecos
28.2
39.2
34.4
1.9x
Table 2:
Machine translation test set BLEU. The decoding speed (last column) is relative to BASE. All models are tested on single TPU v2 accelerator, with batch size 32.
Results. T compares the models' test set BLEU on three machine translation datasets. Overall both Gaussian and arc-cosine variants of REA achieve similar performance to BASE on all three datasets, significantly outperforming Katharopoulos et al. (2020). Differently from the trends in the language modeling experiments, here the gating mechanism does not lead to substantial gains. Notably, all RFA variants decode more than 1.8x faster than BASE.
4.3 LONG TEXT CLASSIFICATION
We further evaluate RFA'S 's accuracy and efficiency when used as text encoders on three NLP tasks from the recently proposed Long Range Arena benchmark (Tay al. 2021). designed evaluate efficient Transformer variants on tasks that require processing long sequences
Experimental setting and datasets. We compare RFA against baselines on the following datasets:
ListOps (LO; Nangia & Bowman 2018) aims to diagnose the capability of modelling hi- erarchically structured data. Given sequence of operations on single-digit integers, the model predicts the solution, also single-digt integer. It is formulated as a 10-way classi- fication. We follow Tay al. (202I) and consider sequences with 500-2,000 symbols.
Character-level text classification with the IMDb movie review dataset MMas al. 201
This binary sentiment classification task.
Character- document retrieval with the ACL Anthology Network (AAN: Radev et al.,
2009) dataset. The model classifies whether there a citation between pair of papers.
To ensure fair comparisons, we implement RFA top of the transformer baseline by Tay et al. (202T), and closely follow their preprocessing. data split, model size, and training procedure. Speed and memory evaluated on the IMDb dataset. For our RFA model, use D = 64 for the IMDb dataset, and D 128 for others. We refer the readers to Tay et al. (2021) for further details.
that RFA outperforms the transformer bascline on two out of the three datasets, achieving the best performance on IMDb with 66% accuracy. Averaging across three datasets, RFA outperforms the transformer by 0.3% accuracy, second only to Zaheer al. (2020) with a 0. % accuracy gap. In terms of time and memory efficiency, RFA is among the strongest. RFA speeds up over the transformer by 1.1-5.3x, varying by sequence length. Importantly, compared to the only two baselines that perform comparably to the baseline transformer model (Tay et al. 2020aa)
Results.
From Table 3 we
Zaheer et al. 2020), RFA has a clear advantage in both speed and memory efficiency, and is the only model that is competitive in both accuracy and efficiency.
htttp::/gitthbb.co ott.:///ooooooooaaaaaaaooo//aaaa
reee-reeeareee massearch/1om

---

Published as a conference paper at ICLR 2021
Accuracy
Speed
Memory
Model
LO
IMDb
AAN
Avg.
IK
2K
3K
4K
1K
2K
3K
4K
Transformer
36.4
64.3
57.5
52.7
1.0
1.0
1.0
1.0
1.00
1.00
1.00
1.00
Wang et al.
2020
35.7
53.9
52.3
47.3
1.2
1.9
3.7
5.5
0.44
0.21
0.18
0.10
Kitacy al.
2020
37.3
56.1
53,4
48.9
0.5
0.4
0.7
0.8
0.56
0.37
0.28
0.24
Tay al.
(0220b
17.1
63.6
59.6
46.8
1.6
3.8
0.55
0.31
0.20
0.16
lay ct al.
(2020a)
37.0
61.7
54.7
51.1
1.1
1.2
2.9
1.4
0.76
0.75
0.74
0.74
Zaheer ct al.
(0220)
36.0
64.0
59.3
53.1
0.9
0.8
1.2
1.1
0.90
0.56
0.40
0.30
Katharopoulos
2020
16.1
65.9
53.1
45.0
1.1
1.9
3.7
5.6
0.44
0.22
0.14
0.11
Choromanski al.
2021
18.0
65.4
53.8
45.7
1.2
1.9
3.8
5.7
0.44
0.22
0.15
0.11
RFA-Gaussian (This work)
36.8
66.0
56.1
53.0
1.1
1.7
3.4
5.3
0.53
0.30
0.21
0.16
Table 3: Accuracy (higher is better) of different models on LO, IMDb, and AAN, along with their speed (higher is better) and peak memory consumption (lower is better) varying sequence lengths (1-4K). Speed and memory are evaluated on the IMDb dataset and relative to the transformer' s. Bold font indicates the best performance in each column, and underlined numbers outperform the transformer accuracy. Transformer's and previous works' numbers are due to Tay et al. (021
Softmax
5000
Decuding Speed: Takers/s
RFA-amoce
7003
RFA- GaAA-Gasssann
6000
3000
Softmax
2000
RFA-arecas
1000
RFA- Gaussia
27
24
Sequence Length
Sequence Length
(a) Speed lengths.
(b) Memory lengths.
Figure 2: Conditional decoding speed (left) and memory overhead (right) varying the output lengths. All models are tested on single TPU v2 accelerator, with greedy decoding and batch size 16.
ANALYSIS
Decoding time and memory varying by sequence length. $3.4 shows that RFA can potentially achieve more significant speedup and memory saving for longer sequences, which we now explore.
We use a simulation conditional generation experiment on to compare RFA'S sequence-to-sequene decoding speed and memory overhead against the baseline's. Here assume the input and output sequences are of the same length. The compared models are of the same size as those described in 84.2 with 6-layer encoders and decoders. Other hyperparameters are summarized Appendix B.2 All models are tested using greedy decoding with the same batch size of 16, TPU v2 accelerator.
From Figures (a) and (b) we observe clear trends. Varying the lengths, both RFA variants achieve consistent decoding speed with nearly-constant memory overhead. In contrast, the baseline decodes slower for longer sequences, taking an increasing amount of memory. Notably, for 2.048-length sequences, RFA decodes around 12x faster than the baseline while using less than 10% of the memory. RFA-arccos slightly outperforms RFA- Gaussian terms of speed and memory efficiency. This is because when using the same D (as do here), the &arcos is half the size of @Gassian- These results suggest that RFA can particularly useful in cequence-to-sequenee tasks with longer sequences, e.g., document-level machine translation (Miculicich et al.) 2018).
Figure 3 Appendix C.I compares the speed and memory consumption in unconditional decoding (e.g., sampling from language model). The overall trends are similar to those in Figure 2
Notes on decoding speed. With a lower memory overhead, REA can use a larger batch size than the baseline. As noted by KKatharopoulos et al. 2020) and Kasai al. 202T). if had used mini-
8

---

Published as a conference paper at ICLR 2021
batches large the hardware allows, RFA could have achieved more significant speed gain. Nonetheless, we control for batch size even though it is not the most favorable setting for RFA, since the conclusion translates better to common applications where a time (e.g.. instantaneous machine translation). For the softmax attention baseline, we follow Ont al. (2018) and cache previously computed query/key/value representations, which significantly improves its decoding speed (over not caching).
generates single sequence
Further analysis results. RFA achieves comparable performance softmax attention. Ap- pendix C.3 empirically shows that this cannot be attributed to RFA learning a good approxima- tion softmax: when we train with one attention but evaluate with the other, the performance is hardly better than randomly- -iitiliiiiiallzzedd untrained models. Yet, an RFA model initialized from pretrained softmax transformer achieves decent training loss after a moderate amount of finetuning steps (Appendix C.4). This suggests some potential applications, e.g., transferring knowledge from a pretrained transformer (e.g., GPT-3: Brown al., 2020) to an RFA model that is more efficient to sample from.
RELATED WORK
One common motivation across the following studies, that shared by this work and the research we have already discussed, is to scale transformers to long sequences. Note that there are plenty orthogonal choices for improving efficiency such as weight sharing (Dehghani al. 2019), quanti- zation (Shen et al. 2020). knowledge distillation (Sanh et al.] (220), and adapters (Houlsby et al.) 2019).
For a detailed overview we refer the reader to Tay et al. (2020e).
Sparse attention patterns. The idea behind these methods is to limit the reception field of attention computation. motivates earlier attempts improving attention's efficiency, and still receives lots interest. The sparse patterns can be set a priori (Liu et al., 2018 8 Qiu et al. [0200] Ho et al. 2020 You et al. 2020, inter alia) or learned from data (Sukhbaatar et al. 2019 Roy et al., 2020 inter alia). For most of these approaches, it is yet to be empirically verified that they are suitable for large-scale Sigueence-t--eqqeene learning; few of them have recorded decoding speed benefits.
Compressed context. Wang et al. 22200) compress the context along the timesteps that the effee- tive sequence length for attention computation is reduced. Another line of work aims to store past context into memory module with limited size (Lee et al. 2019:] Ainslie al.) 2020 Rae al.) 2020 inter alia). so that accessing longer history only moderately increases the overhead. Remi- niscent of RNN language models, RFA attends beyond fixed context window through stateful computation, without increasing time or memory overhead.
7 CONCLUSION
We presented random feature attention (RFA). It views the softmax attention through the lens of kernel methods, and approximates it with random feature methods. With an optional gating mech- anism, RFA provides straightforward way of learning with recency bias. REA'S time and space complexity is linear the sequence length. We use REA as a drop-in substitute for softmax atten- tion in transformer models. On language modeling, machine translation, and long text classification benchmarks, REA achieves comparable or better performance than strong baselines. In the machine translation experiment, RFA decodes twice as fast. Further time and memory efficiency improve- ments can be achieved for longer sequences.
ACKNOWLEDGMENTS
We would like to thank Phil Blunsom, Chris Dyer, Nando de Freitas, Jungo Kasai, Adhiguna Kun- coro, Dianqi Li, Ofir Press, Lianhui Qin, Swabha Swayamdipta, Sam Thomson, the language team at DeepMind and the ARK group the University of Washington for their helpful feedback. We also thank Tay Yi for helping run the Long Range Arena experiments, Richard Tanburn for the advice on implementations, and the anonymous reviewers for their thoughtful comments. This work was supported in part by NSF grant 1562364 and a Google Fellowship. Nikolaos Pappas was supported by the Swiss National Science Foundation under grant number P400P2 18391 1 "UNISON."

---

Published as a conference paper at ICLR 2021
REFERENCES
Joshua Ainslie, Santiago Ontanon, Chris Alberti, Vaclav Cvicek, Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai. Qifan Wang, and Li Yang. ETC: Encoding long and structured inputs in transformers. In Proc. of EMNLP, 2020.
Maximilian Alber, Pieter-Jan Kindermans, Kristof Schüit,, Klaus-Robert Müller, and Fei Sha. An empirical study on the properties of random bases for kernel methods. Proc. of NeurIPS, 2017.
Haim Avron, Vikas Sindhwani, Jiyan Yang, and Michael W. Mahoney. Quasi-Monte Carlo feature maps for shift-invariant kernels. Journal of Machine Learning Research, 17(120):1-38 2016.
Haim Avron, L. Kenneth Clarkson, and P. David and Woodruff. Faster kernel ridge regression using sketching and preconditioning. SIAM J. Matrix Analysis Applications, 2017.
Jimmy Ba, Geoffrey E Hinton, Volodymyr Mnih, Joel Z Leibo, and Catalin lonescu. Using fast weights to attend to the recent past. In Proc. of NeurIPS, 2016.
Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. In Proc. of ICLR, 2019.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In Proc. of ICLR, 2015.
Iz Beltagy, Matthew E. Peters, and Arman Cohan. arXiv: 2004. 05500, 2020.
Longformer: The long-document transformer.
S. Bochner. Harmonic Analysis and the Theory of Probubility. University of California Press, 1955.
Ondfej Bojar, Christian Buck, Christian Federmann, Barry Haddow, Philipp Koehn, Johannes Lev- cling, Christof Monz, Pavel Pecina, Matt Post, Herve Saint-Amand Radu Soricut, Lucia Specia, and Aleš Tamchyna. Findings of the 2014 workshop on statistical machine translation. In Proc. of WMT, 2014.
Tom B. Brown, Benjamin Mann, Nick Ryder. Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. arXiv: 2005. 14165, 2020.
Mauro Cettolo, Jan Niehues, Sebastian Stüker, Luisa Bentivogli, and Marcello Federico. Report on the 11th IWSLT evaluation campaign. In Proc. of /WSLT, 2014.
Kchai Chen, Rui Wang, Masao Utiyama, and Eicchro Sumita. Recurrent positional embedding for neural machine translation. In Proc. of EMNLP, 2019.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. Xrviv: 10509, 2019.
Kyunghyun Cho, Bart van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Hol- ger Schwenk, and Yoshua Bengio. Learning phrase representations using RNN encoder-decoder for statistical machine translation. In Proc. of EMNLP, 2014.
Youngmin Cho and Lawrence K. Saul. Kernel methods for deep learning. Proc. of NeurIPS, 2009.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreca Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J Colwell, and Adrian Weller. Rethinking attention with per- formers. In Proc. of ICLR, 2021.
Djork-Amé Clevert Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (ELUss. In Proc. of ICLR, 2016.
10

---

Published as a conference paper at ICLR 2021
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer-Xe XL: Attentive language models beyond a fixed-length context. In Proc. of ACL, 2019.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. In Proc. ICLR, 2019.
Jacob Devlin, Ming-Wer Weig Chang, Kenton Lee, and Kristina Toutanova. BERT; Pre-training of deep bidirectional transformers for language understanding. In Proc. of NAACL. 2019.
Sergey Edunov, Myle Ott, Michael Auli, David Grangier, and Marc' Aurelio Ranzato. Classical structured prediction losses for sequence to sequence learning. In Proc. of NAACL, 2018.
Yingbo Gao, Christian Herold, Weiyue Wang, and Hermann Ney. Exploring kernel functions in the softmax layer for contextual word classification. In International Workshop on Spoken Language Translation, 2019.
Jie Hao, Xing Wang, Baosong Yang, Longyue Wang, Jinfeng Zhang, and Zhaopeng Tu. Modeling recurrence for transformer. In Proc. of NAACL, 2019.
Geoffrey Hinton, Oriol Vinyals, and Jeffrey Dean. Distilling the knowledge in a neural network. In NeurlPs Deep Learning and Representation Learning Workshop, 2015.
Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans. Axial attention in multidi-
mensional transformers. arXiv: 1912.12180, 2020.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 9(8): 1735-1780, 1997.
Thomas Hofmann, Bemhard Schölkopf, and Alexander J. Smola. Kernel methods in machine learn- ing. Annals of Statistics, 36(3):1171-1220) 2008.
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, An- drea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficient transfer learning for NLP. In Proc. of ICML, 2019.
Jungo Kasai, Nikolaos Pappas, Hao Peng, James Cross, and Noah A. Smith. Deep encoder, shallow decoder: Reevaluating the speed-quality tradeoff in machine translation. In Proc. of ICLR, 2021.
Angelos Katharopoulos, Apoory Vyas, Nikolaos Pappas, and Francois Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In Proc. of ICML, 2020.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Proc. of ICLR, 2015.
Diederik Kingma and Max Welling. Auto-cncoding variational bayes. In Proc. ICLR, 2014.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In Proc. ICLR, 2020.
Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: framework for attention-bag based permutation invariant neural networks. In Proc. of ICMI., 2019.
inrrarationiivvariantt
Shiyang Li, Xiaoyong Jin, Yao Xuan, Xiyou Zhou, Wenhu Chen, Yu-Xiang Wang, and Xifeng Yan. Enhancing the locality and breaking the memory bottleneck of transformer on time series forecasting. In Proc. of NeurIPS, 2019.
J. Liu, Mohammad Saleh, Etienne Pot, Ben Goodrich, Ryan Sepassi, Lukasz Kaiser, and Noam Shazeer. Generating wikipedia by summarizing long sequences. In Proc. of ICLR, 2018.
Andrew L. Maas, Raymond E. Daly, Peter Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proc. of ACI,, 2011.

---

Published as a conference paper at ICLR 2021
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. models. In Proc. of ICLR, 2017.
Pointer sentinel mixture
Stephen Merity, Nitish Shirish Keskar, and Richard Socher. Regularizing and Optimizing LSTM Language Models. In Proc. of ICLR, 2018.
Thomas Miconi, Kenneth Stanley, and Jeff Clune. Differentiable plasticity: training plastic neural networks with backpropagation. Proc. of ICML, 2018.
Lesly Miculicich, Dhananjay Ram, Nikolaos Pappas, and James Henderson. Document-level neural machine translation with hierarchical attention networks. Proc. of EMNLP, 2018.
Abdelrahman Mohamed, Dmytro Okhonko, and Luke Zettlemoyer. Transformers with convolutional context for ASR. arXiv: 1904.11660, 2019.
Nikita Nangia and Samuel Bowman. ListOps: A diagnostic dataset for latent tree learning. In Proc. of NAACL Student Research Workshop, 2018.
Junier Oliva, William Neiswanger, Barnabas Poczos, Eric Xing, Hy Trac, Shirley Ho, and Jeff Schneider. Fast function function regression. In Proc. of AISTATS, 2015.
Myle Ott, Sergey Edunov, David Grangier, and Michael Auli. Scaling neural machine translation. In Proc. of WMT, 2018.
Emilio Parisotto, H. Francis Song, Jack W. Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant M. Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, Matthew M. Botvinick, Nicolas Heess, and Raia Hadsell. Stabilizing transformers for reinforcement learn- ing. In Proc. of ICML, 2020.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In Proc. of ICML, 2018.
Hao Peng, Roy Schwartz, Sam Thomson, and Noah A. Smith. Rational recurrences. Proc. of EMNLP, 2018
Hao Peng, Roy Schwartz, Dianqi Li, and Noah A. Smith. A mixture of heads is better than heads. In Proc. of ACL, 2020.
Matt Post. A call for clarity in reporting BLEU scores. Proc. of WMT, 2018.
Jiezhong Qiu, Hao Ma, Omer Levy, Wen-tau Yih, Sinong Wang, and Jie Tang. Blockwise self- attention for long document understanding. In Findings of EMNLP, 2020.
Dragomir R. Radev, Pradeep Muthukrishnan, and Vahed Qazvinian. The ACL Anthology network. In Proc. of the Workshop on Text and Citation Analysis for Scholarly Digital Libraries, 2009.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners, 2018.
Jack W. Rae. Anna Potapenko, Siddhant M. Jayakumar, Chloe Hillier, and Timothy P. Lillicrap. Compressive transformers for long-range sequence modelling. Proc. of ICLR, 2020.
Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. In Proc. of NeurlPS, 2007.
Ankit Singh Rawat, Jiecao Chen, Felix Xinnan X Yu, Ananda Theertha Suresh, and Sanjiv Kumar. Sampled softmax with random Fourier features. In Proc. of NeurIPS, 2019.
Aurko Roy, Mohammad Taghi Saffar, David Grangier, and Ashish Vaswani. Efficient content-based sparse attention with routing transformers. arXiv: 2003.05997, 2020.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter. arXiv: 1910.01 108, 2020.
12

---

Published as a conference paper at ICLR 2021
J. Schmidhuber. Learning to control fast-weight memories: An alternative to dynamic recurrent networks. Neural Computation, 4(1):131-139 1992.
J. Schmidhuber. Reducing the ratio between learning complexity and number of time varying vari- ables in fully recurrent nets. In Proc. of ICANN, 1993.
Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. In Proc. of ACL, 2016.
Sheng Shen, Zhen Dong, Jiayu Ye, Linjian Ma. Zhewei Yao, Amir Gholami, Michael W. Mahoney, and Kurt Keutzer. Q-BERT: Hessian based ultra low precision quantization of BERT. In Proc. of AAAI, 2020.
Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin. Adaptive attention span in transformers. In Proc. of ACL. 2019.
Yitong Sun. Random Features Methods in Supervised Learning. PhD thesis, The University of Michigan, 2019.
Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng. Synthesizer: Rethinking self-attention in transformer models. arXiv: 2005.00743 2020a.
Yi Tay, Dara Bahri, Liu Yang, Don Metzler, and Da-Cheng Juan. Sparse sinkhorn attention. In Proc. of ICML, 2020b.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey, arXiv: 2009.06732. 2020c.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. In Proc. of ICLR, 2021.
Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philippe Morency, and Ruslan Salakhutdinov. Transformer dissection: An unified understanding for transformer's attention via the lens of kernel. In Proc. of EMNLP, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz Kaiser, and Illia Polosukhin. Attention all you need. In Proc. of NeurIPS, 2017.
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv: 2006.04768, 2020.
Ronald J. Williams and David Zipser. A learning algorithm for continually running fully recurrent neural networks. Newral Computation, 1:270-280, 1989.
Felix Wu, Angela Fan, Alexei Baevski, Yann Dauphin, and Michael Auli. Pay less attention with lightweight and dynamic convolutions. In Proc. of ICLR, 2019.
Zhanghao Wu, Zhijian Liu, Ji Lin, Yujun Lin, and Song Han. Lite transformer with long-short range attention. In Proc. of ICLR, 2020.
Weiqiu You, Simeng Sun, and Mohit Iyyer. Hard-coded Gaussian attention for neural machine translation. In Proc. of ACL, 2020.
Felix Xinnan X Yu, Ananda Theertha Suresh, Krzysztof M Choromanski, Daniel N Holtmann- -ice,, and Sanjiv Kumar. Orthogonal random features. In Proc. of NeurlPS, 2016.
Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed, Big bird: Transformers for longer sequences. arXiv: 2407.14062, 2020.
13

---

Published as a conference paper at ICLR 2021
Appendices
A RANDOM FEATURE ATTENTION IN MORE DETAIL.
A.1 DETAILED COMPUTATION PROCEDURE
Algorithms and describe causal and cross random feature attention's computation procedures.
Algorithm Causal random feature attention.
I: procedure REA-CAUSAL( {q: {q)) (k:) (vi).e))
> a D matrix
P i D-dimensional vector
4:
S. z + 0,
for to N do
qi k; - (q(),, (k;,)
Random feature maps
7:
S 5 S k
k &v
8:
2 + k;
h + q. S/(
10:
end for
11: return (h, }
12:
end procedure
Algorithm 2 Cross random feature attention.
1: procedure RFA-CROSS( {q:) ( (
2: > S is D d matrix
3:
b z is a D-dimensional vector
4:
S,
0, 0
5:
for i to M do
k; ((k;)
Random feature map
7:
S - S + k, 2 v.
8:
z + + k,
9:
end for
10:
for i i= to N do
11:
q. - ②(q.)
> Random feature map
12:
h. +- q, S/
13:
end for
14: return (h;),1
15:
end procedure
A.2 VARIANCE OF RANDOM FOURIER FEATURES
The following result is due to Yu et al. (0066). Using the notation in 2.2
2D (1-e-s)
Var( (x) + +(y)) (y))
(9)
where =||--V sc-yll/A yll/a.
14

---

Published as a conference paper at ICLR 2021
A.3 DERIVATION OF CAUSAL REA
This section presents a detailed derivation of causal RFA as in 33.1 Following Eq. the attended keys and values to the prefix:
but changing
Φ (qt) Eist Φ (k,) Ov;
REA(qr; (kitish (vilise)
(10)
Φ (q) Ejst Φ (kj)
Let S, - Eisi o(k,) 8 Vo and Zz - Eise (kk)); both can be calculated recurrently. Assuming
So = 0 and Zo 0:
S = S + (kt) @ Vt.
2, = 26-1 + (ki), t2
(11)
This completes the derivation of causal REA as in
A.4 RFA WITHOUT NORM-1 CONSTRAINTS
assumes that the queries and keys are unit vectors. This norm-1 constraint not must. Here we present a REA without imposing this constraint. Let C(x) = exp([| [|| /20²). From Eq. we have attn (qt, (k,), {v.})=
exp (q:- k,/0²2
C(k)) (k:) E, C(g:) Φ Φ(k) ① (q)) Σ, Ckk,) 2 8v.
Σ
Σ, exp (q: k//02) V.
(12)
(q:) E, C(k,)
The specific attention computation is similar those in 3.1 In sum, lifting the norm-l constraint brings an additional scalar term C(-).
A.5 RELATING REA-GATE TO SOFTMAX ATTENTION
Drawing inspiration from gated RNNs, "softmax counterpart."
3.2 introduces a gated variant of RFA. Now we study its
k; k, 9i)
II
9j. V v; (1
9:)
II j=i+1
9j.
1, 1.....
(13)
i=++1
h attn (k, Jigu (ilse.)
h, the output at timestep and used for onward computation.
At each step, all prefix keys and values are decayed by gate value before calculating the atten- tion. This implies that the attention computation for q++1 cannot start until that of qe finished. Combined with the linear complexity of softmax normalization, this amounts to quadratic time in sequence length, even for language modeling training.
The above model is less intuitive and more expensive in practice, without the RFA perspective. This shows that REA brings some benefits in developing new attention models.
A.6 DETAILED COMPLEXITY ANALYSIS
considess a sequence- to-eeuueee- model, and breaks down the comparisons to training (with teacher forcing; Williams & Zipser 1989) and autoregressive decoding. Here we assume enough threads to fully parallelize softmax attention across timesteps when the inputs are revealed to the model in full. RFA has lower space complexity, since it never explicitly populates the attention matrices. As for time, RFA trains linear time, and so does the softmax attention: in teacher-foreee forcing training standard transformer decoder parallelizes the attention computation across time steps. The trend of the time comparison differs during decoding: when only one output token is produced at time, RAA decodes linearly in the output length, while softmax attention decodes quadratically.
15

---

Published as a conference paper at ICLR 2021
Time Complexity
Space Complexity
Setting
Model
Encoder
Cross
Causal
Encoder
Cross
Causal
Training w/
softmax
((M)
((M)
O(N)
O(M²) O(MN) O(N2)
teacher forcing
RFA
(M))
O(M) O(N)
O(M) O(M+N) O(N)
softmax
((M)
O(MN) O(N2)
O(M2) O(MN) O(N2)
Decoding
RFA
((M) O(M+M O(N)
O(M) O(M+N O(N)
Table 4: Time and space complexity comparisons between RFA and its softmax counterpart in attentive model, assuming an infinite amount available threads. M and N denote the lengths of the source and target sequences respectively. Teacher forcing train- ing (Williams & Zipser| 1989) and autoregressive decoding are assumed. Blue color indicates the cases where RFA asymptotically outperforms softmax attention.
a sequence-to-sequees
Data
Train
Dev.
Test
Vocab.
WikiText-100
103M
218K
246K
268K
WMT14 EN-DE
4.5M
3K
3K
32K
WMT14 EN-FR
4.5M
3K
3K
32K
IWSLT14 DE-EN
160K
7K
7K
9K/7K
Table 5: Some statistics for the datasets. WikiText-103 split sizes are in number of tokens, while others are in number of instances.
B EXPERIMENTAL DETAILS
Table summarizes some statistics of the datasets used in our experiments. Our implementation is based on JAX
# Random Matrices
50
100
200
BLEU
24.0
25.7
25.8
25.8
Table 6: WMT14 EN-DE development set performance varying the number of random matrices to sample from during training. No beam search or checkpoint averaging is used.
During training, sample a different random projection matrix for each attention head. Prelim- inary experiments suggest this performs better than using the same random projection throughout training (Table (6). Our conjecture that this helps keep the attention heads from "over commit- ting" to any particular random projection (Peng et al. 2020). To avoid the overhead of sampling from Gaussian during training, we this in offline manner. Iee. before training we construct pool of random matrices (typically 200), at each training step we draw from the pool. At test time each attention head uses the same random projection, since no accuracy benefit is observed by using different ones for different test instances.
B.1 LANGUAGE MODELING
We use the fixed sinusoidal position embeddings by Vaswani et al. (2017). All models trained for up to 150K gradient steps using the Adam optimizer (Kingma & Ba 2015). No fg-regularization is used. We apply early stopping based development perplexity. All models trained using 16 TPU v3 accelerators, and tested using a single TPU v2 accelerator.
We compare the models using two model size settings, summarized in Table 7
htttp:://iithub.
16

---

Published as a conference paper at ICLR 2021
Hyperprams.
Small
Big
Layers
16
Heads
Embedding Size
512
1024
Head Size
64
FFN Size
2048
4096
Batch Size
64
64
Learning Rate
[1
10-1,
2.5
10-1, x 10-4)
Warmup Steps
6000
6000
Gradient Clipping Norm
0.25
0.25
Dropout
[0.05, 0.1]
[0.2, 0.25, 0.3]
Random Feature Map Size
64
64
Table 7: Hyperparameters used in the language modeling experiments.
B.2 MACHINE TRANSLATION
WMT14. We use the fixed sinusoidal position embeddings by Vaswani et al. (017). For both EN-DE and EN-FR experiments,
train the models using the Adam (with 31 0.1, 82 0.98, and c = 10-9) optimizer for up to 350K gradient steps. We use a batch size of 1,024 instances for EN-DE, while 4,096 for the much larger EN-FR dataset. The learning rate follows that by Vaswani et al. (017). Early stopping is applied based on development set BLEU. No &2 regularization or gradient clipping used. All models trained using 16 TPU v3 accelerators, and tested using single TPU v2 accelerator. Following standard practice, we average 10 most recent checkpoints at test time. We evaluate the models using SacreBLEU (Post 2018) A beam search with beam size 4 and length penalty 0.6 is used. Other hyperparameters are summarized in Table
Hyperprams.
WMT14
IWSLT14
Layers
6
# Heads
8
Embedding Size
512
512
Head Size
64
64
FFN Size
2048
2048
Warmup Steps
6000
4000
Dropout
0.1
0.3
Cross Attention Feature Map
128
128
Causal Attention Feature Map
64
64
Table Hyperparameters used in the machine translation experiments.
C MORE ANALYSIS RESULTS
C.I MORE RESULTS ON DECODING SPEED AND MEMORY OVERHEAD
Figure 3 compares the RFA'S unconditional decoding speed and memory against the softmax atten- tion. The setting is the same that in e5xceept that here the models do not have encoder. This experiment aims to simulate the applications such as sampling from a language model.
C.2 EFFECT OF RANDOM FEATURE SIZE
This section studies how the size of Φ(.) affects the performance. Table 9 summarize RFA- Gaussian's performance WMT14 EN-DE development set. The model and training the same as that used in 4.2 except random feature size. Recall from 22.2 that the size of Φ0-) is 2D for
17

---

Published as a conference paper at ICLR 2021
1.8
10003
Softmax
Decoding Spoed: Tolenz's
RFA-amoue
RFA-Gausi
8000
Memory:
6000
0.3
Softmax
0.5
RFAwarecns
0.4
RFA- Gauasion
0.2
2
2
24
2
24
Sequence Length
Sequence Longth
(a) Speed lengths.
(b) Memory lengths.
Figure 3: Unconditional decoding speed (left) and memory overhead (right) varying the output lengths, All models tested single TPU v2 accelerator, with greedy decoding and batch size 16.
RFA-Gaussian. When the size of ((-) is too small (32 or 64 for cross attention, 32 for causal at- tention), training does not converge. We observe accuracy improvements by using random features sufficiently large (256 for cross attention and 128 for causal attention); going beyond that, the benefit marginal.
Φ Size
32
64
128
256
512
Φ Size
32
64
128
256
512
BLEU
N/A
N/A
24.9
25.8
26.0
BLEU
N/A
25.3
25.8
25.8
25.6
(a) Varying cross attention & sizes while fixing that of causal attention to be 128.
(b) Varying causal attention & sizes while fixing that of cross attention to be 256.
Table 9: WMT14 EN-DE development set performance of RFA-Gaussian (the size of Φb is 2D; 2.2 varying the random feature sizes. N/A indicates training does not converge. No beam search or checkpoint averaging is used.
C.3 TRAIN AND EVALUATE WITH DIFFERENT ATTENTION FUNCTIONS
RFA achieves comparable performance its softmax counterpart. Does this imply that it learns good approximation to the softmax attention? To answer this question, we consider:
ii RFA-Gaussian model initialized from pretrained roftmax-transformee a
initialized from a pretrained an RFA-GGussie model.
If REA'S good performance be attributed to learning good approximation to softmax, both, without finetunining, should perform similarly to the pretrained models. However, this is nor the case IWSLT14 DE-EN. Both pretrained models achieve more than 35.2 development set BLEU. In contrast, (i) and (ii) respectively get 2.3 and 1.1 BLEU withous finetuning, hardly beating a randonly-initialized untrained model. This result aligns with the observation by Choromanski al.) 202II. and suggests that it is not the case that RFA performs well because it learns to imitate softmax attention's 's outputs.
C.4 KNOWLEDGE TRANSFER FROM SOFTMAX ATTENTION TO RFA
We first supplement the observation Appendix C.3 by finetuning (i) on the same pretraining data. Figure plots the learning curves. It takes RFA roughly 1,500 steps to reach similar training loss to the pretrained model. As baseline, "RFA Reset" resets the multihead attention parameters (i.e., those for query, key, value, and output projections) randomly initialized ones. Its learning curve similar to that of (i), suggesting that the pretrained multihead attention parameters are no more useful RE than randomly initialized ones. To further confirm this observation, "softmax Reset"
18

---

Published as a conference paper at ICLR 2021
10
RFA Rinat
Sofimax Reset
Pretrained
Toaniig
500
10:00
15:0
Number of Firstuning Steps.
Figure 4: Finetuning an RFA-Gaussian model with its parameters initialized from pretrained softmax-tran tranmax-trransfomer.... "Reset" indicates resetting the multihead attention parameters to randomly- initialized ones. The dashed line indicates the training loss of the pretrained model.
resets the multihead attention parameters without changing the attention functions. It converges to the pretraining loss in less than 200 steps.
Takeaway. From the above results IWSLT14, pretrained knowledge in softmax transformer cannot be directly transferred to an RFA model. However, from Figure 4 and a much larger-scale experiment by Choromanski al. (202T). do observe that REA recover the pretraining loss, and the computation cost of finetuning is much less than training a model from scratch. This suggests some potential applications. For example, one might be able initialize RFA language model from softmax transformer pretrained on large-scale data (e.g., GPT-3; Brown et al., (020), and finetune it at a low cost. The outcome would be an RFA model retaining most of the pretraining knowledge, but is much faster and more memory-frie memory-friendly to sample from. We leave such exploration future work.
19


Reviews:
[{"review_id": "UsRcQtQCt_F", "reviewer": "AnonReviewer1", "summary": "", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary\n* This paper proposes a linear time and space attention variant that matches (or exceeds) the accuracy of standard attention while maintaining the speedup of prior work in linear time/space attention.\n* The approach is centered around a linear approximation of softmax attention, and is extended with a gating mechanism similar to a GRU.\n* The accuracy improvements from the gated extension are demonstrated via language modeling on WikiText-103, while the speed-ups are demonstrated on translation.\n\nContributions\n* Demonstrates the importance of the choice of kernel in RFA, as prior results from Katharopoulos et. al [1] underperforms softmax attention, while this work gets comparable performance by changing the kernel.\n* Proposes an extension of RFA with gating that improves accuracy on language modeling, relative to softmax attention.\n\nStrengths\n* The RFA-gated formulation is both more accurate and potentially faster (at least for decoding) than softmax attention, as demonstrated on language modeling.\n* The writing was clear and thorough.\n* The experiments support the claims of improved accuracy via choice of kernel and gating, and preservation of speedups inherited from the linear attention formulation.\n\nWeaknesses\n*  The improved accuracy and speed over the baseline transformer are great, and the experiments serve to nicely illustrate those independently. I would like to see a third application where these two qualities are demonstrated together.\n* The gating mechanism could also be applied to the output of softmax attention, but that comparison is not included. Please correct me if this is incorrect.\n\nRecommendation: Weak Accept\n* Well-written and timely exploration of linear attention.\n* Empirical results demonstrate accuracy improvement over softmax attention, while preserving the linear time complexity.\n* Extension of RFA with a gating mechanism appears to be effective, but I do not believe the claim that it is hard to apply this to softmax attention is valid, leaving the main contribution an exploration of a couple different kernels for linear attention.\n\nQuestions\n* Can you clarify how the RFA defined in section 3.1 differs from the linear attention in Katharopoulos et. al. [1]? What is different other than the choice of phi?\n* Associative reductions such as those described in sections 3.1 and 3.2 can be computed in time logarithmic in sequence length (at the cost of n log n memory consumption) on a parallel device using a binary reduction or the prefix sum trick. Would this result in speed gains, or is the performance of attention in the parallelizable softmax setting already saturated?\n* Are there experiments with a bidirectional gated RFA in the conditional setting, i.e. in the encoder for translation?\n\nNits\n* The formatting of equation 5 is a bit strange, as the start of the equation is inline while the rest is not.\n\n[1] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and Francois Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In Proc. of ICML, 2020.\n\nEdit: I have updated my rating based on author response.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "F1vLmmoggBw", "reviewer": "AnonReviewer3", "summary": "", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper presents a linear time and space attention mechanism based on random features to approximate the softmax. The paper is clearly written and easy to follow. The results are convincing: not chasing SOTA, but comparing to sensible baselines, namely [Baevski & Auli 2019] for language modeling on Wikitext-103, and [Vaswani et al. 2017] for machine translation on WMT14 EN-DE/EN-FR and IWSLT14 DE-EN.\n\nThe difference between theoretical speed-up and experimental speed-up is honestly discussed, and the overhead of the random features is not swept under the rug. However, having an experimental study on the random features dimensions' impact on empirical compute time vs. approximation performance vs. end-task performance would have been a plus.\n\nI had read \"Rethinking Attention with Performers\" when I reviewed this paper, and I originally thought it was the same paper, but (along with notation being different) they start to differ on page 3. Where \"Performers\" goes with positive orthogonal random features (to improve over vanilla RFA), this paper adds a gating mechanism: this adds the possibility to learn some monotically decaying attention over older context, similar to learned receptive fields of attention (as in e.g. [Sukhbaatar et al. 2019]).\n\nOverall, this is a good paper, and I don't see why we should downplay it in light of simultaneous (\"Performers\" got on ArXiV on September 30th, the ICLR deadline was October 2nd) quite similar contribution that the authors took the time to discuss. (It would be even better if they could compare to it in a future version.)", "extra_scores": {}, "extra_text": {}}, {"review_id": "Db7xKJiZVR", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "1. The logic in the introduction is a bit contradictive to me:\n\nSome are able to achieve better asymptotic complexity (citations). while it\nis more challenging to improve on shorter sequences: the additional computation steps required by\nsome approaches can overshadow the time and memory they save.\n\nDoesn't this simply mean that for short sequences there is no such computational burden? \nI think the story starts with pointing out the importance for long-sequence but turns to the topic on short sequence\nwhich is confusing. The need for short sequence acceleration needs to be justified IMO.\n\n2. Following 1, the baseline should be added.\n\nFor a fair comparison, I think the baseline should add those methods as claimed in the introduction\n(Lee et al., 2019; Child et al., 2019; Sukhbaatar et al., 2019; Beltagy et al., 2020, inter alia), (Kitaev et al., 2020; Wang et al.,\n2020; Roy et al., 2020, inter alia) and let us know how badly they performed under the short sequence. In particular we don't know\nif the sacrifice of short sequence time would benefit a lot in long sequences for existing methods. The current experimental baseline can't reflect this.\n\n3. Is the speedup over total computational time or just the attention part?\n\nTo best of my knowledge, under many circumstances in particular for short sequence, attention alone might not be the \nmost time-consuming part of the model. I think it will be helpful for authors to have a complete graph of the computational model used instead of only figure 1 concept graph. Specifically, is there\nany feed-forward computation involved and how many layers of the models used in comparison.\n\n\n4. Introduction of Eq 6,7 is confusing. \n\nUp to eq 5 it's clear whatr's going, but it comes from nowhere to intorduce these 2 modules in eq 6,7. So my understanding is that\nRFA simply refers to the approximation of computing the softmax. So the statement:\n\nfor softmax-attention. The latter is typically used in two different ways in the transformer architecture, each resulting in a different computation for RF\n\nis confusing as the RFA is now redefined. I believe RFA should only refer one thing and I don't think eq(6) and eq(5) leads to the\nsame result. On the other hand, eq 7 should be the same as eq 5. Is this correct?\nIn addition, the notation in (6) looks wrong to me. \\phi(x) as introduced in eq 2 is in R^{2D} but S_{t-1} is in R{D}, not sure what does + mean in this context.\nI couldn't find out where you properly define the meaning of D either.\n\n\n5. Clarification of contribution\n\nEq 6,7 reads like RNN style update but the intuition is lacking. Do you want to claim that this structure design is inspired by RNN and it leads to a better result?\nPut in another way, using RFA in transformer is from Rawat et al., 19 so do you think your major contribution is to design such \na gated usage of RFA?\n\n\n5. Discussion of D\n\nSince RF is not the major contribution, you summarize existing results of FA in sec2.2. I think I'd like to see a discussion of\nsufficient number D analytically or empirically. Could you also cite the convergence bound on this approximation? To me, D looks to be an important efficiency tradeoff. Say sequence length is M\nand feature in d dimension. Original Attention is O(M^2 d). The computation of RFA\nrequires outer product, which is O(D^2d) so overall it's O(M D^2 d), if M is around 64 or 128 (common usage) and D is 64, I actually\ndon't see why RFA could improve 2x. Do you pre-compute and pre-store anything?\n\n\n6. Time analysis on language modeling is not presented. Since it's a efficiency paper, I think it should be complete.\n\n\n\nOverall, I think the paper provides an interesting view of discussion, but there are many flaws in the current version which needs to be corrected before a more serious consideration.  Especially, in terms novelty, the paper is relatively limited as the RF is explored in Rawat et al., 19. So my point 5 is important to answer and I would like to see all the details are clarified in order to make the contribution stronger. \n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "7dbbnudfpV", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "##### Summary:\nThis paper presents Random Kernel Attention which is based on replacing the kernel function in the Linear Attention with random projection kernels. In general, I think the method is novel and quite impactful. Nowadays, some people are still staying away from attention because of its quadratic time and space complexity. To the best of my knowledge, it is the first attention method with linear complexity that can match or even outperforms the conventional attention.\n\n##### Strengths:\n- The method is intuitive and interesting to me.\n- The results are strong. Unlike Linear Attention, the proposed RFA outperforms the original multi-head attention baseline on both LM and MT tasks. This is quite impressive. Based on my experience the Linear Attention with ELU non-linearity can bearly match the performance of the original attention mechanism.\n- The authors provide several in-depth analyses in the appendix. I like the experiments in C.2.\n- It is nice to see that the authors confess that the training is actually increased when using the RFA compared to the original Transformer. Usually, the inference time and memory usage are more important in practice.\n- It is great to see that the authors compare with the baselines that cache the query/key/value representations. Nowadays, some papers avoid it make their speedup look better.\n\n\n##### Weaknesses & suggestions:\n- It is not clear what D is used in the experiments. The authors just vaguely say that they don't observe the improvement by setting it great than 2d. However, it would be better to see plots at least in the appendix. Also, I wonder if it would behave differently with different d. Also, it would be great to see what exactly the number is in the experimental setup to make this paper more reproducible.\n- The arccos feature maps have only D-dimensional features, unlike the Gaussian feature maps which have 2D. It is not clear whether the authors use the same D for both variants or double the D of arccos to keep the feature dimensions the same.\n- After introducing the random projection weights, the number of parameters would increase. It would be better if the number of parameters and the inference speed are both provided in Tables 1 & 2.\n- The authors should clarify that the time complexity in Table 3 is based on the assumption that we have infinite number of threads or GPT/TPU cores that can be scaled up when M is increased. Otherwise, the time complexity of training the softmax model is still O(M^2) because there is a matrix multiplication between matrices of sizes M-by-M and M-by-d. \n\n##### Questions:\n- Based on the experiments, it seems that the Gaussian random feature maps don't really try to approximate \n- How would the gating mechanism perform on the encoder side? Similar to BiLSTM, half of the dimensions can be applied in a backward manner to make it bidirectional.\n- Do you resample the random weights during the time? \n- I wonder if the authors will release their implementation. Based on my quick re-implementation, the proposed RFA doesn't really converge on some other dataset. I believe there might be some differences in how the parameters are initialized which is not clearly described in the paper. Admittedly, there is a chance that I have a bug in the code.\n- Based on the conclusion in C.2 that RFA is not approximating the softmax kernel, would it be better if we just trained those projection matrices instead of fixing them as random matrices?\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: 
This paper proposes an efficient attention mechanism linear in time and space using random features.
The approach has some similarities with the simultaneous ICLR 2021
submission "Rethinking Attention with Performers", with a key difference of a gating
mechanism present in this work, motivated by recency bias. This paper is a
valuable contributions to the efficient attention research topic. The reviewers
appreciate the experiments and the in-depth analysis. I recommend acceptance.

A noteworthy concern brought up in the discussion period has to do with whether the attention mechanism dominates the feed-forward computations in the neural network, and how much this is architecture-specific. The authors provide TPU timings, but I encourage the authors to add a discussion and timings of relative performance of feed-forward vs. attention layers that covers GPU and CPU optimizers as well.
Decision: Accept (Spotlight)