R ETHINKING A TTENTION WITH P ERFORMERS
A BSTRACT
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 attentionkernels, 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 efﬁciently model kernelizable attention mechanisms beyond softmax.
This representational power is crucial to accurately compare softmax with other kernels for the ﬁrst 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 efﬁcient sparse and dense attention methods, showcasing effectiveness of the novel attention-learning paradigm leveraged by Performers.
I NTRODUCTION 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/summarization, time series prediction, generative modeling (e.g. image generation (Parmar et al., 2018)), music generation (Huang et al., 2019), and bioinformatics (Rives et al., 2019; Madani et al., 2020; Ingraham et al., 2019; Elnaggar et al., 2019; Du et al., 2020).
Transformers rely on a trainable attention mechanism that identiﬁes 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 L .
Several solutions have been proposed to address this issue (Beltagy et al., 2020; Gulati et al., 2020; Chan et al., 2020; Child et al., 2019; Bello et al., 2019).
Most approaches restrict the attention mechanism to attend to local neighborhoods (Parmar et al., 2018) or incorporate structural priors on attention such as sparsity (Child et al., 2019), pooling-based compression (Rae et al., 2020) clustering/binning/convolution 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 is used to group together tokens of similar embeddings), sliding windows (Beltagy et al., 2020), or truncated targeting (Chelba et al., 2020).
There is also a long line of research on using dense attention matrices, but deﬁned by low-rank kernels substituting softmax (Katharopoulos et al., 2020;
Shen et al., 2018).
Those methods critically rely on kernels admitting explicit representations as dot-products of ﬁnite positive-feature vectors.
The approaches above do not aim to 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 layers (Child et al., 2019).
Unfortunately, there is a lack of rigorous guarantees for the representation power produced by such methods, and sometimes the validity of sparsity patterns can only be veriﬁed 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 to 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 sufﬁcient.
Approximations based on truncated back-propagation (Dai et al., 2019) are also unable to capture long-distance correlations since the gradients are only propagated inside a 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 ﬁrst 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 on any priors such as 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 we propose.
We believe these methods are of independent interest, contributing to the theory of scalable kernel methods.
Consequently, Performers are the ﬁrst linear architectures fully compatible (via small amounts of ﬁne-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 to efﬁciently model other kernelizable attention mechanisms beyond softmax.
This representational power is crucial to accurately compare softmax with other kernels for the ﬁrst time on large-scale tasks, that are beyond the reach of regular Transformers, and ﬁnd for them optimal attention-kernels.
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 a rich set of tasks ranging from pixel-prediction through text models to protein sequence modeling.
We demonstrate competitive results with other examined efﬁcient 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 as reversible layers (Kitaev et al., 2020) or cluster-based attention (Roy et al., 2020).
FAVOR+ M ECHANISM & P OSITIVE O RTHOGONAL R ANDOM F EATURES
Below we describe in detail the FAVOR+ mechanism - the backbone of the Performer (cid:48) 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 attention-kernels.
2. P RELIMINARIES - REGULAR ATTENTION MECHANISM
Let L be the size of an input sequence of tokens.
Then regular dot-product attention (Vaswani et al., 2017) is a mapping which accepts matrices Q , K ,
V ∈ R L × d 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-directional
Devlin et al., 2018))
dot-product attention has the following form, where A ∈ R L × L is the so-called attention matrix :
[Equation]
Here exp( · ) is applied elementwise, 1 L is the all-ones vector of length L , and diag( · ) is a diagonal matrix with the input vector as the diagonal.
Time and space complexity of computing (1) are O
( L 2 d ) and O ( L 2 + Ld ) respectively, because A has to be stored explicitly.
Hence, in principle, dot-product attention of type (1) is incompatible with end-to-end processing of long sequences.
Bidirectional attention is applied in encoder self-attention and encoder-decoder attention in Seq2Seq architectures.
Another important type of attention is unidirectional dot-product attention which has the form:
[Equation]
where tril( · ) returns the lower-triangular part of the argument matrix including the diagonal.
As discussed in (Vaswani et al., 2017), unidirectional attention is used for autoregressive generative modelling, e.g. as self-attention in generative Transformers as well as the decoder part of Seq2Seq
Transformers.
We will show that attention matrix A can be approximated up to any precision in time
O
( Ld 2 log( d )) .
For comparison, popular methods leveraging sparsity via Locality-Sensitive Hashing (LSH) techniques (Kitaev et al., 2020) have O ( Ld 2 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 preﬁx-sums (all details in the Appendix B.1).
2. G ENERALIZED K ERNELIZABLE A TTENTION
FAVOR+ works for attention blocks using matrices A ∈ R L × L of the form
A ( i, j ) =
K( q (cid:62) i , k (cid:62) j ) , with q i / k j standing for the
i th /j
th query/key row-vector in Q / K and kernel K : R d × R d → R + deﬁned for the (usually randomized) mapping: φ : R d → R r + (for some r > 0 ) as:
[Equation]
We call φ ( u ) a random feature map for u ∈ R d .
For Q (cid:48) , K (cid:48)
∈ R L × r with rows given as φ ( q (cid:62) i ) (cid:62) and φ ( k (cid:62) i ) (cid:62) respectively, Equation 3 leads directly to the efﬁcient attention mechanism of the form:
[Equation]
Here (cid:92) Att ↔ 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
O
( Lrd ) as opposed to O
( L 2 + Ld ) and O ( L 2 d ) of the regular attention (see also Fig. 1).
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 deﬁned in Equation 3, and in particular, can we use it in principle to approximate regular softmax attention ?
(2) How do we implement it robustly in practice, and in particular, can we choose r (cid:28) L for L (cid:29) d to obtain desired space and time complexity gains?
We answer these questions in the next sections.
2. H OW TO AND HOW NOT TO APPROXIMATE SOFTMAX - KERNELS FOR A TTENTION
It turns out that by taking φ of the following form for functions f 1 , ..., f l : R → R , function g : R d → R and deterministic vectors ω i or ω 1 , ..., ω m iid ∼ D for some distribution D ∈ P ( R d ) :
[Equation]
(5) we can model most kernels used in practice.
Furthermore, in most cases D is isotropic (i.e. with pdf function constant on a sphere), usually Gaussian.
For example, by taking h ( x ) = 1 , l = 1 and
D = N (0 , I d ) we obtain estimators of the so-called PNG-kernels (Choromanski et al., 2017)
(e.g. f 1 = sgn corresponds to the angular kernel).
Conﬁgurations: h ( x ) = 1 , l = 2 , f 1 = sin , f 2 = cos correspond to shift-invariant kernels, in particular D = N (0 , I d ) leads to the Gaussian kernel K gauss (Rahimi & Recht, 2007).
The softmax-kernel which deﬁnes regular attention matrix A is given as:
[Equation]
In the above, without loss of generality, we omit √ d -renormalization since we can equivalently renormalize input keys and queries.
Since: SM( x , y ) =
exp( (cid:107) x
[Equation]
2 )K gauss ( x , y ) exp(
[Equation]
y
[Equation]
2 2 ) , based on what we have said, we obtain random feature map unbiased approximation of SM( x , y ) using trigonometric functions with: h ( x ) = exp( (cid:107) x (cid:107) 2 2 ) , l = 2 , f 1 = sin , f 2 = cos .
We call it (cid:100)
SM trig m ( x , y ) .
There is however a caveat there.
The attention module from (1) constructs for each token, a convex combination of value-vectors with coefﬁcients 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 / cos ) leads to unstable behaviours, especially when kernel scores close to 0 (which is the case for many entries of A corresponding to low relevance tokens) are approximated by estimators with large variance in such regions.
This results in abnormal behaviours, e.g. negative-diagonal-values renormalizers D − 1 , and consequently either completely prevents training or leads to sub-optimal models.
We demonstrate empirically that this is what happens for (cid:100) SM trig m and provide detailed theoretical explanations showing that the variance of (cid:100) SM trig m 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 was never proposed.
We propose a robust mechanism in this paper.
Furthermore, the variance of our new unbiased positive random feature map estimator tends to 0 as approximated values tend to 0 (see: Section 3).
Lemma 1 (Positive Random Features (PRFs) for Softmax) .
For x , y ∈ R d , z
=
x
+ y we have:
SM( x , y ) =
[Equation] [Equation]
z ) , (7) where Λ = exp( − (cid:107) x (cid:107) 2 + (cid:107) y (cid:107) 2) and cosh is hyperbolic cosine.
Consequently, softmax-kernel admits a positive random feature map unbiased approximation with h ( x ) =
exp( − (cid:107) x (cid:107) 2 2 ) , l = 1 , f 1 = exp and D = N (0 , I d ) or: h ( x ) = 1 √ 2 exp( − (cid:107) x (cid:107) 2 2 ) , l = 2 , f 1 ( u ) = exp( u ) , f 2 ( u ) =
exp( − u ) and the same D (the latter for further variance reduction).
We call related estimators: (cid:100) SM + m and (cid:100)
SM hyp+ m .
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 outperforms its counterpart.
In Section 4 we further conﬁrm our method’s advantages empirically, using positive features to efﬁciently train softmax-based linear Transformers.
If we replace in (7) ω with √ d ω (cid:107) ω (cid:107) , we obtain the so-called regularized softmax-kernel SMREG which we can approximate in a similar manner, simply changing D = N (0 , I d ) to D = Unif( √ d S d − 1 ) , a distribution corresponding to Haar measure on the sphere of radius √ d in R d , obtaining estimator (cid:92) SMREG + m .
As we show in Section 3, such random features can also be used to accurately approximate regular softmax-kernel.
2. O RTHOGONAL R ANDOM F EATURES (ORF S )
The above constitutes the R+ part of the FAVOR+ method.
It remains to explain the O-part.
To further reduce the variance of the estimator (so that we can use an even smaller number of random features r ), we entangle different random samples ω 1 , ..., ω m 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 a well-known method, yet it turns out that it works particularly well with our introduced PRFs for softmax.
This leads to the ﬁrst theoretical results showing that ORFs can be applied to reduce the variance of softmax/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 ﬁrst exponentially small bounds on large deviations probabilities that are strictly smaller than for non-orthogonal methods.
Positivity of random features plays a key role in these bounds.
The ORF mechanism requires m ≤ d , but this will be the case in all our experiments.
The pseudocode of the entire FAVOR+ algorithm is given in Appendix B.
Our theoretical results are tightly aligned with experiments.
We show in Section 4 that PRFs+ORFs drastically improve accuracy of the approximation of the attention matrix and enable us to reduce r which results in an accurate as well as space and time efﬁcient mechanism which we call FAVOR+.
T HEORETICAL RESULTS
We present here the theory of positive orthogonal random features for softmax-kernel estimation.
All these results can be applied also to the Gaussian kernel, since as explained in the previous section, one can be obtained from the other by renormalization (see: Section 2.3).
All proofs 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:
[Equation]
for independent random samples ω i , and where MSE stands for the mean squared error.
Thus, for SM( x , y ) → 0 we have: MSE( (cid:100)
SM trig m ( x , y ))
→ ∞ and MSE( (cid:100) SM + m ( x , y ))
→ 0 .
Furthermore, the hyperbolic estimator provides additional accuracy improvements that are strictly better than those from (cid:100) SM + 2 m ( x , y ) with twice as many random features.
The next result shows that the regularized softmax-kernel is in practice an accurate proxy of the softmax-kernel in attention.
Theorem 1 (regularized versus softmax-kernel) .
Assume that the L ∞ -norm of the attention matrix for the softmax-kernel satisﬁes: (cid:107) A (cid:107) ∞ ≤ C for some constant C ≥ 1 .
Denote by A reg the corresponding attention matrix for the regularized softmax-kernel.
The following holds:
[Equation]
Furthermore, the latter holds for d ≥ 2 even if the L ∞ -norm condition is not satisﬁed, i.e. the regularized softmax-kernel is a universal lower bound for the softmax-kernel.
Consequently, positive random features for SMREG can be used to approximate the softmax-kernel.
Our next result shows that orthogonality provably reduces mean squared error of the estimation with positive random features for any dimensionality d > 0 and we explicitly provide the gap.
Theorem 2.
If (cid:100) SM ort+ m ( x , y ) stands for the modiﬁcation of (cid:100) SM + m ( x , y ) with orthogonal random features (and thus for m ≤ d ), then the following holds for any d > 0 :
[Equation]
Furthermore, completely analogous result holds for the regularized softmax-kernel SMREG .
For the regularized softmax-kernel, orthogonal features provide additional concentration results - the ﬁrst exponentially small bounds for probabilities of estimators’ tails that are strictly better than for non-orthogonal variants for every d > 0 .
Our next result enables us to explicitly estimate the gap.
Theorem 3.
Let x , y ∈ R d .
The following holds for any a > SMREG( x , y ) and m ≤ d :
[Equation]
where (cid:92) SMREG ort+ m ( x , y ) stands for the modiﬁcation of
[Equation]
SMREG
[Equation]
x
[Equation]
y ) with ORFs,
X = Λ exp( √ d ω (cid:62) (cid:107) ω (cid:107)
( x + y )) , ω ∼ N (0 , I d ) , Λ is as in Lemma 1 and L Z is a Legendre Transform of Z deﬁned as: L Z ( a ) =
sup θ> 0 log( e θa M Z ( θ ) ) for the moment generating function M Z of Z .
We see that ORFs provide exponentially small and sharper bounds for critical regions where the softmax-kernel is small.
Below we show that even for the SM trig mechanism with ORFs, it sufﬁces to take m =
Θ( d log( d )) 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 d of the embeddings, radius R of the ball where all queries/keys live and precision parameter (cid:15) (see:
Appendix F.6 for additional discussion), but does not depend on input sequence length L .
Theorem 4 (uniform convergence for attention approximation) .
Assume that L 2 -norms of queries/keys are upper-bounded by R > 0 .
Deﬁne l =
Rd − 14 and take h ∗ =
exp( l 2 2 ) .
Then for any (cid:15) > 0 , δ = (cid:15) ( h ∗ ) 2 and the number of random projections m =
Θ( dδ 2 log( 4 d 34 R δ ))
the following holds for the attention approximation mechanism leveraging estimators (cid:100)
SM trig with ORFs: (cid:107) (cid:98)
A − A (cid:107) ∞ ≤ (cid:15) with any constant probability, where (cid:98) A approximates the attention matrix A .
E XPERIMENTS
We implemented our setup on top of pre-existing Transformer training code in Jax (Frostig et al., 2018) optimized with just-in-time ( jax.jit ) compilation, and complement our theory with empirical evidence to demonstrate the practicality of FAVOR+ in multiple settings.
Unless explicitly stated, a Performer replaces only the attention component with our method, while all other components are exactly the same as for the regular Transformer.
For shorthand notation, we denote unidirectional/causal modelling as (U) and bidirectional/masked language modelling as (B) .
In terms of baselines, we use other Transformer models for comparison, although some of them are restricted to only one case - e.g. Reformer (Kitaev et al., 2020) is only (U), and Linformer (Wang et al., 2020) is only (B).
Furthermore, we use PG-19 (Rae et al., 2020) 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., 2015)
+ Wikipedia dataset used in BERT (Devlin et al., 2018) and Linformer.
All model and tokenization hyperparameters are shown in Appendix A.
4. C OMPUTATIONAL COSTS
We compared speed-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 ( n heads , n layers , d ff , d ) = (8 , 6 , 2048 , 512) , where d ff 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 memory consumption (since the explicit O ( L 2 ) attention matrix is not stored).
In fact, the Performer achieves nearly optimal speedup and memory efﬁciency 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 efﬁciencies for large L 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. S OFTMAX 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 error than trigonometric sin / cos features.
These two empirically validate the PORF mechanism.
To further improve overall approximation of attention blocks across multiple iterations which further improves training, random samples should be periodically redrawn (Fig. 5, right).
This is a cheap procedure, but can be further optimized (Appendix B.2).
4. S OFTMAX APPROXIMATION ON T RANSFORMERS
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.(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 on a Transformer model (i.e. "Performer-XSOFTMAX"), we demonstrate that:
Backwards compatibility with pretrained models is available as a beneﬁt from softmax approximation, via small ﬁnetuning (required due to error propagation) even for trigonometric
features (Fig. 5, left) on the LM1B dataset (Chelba et al., 2014).
However, when on larger dataset PG-19, 2.
Positive (POS) softmax features (with redrawing) become crucial for achieving performance matching regular
Transformers (Fig. 5, right).
even are
4. M ULTIPLE 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 et al., 2020).
In Fig.
6, the Reformer and Linformer signiﬁcantly drop in accuracy on the protein dataset.
Furthermore, the usefulness of generalized attention is evidenced by Performer-RELU (taking f = ReLU in Equation 5) achieving the highest accuracy in both (U) and (B) cases.
Our proposed softmax approximation is also shown to be tight, achieving the same accuracy as the exact-softmax
Transformer and conﬁrming our theoretical claims from Section 3.
4. L ARGE LENGTH TRAINING - C OMMON 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 ( n heads , d ff , d ) but varying n layers .
Performer/6-layers matches the Reformer/12-layers, while the Performer/12-layers matches the Reformer/24-layers (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 to length L = 8192 from TrEMBL, long enough to model protein interaction networks without the large sequence alignments required by existing methods (Cong et al., 2019).
In this setting, a regular Transformer overloads memory even at a batch size of 1 per chip, by a wide margin.
Thus as a baseline, we were forced to use a signiﬁcantly smaller variant, reducing to ( n heads , n layers , d ff , d ) =
Meanwhile, the Performer trains efﬁciently at a batch size of 8 per chip using the standard (8 , 6 , 2048 , 512) architecture.
We see in Fig.
7 (right subﬁgure) that the smaller Transformer ( n layer = 3 ) is quickly bounded at ≈ 19% , while the Performer is able to train continuously to ≈ 24% .
C ONCLUSION
We presented Performer , a new type of Transformer, relying on our Fast Attention Via positive Orthogonal Random features (FAVOR+) mechanism to signiﬁcantly improve space and time complexity of regular Transformers.
Our mechanism provides to our knowledge the ﬁrst effective unbiased estimation of the original softmax-based Transformer with linear space and time complexity and opens new avenues in the research on Transformers and the role of non-sparsifying attention mechanisms.
B ROADER 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 on biological sequence analysis by enabling the Transformer to be applied to much longer sequences without constraints on the structure of the attention matrix.
The initial application that we consider is the prediction of interactions between proteins on the proteome scale.
Recently published approaches require large evolutionary sequence alignments, a bottleneck for applications to mammalian genomes (Cong et al., 2019).
The potentially broad translational impact of applying these approaches to biological sequences was one of the main motivations of this work.
We believe that modern bioinformatics can immensely beneﬁt 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-speciﬁed interaction properties.
These approaches could be used to augment existing physics-based 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 COemission reduction (Strubell et al., 2019) and lower energy consumption (You et al., 2020), as regular
Transformers require very large computational resources.
Research on Transformers:
We believe that our results can shape research on efﬁcient Transformers architectures, guiding the ﬁeld 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 at the same time, is a gateway to new breakthroughs in bio-informatics, e.g. language modeling for proteins, as we explained in the paper.
Our presented method can be potentially a ﬁrst step.
Backward Compatibility: Our Performer can be used on the top of a 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 a 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 et al., 2018), image processing (Fu et al., 2019), and reinforcement learning/robotics (Tang et al., 2020).
A CKNOWLEDGEMENTS
We thank Nikita Kitaev and Wojciech 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 on this draft.
We lastly thank Yi Tay and Mostafa Dehghani for discussions on comparing baselines.
Valerii Likhosherstov acknowledges support from the Cambridge Trust and DeepMind.
Lucy Colwell acknowledges support from the Simons Foundation.
Adrian Weller acknowledges support from a
Turing AI Fellowship under grant EP/
V025379/1, The Alan Turing Institute under EPSRC grant
EP/N510129/1 and U/B/000074, and the Leverhulme Trust via CFI.
APPENDIX: R ETHINKING A TTENTION WITH P ERFORMERS A H YPERPARAMETERS FOR EXPERIMENTS
This optimal setting (including comparisons to approximate softmax) we use for the Performer is speciﬁed in the Generalized Attention (Subsec.
A.4), and unless speciﬁcally mentioned (e.g. using name "Performer-SOFTMAX"), "Performer" refers to using this generalized attention setting.
A. M ETRICS
We report the following evaluation metrics:
Accuracy : For unidirectional models, we measure the accuracy on next-token prediction, averaged across all sequence positions in the dataset.
For bidirectional models, we mask each token with 15% probability (same as (Devlin et al., 2018)) and measure accuracy across the masked positions.
Perplexity :
For unidirectional models, we measure perplexity across all sequence positions in the dataset.
For bidirectional models, similar to the accuracy case, we measure perplexity across the masked positions.
Bits Per Dimension/Character (BPD/BPC) :
This calculated by loss divided by ln(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. PG-19 P REPROCESSING
The PG-19 dataset (Rae et al., 2020) is presented as a challenging long range text modeling task.
It consists of out-of-copyright Project Gutenberg books published before 1919.
It does not have a ﬁxed vocabulary size, instead opting for any tokenization which can model an 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 to number of tokens in the original dataset.
The original dataset token count per split is: train=1973136207, validation=3007061, test=6966499.
Our sentencepiece tokenization yields the following token counts per split: train=3084760726, valid=4656945, and test=10699704.
This gives log likelihood multipliers of train=1.5634, valid=1.5487, test=1.5359 per split before computing perplexity, which is equal to exp( log likelihood multiplier ∗ loss ) .
Preprocessing for TrEMBL is extensively explained in Appendix C.
A. T RAINING H YPERPARAMETERS
Unless speciﬁcally stated, all Performer + Transformer runs by default used 0 .
5 grad clip, 0 .
1 weight decay, 0 .
1 dropout, 10 − 3 ﬁxed learning rate with Adam hyperparameters ( β 1 = 0 .
9 , β 2 = 0 .
98 , (cid:15) =
10 − 9 ) , with batch size maximized (until TPU memory overload) for a speciﬁc model.
All 36-layer protein experiments used the same amount of compute (i.e. 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. A PPROXIMATE S OFTMAX A TTENTION D EFAULT V ALUES
The optimal values, set to default parameters 1 , are: renormalize_attention =
True, numerical stabilizer = 10 − 6 , number of features = 256, ortho_features =
True, ortho_scaling = 0.0.
[Table]
A. G ENERALIZED A TTENTION D EFAULT V ALUES
The optimal values, set to default parameters 2 , are: renormalize_attention =
True, numerical stabilizer = 0.0, number of features = 256, kernel = ReLU, kernel_epsilon = 10 − 3 .
A. R EFORMER D EFAULT V ALUES
For the Reformer, we used the same hyperparameters as mentioned for protein experiments, without gradient clipping, while using the defaults 3 (which instead use learning rate decay) for ImageNet-64.
In both cases, the Reformer used the same default LSH attention parameters.
A. L INFORMER D EFAULT V ALUES
Using our standard pipeline as mentioned above, we replaced the attention function with the Linformer variant via Jax, with δ = 10 − 6 , k = 600 (same notation used in the paper (Wang et al., 2020)), where δ is the exponent in a renormalization procedure using e − δ as a multiplier in order to approximate softmax, while k is the dimension of the projections of the Q and K matrices.
As a sanity check, we found that our Linformer implementation in Jax correctly approximated exact softmax’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 et al., 2020), as:
• We use k = 600 , which is more than twice than the default k = 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 .
2 https://github.com/google-research/google-research/blob/master/ performer/fast_attention 3 https://github.com/google/trax/blob/master/trax/supervised/configs/
B M AIN A LGORITHM : FAVOR+
We outline the main algorithm for FAVOR+ formally:
Algorithm 1: FAVOR+ (bidirectional or unidirectional).
Input : Q , K ,
V ∈ R
L × d , isBidirectional - binary ﬂag.
Result: (cid:100)
Att ↔ ( Q , K , V ) ∈
R L × L if isBidirectional , (cid:100)
Att → ( Q , K , V ) ∈ R L × L otherwise.
Compute Q (cid:48) and K (cid:48) as described in Section 2.2 and Section 2.3 and take C :=
[ V 1 L ] ; if isBidirectional then
[Equation]
else
Compute G and its preﬁx-sum tensor G PS according to (11);
[Equation]
end [Buf
[Equation]
return diag(buf 4 ) − 1 Buf 3 ;
B. U NIDIRECTIONAL C ASE AND P REFIX S UMS
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.e. all its entries not in the lower-triangular part (which contains the diagonal) are zeroed (see also Fig.
the outer products of random features corresponding to keys with value-vectors.
At each given iteration of the preﬁx-sum algorithm, a random feature vector corresponding to a query is multiplied by the most recent preﬁx-sum (obtained by summing all outer-products corresponding to preceding tokens) to obtain a 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 tril( Q (cid:48)
( K (cid:48) )
(cid:62) ) C without constructing and storing the L × L -sized matrix tril( Q (cid:48)
( K (cid:48) )
(cid:62) ) explicitly, where C
=
[ V
[Equation]
] ∈ R L × ( d +1) .
In order to do so, observe that ∀ 1 ≤ i ≤ L :
[Equation]
) where G , G PS ∈ R L × M × ( d +1) are 3d-tensors.
Each slice G PS: ,l,p is therefore a result of a preﬁx-sum (or cumulative-sum) operation applied to G : ,l,p : G PS i,l,p = (cid:80) i j =1 G i,l,p .
An efﬁcient algorithm to compute the preﬁx-sum of L elements takes O ( L ) total steps and O (log L ) time when computed in parallel (Ladner & Fischer, 1980;
Cormen et al., 2009).
See Algorithm 1 for the whole approach.
B. O RTHOGONAL R ANDOM F EATURES - E XTENSIONS
As mentioned in the main text, for isotropic Ω (true for most practical applications, including regular attention), instead of sampling ω
i independently, we can use orthogonal random features (ORF)
(Yu et al., 2016; Choromanski et al., 2017; 2018b): these maintain the marginal distributions of samples ω i 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 et al., 2016).
ORFs were introduced to reduce the variance of Monte Carlo estimators (Yu et al., 2016; Choromanski et al., 2017; 2018b; 2019a;
Rowland et al., 2019; Choromanski et al., 2018a; 2019b) 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 brieﬂy review the most efﬁcient 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 ω -samples (with different rows corresponding to different samples) in O ( md ) space.
Provides algorithm for computing Wx in O ( md ) time for any x ∈ R d .
Gives unbiased estimation.
Requires one-time O ( md 2 ) preprocessing (Gram-Schmidt orthogonalization).
(2) Hadamard/Givens ORFs [H/G-ORFs]: Applies random Hadamard (Choromanski et al., 2017) or Givens matrices (Choromanski et al., 2019b).
Encodes matrix W in O ( m ) or O ( m log( d )) space.
Provides algorithm for computing Wx in O ( m log( d )) time for any x ∈ R d .
Gives small bias (tending to 0 with d → ∞ ).
B. T IME AND S PACE C OMPLEXITY - D ETAILED A NALYSIS
We see that a variant of bidirectional FAVOR+ using iid samples or R-ORFs has O ( md + Ld + mL ) space complexity as opposed to Θ( L 2 + Ld ) space complexity of the baseline.
Unidirectional
FAVOR+ using fast preﬁx-sum pre-computation in parallel (Ladner & Fischer, 1980; Cormen et al., 2009) has O ( mLd ) space complexity to store G PS which can be reduced to O ( md + Ld + mL ) by running a simple (though non-parallel in L ) aggregation of G PS i, : , : without storing the whole tensor
G PS in memory.
From Subsec.
B.2, we know that if instead we use G-ORFs, then space complexity is reduced to O ( m log( d ) + Ld + mL ) and if the H-ORFs mechanism is used, then space is further reduced to O ( m + Ld + mL )
=
O ( Ld + mL ) .
Thus for m, d (cid:28) L all our variants provide substantial space complexity improvements since they do not need to store the attention matrix explicitly.
The time complexity of Algorithm 1 is O
( Lmd ) (note that constructing Q (cid:48) and K (cid:48) can be done in time
O ( Lmd ) ).
Note that the time complexity of our method is much lower than O ( L 2 d ) of the baseline for L (cid:29) m .
As explained in Subsec.
B.2, the R-ORF mechanism incurs an extra one-time O ( md 2 ) cost (negligible compared to the O ( Lmd ) term for L (cid:29) d ).
H-ORFs or G-ORFs do not have this cost, and when
FAVOR+ uses them, computing Q (cid:48) and K (cid:48) can be conducted in time O
( L log( m ) d ) as opposed to O ( Lmd ) (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 m allows a trade-off between computational complexity and the level of approximation: bigger m results in higher computation costs, but also in a lower variance of the estimate of A .
In the theoretical section from the main body we showed that in practice we can take
M = Θ( d log( d )) .
Observe that the FAVOR+ algorithm is highly-parallelizable, and beneﬁts from fast matrix multiplication and broadcasted operations on GPUs or TPUs.
C E XPERIMENTAL D ETAILS FOR P ROTEIN M ODELING T ASKS C
T R EMBL D ATASET
[Table]
We used the TrEMBL dataset 4 , which contains 139,394,261 sequences of which 106,030,080 are unique.
While the training dataset appears smaller than the one used in Madani et al.
(Madani et al., 2020), we argue that it includes most of the relevant sequences.
Speciﬁcally, the TrEMBL dataset consists of the subset of UniProtKB sequences that have been computationally analyzed but not manually curated, and accounts for ≈ 99 .
5% of the total number of sequences in the UniProtKB dataset 5 .
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 for valuation, and an IID split, where the remaining protein sequences are split randomly into train, valid, and test tests.
We held-out the following protein families (PF18369, PF04680, PF17988, PF12325, PF03272, PF03938,
PF17724, PF10696, PF11968, PF04153, PF06173, PF12378, PF04420, PF10841, PF06917, PF03492,
PF06905, PF15340, PF17055, PF05318), which resulted in 29,696 OOD sequences.
We note that, due to deduplication and potential TrEMBL version mismatch, our OOD-Test set does not 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 1.
In the standard sequence modeling task, given the length statistics that are reported in the table, we clip single sequences to maximum length L = 1024 , which results in few sequences being truncated signiﬁcantly.
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 L = 8192 .
[Figure]
A random baseline, 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 amino acids in our dataset may be far from uniform, so we also consider an 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 use 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 ﬁgure shows only the 20 standard amino acids, colored by their class, for comparison with the visualization on the TrEMBL web page 6 .
[Figure]
Table 2 contains the results on the single protein sequence modeling task ( L = 1024 ).
We report accuracy and perplexity as deﬁned in Appendix A:
[Table]
C. A TTENTION M ATRIX I LLUSTRATION
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 to 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 modiﬁed ( Q, K, V ◦ ) triple, where V ◦ 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 is equal to the sequence length, and that a non-zero output on a dimension i can only arise from a non-zero attention weight to the i th 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_BOVIN protein sequence 7 , one of the most extensively studied globular proteins, which contains 100 amino acids.
In Figure 10, we show the attention matrices for the ﬁrst 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 ﬁxed positions.
These patterns are consistent with the patterns found in Transformer models trained on natural language
(Kovaleva et al., 2019).
In Figure 12 we highlight these attention patterns by focusing on the ﬁrsttokens, and in Figure 11, we illustrate in more detail two attention heads.
Amino acid similarity.
Furthermore, we analyze the amino-acid similarity matrix estimated from the attention matrices produced by the Performer model, as described in Vig et al.
(Vig et al., 2020).
We aggregate the attention matrix across 800 sequences.
The resulting similarity matrix is illustrated in Figure 13.
Note that the Performer recognizes highly similar amino acid pairs such as (D, E) and (F, Y).
[Figure]
D E XTENDED APPROXIMATION RESULTS D. B ACKWARDS C OMPATIBILITY - E RROR P ROPAGATION
Although mentioned previously (Sec. 4.2) that the Performer with additional ﬁnetuning is backwards compatible with the Transformer, we demonstrate below in Fig.
14 that error propagation due to nonattention components of the Transformer is one of the primary reasons that pretrained Transformer weights cannot be immediately used for inference on the corresponding Performer.
D. A PPROXIMATE S OFTMAX - E XTENDED P ROPERTIES
We show the following properties of our softmax approximation, in Fig. 15:
Redrawing: While the beneﬁts of redrawing features was shown in Subsec.
4.3 of the main body of the paper, we also its beneﬁts when there are multiple layers with large scale (16x16 TPU-v2) training.
Unidirectional: While we have shown on TrEMBL that Performer with generalized ReLU attention outperforms softmax, we also show that approximate softmax attention can still be a solid choice, for example on ImageNet64 (U).
After 100K steps of training, the Performer-ReLU, Performer-Softmax, and Performer-Softmax (SMREG) variants achieve respectively, 3.67, 3.69, 3.67 BPD.
Instability of Trigonometric Features: We see the full view of the unstable training curve when using Trigonometric softmax.
D. G ENERALIZED A TTENTION
We investigated Generalized Attention mechanisms (mentioned in Sec. 2.2) on TrEMBL
when
L = 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 variables in FAVOR, we compared several kernels and also renormalization on/off (Fig. 16 and
Fig. 17), where Renormalize corresponds to applying D − 1 operator in attention, as for the standard mechanism, though we noticed that disabling it does not necessarily hurt accuracy) to produce the best training conﬁguration 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 use 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).
emphasize the highest accuracy runs but also show the NaN issues with certain kernels
D. C OMPARISON WITH L INEAR T RANSFORMER
We use the attention implementation of the Linear Transformer from (Katharopoulos et al., 2020), which mainly involves setting our feature map φ ( x ) = elu ( x ) + 1 , where elu ( x ) is the shifted-eLU function from (Clevert et al., 2016).
For the sake of fairness 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
ReLU nonlinearity.
We also used the exact same training hyperparameters as Performer-ReLU on our exact ProGen setting from Fig.
Ultimately, we empirically found that the Linear Transformer possessed numerical instability during training via unstable training curves, ultimately stopping training by producing exploding gradients (NaNs) (Fig. 18).
E C OMPUTATION COSTS - E XTENDED RESULTS
In this subsection, we empirically measure computational costs in terms wall clock time on forward and backward passes for three scenarios in Fig.
Performer, with varying number of layers.
We show that our method can scale up to (but not necessarily limited to) even 20 layers.
Attention time complexities when comparing standard attention (from Transformer) and FAVOR (from Performer).
Note that the maximum memory size here is not reﬂective of the maximum memory size in an actual model (shown below), as this benchmark requires computing explicit tensors (causing memory increases) in Jax, while a model does not.
Time complexities when comparing the Transformer and Performer models.
"X" (OPT) denotes the maximum possible speedup achievable, when attention simply returns the Vvector, showing that the Performer is nearly optimal.
We see that the maximum possible power of 2 length allowed on a V100 GPU (16GB) is 2 15 = 32768 using regular dimensions.
Since some of the computational bottleneck in the Transformer may originate from the extra feed-forward layers (Kitaev et al., 2020), we also benchmark the “Small" version, i.e.
( n heads , n layers , d ff , 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 ( n heads , n layers , d ff , d ) = (8 , 6 , 2048 , 512) .
F T HEORETICAL RESULTS
We provide here the proofs of all theoretical results presented in the paper.
F. P ROOF OF L EMMA
Proof.
We ﬁrst deduce that for any a , b ∈ R d
[Equation]
Next
[Equation]
let w ∈ R d .
We use the fact that
[Equation]
for any c ∈ R d and derive:
[Equation]
That completes the proof of the ﬁrst part of the lemma.
An identity involving hyperbolic cosine function is implied by the fact that for every u ∈ R d and ω ∼ N (0 , I d ) the following is true:
[Equation]
The cancellation of the odd moments E
[( ω (cid:62) u ) 2 i +1 ] follows directly from the fact that ω is taken from the isotropic distribution (i.e. distribution with pdf function constant on each sphere).
That completes the proof.
F. P ROOF OF L EMMA
Proof.
Denote: z = x + y and ∆ = x − y .
Note that by using standard trigonometric 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 (0 , I d ) :
[Equation]
Using the fact that (see: Lemma 1 in (Yu et al., 2016); note that in that lemma they use notation: z for what we denote as: (cid:107) ∆ (cid:107) ):
[Equation]
(14) we obtain:
[Equation]
which completes the ﬁrst part of the proof.
To obtain the formula for: MSE( (cid:100) SM + m ( x , y )) notice ﬁrst that:
[Equation] [Equation]
The above immediately follows from the fact that positive random feature maps provide unbiased estimation of the softmax-kernel, thus the following is true:
[Equation]
Therefore we obtain:
[Equation]
where the last equality follows from Equation
[Equation]
Therefore we have:
[Equation]
Finally,
[Equation]
In the chain of equalities above we used the fact that random variables exp( ω (cid:62) z ) and exp( − ω (cid:62) z ) have the same distribution.
This is true since ω and − ω have the same distribution ( ω is Gaussian).
That completes the proof.
F. P ROOF OF T HEOREM
Proof.
Let x ,
y ∈ R
d be respectively a query/key.
Note that from the deﬁnition of SMREG( x , y ) we have for z = x + y :
[Equation] [Equation]
where e 1 def = (1 , 0 , ..., 0) (cid:62) ∈ R d .
To obtain the above we used the fact that N (0 , I d ) is isotropic (that in particular implies zeroing of the even terms in the Taylor expansion).
Let us denote: A ( k, d ) def = E ω ∼N
[Equation]
I
[Equation] [Equation]
e
[Equation]
2 k ] .
It turns out that:
[Equation]
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:
Lemma 3.
Expression A (2 k, d ) satisﬁes the following for k ∈ N :
[Equation]
Proof.
Note ﬁrst that for d ≥ 2 the density function p d ( θ ) of the angle between a vector r ∈ R d chosen uniformly at random from the unit sphere and e 1 is given by the following formula:
[Equation]
Let us denote: F ( k, d ) def = (cid:82)
[Equation]
cos k ( θ ) sin d ( θ ) dθ .
Using partial integration
[Equation]
we get:
[Equation]
Thus we conclude that: F ( k, d ) =
k − 1 d +1 F ( k − 2 , d + 2) .
Therefore we have:
[Equation]
We again conduct partial integration and get:
[Equation]
Therefore we conclude that:
[Equation]
which completes the proof
[Equation]
Applying the above lemma
[Equation]
we get:
[Equation]
where w = (cid:107) z (cid:107) 2 2 and f ( k, d ) =
d k ( d +2 k − 2)
( d +2 k − 4) · ... · d .
Thus we obtain:
[Equation] [Equation]
Note ﬁrst that for k ≥ 1 we have: f ( k, d ) ≤ 1 , thus:
[Equation]
[Equation]
We also have for l = d 13 :
[Equation]
where Po( w ) stands for the random variable of Poisson distribution with parameter w .
Therefore we get for t
= ln( lw ) :
[Equation] [Equation]
where the last equality is implied by the formula for the Laplace Transform for the Poisson random variable:
[Equation]
− 1)) .
[Equation]
Notice that: w = (cid:107) z (cid:107) 2=
[Equation]
≤
[Equation]
C ) .
We conclude that:
[Equation]
That completes the proof.
F. P ROOFS OF T HEOREM 2,T HEOREM 3 & B EAUTIFUL F UNCTIONS
We will provide here much more general theoretical results which will imply Theorem 3 and Theorem 2.
We need the following deﬁnition:
Deﬁnition 1.
We say that function F : R n → R is beautiful if F can be expressed as:
[Equation]
for a probabilistic isotropic distribution Ω , and where g : R → R is an entire function with nonnegative power-series coefﬁcients (i.e. g ( x )
=
(cid:80) ∞
i =
a i x i for every x ∈ R and with a i ≥ 0 for i = 0 , 1 , ... ).
In the formula above we assume that the expectation on the RHS exists.
Interestingly, beautiful functions can be used to deﬁne softmax and consequently, Gaussian kernels (both standard and regularized), leading to our PRF mechanism presented in the main body of the paper, as we explain below.
Remark 1.
If one takes Ω = N (0 , I d ) (note that N (0 , I d ) is isotropic) and g : x → exp( x ) (such g is clearly entire with nonnegative power-series coefﬁcient) then the following is true for z = x
+ y :
[Equation]
Similarly: SMREG( x , y ) =
exp( − (cid:107)
[Equation]
+
[Equation]
(cid:107)
[Equation]
) F Ω reg ,g ( z ) , where Ω reg stands for the distribution corresponding to Haar measure on the sphere of radius √ d (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.
We will consider two estimators of the beautiful functions from Deﬁnition 1 that directly lead (through Remark 1) to: PRF-based approximation of the softmax-kernel and its enhanced version with orthogonal features.
Standard Monte Carlo estimator samples independently ω iid1 , ..., ω iid m iid ∼ Ω , where m stands for the number of samples and then computes:
[Equation]
Orthogonal Monte Carlo estimator samples ω ort1 , ..., ω ort m ( m ≤ d ) in such a way that marginally we have:
[Equation]
but (
[Equation]
0 for i (cid:54) =
j (such an orthogonal ensemble can be always created if Ω is isotropic, as we already mentioned in the main body of the paper).
We deﬁne:
[Equation]
F.4. O RTHOGONALITY UNIVERSALLY IMPROVES CONCENTRATION
Denote by M Z ( θ ) =
E [ e θZ ] a moment generating function of the random variable Z .
Note ﬁrst that estimators of beautiful functions based on standard Monte Carlo procedure using independent vectors ω iid i guarantee strong concentration bounds since independent ω
i s provide a way to obtain exponentially small upper bounds on failure probabilities through moment generating functions.
We summarize this classic observation which is a standard application of Markov’s Inequality below.
Lemma 4.
Consider an estimator (cid:98)
F iid m ( z ) of the beautiful function F evaluated at z .
Then the following holds for any a > F ( z ) :
[Equation] [Equation]
where X = g ( w (cid:62) z ) , w ∼ D and L Z stands for a Legendre Transform of the random variable
Z deﬁned as: L Z ( a ) =
sup θ> 0 log( e θa M Z ( θ ) ) .
Furthermore, L X ( a ) > 0 .
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 5 (orthogonality provides smaller tails) .
If F Ω ,g is a beautiful function then the following holds for m ≤ d , X as in Lemma 4 and any a > F ( z ) :
[Equation]
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 d , not just asymptotically for d large enough.
We also obtain similar result regarding mean squared errors (MSEs) of the considered estimators:
Theorem 6.
If F Ω ,g is a beautiful function then the following holds for m ≤ d :
[Equation]
As before, an orthogonal estimator leads to better concentration results and as before, this is the case for any d > 0 , not only asymptotically for large enough d .
Note that from what we have said above, Theorem 2 and Theorem 3 follow immediately from
Theorem 6 and Theorem 5 respectively.
Thus in the remainder of this section we will prove Theorem 6 and Theorem 5.
F.4. P ROOF OF T HEOREM
Proof.
Note that by the analogous application of Markov’s Inequality as in Lemma 4, we get:
[Equation]
where we have: X ort i =
g
(( ω ort i ) (cid:62) z ) .
We see that it sufﬁces to show that for any θ > 0 the following holds: E [ e θ ( X ort1 + ... + X ort m ) ] < E [ e θ ( X iid1 + ...
+ X iid m ) ] .
We have:
[Equation]
) where S j = { ( j 1 , ..., j m ) ∈ N × ...
× N : j 1 , ..., j m ≥ 0 , j 1 + ... + j m = j } and for some positive constants c ( j 1 , ..., j m ) .
Thus we have:
[Equation]
Similarly, we get:
[Equation]
Therefore we get:
[Equation]
Note ﬁrst that using the fact that f is entire, we can rewrite each X ort i as:
[Equation]
48) where f ( x ) = (cid:80) ∞ s =
a s x s and a 0 , a 1 , ... ≥ 0 .
Similarly,
[Equation]
By plugging in the above formulae for X ort i and X iid i int the formula for ∆ and expanding powerexpressions, we obtain:
[Equation]
for some ordered subsets of indices (with potentially repeating entries) D ( j 1 , ..., j m ) (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
(cid:98) ∆( d 1 , ..., d m ) deﬁned as:
[Equation]
Our next goal is to re-write the formula for (cid:98) ∆( d 1 , ..., d m ) .
Denote:
[Equation]
Observe that Y has the same distribution as Y (cid:48) deﬁned as:
[Equation]
where g is a Gaussian vector taken from the N (0 , I d ) distribution, independently from:
[Equation]
This comes from the fact that for a ﬁxed z one can think about the set:
[Equation]
a random rotation of the system of m canonical basis vectors: e 1 , ..., e m .
Thus instead of applying a random rotation to: e 1 , ..., e m , one can equivalently randomly rotate vector z .
Randomly rotated vector z has the same distribution as:
[Equation]
Now note that lengths of vectors ω ort 1 , ..., ω ort m are chosen independently.
Therefore we obtain:
[Equation]
where v ∼ g
[Equation]
g
[Equation]
Denote g = ( g 1 , ..., g d ) (cid:62) .
Thus we obtain:
[Equation] [Equation]
Now let us focus on the second expression from the formula on (cid:98) ∆( d 1 , ..., d m ) .
We have:
[Equation]
where the ﬁrst equality comes from the fact that different ω iid i 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 s ∈ N + such that s ≤ n and every k 1 , ..., k s ∈ N + the following holds:
[Equation]
Proof.
Take r = g (cid:107) g
(cid:107) 2 (cid:107) ˜
g (cid:107) 2 , where ˜ g is an independent copy of g .
Note that r ∼ g .
We have:
[Equation]
where the ﬁrst equality comes from the independence of different elements of z = ( z 1 , ..., z n ) (cid:62) and the second equality is implied by the fact that ˜ g is independent from g .
Therefore we have:
[Equation]
That completes the proof since z ∼ g and ˜ g ∼ g .
Note that by Lemma
[Equation]
we can rewrite the right expression from the formula on (cid:98) ∆( d 1 , ..., d m ) as:
[Equation]
The left expression from the formula on (cid:98) ∆( d 1 , ..., d m ) can be rewritten as:
[Equation]
Since marginal distributions of ω ort i and ω iid
i are the same, we can rewrite (cid:98) ∆( d 1 , ..., d n ) as
[Equation]
where τ ( d 1 , ..., d m ) is deﬁned as:
[Equation]
We need now few observations regarding (cid:98) ∆( d 1 , ..., d m ) .
Note ﬁrsr that since odd moments of the Gaussian scalar distribution N (0 , 1) are zero, (cid:98) ∆( d 1 , ..., d m ) is zero if at least of of d i is odd.
Furthermore, (cid:92) ∆( d 1 , ..., d m ) is trivially zero if all but at most one d i are zero.
With our new notation
[Equation]
∆ can be rewritten as:
[Equation]
Note also that we have:
[Equation]
j 1 ,...
[Equation]
L ( d 1 , ..., d m
[Equation]
Therefore (see: our observations on (cid:98) ∆( d 1 , ..., d m ) ) to complete the proof it sufﬁces to show that: τ ( d 1 , ..., d m ) ≤ dd +2 if at least two: d i , d j for i (cid:54) =
j are nonzero and all d i are even.
Lemma 6.
The following holds if for some i (cid:54) =
j we have: d i , d j > 0 and all d i are even:
[Equation]
Proof.
Note that τ ( d 1 , ..., d m ) can be rewritten as:
[Equation]
where µ d ( j ) stands for the j th moment of the χ -distribution with d degrees of freedom.
Note that µ d ( j ) =
2 j 2 Γ( d + j 2 ) Γ( d 2 ) , where Γ is the so-called Gamma-function .
Using the fact that: Γ( n ) =
( n − 1)!
and Γ( n + 12 ) =
(2 n − 1)!!
2 n √ π for n ∈ N + , it is easy to see that for a ﬁxed d , the RHS of the Equality 67 is maximized when d i =
d j = 2 and d k = 0 for some i (cid:54) =
j and k / ∈ { i, j } .
Furthermore, straightforward calculations show that in that case the value of the RHS from Equality 67 is dd +2 .
That completes the proof of the Lemma, and consequently, the proof of the entire Theorem.
F.4. P ROOF OF T HEOREM
Proof.
We will use the notation from the proof of Theorem 5.
Since both estimators: (cid:98) F ort m ( z ) and (cid:98) F iid m ( z ) are unbiased, we have: MSE( (cid:98) F ort m ( z )) =
Var( (cid:98) F ort m ( z )) and MSE( (cid:98) F iid m ( z )) =
Var( (cid:98)
F iid m ( z )) .
We have:
[Equation]
Similarly,
[Equation]
We have:
[Equation]
Similarly, we get:
[Equation]
Therefore
[Equation]
since marginal distributions of X iid i and X ort i are the same, we have:
[Equation]
Plugging in the formula for X ort i and X iid i from Equation 48 and Equation 49, and using our analysis from the proof of Theorem 3 we obtain:
[Equation]
for ω ∼ Ω and r ∼ N (0 , 1) .
Thus, using Lemma
[Equation]
we get:
[Equation]
That completes the proof
[Equation]
F. P ROOF OF T HEOREM
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 case 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 m opt , 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 m opt =
Θ( d log( d )) , then with O ( Ld 2 log( d )) -time, we can approximate A up to any precision, regardless of the number of tokens L .
In order to provide those guarantees, we leverage recent research on 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) .
Deﬁne entries of the attention matrix
A as follows:
A i,j = g ( q (cid:62) i )K( 1 d 14 q (cid:62) i , 1 d 14 k (cid:62) j ) h ( k (cid:62) j ) for some g, h : R d → R and where K is a radial basis function (RBF) kernel (Choromanski et al., 2018b) with corresponding spectral distribution Ω (e.g. Gaussian kernel for which Ω = N (0 , I d ) ).
Assume that the rows of matrices Q and K are taken from a ball B ( R ) of radius R , centered at 0 (i.e. norms of queries and keys are upperbounded by R ).
Deﬁne l
=
Rd − 14 and take g ∗ = max x ∈ B ( l ) | g ( x ) | and h ∗ = max x ∈ B ( l ) | h ( x ) | .
Then for any (cid:15) > 0 , δ = (cid:15) g ∗ h ∗ and the number of random projections m = Ω( dδ 2 log( 4 σR δd 14 )) for σ = E ω ∼ Ω [ ω (cid:62) ω ] the following holds: (cid:107) (cid:98) A − A (cid:107) ∞ ≤ (cid:15) with any constant probability, where (cid:98) A approximates generalized attention matrix via orthogonal trigonometric random features.
The result holds in particular for regular softmax-attention for which K is a Gaussian kernel and g ( x )
= h ( x )
=
exp( (cid:107) x (cid:107) 2 2 ) .
In that case m opt = Ω( dδ 2 log( 4 d 34 R δ ))
since σ = d .
Proof.
Let D Q be a diagonal matrix with entries of the form: g ( q (cid:62) i ) and let D K be a diagonal matrix with entries of the form: h ( k (cid:62) i ) .
Denote B =
[K( 1 d 14 q (cid:62) i , 1 d 14 k (cid:62) j )] i,j ∈ R L × L .
Denote by (cid:98) A and approximation of the attention matrix obtained from trigonometric orthogonal random features and by (cid:98) 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
f i satisfy f 1 ( x ) = sin( x ) ,
f 2 ( x ) = cos( x ) (thus in particular are bounded).
Also, it is not hard to observe (see for instance analysis in Claim 1 from (Rahimi & Recht, 2007)) that we can take:
L f = 1 (for L f as in Theorem 3 from (Lin et al., 2020)).
Using Theorem 3 from (Lin et al., 2020), we conclude that:
[Equation]
with any constant probability as long as m = Ω( dδ 2 ) log( σ · diam( M ) δ ) , where σ = E [ ω (cid:62) ω ] and M is the diameter of the smallest ball M containing all vectors of the form z
= Q i d 14 − K j d 14 .
Since (cid:107) Q i (cid:107) 2 , (cid:107) K j (cid:107) 2 ≤ R , we conclude that (cid:107) z (cid:107) 2 ≤ 2 R d 14 and thus one can take diam( M )
=
4 R d 14 .
We have:
[Equation]
Taking δ = (cid:15) g ∗ h ∗ completes the proof.
F. D ISCUSSION OF T HEOREM
As a consequence of Theorem 4, the number m of random projections required to approximate the attention matrix within (cid:15) error is a function of data dimensionality d , the parameter (cid:15) and the radius
R of the ball within which the queries and keys live:
[Equation]
The dependence on d and (cid:15) is fairly easy to understand: with a larger dimensionality d we need more random projeections (on the order of magnitude d log( d ) ) to get an approximation within (cid:15) error.
The dependence on R means that the length of queries and keys cannot grow at a ﬁxed m if we want to retain the quality of the approximation.
In particular, this means that FAVOR cannot approximate hard attention on sequences of unlimited length with a ﬁxed m .
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 our experiments, this limitation does not manifest itself in practice at the lengths we experimented with.