Rethinking Attention with Performers
Authors: 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 J Colwell
Adrian Weller
Venue: ICLR 2021 Oral (2021)
Keywords: performer
transformer
attention
softmax
approximation
linear
bert
bidirectional
unidirectional
orthogonal
random
features
FAVOR
kernel
generalized
sparsity
reformer
linformer
protein
trembl
uniprot

Abstract: We introduce Performers, Transformer architectures which can estimate regular (softmax) full-rank-attention Transformers with provable accuracy, but using only linear (as opposed to quadratic) space and time complexity, without relying on any priors such as sparsity or low-rankness. To approximate softmax attention-kernels, Performers use a novel Fast Attention Via positive Orthogonal Random features approach (FAVOR+), which may be of independent interest for scalable kernel methods. FAVOR+ can also be used to efficiently model kernelizable attention mechanisms beyond softmax. This representational power is crucial to accurately compare softmax with other kernels for the first time on large-scale tasks, beyond the reach of regular Transformers, and investigate optimal attention-kernels. Performers are linear architectures fully compatible with regular Transformers and with strong theoretical guarantees: unbiased or nearly-unbiased estimation of the attention matrix, uniform convergence and low  estimation variance. We tested Performers on a rich set of tasks stretching from pixel-prediction through text models to protein sequence modeling. We demonstrate competitive results with other examined efficient sparse and dense attention methods, showcasing effectiveness of the novel attention-learning paradigm leveraged by Performers. 

Published as a conference paper at ICLR 2021
RETHINKING ATTENTION WITH PERFORMERS
Krzysztof Choromanski",, Valerii Likhosherstov David Dohan "1 Xingyou Song*!
Andreea Gane*1, Tamas Sarlos"1, Peter Hawkins*) .Jared Davis Afroz Mohiuddin' Lukasz Kaiser', David Belanger', Luey Colwell Adrian Weller
'Google" University of Cambridge DDeppiind *Alan Turing Institute
ABSTRACT
We introduce Performers, Transformer architectures which can estimate regular (softmax) bull-rank-attentin Transformers with provable accuracy, but using only linear (as opposed to quadratic) space and time complexity, without relying on any priors such as sparsity or low-rankness. To approximate softmax attention- kernels, Performers use novel Fast Attention Via positive Orthogonal Random features approach (FAVOR+), which may be of independent interest for scalable kernel methods. FAVOR+ also be used efficiently model kernliizble attention mechanisms beyond softmax. This representational power is crucial to accurately compare softmax with other kernels for the first time large-scale tasks, beyond the reach of regular Transformers, and investigate optimal attention- -ernels.. Performers are linear architectures fully compatible with regular Transformers and with strong theoretical guarantees: unbiased or nearly-unbiased estimation of the attention matrix, uniform convergence and low estimation variance. We tested Performers on a rich set of tasks stretching from pixel-prediction through text models to protein sequence modeling. We demonstrate competitive results with other examined efficient sparse and dense attention methods, showcasing effectiveness of the novel attention-learning paradigm leveraged by Performers.
INTRODUCTION AND RELATED WORK
Transformers (Vaswani et al. 2017 Dehghani et al. 2019 are powerful neural network architectures that have become SOTA in several areas of machine learning including natural language processing (NLP) (e.g. speech recognition (Luo et al. 2020)), neural machine translation (NMT) (Chen et al. 2018), document generation/ caairtiionsummmaizzaiii time series prediction, generative modeling (c.g. image generation (Parmar et al. 2018)), music generation (Huang et 2019), and bioinformatics (Rives)
ct al. 2019 Madani et al. 2020 Ingraham et al. 2019; Elnaggar et al. 2019 Du et al. 2020).
Transformers rely trainable attention mechanism that identifies complex dependencies between the elements of each input sequence. Unfortunately, the regular Transformer scales quadratically with the number of tokens L in the input sequence, which is prohibitively expensive for large L and precludes its usage in settings with limited computational resources even for moderate values of T. Several solutions have been proposed to address this issue (Beltagy al., 20200] Gulati et al.)
2020: Chan et al.) 2020 Child et al. [2019] Bello et al.) 2019)). Most approaches restrict the attention mechanism attend to local neighborhoods (Parmar al. 2018) or incorporate structural priors on attention such as sparsity (Child et al. 2019). pooling based compression (Rae et al. 2020)
buanerrnn/binning/oooootttoooooo techniques (e.g. (Roy et al., 2020) which applies k-means clustering to learn dynamic sparse attention regions, or (Kitaev et al. 2020), where locality sensitive hashing used to group together tokens of similar embeddings), sliding windows (Beltagy et al.) 2020),
or truncated targeting (Chelba et al. 2020). There is also long line of research on using dense attention matrices, but defined by low-rank kernels substituting softmax (Katharopoulos et al. 2020 Shen et al. 2018). Those methods critically rely kernels admitting explicit representations as dot-products of finite positive-feature vectors.
The approaches above do not aim approximate regular attention, but rather propose simpler and more tractable attention mechanisms, often by incorporating additional constraints (e.g. identical query and key sets as in (Kitaev et al. 2020)), or by trading regular with sparse attention using more
'Equal contribution. Correspondence to (kchoro, dcolwell 11 :@google com.
Code for Transformer models on protein data can be found in github.bhb... , com/ googloogleeeeoooooo wacooggg--ereeeeee research
goog bboggee--eeeeeccttttte/////rrsae
goog hhtpoeeeeeeeeeeeeoooooooeeeeeeer
and Performer code can be found in g siihhu..... con/
Google Al Blog: https:/ //
ai http
htm

---

Published as a conference paper at ICLR 2021
layers (Child et al. 2019). Unfortunately, there is lack of rigorous guarantees for the representation power produced by such methods, and sometimes the validity of sparsity patterns can only be verified empirically through trial and error by constructing special GPU operations (e.g. either writing C++ CUDA kernels (Child et al. 2019) or using TVMs (Beltagy et al. 2020)). Other techniques which aim reduce Transformers' space complexity include reversible residual layers allowing one-time activation storage in training (Kitaev et al. 2020)) and shared attention weights (Xiao et al. 2019). These constraints may impede application to long-sequence problems, where approximations of the attention mechanism are not sufficient. Approximations based on truncated back-prop propagation (Dai et al., 2 are also unable to capture Long-distance correlations since the gradients are only propagated inside localized window. Other methods propose biased estimation of regular attention but only in the non-causal setting and with large mean squared error (Wang et al. 2020).
In response, we introduce the first Transformer architectures, Performers, capable of provably accurate and practical estimation of regular (softmax) full-rank attention, but of only linear space and time complexity and not relying any priors such sparsity or low-rankness. Performers use the Fast Attention Via positive Orthogonal Random features (FAVOR+) mechanism, leveraging new methods for approximating softmax and Gaussian kernels, which propose. We believe these methods are of independent interest, contributing to the theory of scalable kernel methods. Consequently, Performers are the first linear architectures fully compatible (via small amounts of fine-tuning) with regular Transformers, providing strong theoretical guarantees: unbiased or nearly-unbiased estimation of the attention matrix, uniform convergence and lower variance of the approximation.
FAVOR+ can be also applied efficiently model other kernelizable attention mechanisms beyond softmax. This representational power is crucial to accurately compare softmax with other kernels for the first time on large-scale tasks, that are beyond the reach of regular Transformers and find for them optimal
irrention-kennhte FAVOR+ can also be applied beyond the Transformer scope as a
more scalable replacement for regular attention, which itself has a wide variety of uses in computer vision (Fu et al. 2019), reinforcement learning (Zambaldi et al., 2019). training with softmax cross entropy loss, and even combinatorial optimization (Vinyals et al. 2015).
We test Performers on rich set of tasks ranging from pixel-prediction through text models to protein sequence modeling. We demonstrate competitive results with other examined efficient sparse and dense attention methods, showcasing the effectiveness of the novel attention learning paradigm leveraged by Performers. We emphasize that in principle, FAVOR+ can also be combined with other techniques, such reversible layers (Kitaev al. 2020) cluster-based attention (Roy et al. 2020).
FAVOR+ MECHANISM & POSITIVE ORTHOGONAL RANDOM FEATURES
Below we describe in detail the FAVOR+ mechanism- - the backbone of the Performer's s architecture. We introduce a new method for estimating softmax (and Gaussian) kernels with positive orthogonal random features which FAVOR+ leverages for the robust and unbiased estimation of regular (softmax) attention and show how FAVOR+ can be applied for other iitention-kenells
2.1 PRELIMINARIES REGULAR ATTENTION MECHANISM
Let L be the size of an input sequence of tokens. Then regular dot-product attention (Vaswan et al. 2017) is a mapping which accepts matrices Q K, V E R gxdd as input where d' is the hidden dimension (dimension of the latent representation). Matrices Q, K, V are intermediate representations of the input and their rows can be interpreted as queries, keys and values of the continuous dictionary data structure respectively. Bidirectional (or non con-directional (Devlin et al. 2018) dot-product attention has the following form, where A E IR 1. I. is the so-called attention matrix:
Att V)
= DV, A eXK /Vd). D (A1l).
(1)
Here exp(:) is applied elementwise, 1, is the all-ones vector of length L. and diag(-) is diagonal matrix with the input vector as the diagonal. Time and space complexity of computing 11 are (T.2d) and ( + Ld) respectively, because A has to be stored explicitly. Hence, in principle, dot-producct attention of type (I) is incompatible with end-to-end processing of long sequences. Bidirectional attention applied in encoder self-attention and encoder-decoder attention in Seq2Seq architectures.
Another important type of attention is unidirectional dot-product atlention which has the form:
Att,, K, V) D --VAV, Ā tril(A), D :diag( A A11);
(2)
2

---

Published as a conference paper at ICLR 2021
where tril(-) returns the lower- trianglla part of the argument matrix including the diagonal. As discussed in /Vaswani et al. (017). unidirectional attention is used for autoregressive generative modelling, e.g. as self-attention in generative Transformers well the decoder part of Seq2Seq Transformers.
We will show that attention matrix A can be approximated up to any precision in time O log (d)) For comparison, popular methods leveraging sparsity via Locality-Sensitiv Hashing (LSH) tech- niques (Kitaev al. 2020) have OLde log L) time complexity. In the main body of the paper we will describe FAVOR+ for bidirectional attention. Completely analogous results can be obtained for the unidirectional variant via the mechanism of prefix-sums (all details in the Appendix B.I
2.2 GENERALIZED KERNELIZABLE ATTENTION
FAVOR+ works for attention blocks using matrices A E R I. of the form A(i,)) = K(q kj) ) with qt/k, standing for the tth 1jth query/key row-vector in Q/K and kernel K : IRd IRd + R+ defined for the (usually randomized) mapping: Rd RT (for some r' >> 0) as:
K(x,y) = ((x(x) p(y)].
(3)
We call (bu) random feature map for E Rd. For Q', K' E R4xr with rows given (qq and o(k,) respectively, Equation 3 leads directly to the efficient attention mechanism of the form:
Att V)
V)), δ diag( QQ Q'(((') ((K))(1))) 1z)).
Here At stands for the approximate attention and brackets indicate the order of computations. It is easy to see that such a mechanism is characterized by space complexity O(Lr + Ld + rd) and time complexity (Trd) as opposed to ( + Ld) and (II,2 of the regular attention (see also Fig. 1..
(
((rd)
batch dlin
O(rrd
E
x
L L
L
L X d
P
L
L d
KK')
T
A
attention nechanise
Figure I: Approximation the regular attention mechanism AV (before D -renomallization) via (random)
feature maps. Dashed blocks indicate order of computation with corresponding time complexities attached.
The above scheme constitutes the FA-part of the FAVOR+ mechanism. The remaining OR+ part answers the following questions: (1) How expressive is the attention model defined in Equation 3 and in particular, can we use in principle to approximate regular softmax attention ? (2) How do we implement it robustly in practice, and in particular, can we choose < L for L >> d to obtain desired space and time complexity gains? We answer these questions in the next sections.
2.3 HOW TO AND HOW NOT TO APPROXIMATE SOFTMAX-KERNHII FOR ATTENTION
turns out that by taking & of the following form for functions fi,..., ft R > R, function
g Rdd + R and deterministic vectors W; or wI..... Wm iid for some distribution D E P(Rd):
h(x)
o(x)
Vm (fi (fi(WT x)... f1(m)))) flw x)..., .... f((wTT)))
(5)
we can model most kernels used in practice. Furthermore, in most cases D is isotropic (i.e. with pff function constant on a sphere), usually Gaussian. For example, by taking h(x) = 11, = and
D = ~N L2) we obtain estimators of the so-called PNG Kernell (Choromanski et al. 2017 (e.g. fi sgn corresponds to the angular kernel). Configurations: h(x) 1, = 2, f1 sin, fo
cos correspond to thift-invariin kernels, in particular D = (O I2) leads to the Gaussian kernel K Kgauss (Rahimi & Recht 2007). The voftmax-kered which defines regular attention matrix A is given as:
3

---

Published as a conference paper at ICLR 2021
SM(x, y) def exp(x y).
(6)
In the above, without loss of generality, we omit Vd-renormalization since we can equivalently
renormalize input keys and queries. Since: (M y) exp( |x]|_ gx exp( ly|l based
on what we have said, we obtain random feature map unbiased approximation of SM y) using
trigonometric functions with: h(x) exp( | ), 2, fi = sin, /2 cos. We call it SM trig
m (x
There is however a caveat there. The attention module from ①T) constructs for each token, a convex combination of value-vectors with coefficients given as corresponding renormalized kernel scores. That is why kernels producing non-negative scores are used. Applying random feature maps with potentially negative dimension-values (sin /oo) leads to unstable behaviours, especially when kernel scores close to (which is the case for many entries A corresponding to low relevance tokens) are approximated by estimators with large variance in such regions. This results in abnormal behaviours, e.g. wereeedignnn---eee
renormalizers D and consequently either completely prevents training or leads to sub-optimal models. We demonstrate empirically that this is what happens for trig
srig SM,, and provide detailed theoretical explanations showing that the variance of SM,, is large as approximated values tend to 0 (see: Section 3). This is one of the main reasons why the robust random feature map mechanism for approximating regular softmax attention
never proposed.
We propose a robust mechanism in this paper. Furthermore, the variance of our new unbiased positive random feature map estimator tends to approximated values tend to (see: Section 3
Lemma (Positive Random Features (PRFs) for Softmax). For x, E R = x x we have:
|x112
exp W - ( cosh(w T z), (7)
SM(X, y) (w--A((,LL)) exp w
|xx| |xx|++ +xl'+>>>l') and cosh is hyperbolic cosine. Consequently, seftmax-kerue admits
where A exp(
positive random randomfeature feature map unbiased approximation with h(xx) = exp(- |ل] 1, fi = exp and = (0, or: h(x) - -/222222000 exp( |x]]" 2, fi (w) = exp(u), f2(u) = exp(-u) and the hyp+ same D (the latter for further variance reduction). We call related estimators: SM,, and SM,,
Figure 2: Left: Symmetrized (around origin) utility function r (defined as the ratio of the mean squared errors (MSEs) of estimators built on: trigonometric and positive random features) as function of the angle & (in radians) between input feature vectors and their lengths 7. Larger values indicate regions of (ś, []-space with hetter performance of positive random features. We see that for critical regions with ø large enough (small enough softmax-kent presented for domain [-x, T| x [-2,2] 2]. Right: The slice of function for fixed = 1 and varying angle a. Right Upper Corner:
values) our method is arbitrarily more accurate than trigonometric random features. Plot
Comparison of the MSEs of both the estimators in a low softmax-kerne value region.
In Fig. 2 we visualize the advantages of positive versus standard trigonometric random features. In critical regions, where kernel values are small and need careful approximation, our method outper- forms its counterpart. In Section we further confirm our method's 's advantages empirically, using positive features to efficiently train softmax-based linear Transformers. If we replace in (7) w with we obtain the so-called regularized softmax-kernel SMREG which we can approximate in
vd & | a similar manner, simply changing D = (O, La) to D = Unif(Vds"* )), a distribution correspond-
ing to Haar measure on the sphere radius /d in Rdd obtaining estimator SMREG As we show in Section such random features also be used to accurately approximate regular softmax-kernel.
4

---

Published as a conference paper at ICLR 2021
2.4 ORTHOGONAL RANDOM FEATURES (ORFs)
The above constitutes the R+ part of the FAVOR+ method. remains to explain the O-part. To further reduce the variance of the estimator (so that features r), we entangle different random samples W1, Wm to be exactly orthogonal. This can be done while maintaining unbiasedness whenever isotropic distributions D are used (i.e. in particular in all kernels we considered so far) by the standard Gram-Schmidt orthogonalization procedure (see (Choromanski et al. 2017)) for details). ORFs is well-known method, yet turns that it works particularly well with our introduced PRFs for softmax. This leads to the first theoretical results showing that ORFs can be applied reduce the variance of noatmax/Gaussian kernel estimators for any dimensionality d rather than just asymptotically for large enough d (as is the case for previous methods, see: next section) and leads to the first exponentially small bounds on large deviations probabilities that are strictly smaller than for non-orthogonal methods. Positivity of random features plays key role in these bounds. The ORF mechanism requires m S but this will be the case in all our experiments. The pseudocode of the entire FAVOR+ algorithm is given in Appendix B
use even smaller number of random
Our theoretical results are tightly aligned with experiments. We show Section 4 that PREs+ORFs drastically improve accuracy of the approximation of the attention matrix and enable us reduce which results in an accurate as well as space and time efficient mechanism which we call FAVOR+.
3 THEORETICAL RESULTS
We present here the theory of positive orthogonal random features for softmax-kernn estimation. All these results be applied also to the Gaussian kernel, since explained in the previous section, one can be obtained from the other by renormalization (see: Section 2.3). All proof's and additional more general theoretical results with a discussion are given in the Appendix.
Lemma 2 (positive (hyperbolic) versus trigonometric random features). The following is true:
MSE(SM,, trig (x,
1
exp([Ix + 11) (x, - exp(-lxx y1
2m
1
(x, y) exp(- MSF(SM, m (x, = (
(8)
MSE hyp+
m (x, = 5 expp - +SS)) y)),
for independent random samples Wa and where MSE stands for the mean squared error.
trig
Thus, for SM(x,y + 0 we have: MSE(SM,,, (x, oo and MSE(SM, (x,y)) 0. Further- more, the hyperbolic estimator provides additional accuracy improvements that are strictly better than those from SM (x, y) with twice as many random features. The next result shows that the regularized softmax -ernel is in practice accurate proxy of the softmax kernel in attention.
Theorem 1 (regularized versus softmax-kernel) Assume that the Los -nom of the attention matrix for the softmax-kett kernel satisfies: ||ll bo C for some constant C 2 1. Denote by A reg the corresponding attention matrix for the regularized softmax-kerne The following holds:
A reg (i,j)
+ 2 2
(o $1
inf 1.J
and sup
A
dt
dt
A(i.j)
Furthermore, the latter holds for d 2 even if the Loe condition not satisfied, i.e. the regularized softmax-kerne kernel is universal lower bound for the pofimax-kerned
Consequently, positive random features for SMREG can be used approximate the coftmax-kerne Our next result shows that orthogonality provably reduces mean squared error of the estimation with positive random features for any dimensionality
and we explicitly provide the gap.
ort- Theorem 2. If SM, (x. y) stands for the modification of SM y) with orthogonal random features (and thus for S d), then the following holds for any d >> 0:
MSE(SM,, m (x.y) < MSM (x.ss) - 1000-(111 (1-1) d+2
ort.
_SM
2 SM y).
(10)
Furthermore, completely analogous result holds for the regularized seftmax-kernel SMREG.
5

---

Published as a conference paper at ICLR 2021
For the regularized deftmax-kerneds orthogonal features provide additional concentration results the first exponentially small bounds for probabilities of estimators' tails that are strictly better than for non-orthogonal variants for every d > O. Our next result enables us to explicitly estimate the gap.
Theorem 3. Xerx,y E Rd. The following holds for any > SMREG(X. y) and m S d:
+ PISMREG, (x,y) > a) I exp( exp(-mEx(a)), PISMREG,. m
ort+ (x,y) > a) S d+2 expl exp(-m=x(a))
where SMREG,, in oct+ (x,)) stands for the modification of SMREG (x, with ORFs, X =
A exp(vá iillaa (x y)), so (O, Ia). A is as in Lemma and Lz Legendre Transform of Z defined as: La(a)= = supa>o log( M2(0)) for the moment generating function Mz of Z.
We see that ORFs provide exponentially small and sharper bounds for critical regions where the softmax-kerne is small. Below we show that even for the SM mechanism with ORFs, it suffices to take m O (og())) random projections to accurately approximate the attention matrix (thus if not attention renormalization, PRFs would not be needed). In general, m depends on the dimensionality of the embeddings, radius the ball where all queries/keys live and precision parameter (see: Appendix F for additional discussion), but does not depend on input sequence length L.
Theorem (uniform convergence for attention approximation). Assume that L2-norms of queries/keys are upper-bounded by R > O. Define = Rd and take h² =exp(+) (.) Then for any E > 0, = (hčy and the number of random projections m = O log( ( the fol- lowing holds for the attention approximation mechanism leveraging estimators SM trig with ORFs: || All Se e with any constant probability, where approximates the attention matrix A.
EXPERIMENTS
We implemented our setup on top of pre-existing Transformer training code in Jax (Frostig et al., 2018) optimized with just-in-time (jax. jitt compilation, and complement our theory with em- pirical evidence to demonstrate the practicality of FAVOR+ multiple settings. Unless explicitly stated, a Performer replaces only the attention component with our method, while all other com- ponents are exactly the same as for the regular Transformer. For shorthand notation, we denote unidirectional/causal modelling as (U) and bidirectional/masked language modelling (B).
In terms of baselines, we use other Transformer models for comparison, although some of them restricted to only one case c.g. Reformer (Kitaev al., (0200) is only (U), and Linformer (Wang et al. (020) is only (B). Furthermore, we use PG-19 (Rac et al. 20200 as an alternative (B) pretraining benchmark, as it is made for long-length sequence training compared to the (now publicly unavailable) BookCorpus (Zhu et al. 2013) + Wikipedia dataset used in BERT (Devlin et al. 2018
and Linformer. All model and tokenization hyperparameters are shown in Appendix A
Model Forward Timing (Regular)
Model Backward Timing (Regular)
Batch-1
Baich-2
Batchie
(ogy(7 (sec)
Batchen
Performer
-6
*
0PT.BATCH--
-5
-8
-10
-7
10
12
14
16
18
8
10
12
14
(o):(L)
(og:(L)
Figure 3: Comparison of Transformer and Performer in terms of forward and backward pass speed and maximum L allowed. "X" (OPT) denotes the maximum possible speedup achievable. when attention simply returns the V-matrix. Plots shown up to when model produces an out of memory error on a V100 GPU with 16GB. Vocabulary size used was 256. Best in color.
4.1 COMPUTATIONAL COSTS
We compared speed-wist wise the backward pass of the Transformer and the Performer in (B) setting, as it is one of the main computational bottlenecks during training, when using the regular default size (Mreads: Mayerss drf, d) = (8, 6, 2048, 512). where y denotes the width of the MLP layers. We observed (Fig. 3) that in terms of L., the Performer reaches nearly linear time and sub-quadratic

---

Published as a conference paper at ICLR 2021
memory consumption (since the explicit O attention matrix is not stored). In fact, the Performer achieves nearly optimal speedup and memory efficiency possible, depicted by the "X"-line when attention is replaced with the "identity function" simply returning the V-matrix. The combination of both memory and backward pass efficiencies for large I. allows respectively, large batch training and lower wall clock time per gradient step. Extensive additional results are demonstrated in Appendix E by varying layers, raw attention, and architecture sizes.
4.2 SOFTMAX ATTENTION APPROXIMATION ERROR
We further examined the approximation error via FAVOR+ in Fig. We demonstrate that 1.
Orthogonal features produce lower error than unstructured (IID) features, 2. Positive features produce lower emrrr than trigonometric sin/cos features. These two empirically validate the PORF mechanism.
Attention Output (Orthogonal vs IID)
Attention Output (Sin/Cos Positive)
10c
104
Sir/Cos
10a
Orthogonal
10a
Positive
10.
Output ASE
Output MSE
10a
10a
102
104
10a
10a
109
10-1
10-2
10
25
50
75
100
125
150
175
200
100
125
150
175
200
Features
Num. Features
Figure 4: MSE of the approximation output when comparing Orthogonal vs IID features and trigonometric sin/cos vs positive features. We took L = 4096, d = 16, and varied the number of random samples m. Standard deviations shown across 15 samples of appropriately normalized random matrix input data.
To further improve overall approximation of attention blocks across multiple iterations which further improves training, random samples should be periodically redrawn (Fig. [ right). This a cheap procedure, but can be further optimized (Appendix B.2.
4.3 SOFTMAX APPROXIMATION ON TRANSFORMERS
Even if the approximation of the attention mechanism is tight, small errors can easily propagate throughout multiple Transformer layers (e-g. MLPs, multiple heads), as we show in Fig. [4] (Appendix). In other words, the model's Lipschitz constant can easily scale up small attention approximation error, which means that very tight approximations may sometimes be needed. Thus, when applying FAVOR(+)'s softmax approximations SOFTMAX"), we demonstrate that:
Transformer model (i.e. "Performer-X.
1. Backwards compatibility with pretrained models is available benefit from softmax approxima- tion, via small finetuning (required due to error propagation) even for trigonometric features (Fig. 5 left) the LMIB dataset (Chelba et al., 2014). However, when larger dataset PG-19, 2. Positive (POS) softmax features (with redrawing) become crucial for achieving performance matching regular Transformers (Fig. right).
LM1B (B) Backward Compatibility (L 512)
PG-19 (B), Layer Regular Model (L (L= 1024)
60K
50K
Lifemmer PEDRAA
Pefferec POS SCF MAX NO RECHUM neturee Sor Thuro
Accuracy
Perplexity
40K
30K
Performer-TRIG TRIG-SOFTMAX (Finetune)
Transformer
20K
10K
0.00
0K
10K
20K
30K
40K
SOK
63K
DK
20K
40K
BOK
B0K
100K
Gradient Steps
Steps
Figure 5: We transferred the original pretrained Transformer's 's weights into the Performer, which produces an initial non-zero 0.07 accuracy (dotted orange line), but quickly recovers accuracy in small fraction of the original number of gradient steps. However on PG-19, Trigonometric (TRIG) softmax approximation becomes highly unstable (full curve in Appendix [D.2). while positive features (POS) (without redrawing) and Linformer (which also approximates softmax) even with redrawn projections, plateau at the same perplexity. Positive softmax with feature redrawing is necessary to match the Transformer, with SMREG (regularization from Sec. 3 allowing faster convergence. Additional ablation studies over many attention kernels, showing also that trigonometric random features lead even to NaN values training are given Appendix D.3
7

---

Published as a conference paper at ICLR 2021
4.4 MULTIPLE LAYER TRAINING FOR PROTEINS
We further benchmark the Performer on both (U) and (B) cases by training a 36-layer model using protein sequences from the Jan. 2019 release of TrEMBL (Consortium) 2019), similar to (Madani ct al., 2020). In Fig. 6 the Reformer and Linformer signifccntll drop in accuracy on the protein dataset. Furthermore, the usefulness of generalized attention is evidenced by Performer-RELU (taking
= R in Equation 5 achieving the highest accuracy in both (U) and (B) cases. Our proposed softmax approximation also shown to be tight, achieving the same accuracy as the exact-softmax Transformer and confirming our theoretical claims from Section 3
TrEMBL, 36 Layers (U) (L 1024)
TrEMBL, 36 Layers (B) (L - 1024)
0.35
0.35
0.30
Accuracy
Accuracy
0.25
0:25
0.20
0.20
Performer REEL (8)
Perfomer RELU (U)
Performer SOFTMAX (B)
0.15
Transformr (U)
Transhormer (B)
Retormer (U]
Linformer
0.10
0.10
OK
50K
100K
0K
100K
200K
300K
Stops
Steps
Figure 6: Train = Dashed, Validation = Solid. For TrEMBL, we used the exact same model parameters (Rheads, Players, dz d) = (8, 36, 1024, 512) from (Madani ct al. (020) for all runs. For fairesss, all TrEMBL. experiments used 16x16 TVU-s's. Batch sizes were maximized for each separate run given the compute constraints. Hyperparameters can found in Appendix A Extended results including dataset statistics, out of distribution evaluations, and visualizations, can found in Appendix C
4.5 LARGE LENGTH TRAINING COMMON DATASETS
On the standard (U) ImageNet64 benchmark from Parmar et al.) 2018 with L 12288 which is unfeasible for regular Transformers, we set all models to use the same [Nheads: dr d) but varying Mlayers- Performer/6-layerr matches the Reformer/1 2-layers, while the Berformer/12-layerr matches the Reformeer244layerrs (Fig. 7 left). Depending on hardware (TPU or GPU), we also found that the Performer can be 2x faster than the Reformer via Jax optimizations for the (U) setting.
For a proof of principle study, we also create an initial protein benchmark for predicting interactions among groups of proteins by concatenating protein sequences length L 8192 from TrEMBL. long enough to model protein interaction networks without the large sequence alignments required by existing methods (Cong al. 2019p. this setting, a regular Transformer overloads memory even at batch size of 1 per chip, by wide margin. Thus as baseline, we were forced to use significantly smaller variant, reducing to (Nrends: Riagers: drs, d) = (8, 21,2,3). 256, 256). Meanwhile, the Per- batch size of 8 per chip using the standard (8, 6, 2048, 512) architecture. We see Fig. 7 (right subfigure) that the smaller Transformer (Miager 3) is quickly bounded at 19%, while the Performer is able to train continuously to z 24%.
former trains efficiently
ImageNet64 12288)
TrEMBL-Concat (B), 8192)
Pertomer RELU
0.24
Pertmee RELD 321.
0.22
4.2
Reefomen 144
24L
0.20
4.0
0.18
3.8
meformert Lapers
0.16
Transfomer. (d:errdd:200)
3.6
Tranofomer, 2 Lxjer (du2500
0.14
Transfomer, 3 Layers (d=250)
3.4 OK
25K
50K
75K
100K
0K
100K
200K
300K
Steps
Steps
Figure 7: Train = Dashed, Validation = Solid. For ImageNet64, all models used the standard (Matadda d d) (8, 2028048555 512). We further show that our positive softmax approximation achieves the same performance as ReI.U in Appendix [D.2] For concatenated TrEMBI,, we varied Mayyers E {1, 2,33 for the smaller Transformer. Hyperparameters can be found in Appendix
CONCLUSION
We presented Performer, a new type of Transformer, relying on our Fast Attention Via positive Or- thogonal Random features (FAVOR+) mechanism to significantly improve space and time complexity of regular' Transformers. Our mechanism provides to our knowledge the first effective unbiased esti- mation of the original bbased Transformer with linear space and time complexity and opens new avenues in the research on Transformers and the role of non non-sparsifying attention mechanisms,
8

---

Published as a conference paper at ICLR 2021
6 BROADER IMPACT
We believe that the presented algorithm can be impactful in various ways:
Biology and Medicine: Our method has the potential to directly impact research biological sequence analysis by enabling the Transformer to be applied to much longer sequences without constraints the structure of the attention matrix. The initial application that consider is the prediction of interactions between proteins on the proteome scale. Recently published approaches require large evolutionary sequence alignments, bottleneck for applications mammalian genomes (Cong et al. 2019). The potentially broad translational impact of applying these approaches to biolog- ical sequences was of the main motivations of this work. We believe that modern bioinformatics can immensely benefit from new machine learning techniques with Transformers being among the most promising. Scaling up these methods to train faster more accurate language models opens the door to the ability to design sets of molecules with pre-specified interaction properties. These approaches could be used to augment existing physics- basedd design strategies that are of critical importance for example in the development of new nanoparticle vaccines (Marcandalli et al. 2019).
Environment: As we have shown, Performers with FAVOR+ are characterized by much lower compute costs and substantially lower space complexity which can be directly translated to CO2 emission reduction (Strubell et al., 2019) and lower energy consumption (You al. 2020). as regular Transformers require very large computational resources.
Research on Transformers: We believe that our results can shape research on efficient Transformers architectures, guiding the field towards methods with strong mathematical foundations. Our research may also hopefully extend Transformers also beyond their standard scope (e.g. by considering the Generalized Attention mechanism and connections with kernels). Exploring scalable Transformer architectures that can handle L, of the order of magnitude few thousands and more, preserving accuracy of the baseline the same time, is gateway new breakthroughs in bio-informatics, e.g. language modeling for proteins, as we explained in the paper. Our presented method can be potentially first step.
Backward Compatibility: Our Performer can be used on the top of regular pre-trained Transformer as opposed to other Transformer variants. Even if up-training is not required, FAVOR+ can still be used for fast inference with no loss of accuracy. We think about this backward compatibility as very important additional feature of the presented techniques that might be particularly attractive for practitioners.
Attention Beyond Transformers: Finally, FAVOR+ can be applied to approximate exact attention also outside the scope of Transformers. This opens a large volume of new potential applications including: hierarchical attention networks (HANS) (Yang et al., 2016), graph attention networks (Velickovic al. 2018), image processing (Fu et al. 2019). and reinforcement learning/robotics
Tang et al. 2020).
7 ACKNOWLEDGEMENTS
We thank Nikita Kitaev and Wojeiech Gajewski for multiple discussions on the Reformer, and also thank Aurko Roy and Ashish Vaswani for multiple discussions on the Routing Transformer. We further thank Joshua Meier, John Platt, and Tom Weingarten for many fruitful discussions on biological data and useful comments this draft. We lastly thank Yi Tay and Mostafa Dehghani for discussions on comparing baselines.
Valerii Likhosherstov acknowledger support from the Cambridge Trust and DeepMind. Lucy Colwell acknowledges support from the Simons Foundation. Adrian Weller acknowledges support from a Turing Al Fellowship under grant EP/V025379/1, The Alan Turing Institute under EPSRC grant EP/N510 101 29/1 and U/B/000074, and the Leverhulme Trust via CFI.
9

---

Published as a conference paper at ICLR 2021
REFERENCES
Irwan Bello, Barret Zoph, Ashish Vaswani, Jonathon Shlens, and Quoe Le. Attention augmented abs/
convolutional networks. CoRR, abs/1904.09925,, 2019. URL
1904. 09925
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. CoRR, abs/2004.05150, 2020. URL/tttp:: petecps://arii... barrgg/bb//004/4 2004. 05150
William Chan, Chitwan Saharia, Geoffrey E. Hinton, Mohammad Norouzi, and Navdeep Jaitly. Imputer: Sequence modelling via imputation and dynamic programming. CoRR, abs/2002.08922 2020. URLhhtts:///axxi
abs /202 08926
Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillipp Koehn, and Tony Robinson. One billion word benchmark for measuring progress in statistical language modeling. In INTERSPEECH 2014, 15th Annual Conference of the International Speech Communication Association, Singapore, September 14-18, 2014, pp. 2635-2639, 2014.
Ciprian Chelba, Mia Xu Chen, Ankur Bapna, and Noam Shazeer. Faster transformer decoding: N-gram masked self-attention. CoRR, abs/2001.04589, 2020. URL https:://axii..orrtt
abbs/201. 04589
Mia Xu Chen, Orhan Firat, Ankur Bapna, Melvin Johnson, Wolfgang Macherey, George F. Foster, Llion Jones, Mike Schuster, Noam Shazeer, Niki Parmar, Ashish Vaswani. Jakob Uszkoreit, Lukasz Kaiser, Zhifeng Chen, Yonghui Wu, and Macduff Hughes. The best of both worlds: Combining recent advances in neural machine translation. Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics, ACL 2018, Melbourne, Australia, July 15-20, 2018, Volume Long Papers, pp. 76-86, Association for Computational Linguistics, 2018, doi: 10.18653/v 1/P18- 1008. UtLhttss: www http://ww..bbb....aaaaaoooolalaa
Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. CoRR, 0bbs/1904.11 10509, 2019.
Krzysztof Choromanski, Carlton Downey, and Byron Boots. Initialization matters: Orthogonal predic- tive state recurrent neural networks. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. 3bwbbb)
OpenReview.net, 2018a. URL
httttpp:///pperrvvieee wwtpp//wpweeee..eeeoeeeee net htt///rut///reeeee
Krzysztof Choromanski. Mark Rowland, Tamás Sarlós, Vikas Sindhwani, Richard E. Turner, and Adrian Weller. The geometry of random features. In International Conference on Artificial Intelligence and Statistics, AISTATS 2018, 9-11 April 2018, Playa Blanea, Lanzarote, Canary Islands, Spain, volume 84 of Proceedings of Machine Learning Research, pp. 1-9, PMLR, 2018b. URL httpLtttpp//rrreeeinggg.ou mlr. http:///////////oooaeeeteeeeeeee
html
Krzysztof Choromanski, Aldo Pacchiano, Jeffrey Pennington, and Yunhao Tang. KAMA-NNS: now-dimensional rotation based neural networks. In The 22nd International Conference on Artificial Intelligence and Statistics, AISTATS 2019, 16-18 April 2019, Naha, Okinawa, Japan, volume 89 of Proceedings of Machine Learning Research, pp. 236 245.2.5.12 PMLR, 2019a. URL http://proceed: ngs mmm www.ppees/veee wttp://or.aaas......m 9a html
Krzysztof Choromanski, Mark Rowland, Wenyu Chen, and Adrian Weller. Unifying orthog- onal Monte Carlo methods. In Proceedings of the 36th International Conference on Ma- chine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pp. 1212- PMLR, 2019b. URL http: //proceedings. mlr press /v7 htt.//oorrmacccec html
Krzysztof Marcin Choromanski, Mark Rowland, and Adrian Weller. The unreasonable effectiveness of structured random orthogonal embeddings. In Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, 4-9 December 2017, Long Beach, CA, USA, pp. 219-228, 2017.
10

---

Published as a conference paper at ICLR 2021
Djork-Arn Clevert Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). In 4th International Conference on Learning Represen- tations, ICLR 2016, San Juan, Puerto Rico May 2-4, 2016, Conference Track Proceedings, 2016. URL
absp///llll1111/ 07289
Qian Cong, Ivan Anishchenko, Sergey Ovchinnikov, and David Baker. Protein interaction networks revealed by proteome coevolution. Science, 365(6449) 2019.
UniProt Consortium. Uniprot: worldwide hub of protein knowledge. Nucleic acids research, 47 (D1): D515, 2019.
Thomas H. Cormen, Charles E. Leiserson, Ronald L. Rivest, and Clifford Stein. Introduction to Algorithmss, 3rd Edition. MIT Press, 2009. ISBN 978-0-262-003884-8 URL http::/mitpreas a gor t.mss
mi www.dd//ooos/n.t..too
Zihang Dai, Zhilin Yang, Yiming Yang, William W. Cohen, Jaime Carbonell, Quoe V. Le, and Ruslan Salakhutdinov. Tr Language modeling with longer-term dependency, 2019. URL
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. In 7th International Conference Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL hnttps://operreviee
net/Eorum? Ho-
Jacob Devlin, Ming-Wet Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. CoRR, abs/1810.04805, 2018. URL
attp:://axii: wwg. abs 1810 04805
Yilun Du, Joshua Meier, Jerry Ma. Rob Fergus, and Alexander Rives. Energy-based models for atomic-resolution protein conformations. arXiv preprint arXiv:2004.13167 2020.
Ahmed Elnaggar, Michael Heinzinger, Christian Dallago, and Burkhard Rost. End-to-end multitask learning, from protein language to protein features without alignments. bioRxiv, pp. 864405, 2019.
Roy Frostig, Matthew Johnson, and Chris Leary. Compiling machine learning programs via high- level tracing. In Conference on Machine Learning and Systems 2018, 2018. URL http::/tt/// www sysml. wttpcc//cccc///////
pdf
Jun Fu, Jing Liu, Haijie Tian, Yong Li, Yongjun Bao, Zhiwei Fang, and Hanqing Lu. Dual attention network for scene segmentation. In IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2019, Long Beach, CA, USA, June 16-20, 2019, pp. 3146-3154 2019.
Anmol Gulati, James Qin, Chang-Cheng Chiu, Niki Parmar, Yu Zhang, Jiahui Yu, Wei Han, Shibo Wang, Zhengdong Zhang. Yonghui Wu, and Ruoming Pang. Conformer: Convolution-augmente transformer for speech recognition, 2020.
Cheng-Zhi Zhi Anna Huang, Ashish Vaswani, Jakob Uszkoreit, Ian Simon, Curtis Hawthorne, Noam Shazeer, Andrew M. Dai, Matthew D. Hoffman, Monica Dinculescu, and Douglas Eck. Music transformer: Generating music with long-term structure. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. SRRLstttps//opeerriiiee htt.opprnrreeeeeeeeeeeeeereeoeee foreet//ooommaam idd rer.:emmmmirrJeeeeeheeaeae
John Ingraham, Vikas Garg, Regina Barzilay, and Tommi Jaakkola. Generative models for graph- based protein design. In Advances in Neural Information Processing Systems, pp. 15794-15805, 2019.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. CoRR, abs/2006.16236 2020. URL
https://arxxv.oog////s//// 2606.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In 8th International Conference Learning Representations, ICLR 2020, Addis Ababa, Ethiopia. April 26-30, 2020. OpenReview.net, 2020. URL http:t/tp/:://eeeeeeeew net ffo.mmiddmd FkgN

---

Published as a conference paper at ICLR 2021
Olga Kovaleva, Alexey Romanov, Anna Rogers, and Anna Rumshisky. Revealing the dark secrets of bert. arXiv preprint arXiv: 1aKV::9908.08993 2019.
Taku Kudo and John Richardson. Sentencepiece: A simple and language independent subword
tokenizer and detokenizer for neural text processing. CoRR, tos//1808.663 2018. http
//arxiv. www. abs 1808 06226
Richard E. Ladner and Michael J. Fischer. Parallel prefix computation. J. ACM, 2784)8831-838
October 1980. ISSN 0004-5411, doi: 10.1145/322217.522222. URL.:ttp:://o1..oo
1145/322217 322232
Han Lin, Haoxian Chen, Tianyi Zhang, Clément Laroche, and Krzysztof Choromanski. Demystifying orthogonal Monte Carlo and beyond. CoRR, abs/2005.13590, 2020.
Haoneng Luo, Shiliang Zhang, Ming Lei, and Lei Xie. Simplified self-attention for transformer-based end-to-end speech recognition. CoRR, 10s/20005.1063,, 2020. URL https://aaxxvv.org
2aas/2005.1444
Ali Madani, Bryan McCann, Nikhil Naik, Nitish Shirish Keskar, Namrata Anand, Raphael R. Eguchi, Po-Ssu Huang, and Richard Socher. Progen: Language modeling for protein generation. CoRR, abs/2004.03497, 2020. URL httppp/////aavv..rrr//////////
03497
Jessica Marcandalli, Brooke Fiala, Sebastian Ols, Michela Perotti, Willem de van der Schueren, Joost Snijder, Edgar Hodge, Mark Benhaim, Rashmi Ravichandran, Lauren Carter, et al. Induction of potent neutralizing antibody responses by a designed protein nanoparticle vaccine for respiratory syncytial virus. Cell, 176(6):1420-144) 1431, 2019.
Nikita Nangia and Samuel R. Bowman. Listops: diagnostic dataset for latent tree learning. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics, NAACI-HIT 2018, New Orleans, Louisiana, USA, June 2-4, 2018, URL https
Student Research Workshop, pp. 92-99, 2018. doi: 10. 18653/v1/n18-4400
dod //doi.org/ 1.O.18853/v111
4013
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. Proceedings of the 35th International Conference on Machine Learning, ICML. 2018, Stockholnsmässan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 4052 4061. PMLR, 2018. URL httt:://proceedinggs mlr. prwss / vv0 2:000parmmn wwwppaamaallee www.mmarlaa..a....
Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, Chloe Hillier, and Timothy P. Lillicrap. Com- pressive transformers for long-range sequence modelling. In International Conference on Learning Representations, 2020. URL /ttp:::////e////eeeeeeee net www.mmrrummmmmm 21-
Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. In Advances in Neural Information Processing Systems 20, Proceedings of the Twenty-First Annual Conference on Neural Information Processing Systems, Vancouver, British Columbia, Canada, December 3-6, 2007, pp. 1177-1184. Curran Associates, Inc., 2007. URL httttttttp/pppprttttt nips. cc/
Alexander Rives, Siddharth Goyal, Joshua Meier, Demi Guo, Myle Out, C. Zitnick, Jerry Ma, and Rob Fergus. Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. bioArxiv, 04 2019. doi: 10. 101//2002..
Mark Rowland, Jiri Hron, Yunhao Tang, Krzysztof Choromanski, Tamás Sarlós, and Adrian Weller. Orthogonal estimation of W Wasserstein distances. In The 22nd International Conference on Artificial Intelligence and Statistics, AISTATS 2019, 16-18 April 2019, Naha, Okinawa, Japan, volume 89 of Proceedings of Machine Learning Research, pp. 186-195. PMLR, 2019. URL http://
proceedings mlr. www.sppesssssseee v89/ http///owaaanlndd. html
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. CoRR, abs/2003.05997, 2020. URL https: ://arxiv
orrg/bbs 2003. 05997
12

---

Published as a conference paper at ICLR 2021
Zhuoran Shen, Mingyuan Zhang, Shuai Yi, Junjie Yan, and Haiyu Zhao. Factorized attention: Self-attention with linear complexities. CoRR, abs/1812.01243, 2018. URL http://axxiv
org/ arr//bs/Lae 1812. 01243
Emma Strubell, Ananya Ganesh, and Andrew McCallum. Energy and policy considcrations for deep learning in NLP. CoRR, abs/1906.02243, 2019. URL http:://axxi...oot org/abs///aoooooooooo 1argg/abs19900 /1906
02243
Yujin Tang, Duong Nguyen, and David Ha. Neuroevolution of sef-itteppeeaaeeeeee agents. CoRR, abs/2003.08 16,, 2020. URL httpp//s///exxi...rrr/////222..0
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. 2021.
Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philippe Morency, and Ruslan Salakhut- dinov. Transformer dissection: An unified understanding for transformer's attention via the lens of kernel. In Proceedings of the 2019 Conference on Empirical Methods Natural Lan- guage Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp. 4335-4344, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems 30, pp. 5998-6008
Curran Associates, Inc., 2017. URL (ttpttt:1/ /tt.ttp::/pppest papers you pouu- pd
nips. wttpscc/ppppr///.// 1araper/788 /718 81 -
Petar Velickovic, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pictro Liō, and Yoshua Bengio. Graph attention networks. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018. Utthtttps:/ /ttp///eeeeeeeeeeeeeeeeeeee
Jesse Vig. A multiscale visualization of attention in the transformer model. arXiv preprint arXiv:1906.05714. 2019.
Jesse Vig and Yonatan Belinkov. Analyzing the structure of attention in transformer language model. CoRR, 042844, 2019. http:/tpp///xxi..pttttt/t//
04084444.4
Jesse Vig, Ali Madani, Lav R. Varshney, Caiming Xiong, Richard Socher, and Nazneen Fatema Rajani. Bertology meets biology: Interpreting attention in protein language models. CoRR, abs/2006. 15522, 2020. htRl:htts:://aas http://.orgga/aa////////e/se//// abtp/////...... 122.2220666.222222222222222
Oriol Vinyals. Meire Fortunato, and Navdeep Jaitly. Pointer networks. In Advances in Neural Information Processing Systems 28: Annual Conference on Neural Information Processing Systems 2015, December 7-12, 2015, Montreal, Quebec, Canada, pp.
2015.
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. CoRR, abs/2006.04766 04768, 2020. URL 04768
Tong Xiao, Yinqiao Li, Jingbo Zhu, Zhengtao Yu, and Tongran Liu. Sharing attention weights for fast transformer. In Proceedings of the Twenty-Eighth International Joint Conference on Artificial Intelligence, IJCAI 2019, Macao, China, August 10-16, 2019, pp. 2292-5298 ijcai.org. 2019. doi: 10.24963/jjcai.2011 20199//33..0 /tt
Zichao Yang, Diyi Yang. Chris Dyer, Xiaodong He, Alexander J. Smola, and Eduard H. Hovy. Hierarchical attention networks for document classification. In NAACL HLT 2016, The 2016 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, San Diego California, USA, June 12-17, 2016, pp. 1480-1499 1489. The Association for Computational Linguistics, 2016. doi: 10.18653/w1/n166 174. URL https
7/doi .wl.oorg//ie 10. 18653/11/ 1-17
13

---

Published as a conference paper at ICLR 2021
Haoran You, Chaojian Li. Pengfei Xu, Yonggan Fu, Yue Wang, Xiaohan Chen, Richard G. Baraniuk,
Zhangyang Wang, and Yingyan Lin. Drawing early-bird tickets: Toward more efficient training of deep networks. In International Conference Learning Representations, 2020. URL https:
//openreviee new. net o (ot.rrmmmmd-BB.ssrrrss....
Felix X. Yu, Ananda Theertha Suresh, Krzysztof Marcin Choromanski, Daniel N. Holtmann-Rice, and Sanjiv Kumar. Orthogonal random features. In Advances in Neural Information Processing Systems 29: Annual Conference on Neural Information Processing Systems 2016, December 5-10,
2016, Barcelona, Spain, pp. 1975-1983, 2016.
Vinícius Flores Zambaldi, David Raposo. Adam Santoro, Victor Bapst, Yujia Li, Igor Babuschkin, Karl Tuyls, David P. Reichert, Timothy P. Lillicrap, Edward Lockhart, Murray Shanahan, Victoria Langston, Razvan Pascanu, Matthew Botvinick, Oriol Vinyals, and Peter W. Battaglia. Deep reinforcement learning with relational inductive biases. In 7th International Conference Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019, 2019.
Yukun Zhu, Ryan Kiros, Richard S. 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 2015 IEEE International Conference on Computer Vision, ICCV 2015, Santiago, Chile, December 7-13, 2015, pp. 19-27, 2015. doi: 10.1 109/ICCV.205.11. URL
https: //do:://do.... www.rrgg.................. 111100111111CCCC 2015..11....111
14

---

Published as a conference paper at ICLR 2021
APPENDIX: RETHINKING ATTENTION WITH PERFORMERS
A HYPERPARAMETERS FOR EXPERIMENTS
This optimal setting (including comparisons to approximate softmax) we use for the Performer specified in the Generalized Attention (Subsec. [A.4), and unless specifically mentioned (e-g. using name "PefOO) "Performer' refers to using this generalized attention setting.
A.1 METRICS
report the following evaluation metrics:
For unidirectional models, we measure the accuracy on next-token prediction, averaged across all sequence positions in the dataset. For bidirectional models, mask each token with 15% probability (same as (Devlin et al. 2018 ) and measure accuracy across the masked positions.
Accuracy:
2. Perplexity: For unidirectional models, we measure perplexity across all sequence positions in the dataset. For bidirectional models, similar to the accuracy case, measure perplexity across the masked positions.
3. Bits Per DimensionnChaarader (BPD/BPC): This calculated by loss divided by In(2). We used the full evaluation dataset for TrEMBL in the plots in the main section, while for other datasets such as ImageNet64 and PG-19 which have very large evaluation dataset sizes, we used random batches (>2048 samples) for plotting curves.
A.1.1 PG-19 PREPROCESSING
The PG-19 dataset (Rae al. 2020) is presented as challenging long range text modeling task. consists of out-of-copyright Project Gutenberg books published before 1919. It does not have a fixed vocabulary size, instead opting for any tokenization which model arbitrary string of text. We use a unigram SentencePiece vocabulary (Kudo & Richardson 2018) with 32768 tokens, which maintains whitespace and is completely invertible to the original book text. Perplexities are calculated as the average log-likelihood per token, multiplied by the ratio of the sentencepiece tokenization number of tokens the original dataset. The original dataset token count per split is: train=19731 136207,, validation=30077611, test=6966499. Our sentencepiece tokenization yields the following token counts per split: train=30847767766 valid=4656945, and test= 10699704. This gives log likelihood multipliers of train=1. 1.6644, valid=1.5487. test=1.5359 per split before computing perplexity, which equal to exp(log likelihood multiplier loss).
Preprocessing for TrEMBL is extensively explained in Appendix
A.2 TRAINING HYPERPARAMETERS
Unless specifically stated, all Performer Transformer runs by default used 0.5 grad clip, 0.1 weight decay, 0.1 dropout, 10-3 3 fixed learning rate with Adam hyperparameters (8) = 0.9, .2 0.98, 10 9), with batch size maximized (until TPU memory overload) for specific model.
All 36-layer protein experiments used the same amount compute (i.c. 16x16 TPU-v2, 8GB per chip). For concatenated experiments, 16x16 TPU-v2's were also used for the Performer, while 8x8's were used for the 1-3 layer (d = 256) Transformer models (using 16x16 did not make a difference in accuracy).
Note that Performers are using the same training hyperparameters as Transformers, yet achieving competitive results this shows that FAVOR can act as a simple drop-in without needing much tuning.
A.3 APPROXIMATE SOFTMAX ATTENTION DEFAULT VALUES
The optimal values, to default parameter f
attention amormalize_attenntion True, numerical stabilizer 10-6 number of features 256, attho_featurs = True, ortho_scaling = 0.0.
are: renormalize
rese://m/oooooeeeeeeeeeeeeeeee rahoog http://swwrrcaa..bbbb/aaaaeeee serformer/Faat_attenttt
ob/
15

---

Published as a conference paper at ICLR 2021
A.4
GENERALIZED ATTENTION DEFAULT VALUES
The optimal values, set to default parameter are: renormalize nnormalize_atteetioonn True, numerical stabilizer 0.0, number of features = 256, kernel ReLU, dernel_epsilon = 10-3.
A.5 REFORMER DEFAULT VALUES
For the Reformer, we used the same hyperparameters as mentioned for protein experiments, without gradient clipping, while using the defaults (which instead use learning rate decay) for ImageNet-64, In both cases, the Reformer used the same default LSH attention parameters.
A.6 LINFORMER DEFAULT VALUES
Using our standard pipeline as mentioned above, we replaced the attention function with the Linformer variant via Jax, with 8 10 6 k 600 (same notation used in the paper (Wang et al. 2020)), where is the exponent in renormalization procedure using e- as a multiplier in order to approximate softmax, while &o is the dimension of the projections of the Q and K matrices. As sanity check, we found that our Linformer implementation in Jax correctly approximated exact softmax's 's output within 0.02 error for all entries.
Note that for rigorous comparisons, our Linformer hyperparameters are even stronger than the defaults found in (Wang al. 2020 as:
We use 600, which is more than twice than the default 256 from the paper, and also twice than our default m 256 number of features.
We also use redrawing, which avoids "unlucky" projections on Q and K.
http:////thaaoooooee..rrrr//o/ar refoomer_imgenettl..g
p
16

---

Published as a conference paper at ICLR 2021
B MAIN ALGORITHM: FAVOR+
We outline the main algorithm for +AVOR+ formally:
Algorithm 1: +AVOR+ (bidirectional or unidirectional).
Input: : Q, K, V Rixdd isBidirectional binary flag.
Result: Atte (Q, K. V) E R isBidircctinall Att Q K. V) R4XL otherwise.
Compute Q' and K' as described Section 2.2 and Section 2.3 and take [V 1r];
if isBidirectional then
Buf == (K') ( e [Mx(++1, Buf == Q'Bufi E R(x(d+1); else
Compute G and its prefix-sum tensor G according to
Buf G Qi
GL Q1] E Llx(d+1);
end
[Buf buf4] == Bufg, Bufs 12xdd buf4 R4;
return (iag(buls)="Bufe
B.
UNIDIRECTIONAL CASE AND PREFIX SUMS
We explain how our analysis from Section 2.2 can be extended to the unidirectional mechanism
in this section. Notice that this time attention matrix A is masked, i.c. all its entries not in the oower-trianullat part (which contains the diagonal) are zeroed (see also Fig. D
Q'
E
A
V
(K')
outer
element of the proftx-sum
Figure 8: Visual representation of the prefix sum algorithm for unidirectional attention. For clarity, we omit attention normalization in this visualization. The algorithm keeps the prefix-sum which is a matrix obtained by summing the outer products of random features corresponding to keys with value-vectors. At each given iteration of the prefix-sum algorithm, random feature vector corresponding query is multiplied the most recent prefix-sun sum (obtained by summing outer- products corresponding preceding tokens) to obtain new row of the matrix AV which is output by the attention mechanism.
For the unidirectional case, our analysis is similar as for the bidirectional case, but this time our goal is to compute (ril(Q' ((K(((((")))) T C without constructing and storing the L-sized matrix (ril(Q' (il(QQ((K) (K')T
(K')
explicitly, where C [V 1z] E Lx++1), In order to do so, observe that S L.
(tril(Q' (KK) GP x Q GP
EG j=1
G... = E Mx(d+1) (1I)
where G, GPS (L++) 3d-tensors. Each slice GPS. is therefore result of a prefix-sum
(or cumulative-sum) operation applied to G.I.p: G.I.p. E₁=1 Gi An efficient algorithm to compute the prefix-sum of I. elements takes O(L) total steps and Ollg L) time when computed in parallel (Ladner & Fischer 1980 Cormen et al. 2009). See Algorithm for the whole approach.
B.2 ORTHOGONAL RANDOM FEATURES EXTENSIONS
As mentioned in the main text, for isotropic Ω (true for most practical applications, including regular attention), instead of sampling w; independently, can orthogonal random features (ORF) Yu
17

---

Published as a conference paper at ICLR 2021
Choromanski et al. 2017: 201 8b): these maintain the marginal distributions of samples Wi while enforcing that different samples are orthogonal. If we need m > d, ORFs still can be used locally within each d d block of W Yu al. 2016).
et al.
2016
ORFs were introduced to reduce the variance of Monte Carlo estimators Yu et al. 2016 Choromanski
al. 2017 2018b 2019a Rowland et al. 2019| Choromanski et al. 201 201 9b) and we showed in the theoretical and experimental sections from the main body that they do indeed lead to more accurate approximations and substantially better downstream results. There exist several variants of the ORF-mechanism and in the main body we discussed only the base one (that we refer to here as regular). Below we briefly review the most efficient ORF mechanisms (based on their strengths and costs) to present the most complete picture.
(1) Regular ORFs [R-ORFs]:
Applies Gaussian orthogonal matrices (Yu et al., 2016). Encodes matrix W of wasmpless (with different rows corresponding to different samples) in ((m)) space. Provides algorithm for computing Wx in O(m)) time for any x RRd. Gives unbiased estimation. Requires one-time ((md²) preprocessing (Gram-Schmidt orthogonalization).
(2) Hadamard/Givens ORFs [H/G-ORFs]: Applies random Hadamard (Choromanski al., 2017) or Givens matrices (Choromanski et al., 2019b). Encodes matrix W in O(m) O(m log(d) (d)) space. Provides algorithm for computing Wx in ((m (og()) time for any x E RR.. Gives small bias (tending to 0 with d + oo).
B.3 TIME AND SPACE COMPLEXITY DETAILED ANALYSIS
We see that variant of bidirectional FAVOR+ using iid samples or R-ORF's has O(md + Ld mL) space complexity as opposed to O + Ld) space complexity of the baseline. Unidirectional FAVOR+ using fast prefix-sum pre-oomputation in parallel (Ladner & Fischer 1980; Cormen et al.) 2009) has (mL)) space complexity to store GS which can be reduced to D(mdd Ld mL) by running simple (though non-parallel in T.) aggregation of GPS without storing the whole tensor k,1!
GPS in memory. From Subsee. B.2. we know that if instead we use G-ORFs, then space complexity is reduced to (m log(d) + La + m) and if the H-ORFs mechanism is used, then space is further reduced to O(m+ mL) = D(Ld+mL) Thus for m, < L all our variants provide substantial space complexity improvements since they do not need store the attention matrix explicitly.
The time complexity of Algorithm is (Lmd) (note that constructing Q' and K' can be done in time ((Lmd)). Note that the time complexity of our method is much lower than OLL2d) of the baseline for >> m.
As explained in Subsec. B.2 the R-ORF mechanism incurs extra one-time ((md)) cost (negligible compared to the O term for T. >> d). H-ORFs or G-ORFs do not have this cost, and when FAVOR+ uses them, computing Q' and K' be conducted in time OL (og())d) as opposed to ((mmd) (see: Subsec. B.2). Thus even though H/G-ORFs do not change the asymptotic time complexity, they improve the constant factor from the leading term. This might play an important role in training very large models.
The number of random features allows trade-off between computational complexity and the level of approximation: bigger m results in higher computation costs, but also in lower variance of the estimate of A. In the theoretical section from the main body we showed that in practice we can take M = O (d log(d)).
Observe that the FAVOR+ algorithm is Sighly-parallelizabbe and benefits from fast matrix multiplica- tion and broadcasted operations on GPUs or TPUs.
18

---

Published as a conference paper at ICLR 2021
C EXPERIMENTAL DETAILS FOR PROTEIN MODELING TASKS
C.1 TREMBL DATASET
Length Statistics
Dataset
Set Name
Count
Min
Max
Mean
STD
Median
Train
104,863,744
74,488
353.09
311.16
289.00
Valid
102,400
11,274
353.62
307.42
289.00
TrEMBL
Test
1,033,216
32,278
353.96
312.23
289.00
OOD
29,696
24
4,208
330.96
269.86
200.00
TrEMBL
Train 4,532,224
8,192
8,192
8, 192
0
8, 1,192
(concat)
Valid
4,096
Table Statistics for the TrEMBL single sequence and the long sequence task.
We used the TrEMBL datase which contains 139,394,261 sequences of which 1 06,000,080 are unique. While the training dataset appears smaller than the one used in Madani et (Madani et al.) 2020), we argue that it includes most of the relevant sequences. Specifically, the TrEMBL dataset consists of the subset of UniProtKB sequences that have been computationally analyzed but not manually curated, and accounts for Ra 99.5% of the total number of sequences in the UniProtKB d
Following the methodology described in Madani et al. Madani et al. 2020), we used both an OOD-Test set, where a selected subset of Pfam families are held-out valuation, and an IID split, where the remaining protein sequences split randomly into train. valid, and test tests. We held-out the following protein families (PF18369, PF04680, PF17988, PF12325, PF03272, PF03938, PF17724, PF10696, 1968, PF04153, PF06173, PF12378, PF04420. PF10841, PF06917, PF03492, PF06905, PF15340, PF1 7055, PF05318), which resulted in 29.696 OOD sequences. We note that, due deduplication and potential THEMBL version mismatch, our OOD-Test set does match exactly the one in Madani et al. (Madani et al.) 2020). We also note that this OOD-Test selection methodology does not guarantee that the evaluation sequences are within a minimum distance from the sequences used during training. In future work, we will include rigorous distance based splits.
The statistics for the resulting dataset splits are reported in Table In the standard sequence modeling task, given the length statistics that are reported in the table, we clip single sequences to maximum length 1024, which results in few sequences being truncated significantly.
In the long sequence task, the training and validation sets are obtained by concatenating the sequences, separated by an end-of-sequence token, and grouping the resulting chain into non-overlapping sequences of length I, = 8192.
C.2 EMPIRICAL BASELINE
1333
LS5
007
0007
0005
202
Tymeen
Tyars
Figure 9: Visualization of the estimated empirical distribution for the 20 standard amino acids, colored by their class. Note the consistency with the statistics on the TrEMBL web page.
random bascline, with uniform probability across all the vocabulary tokens at every position, has accuracy 5% (when including only the 20 standard amino acids) and 4% (when also including the 5 anomalous amino acids (Consortium 2019).. However, the empirical frequencies of the various
https://www.uniproo.orr/////
19

---

Published as a conference paper at ICLR 2021
amino acids in our dataset may be far from uniform, also consider empirical baseline where the amino acid probabilities are proportional to their empirical frequencies in the training set.
Figure 9 shows the estimated empirical distribution. We both the standard and anomalous amino acids, and we crop sequences to length 1024 to match the data processing performed for the Transformer models. The figure shows only the 20 standard amino acids, colored by their class, for comparison with the visualization on the TrEMBL web page
C.3
TABULAR RESULTS
Table 2
contains the results on the single protein sequence modeling task (L 1024). We report accuracy and perplexity defined in Appendix A
Model Type
Set Name
Model
Accuracy
Perplexity
Empirical Baseline
9.92
17.80
Test
Transformer
30.80
9.37
Performer (generalized)
31.58
9.17
UNI
Empirical Baseline
9.07
17.93
OOD
Transformer
19.70
13.20
Performer (generalized)
18.44
13.63
Transformer
33.32
9.22
Test
Performer (generalized)
36.09
8.36
Performer (softmax)
33.00
9.24
BID
Transformer
25.07
12.09
OOD
Performer (generalized)
24.10
12.26
Performer (softmax)
23.48
12.41
Table 2: Results on single protein sequence modeling (L = 1024). We note that the empirical baseline results are applicable to both the unidirectional (UNI) and bidirectional (BID) models.
C.4 ATTENTION MATRIX ILLUSTRATION
In this section we illustrate the attention matrices produced by a Performer model. We focus on the bidirectional case and choose one Performer model trained on the standard single-sequence TrEMBL task for over 500K steps. The same analysis can be applied to unidirectional Performers as well.
We note that while the Transformer model instantiates the attention matrix in order compute the attention output that incorporates the (queries Q, keys K, values V/) triplet (see Eq. 1 in the main paper), the FAVOR mechanism returns the attention output directly (see Algorithm 1 To account for this discrepancy. we extract the attention matrices by applying each attention mechanism twice: once on each original (Q, K, V) triple to obtain the attention output, and once on a modified (Q, K, Vo) triple, where Vo contains one-hot indicators for each position index, to obtain the attention matrix. The choice of v ensures that the dimension of the attention output equal the sequence length, and that a non-zero output on a dimension i can only arise from a non-zero attention weight to the 7th sequence position. Indeed, in the Transformer case, when comparing the output of this procedure with the instantiated attention matrix, the outputs match.
Attention matrix example. We start by visualizing the attention matrix for an individual protein sequence. We use the BPT1 BOOBINIINN protein sequence one of the most extensively studied globular proteins, which contains 100 amino acids. In Figure 10 we show the attention matrices for the first 4 layers. Note that many heads show a diagonal pattern, where each node attends to its neighbors, and some heads show a vertical pattern, where each head attends to the same fixed positions. These patterns consistent with the patterns found in Transformer models trained on natural language
20

---

Published as a conference paper at ICLR 2021
highlight these attention patterns by focusing the first 25 tokens, and in Figure 11 we illustrate in more detail two attention heads.
(Kovaleva et al. 2019). In Figure 12
Amino acid similarity. Furthermore, the attention matrices produced by the Performer model, as described Vig et al. (Vig et al. 2020). We aggregate the attention matrix across 800 sequences. The resulting similarity matrix is illustrated in Figure I3 Note that the Performer recognizes highly similar amino acid pairs such as (D, E) and (F, Y).
analyze the amino-acid similarity matrix estimated from
We show the attention matrices for the first layers and all 8 heads (cach row is layer, each column head index, each cell contains the attention matrix across the entire BPT1 BPTI_BOVIN that many heads vertical pattern, where each head attends to the same fixed positions.
Figure 10:
BOVIN protein sequence). Note diagonal pattern, where each node attends to its neighbors, and some heads show
145
230
d15
Figure 11:
We illustrate in more detail attention heads. The sub-figures correspond respectively (1)
Head 1-2 (second layer, third head), Head (fifth layer. second head). Note the block attention in Head 1-2 and the vertical attention (to the start token ('M') and the 85th token ('C")) )) in Head 4-1.
21

---

Published as a conference paper at ICLR 2021
Figure 12: We highlight the attention patterns by restricting attention to the first 25 tokens (note that we do
not renormalize the attention these tokens). The illustration based Vig al. (Vig 2019 Vig & Belinkov 2019). Note that, similar to prior work on protein Transformers (Madani 2020). the attention matrices include both local and global patterns.
1.0
0.8
0.6
0.4
0.2
0.0
0.0
ACDE G I
A
Figure 13: Amino acid similarity matrix estimated from attention matrices aggregated across small subset of sequences, as described in Vig et al. (Vig et al. (2020). The sub-figures correspond respectively to: (1) the normalized BL.OSUM matrix, (2) the amino acid similarity estimated via trained Performer model. Note that the Performer recognizes highly similar amino acid pairs such (D, E) and (F, Y).
22

---

Published as a conference paper at ICLR 2021
D EXTENDED APPROXIMATION AND COMPARISON RESULTS
D.1 BACKWARDS COMPATIBILITY ERROR PROPAGATION
Although mentioned previously (Sec. that the Performer with additional finetuning is backwards compatible with the Transformer, demonstrate below in Fig. [4 that error propagation due to non- attention components the Transformer is one of the primary reasons that pretrained Transformer weights cannot be immediately used for inference the corresponding Performer.
Performer Approx. Error
0.35
0.30
Output MSE
0.25
0.20
0.15
0.10
0.05
1
2
3
4
5
6
8
Num. Layers
Figure 14: Output approximation errors between a vanilla Transformer and a Performer (with orthogonal features) for varying numbers of layers.
D.2 APPROXIMATE SOFTMAX EXTENDED PROPERTIES
We show the following properties of softmax approximation, in Fig. 15
Redrawing: While the benefits of redrawing features was shown in Subsec. 4.3 of the main body of the paper, we also its benefits when there multiple layers with large scale (16x16 TPU-v2) training.
While have shown TrEMBL that Performer with generalized ReLU attention outperforms softmax, we also show that approximate softmax attention can still be solid choice, for example on ImageNet64 (U), After 100K steps of training, the Performer-Rele ReLU, Serformer-Softtmaa and Performer-Softmax (SMREG) variants achieve respectively, 3.67, 3.69, 3.67 BPD.
Unidirectional:
Instability of Trigonometric Features: We see the full view of the unstable training curve when using Trigonometric softmax.
TrEMBL. (B) (L= 1024)
ImageNet64, Layer Regular Model
120K
PG-19 (B). Layer Regular Model fos
Performen
Performe POS SOTTHON SMR2S
200K
80K
5.0
Perpleaity
40K
Tardomer
Performen SOFTHAX No Sodron Sted
OK
100K
200K 200K
3.5gx
75K
50K
75K
300K
0K
20%
40KC 60K
D0k
100K
Figure 15: Best viewed zoomed in. Left: The importance of redrawing features. If redrawing is not used, an "unlucky" set of random features may cause training degradation, shown by the early-stoppe curve with Seed 1, while 'lucky' set of random features may cause no issue, shown by the curve with Seed 2. Redrawing allows the training to correct itself, seen at the black vertical line. Middle: Using the same 8x8 TPU-v2 compute and same 6-layer standard model, approximate softmax with positive features achieves the same result as generalized ReLU attention. Right: Zoomed out view of right subfigure of Fig. showing that Trigonometric softmax causes very unstable training behaviors.
D.3 GENERALIZED ATTENTION
We investigated Generalized Attention mechanisms (mentioned in Sec. 2.2 TrEMBL when = 512 for various kernel functions. This is similar to (Tsai et al. 2019) which also experiments with various attention kernels for natural language. Using hyperparameter sweeps across multiple
23

---

Published as a conference paper at ICLR 2021
variables in FAVOR, we compared several kernels and also renormalization on/off (Fig. [6 and Fig. 17). where Renormalize corresponds to applying D operator in attention, as for the standard mechanism, though we noticed that disabling it does not necessarily hurt accuracy) to produce the best training configuration for the Performer. We note that the effective batch size slightly affects the rankings (as shown by the difference between 2x2 and 4x4 TPU runs) we by default the generalized ReLU kernel with other default hyperparameters shown in Appendix [A] as we observed that they are empirically optimal for large batch size runs (i.e. 8x8 or 16x16 TPU's).
Kernels
Renormalization
Orthogonality
agrraic
0.000001
Orhe
0.00001
ID
0.003
0.0001
0.26
Accuracy
Drn
0.2
0.1
10a 102
10 10 10 102
10 10f 10a
10a 102
102t
10
10 10
Gradient Steps
Figure 16: To emphasize the highest accuracy runs but also show the NaN issues with certain kernels which caused runs to stop early, we set both x and y axes to be log-scale. We tested kernels defined by different functions f (see: Sec. 2.2p: sigmoid, exponential, ReLU, absolute, gelu, cosine (original softmax approximation), tanh, and identity. All training runs were performed on 2x2 TPU-v2's, s, 128 batch size per device.
Kernels
Renormalization
Orthogonality
0.294
Alg mo
0,900001
Ortho
exp
0.001
ID
0.284
qeuu
20th
Accuracy
0.2
0.01
10a
102
10
10a
10a
10a
10.
10
10a
10 102
10a
10*
10a
10a
Gradient Steps
Figure 17: We also performed similar setup as Fig. 16 for 4x4 TPU-v2's.
D.4 COMPARISON WITH LINEAR TRANSFORMER
We use the attention implementation of the Linear Transformer from (Katharopoulos al. 2020). which mainly involves setting our feature map ((x) = elu(x) + 1, where elu (x) is the shifted-eLU
function from (Clevert al. 2016).
TrEMBL, 36 Layers (U) (L- = 1024)
TrEMBL, 36 Layers (B) (L = 1024)
Performer-BB RELU
0.35
Transfermer [U]
0.25
Liner Transformer (U). Seeds
0.30
Accurscy
0.25
0.20
0.20
Performer- RCUU (8)
0.15
Performer SOP IMAX (8)
Trensformer (B)
0.10
Liwwr Transformen (5)
OK
10K
25K
50K
75K
100K
125K
Steps
Steps
Figure 18: Left: In the unidirectional 36-ProGen setting, we ran 3 seeds of the Linear Transformer, and found that all 3 seeds produced exploding gradients very early on, stopping the training run. Right: The Linear Transformer in the bidirectional setting also produced an exploding gradient in the middle of training, near 125K steps. Exploding gradients can be evidenced by the sharp drop in train accuracy right before a NaN error.
For the sake of faimess and to prevent confounding results, while (Katharopoulos et al. 2020) also uses the GeLU nonlinearity for the MLPs in the Linear Transformer, we instead use the original ReL.U nonlinearity. We also used the exact same training hyperparameters as Performer-Rel.. on
24

---

Published as a conference paper at ICLR 2021
exact ProGen setting from Fig. 6 Ultimately, empirically found that the Linear Transformer possessed numerical instability during training via unstable training curves, ultimately stopping training by producing exploding gradients (NaNs) (Fig. I8..
D.5 LONG RANGE ARENA
Performers compared against many additional (scalable and not scalable) methods not included in our paper: Local Attention, Sparse Attention, Longformer, Sinkhorn Transformer, Synthesizer; Big Bird and the aforementioned Linear Transformer on challenging long range context tasks in the Long Range Arena (Tay et al., 2021), with Fig. 19 displaying the original paper's results. Performers obtain the largest LRA (Long Range Arena) score among all tested scalable Transformers methods (which we define by having speed of > 100 examples/sec).
Tasks used for comparison include: (1) longer variation of the standard ListOps task proposed in (Nangia & Bowman, (0188) (2) byte-level text classification using real-world data, (3) byte-level document retrieval, (4) image classification on sequences of pixels, and (5) Pathfinder task (long- range spatial dependency problem). In the Long Range Arena paper, the authors found that all models do not learn anything on Path-X task (denoted by FAIL), contrary to the Pathfinder task, which shows that increasing the sequence length can cause seriously difficulties for model training.
LieOpp
Trat
Reriewl
Image
Pattfindee
Path-X
Avg
Transformer
36.37
64.27
$7.46
71.40
$4.30
56
Local Amention
15:32
$2.98
53.38
11:05
05:55
EAIL
16.06
Big Transformer
Sparse Trans.
17.07
63.58
59.59
44.34
71.71
FAIL
51.24
Lenffeeeer
$6.90
Retormer
Sinthom Thuns.
33,67
37.37
61.30
56.10
$3.83
41.23
68.50
07.45
FAIL
FAIL
50.7
$1.39
Synthesizer
Syaleoizer
Higttrd
36.05
36.99
61.68
64.12
$4.67
41.61
69.45
24.87
FAIL
FAIL
52.58
58.01
LRA Score
Linformer
Performer
Perfomeer
77.05
Reformer
Sinkham
Taak Avg (845
(9.)))
61 (4.6)
55 (2.6)
41 (1.80
72.537)
PAIL
52 (2.4)
50
Linear Transformer
Stepe per saain
Peak Menory
Usage: (GBB)
Model
1K
2K
3K
4X
28
3K
4K
Travforrer
2.3
Local Aneniion
22
(34(3-7
7.4 D5)
(5.3))
0.42
0.76
1.05
1.37
Local Attention
Linfous
(3.22.2)
9.1 (1.5k)
8.5 (3. Tx)
(5.5k)
0.37
0.55
0.99
0.39
46
Refoemer
2.400.333
2.2.0..xl
1.5 (0.7x)
1.108.0)
0.99
1.53
2.26
Synthesier
8.7(1.30)
(3.30
1900.00
BigBidd
0..30.3
3.9 (0.8k)
1.7 (1.24)
(.5..
0.77
1.40
2.18
2.58
44
lince Trmes
91 (1-15)
(1.0k
5,6 (3.7K)
(.
0.37
0.57
0.80
1,01
100
250
330
Pertomeer
9.5 20)
9/4 (3.0k
8.7 03 80
80
0.37
0.59
0.82
1.06
Speed (examples per sec)
Figure 19: Upper Table: Results on Long-RRange Arena benchmark. Best model is in boldface and
second best is underlined. Lower Table: Benchmark results of all X-former models with consistent batch size of 32 across all models. The authors report relative speed increase/decrease in comparison with the vanilla Transformer in brackets besides the steps per second. Memory usage refers to per device memory usage across each TPU device. Benchmarks are run on 4x4 TPU-v3 chips. Right Fig: Performance (y-axis), speed (x-axis), and memory footprint (size of the circles) of different models.
25

---

Published as a conference paper at ICLR 2021
E COMPUTATION COSTS EXTENDED RESULTS
In this subsection, we empirically measure computational costs in terms wall clock time on forward and backward passes for three scenarios in Fig. 20
1. Performer, with varying number of layers. We show that our method can scale up to (but not necessarily limited to) even 20 layers.
2. Attention time complexities when comparing standard attention (from Transformer) and FAVOR (from Performer). Note that the maximum memory size here is not reflective of the maximum memory size an actual model (shown below), as this benchmark requires computing explicit tensors (causing memory increases) in Jax, while a model does not.
3. Time complexities when comparing the Transformer and Performer models. "X" (OPT) denotes the maximum possible speedup achievable, when attention simply returns the V- vector, showing that the Performer is nearly optimal. We see that the maximum possible power of length allowed on V100 GPU (I6GB) 215 32768 using regular dimensions.
Since some of the computational bottleneck in the Transformer may originate from the ex- tra feed-forward layers (Kitaev et al., 2020), we also benchmark the "Small" version, ie. (Nheads Rayeras djf, d) (1, 6, 64, 64) as well, when the attention component is the dominant source of computation and memory. We remind the reader that the "Regular" version consists of (Theads; Mayers, df,, d) (8, 6, 2048, 512).
Performer Forward Timing (Small)
Performer Backward Timing (Small)
logy() (sec)
-2
Layars
Layors
Shou Lind
-10
-10
12
-12
8
10
12
14
16
18
20
22
8
10
12
14
16
18
log,lL)
(0g,((
Performer Fonward Timing (Regular)
Performer Backward Timing (Regular)
Layers
Lavers
Layers
-2
24 Layora
Layers
-2
Meggi7)
1-Sapa
-8
-8
-10
-10
10
12
16
18
10
12
16
(ogs(()
logs(L)
Attention Forward Timing (Small)
Attention Backward Timing (Small)
Batch=1
Batch-2
1oge/73 (sec)
Betch=4
Batch-B
-6
Transformer
-8
Performer
-8
-10
-10
-12
-12
20
15.0
22.5
logy(L)
log.1L]
Attention Forward Timing (Regular)
Attention Backward Timing (Regular)
-2
Batch-1
Batch=2
togut)) (sec)
Batch=4
Transtormar
-6
Performer
-8
-10
-12
logat11
(00x(L)
Figure 20: Captions (1) and (2) for each 2x2 subfigure mentioned above.
26

---

Published as a conference paper at ICLR 2021
Model Forward Timing (Small)
Model Backward Timing (Smail)
Dwich-1
Batch-2
Batch-4
Megyth (sec)
Fransformor
Performer
*
OPT 1PT1881
-8
-6
-10
-8
6
12
14
18
20
10
14
16
18
logs(L)
(ogs(L)
Model Forward Timing (Regular)
Model Backward Timing (Regular)
Barch-1
Beichn?
-2
Batche4
(ogy(7) (sec)
-4
Trandformer
Performer
OPT.BATCH-1
-6
-10
18
12
Figure 21: Caption (3) for this 2x2 subfigure mentioned above.
27

---

Published as a conference paper at ICLR 2021
F THEORETICAL RESULTS
We provide here the proofs of all theoretical results presented the paper.
F.1 PROOF OF LEMMA !
Proof. We first deduce that for any b R
yM exp(x* y) exp([11ze|1/2) exp([1æ + y1l exp(-1ly11//2)
Next, let E TRd. We the fact that
(2x) (2x)- 4/2 / expp(-11w- cll2/2)dw
for any E Ied and derive:
exp(1 x + y1ll = (2x) exp( (1x + yy1l/22 / exp(-Ilw - (x + y) /2)
(2/3-
exp(-1|w 2/2 w TT + y) || + y11 + 1| + u11 1//222
= (2x) (2x)-4/2
4/2 = (2x) 4/2
exp(-1|w||| + wT (x + TT(x
y))
exp(w" x) exp(w y)
= lexp: x))
exp(w" y)].
That completes the proof of the first part of the lemma. An identity involving hyperbolic cosine function is implied by the fact that for every u E Fd and w -N(0, La) the following is true:
[B[( 1)31
1 [[ 1)22) E[( B[(-w' w Tn)ay
E[exp(w T u)]
(12)
(2))! 2
(2i)!
The cancellation the odd moments [[(w" T u) 1 follows directly from the fact that w is taken from the isotropic distribution (i.e. distribution with pdf function constant each sphere). That completes the proof.
F.2 PROOF OF LEMMA 2
Proof. Denote:
and A x y. Note that by using standard trigonometrie identities (and the fact that the variance of the sum of independent random variables is the sum of variances of those random variables), we can get the following for (N La):
trig (x. = exp exp[[x]] |x|| + + Var(c A)).
1
MSE
(13)
Using the fact that (see: Lemma 1 in Yu et al. 2016;; note that in that lemma they use notation:
for what we denote as: |
Var(cos(w T A)) 71 exp(-11411)1)
(14)
we obtain:
trig (x) MSE(SM
1
exp( ||
[xp[[|x|]
+ +|y (1 exp exp(-1AA13))
2m
(15)
1
exp
exp([lz||2)SM)))
| (x, exp( exp(-1AA))))
2m
which completes the first part of the proof. To obtain the formula for: MSE(SM, (x, xx) notice first that:
(exp)) T 2)] ] = exp( | 2 )
(16)
28

---

Published as a conference paper at ICLR 2021
The above immediately follows from the fact that positive random feature maps provide unbiased estimation of the softmax -eenal,eerne,, thus the following is true:
SM((x y) = exp( |x| + 1lyl?
(ox, z)].
(17)
2
Therefore we obtain:
MSE (x, y)) m
Varl))) [z])
1 exp( (||x + )lyll (2 T 2)] (E/cxp(w> "z)))²)=
(18)
m
exp( (||x + |y
- exp(1 z )),
where the last equality follows from Equation 16 Therefore we have:
1
MSE(SM (x,
(|| + [yyl13) exp(Ila|) exp(llx))))exp((1131)
m
(19)
(x)(
m
Finally,
hyp+ MSE(SM,, (x) = 4m
1
exp( ||x| + 1yl )²( (Var (T)) 2)) + Var (exp( -w 22))+
exp( ||x|2 + || Var 12))+
2Cov(exp(w T 2)), exp(-xe 2))))
4m
2Cov 2Cov(e T 2)), exp( -ww 2))))) 2m
(||[=(1|x|' + [1x||
(Var(exp(w) T 2)) + 1 - (E(exp((w) T 2)))²) = 2m
ppp(-(1|[] + 11yl13))
(exp(2)|z(()) - exp(11z]2) + exp((lz|()) 2m
exp( (|||= +
- 1)a
1 (1 exp(-1<> M (x. y))-
(20)
In the chain of equalities above we used the fact that random variables exp(w z) and exp( z) have the same distribution. This is true since and -w have the same distribution (w is Gaussian). That completes the proof.
F.3 PROOF OF THEOREM
Proof. Let Rd be respectively query/key. Note that from the definition of SMREG(X.y) we have for z = x yy:
|x||
1
2
20
SMREG(X, y) = exp(
|z|| 2k & B
e )2k
(21)
2k)
|w]|2
where def (1.0...., T) E Rdd To obtain the above we used the fact that W/ Ia) is isotropic (that particular implies zeroing of the even terms in the Taylor expansion).
Let us denote: A(k, d) EL ( e1 26 It turns out that:
(2k 1)!!
A(2k, d) = (d + 2k:
(22)
2) 2k
d
The proof of that fact can be found in the supplement of (Choromanski et al. 2018b), yet we provide it below for completeness and the convenience of the Reader:
29

---

Published as a conference paper at ICLR 2021
Lemma 3. Expression A(2k, d) satisfies the following k E N:
(2k: - -1)! 1)!!
A(2k, d) (d + 2k 2) 2 4) ... .d.-d'
(23)
Proof. Note first that for d 2 2 the density function pa( of the angle between a vector E R
chosen uniformly at random from the unit sphere and e1 is given by the following formula:
sin sind~2(0)
Pa(O) = So sin d- 2(0) d0
(24)
Let us denote: F(k,d) def So cos sin si (0)d0. Using partial integration, get:
[" cos co5)(0) (0) (O) { cos (0) (0)sinn (O)sn((O)))
cos*~1(0) sin (0)l J/- sin(0)((k cos * (0) sin sin sin (0)+
(25)
d cos (0) sin (0))
Thus conclude that: F(k,d) d)
k-1 J+I F(k 2, + 2). Therefore we have:
(d + () + 3) .... (d + 2k 1) /.
(2k - 1)!!
F(2k, d)
sin 2g (0)c0.
(26)
again conduct partial integration and get:
sin d (0)d0 (8))(1())
(27)
d
Jo sin sind-2 d-2 (0)d0 d-1 1
/- sin (0)d0
d
Therefore conclude that:
(2k 1)!!
d+ 2k: - 3 d 2h
A(2k,d) 2-30- 2-244
(d. 1). ..... ..... (d+2k 3) d 2h: - 2d + 2h
(28)
(2k: 1)!!
(d 2k
2) 2K 4) d'
which completes the proof.
Applying the above lemma, get:
|x|
2 + yy
2
(2k: 1)!!
AMREG((,y)) y) = eep((-
|z|] 2h dh
2
(2k))
(d + 2k 2)( + 2k
4) d
k=
||x|| 2 ||y|1
2 #
= exxp(
f(k, d),
2
k!
t 0
(29)
where w |x] and f(k, d)
d (d42k d)...dd
Thus obtain:
SMREG((xy) SM(x,y) -
1
k! f(k d).
(30)
k=0
Note first that for > we have: f(k, d) < 1, thus:
SMRFG(x.y) SM(x.y).
(31)
30

---

Published as a conference paper at ICLR 2021
We also have for 1 = d3:
SMREG(X,))
w
0c
w
Σ k=0
[(k, d) - Σ
[(k d)
yM
h!!
k!
k=1+1
w
wh k!
w
(32)
f(1. d].
+
f(k, d) 2 f(1, ((1,d)(1 - -
k!
k!
k=
km1+1
R=1+1
f(1,d)(1 P(Po( (]),
where Po(w) stands for the random variable of Poisson distribution with parameter w. Therefore we get for = In( )
SMREG(X, y)
21
2
) PP(P(()) > 1] 2
SM(x,y)
d
21 - [] )) P((Po(w) 2 21)
exp(IIn(1
d
20
(2
exp
(1 - Pleeep(PPo(())) U1) 2 11)
(33)
2
1
cxp(
+ 0
d3
di
2
1
exp(
+
() cxp( 1(t 1))),
d
d3
where the last equality is implied by the formula for the Laplace Transform for the Poisson random variable:
E[exp((Po(w)) = exp(x(exp(() - -)).
(34)
Notice that: lell (M(M (2ak 5 2nnn.C). We conclude that:
SMREG(X. y)
2 dt
1 + ol d$
)) 1111 0- 2e ln(()()))))))) ln(c))==)) =1 )
db
2 dt
+ ol 1
2 (1
(35)
yM(x,y) That completes the proof.
dt
F.4 PROOFS OF THEOREM 2 THEOREM 3 & BEAUTIFUL FUNCTIONS
We will provide here much more general theoretical results which will imply Theorem We need the following definition:
3 and Theorem
2
Definition 1. We say that function F: 1TT + IR is beautiful if F can be expressed as:
Fa. = Eunnil(w T 2)],
(36)
for a probabilistic isotropic distribution Q, and where g R - IR is entire function with non- negative power-series coefficients (i.e. g(x) = Dieo a,x" for every T E & and with a; Z 0 foror i= 0, I...). In the formula above we assume that the expectation the RHS exists.
Interestingly, beautiful functions can be used to define softmax and consequently, Gaussian kernels (both standard and regularized), leading our PRF mechanism presented in the main body of the paper, as we explain below.
Remark 1. If one takes 92 = (0, La (note that (0 Ia) is isotropic) and x + exp(x) (such is clearly entire with nonnegative power-series coefficient) then the following is true for z x + y:
SM(x,y) = exp( || + 11yl2 (z).(z).
(37)
2
Similarly: SMREG(x,y) = exp(- ||x|| + w1l
) Forez.g(2 where Droy stands for the distribution corresponding Haar measure the sphere of radius v (which is clearly isotropic). Therefore general concentration results for Monte Carlo estimators of beautiful functions immediately imply corresponding results for the (standard and regularized) softmax (and thus also Gaussian) kernel.
31

---

Published as a conference paper at ICLR 2021
We will consider two estimators of the beautiful functions from Definition that directly lead (through R to: PRF-based approximation of the softmax- and its enhanced version with orthogonal features. Standard Monte Carlo estimator samples independently wiid .... lid lid where m stands for the number of samples and then computes:
12
Fild (z)
def 1
M7 Σ )((!!l))))) Tz).
(38)
Orthogonal Monte Carlo estimator samples witt ...., "..,, w (m S d) in such way that marginally with have: wi ~ S2, but T = O for o j (such an orthogonal ensemble can be always created
is isotropic, we already mentioned in the main body of the paper). We define:
port
def
(((
(39)
F.4.1 ORTHIOGNNALITY UNIVERSALLY IMPROVES CONCENTRATION
Denote by M2) = | a moment generating function of the random variable Z. Note first that estimators of beautiful functions based on standard Monte Carlo procedure using independent vectors with guarantee strong concentration bounds since independent wiS provide a way to obtain exponentially small upper bounds on failure probabilities through moment generating functions. We summarize this classic observation which is standard application of Markov's 's Inequality below.
Lemma 4. Consider an estimator Fud (z) of the beautiful function F evaluated at z. Then the following holds for any a >> F(z):
P Filld (z) > a] S emCx(())
(40)
where X g(( 2). DD and Lz stands for Legendre Transform of the random variable Z d0 defined as: Lz((a) = supa>o log( Mz(3j) Furthermore, Lx (a) >
The above result provides us with exponentially small (in Legendre Transform) upper bounds on tail probabilities for the standard estimator. Below we provide our two main theoretical results.
Theorem (orthogonality provides smaller tails). If Fn.9 is a beautiful function then the following holds for m < d, X in Lemma 4 and any F(z):
d
Fot P
> al S mmxx(a)
(41)
d 2
This result shows that features obtained from the ensembles of pairwise orthogonal random vectors provide exponentially small bounds on tail probabilities and that these bounds are strictly better than for estimators using unstructured features. Furthermore, the result is universal, i.e. holds for any dimensionality not just asymptotically for d large enough.
We also obtain similar result regarding mean squared enrors (MSEs) of the considered estimators:
Theorem 6. If Fo.g is a function then the following holds for m S d:
1 2
MSE Fort (2)) < MSE Flid (z)) (1
F2 5. (z).
(42)
m
d+2
As before, orthogonal estimator leads to better concentration results and as before, this is the for any > O, not only asymptotically for large enough d.
Note that from what have said above, Theorem 2 and Theorem follow immediately from Theorem and Theorem respectively.
Thus in the remainder of this section will prove Theorem and Theorem
32

---

Published as a conference paper at ICLR 2021
F.4.2 PROOF OF THEOREM 5
Proof. Note that by the analogous application of Markov's Inequality as Lemma 4 we get:
a) v Ex
+... Xm)
Pl
Port (2)t(()) m (z))
(43)
emm
where we have: X ((wfrt) 2).
T2). We see that it suffices to show that for any 0 >> 0 the following
+..+ holds: E[c"(x]" +.. xxx)) < E[c X1 +... xmx 1.
We have:
Xont
O +..
E
X
X x
j!
0
2m
(44)
X
Porty p|
=
Σ
B Σ
dh Im
(Xort)
j=0 j!
j=0
1=1
(m JES,
where S ((i
EN ... x N: J
Im 2 0,11 ... +........ im = 5} and for some positive
constants c(ji...,Ja
Thus we have:
Ele [(PP' X
Σ
cti Jm. Dn (Xort))m)
(45)
j
1=0 WESndd))ES
Similarly, we get:
0
Ele 1xx +
xtid
Σ
Σ
cii,, Jm) ... (Xidd jmm.
(46)
j!
m
1=0
Therefore we get:
A O 0(x +x.
x
E(e E(c +... X
9(xpm
Σ
(m) (E((x (Xidd)m)) - B[(xxregá ............. (X pim]
j!
== (J
(47)
Note first that using the fact that f entire, we can rewrite each X as:
X ( T2)"
(48)
where f(x) Ex U,t" and ap al 2 0. Similarly,
iid Σ TT)) Tz)"
(49)
0
By plugging in the above formulae for Xort and Xiid int the formula for A and expanding power- expressions, we obtain:
Σ
cm......JJm)
Σ
Aldi,..,, dmn);
(50)
for some ordered subsets of indices (with potentially repeating entries) D(ji..... Jm) (exact formula for those can be given but we do not need it to complete the proof and since it is technical, it would unnecessarily complicate the proof so we skip it) and (ddi,.., dmm.) defined as:
Å A(di,..., dm,) B(((w]]d) Tz)d= 22)......
.... ((w)l) Tt) 2) dm ] E[((wprt)T Tz)a ......... ((w [ T2) (51)
33

---

Published as a conference paper at ICLR 2021
Our next goal is to re-write the formula for Á (d1, .., dm). Denote:
= ((w)t) 2) ((wm)) 2)dm.
(52)
Observe that Y has the same distribution as Y' defined as:
g == (e1 Igll
12
(e, [2112)mm (1|uptt1/2 11) "id..... ".....///// (12m)1 1l2)dmm
(53)
where g is Gaussian vector taken from the ((O,, Ia) distribution, independently from:
| 12,,.,, || [
This comes from the fact that for fixed z one can think about the set:
Co
random rotation of the system of m canonical basis vectors: e1. Cm- Thus instead of applying random rotation to: c1,..., Gm: can equivalently randomly rotate vector Randomly rotated vector z has the same distribution as: Tell ||z||2
Now note that lengths of vectors wirt wott chosen independently.
Therefore we obtain:
[[((w zz) ((w (
=)]u]=
(54)
[[(1)97*1122 1[ [12)dm]
E[(e] v) ..... (e,n () ||11/ t.tt
where v Iialia"
Denote
(g), gd) T. Thus we obtain:
E (ww T z2)d
((w T 2z)dm]
[ [[([))r1122)) ..... [ (w [l2)dm]
|| +
9
ym
(55)
t..
9 +
Id
Now let us focus on the second expression from the formula on Â(d1, dmn). We have:
m
[[((wi T T2)d
(ww Lid
2) dm
[[((w lidy T2)d]=
(56)
d,
Hid [|2)
[.:]..... (i [1] [12)dm] ||z 110 t..
EE
g
+
ga
where the first equality comes from the fact that different wildds s are independent and the second one is implied by the analogous analysis to the one conducted above.
We will need the following lemma:
Lemma 5. For every E N such that S n. and every ki, Kg N4 the following holds:
g g
11 B(g
(57)
+ + 9d
E Bl +
Fd
Proof.
Take
| a
where g an independent copy of g. Note that g. We have:
B
E
E E tko.
(58)
0
where the first equality comes from the independence of different elements (2), 2n) and the second equality is implied by the fact that g independent from g.
34

---

Published as a conference paper at ICLR 2021
Therefore have:
91
B[r
E
(59)
2ki
E
gi
+
9a
That completes the proof since z ~ g and g
g.
Note that by Lemma we can rewrite the right expression from the formula on (dd1.... dm) as:
E3
[[[[4] 1 [2) [ (2)
||z|| [d + +m
(60)
El V gf + + 9
dm
The left expression from the formula Â(d1, dm) can be rewritten
Ldd dm) [[[[[[] l12) |2) I E3 [[1w (12) [2)]""]
|| +dm
Elg d+]
(61)
Iá Am
BlVg5 gi ... 9á Elvgi gi
Since marginal distributions of woth and wiln are the same, we can rewrite (dd dn) as:
Í(d dm) L(di,..,dm)L ....)))),,
(62)
where (dd..... dm) is defined as:
d
T(d1, dm) E El
92
(63)
E of ... 9d
We need now few observations regarding A(d1, ..., dm). Note firsr that since odd moments of the Gaussian scalar distribution (0, 1) are zero, Add....,,m, ) is zero if at least of of d; is odd. Furthermore, (dd....., dm)) is trivially zero if but at most one d2 are zero. With our new notation, A can be rewritten as:
0
Σ
c(fi,..,, jm)
Σ
"mi [(.,, dm)),
j!
j=0 (Esssm)ES
(m)...)))
(64)
Note also that we have:
e ((xii + + Xiid)
Σ
05
Σ
cmi
Σ
L(d2, i(d,,....d dm). (65)
j!
Therefore (see: our observations on A(d1, dm))) to complete the proof it suffices to show that: T(h,,.,. dm) S a+2 if at least two: di d, for + j are nonzero and all d, are even.
Lemma 6. The following holds for some # j we have: di, dj and all d, are even:
d
Tddi,.., d
(66)
Proof. Note that 7(d1..... dm) can be rewritten as:
1 Ma (d.) Ha M d.)
7(1, ..., dm))
(67)
where (a( stands for the 8th moment of the x-distribution with degrees of freedom. Note that where I' is the callel Gamma-finction.
ma(1) = 2 I( das)
T(
35

---

Published as a conference paper at ICLR 2021
that Using the fact that: [(n) (n 1)! and [ + = (2n-1)!! V for E N+, it easy for fixed d, the RHS of the Equality 67 is maximized when d; = dj = and dy = 0 for some i * j and & {i Furthermore, straightforward calculations show that in that case the value of the RHS That completes the proof of the Lemma, and consequently, the proof of the
from Equality 67 is d4-2 entire Theorem.
d'
F.4.3 PROOF OF THEOREM 6
Proof. We will use the notation from the proof of Theorem 5 Since both estimators: Fort (z) and
Fiidd (z) are unbiased, we have: MSE(FNM Fort (z)) = Var Fort (z)) and MSE(F Fiid (z)) = Var ( Fiid (z)).
Var(Fort
We have:
Var( Fiid (2)) = BB( (z) E P (2))) [ B[(Fm (2)) (2))) F2(2).
(68)
Similarly,
Var( Foor Var(Fort (2)) (2)) = B[( (2))² FF2(z)
(69)
We have:
m
BI( Fid (2))
D m
[[xiid)
Σ
xix
(70)
m
i
Similarly, we get:
m
[ Port (z))²]
[
+ m
E[x Xort].
(71)
m
Therefore, since marginal distributions of X and X are the same, we have:
MSE MSE(Finde (z)) MSE( Froet (z)
m
[E X X B[x x
2
m
(72)
(1
(E[X iid
x
- E[x x=
Plugging in the formula for X and X, from Equation 48 and Equation 49 and using analysis from the proof of Theorem we obtain:
MSE( Fild (z)) MSE Frort (z)) =
(z))
aran || = [| [1
t,a=0
(73)
E
E 1
(1 7(t. u)).
EIy o +
g
E
gi +
+ 93
for w Ω and (0 1). Thus, using Lemma 6 get:
1
2
8
MSE(F(i(z)) (z)) MSE(F (z)) 2 (1
a:a.||z|11
1 " B[| ||||:|/||BB|Bpp [|][]
m
+
(74)
E
o
+
+ 9 1E 9f + + 9 nd
2
R((r')
= (1 2
at ||
= (1
F2
F0. (z).
B oी
+ 93
d + 2
That completes the proof.
36

---

Published as a conference paper at ICLR 2021
F.5 PROOF OF THEOREM
We showed in the main body of the paper that in contrast to other methods approximating the attention matrix A, our algorithm provides strong concentration guarantees. This is the also for trigonometric random features, yet, as discussed in the main body of the paper, due to attention renormalization and higher variance of the estimation of small entries of the attention matrix, trigonometric mechanism is sub-optimal. We show here that Mopt the optimal number of random projections for the trigonometric orthogonal mechanism for accurate estimation of the attention matrix does not depend on L but only on d. In fact, we prove that if we take Mopt O log(d)) (d)) then with O Ldd log(d))-timme we can approximate A up to any precision, regardless of the number of tokens L. In order provide those guarantees, we leverage recent research the theory of negative dependence for ORFs (Lin et al. 2020).
We prove the more general version of Theorem 4 from the main body of the paper:
Theorem 7 (Uniform convergence for the trigonometric mechanism). Define entries of the attention matrix A as follows: As,j = g(q )K{ q. kj (h(k)) for some g. h R + and where K is radial basis function (RBF) kernel (Choromanski et al. 2018b) with corresponding spectral distribution (e.g. Gaussian kernel for which Ω = N Ia)). Assume that the rows of matrices Q and K taken from a ball (R) of radius R, centered at (i.e. norms of queries and keys are upper- bounded by R). Define
di
Rd and take g" maxxe (g(x)) and h" maxxe [h(x)). Then for any c > 0, g- and the number of random projections m = S2( d. ga log( doll)) for 5d
= Ew [w] the following holds: ||A A S with any constant probubility, where A approximates generalized attention matrix via orthogonal trigonometric random features.
The result holds in particular for regular softmax-attennion for which K is a Gaussian kernel and g(x) 1(x) exp( [xx]2)) In that case Mopt = 2 fz log( 4d R) since a d.
Proof. Let DQ be diagonal matrix with entries of the form: g(q and let Dx be a diagonal matrix with entries of the form: h A(,T). Denote B [K(-) d q.' / kk,] ]a, BL*X, Denote by Å and approximation of the attention matrix obtained from trigonometric orthogonal random features and by B an approximation of matrix B that those random features provide. We rely on Theorem 3 from (Lin et al. 2020). Note that we can apply it in our case, since for RBF kernels the corresponding functions satisfy fi sin(x), fz(x)) cos(x) (thus in particular are bounded). Also, it is not hard to observe (see for instance analysis in Claim from (Rahimí & Recht 2007) that we can take:
L (for L, in Theorem from (Lin et al. 2020)). Using Theorem from (Lin al. 2020).
we conclude that:
B
B|| 0000 5
(75)
with any constant probability as long as m = Q g4) log( (=diamm/M) where a = Efw T w
and M
the diameter of the smallest ball M containing all vectors of the form z Q.
K
Since
dt
1Q:12, |K1 S R, we conclude that ||z||2 S
2R and thus one can take diam(M) = aà
48
We have:
db
/ S 88" " A | [B B)DK|l> ||Dqll>>|BB
(76)
Taking == h- completes the proof.
F.6 DISCUSSION OF THEOREM
As a consequence of Theorem the number m of random projections required to approximate the attention matrix within error is function of data dimensionality d, the parameter and the radius R of the ball within which the queries and keys live:
m v (c, d, R).
The dependence on d and is fairly easy to understand: with a larger dimensionality d we need more random projecctions (on the order of magnitude d logg( (d) to get approximation within error. The dependence on R means that the length of queries and keys cannot grow at a fixed m we want to retain the quality of the approximation. In particular, this means that FAVOR cannot approximate
37

---

Published as a conference paper at ICLR 2021
hard attention sequences of unlimited length with fixed When the sequence length increases, even the standard attention requires longer and longer vectors to make the softmax concentrated enough to pick single elements. Nevertheless, as seen in experiments, this limitation does not manifest itself in practice at the lengths we experimented with.
38


Reviews:
[{"review_id": "hHGju8kG9_0", "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": "This is a solid paper that presents a computationally less expensive, unbiased, low-variance estimator of the Transformer architecture. \n\n**Strengths**:\n1. The authors provide mathematical guarantees for the suggested estimator.\n2. The estimator (called FAVOR+) seems a more scalable replacement for regular attention.\n3. Theory is empirically verified in a variety of experimental settings.\n\n**Question**: \nIn Introduction, you mention that your method does not rely on low-rankness, but doesn't your equation (4) as well as Fig. 1 assume that the attention matrix $\\mathbf{A}$ allows low-rank decomposition? (I suppose $r<d$)\n\n**Suggestion**:\nWill it be possible to compare a pre-trained Performer-based encoder (like BERT) against its competitors (Linformer, Reformer) on a [GLUE benchmark](https://openreview.net/pdf?id=rJ4km2R5t7) or on [probing tasks](https://openreview.net/forum?id=SJzSgnRcKX)? I see that you compare them (Per-/Lin-/Re-former) against the original Transformer on the pre-training task itself (Fig. 5), but I am wondering how well the pre-training performance transfers to downstream tasks in each case.\n\n**Minor issues**:\n1. In the proof of Lemma 1 you first prove that $\\exp(\\left\\lVert\\mathbf{x}+\\mathbf{y}\\right\\rVert^2/2)=\\mathbb{E}_{\\omega\\sim\\mathcal{N}(\\mathbf{0},\\mathbf{I})}[\\exp(\\omega^\\top\\mathbf{x})\\cdot\\exp(\\omega^\\top\\mathbf{y})]$, however this is simply an m.g.f. of $\\omega$ at $\\mathbf{x}+\\mathbf{y}$, and thus IMHO does not need to be redirived. (I assume that the formula for the m.g.f. of a multivariate Gaussian random vector is well-known).\n2. Similarly, eq. (16) is the m.g.f. of $\\omega\\sim\\mathcal{N}(\\mathbf{0},\\mathbf{I})$, it does not have to \"immediately follow from\" some other fact.\n\n**Typos**:\n1. In eq. (18), a norm is missing in the last term: $\\mathbf{z}^2$ -> $\\left\\lVert\\mathbf{z}\\right\\rVert^2$.\n2. In the text after eq. (18): inequality -> equality\n\n**Update after the author's response**: The authors have answered my questions during the rebuttal period and I am satisfied with the response. Hence updating my score: 7 $\\to$ 8.", "extra_scores": {}, "extra_text": {}}, {"review_id": "TRWTizOqxrD", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "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": "This paper proposes a set of techniques called Fast Attention Via positive Orthogonal Random features (FAVOR+) to approximate softmax self attention in Transformers and achieve better space and time complexity when the sequence length is much higher than feature dimensions ($L\\gg d$). The resulting architecture, Performers, is provably and practically accurate in estimating regular full-rank attention without relying on any priors such as sparsity or low-rankness. It can also be applied to efficiently model other kernalizable attention mechanisms beyond softmax, achieving better empirical results than regular Transformers on some datasets with such strong representation power. Performers are tested on a rich set of tasks including pixel-prediction, language modeling and protein sequence modeling, and demonstrated competitive results with other examined efficient sparse and dense attention models. I think this paper is a solid step towards more efficient Transformers for practical long-sequence data, and should be accepted. Below I raise two potential questions and look forward to the solutions in the future. \n\n1. For Performers, speedup can only be achieved when $L\\gg d$, which means Performers might have less advantages on wider Transformers. The time complexity of Performers for approximating the attention matrix is $O(Ld^2\\log(d))$, while a regular Transformer has a time complexity of $O(L^2d)$ to compute the attention matrix. From Figure 3, it looks like for a model similar to BERT-Base, Performers is only faster when the sequence length is larger than 512. I wonder how wide Performers can be while  preserving faster speed than Transformers with the same size on the current datasets, and how their test scores compare. \n\n2. It looks like Performers do not converge as fast as Transformers on larger language modeling datasets like PG-19. I wonder whether similar phenomenon will happen in other domains when the dataset is larger, and feel quite curious about the cause.", "extra_scores": {}, "extra_text": {}}, {"review_id": "PaQjUjJCTJ0", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "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": "### Summary\n\nThe authors propose to use the kernel feature map self-attention formulation introduced in [1] to efficiently approximate the softmax attention. The main contribution of the paper lies in the proposed _positive random features_ that can approximate softmax with a strictly positive feature map without which the training is unstable. The authors also show that an approximation of softmax is not necessary for good performance and actually use ReLU random features to achieve their best results when training from scratch.\n\n### Strengths\n\n- The paper deals with a very pressing and important issue, that of the scalability of self-attention.\n- The positive random features are also useful outside of the context of self-attention for efficiently approximating softmax.\n- The experimental results provide strong evidence about the performance of training transformers with kernel feature-maps.\n\n### Weaknesses\n\nThe biggest weakness of the paper in my opinion is the lack of comparison with a simple feature-map as proposed in [1]. Since the authors also use the ReLU random features, we establish that approximating softmax is not necessary for good performance.\n\n1. What would the performance be if a simple deterministic feature map was used?\n2. What would it be if the computational cost was equalized either by adding more layers or by increasing the dimensionality of the queries and keys?\n\nThe second weakness of the paper concerns the evaluation of the practical softmax approximation capabilities. I find the theoretical results interesting and important but I would like more experimental evidence. Without fine-tuning, the authors provide evidence that the approximation does not work (Fig 5).\n\n1. What would happen, for instance, in a toy task where the Lipschitz constant of the transformer layers was kept low? How big would the feature map need to be in order for the approximation to work in such a simple case?\n2. What is being approximated in Fig 4? Is it a randomly initialized attention? What is the rank of the attention matrix?\n3. How good would the approximation be for an attention matrix that is almost full rank and how many features would we need then?\n\n### Reasons for my recommendation\n\nI am recommending acceptance because I believe that the positive random features is an important contribution both for transformers and for kernel approximation. In addition, the experimental results are impressive and show that fast kernelized attention indeed works in practice. My only reservation for a higher score is, as mentioned in the weaknesses section, the lack of comparison with simpler feature maps under equalized computation time.\n\n[1]: Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention", "extra_scores": {}, "extra_text": {}}, {"review_id": "MAyJ4XGUF6K", "reviewer": "AnonReviewer3", "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": "##########################################################################\n\nSummary:\n \n\nThe paper proposed a theoretically grounded O(N) approximation of the softmax attention. The key idea is to interpret attention as a kernel function and construct the random feature projection that can reproduce this kernel. It is highly non-trivial to derive a feature mapping that can accurately approximate the softmax kernel. To better approximate the softmax kernel, the author proposed some important design choices, all of which are supported by theoretical and empirical evidences. The author showed that 1) adopting non-negative random features is very essential to the approximation and the proposed Positive Random Features (PRF) can effectively reduce the variance when the attention values are small, 2) drawing orthogonal random matrices can further reduce the variance of the approximation, 3) the final proposed Performer model runs faster, takes less memory, and has better performance than other O(N) and O(N logN) attention methods.\n\n\n##########################################################################\n\nReasons for score: \n \n\nThe paper is very well-written and should be accepted. This is an important landmark in the research about O(N) attention. The design of the random feature mapping is reasonable and theoretical analysis is convincing. Experiments show that Performer is better than the other O(N) attention methods and also other efficient attention methods.\n \n\n##########################################################################\n\nPros: \n \n\n1. The paper gives a provable O(N) approximation of the softmax attention. The method works without assumptions on the structure of the attention map (like sparsity). The theoretical proofs provide good insights on how to design a good O(N)-complexity approximation to the  attention mechanism.\n\n2. Apart from approximating the softmax attention, the proposed FAVOR+ method can be utilized to approximate other attention kernels. In fact, the author has experimented with Performer-ReLU, which outperforms Performer in some experiments. This provides the insight that softmax attention may not be the best choice.\n\n3. The author conducted very comprehensive ablation studies on different components of the proposed method. This includes: 1) effectiveness of using the positive features, 2) drawing orthogonal random samples, 3) redrawing the random samples\n \n\n##########################################################################\n\nCons: \n \n\n1. From Figure 5 and also Figure 15, periodic redrawing is quite essential. However, the author has not mentioned about the implementation details on how they redraw the random samples and how to choose the period. For me, I feel that this hyper-parameter should be important because the model may need to ensure that each group of samples has been trained for a sufficient amount of time.\n\n2. Performer-ReLU has replaced the attention kernel and can sometimes be better than the softmax attention. Thus, I feel the author may also want to compare with the linear attention method in ((Katharopoulos et al., 2020) \"Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention\".\n\n\n##########################################################################\n\nQuestions during rebuttal period: \n \n\nPlease address and clarify the cons above \n \n\n#########################################################################\n\nTypos: \n\n(1) Page 5, after \"than those from SM_{2m}(x, y)\", there is an additional right bracket.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This is a solid paper that proposes a new method for approximating softmax attention in transformer architectures that scales linearly with the size of the sequence. Even though linear architectures have been proposed before using a similar idea (Katharopoulos et al 2020), this paper provides a better solution along with theoretical analysis and makes a rigorous empirical comparison against other methods. All reviewers agree that this is a strong paper that should be accepted. I suggest citing the recent paper https://arxiv.org/abs/2011.04006 (Long Range Arena, mentioned in the discussion) which provides further comparisons on long-range benchmarks, including the method presented in this paper and Katharopoulos et al 2020, along with a detailed discussion of the differences between the two methods.
Decision: Accept (Oral)