Y OU O NLY S AMPLE (A LMOST ) O NCE : L INEAR C OST S ELF -A TTENTION V IA B ERNOULLI S AMPLING
Anonymous authors
Paper under double-blind review
A BSTRACT
Transformer-based models have come to dominate the landscape in a wide range of natural language processing (NLP) applications.
The heart of the transformer model is the self-attention mechanism, which captures the interactions of token pairs in the input sequences and consequently, depends quadratically on the input sequence length.
It is known that training such models on longer sequences is quite expensive, and often, prohibitively so.
We show that a Bernoulli sampling attention mechanism based on Locality Sensitive Hashing (LSH), decreases the quadratic complexity to linear.
We bypass the quadratic cost by considering selfattention as a sum of individual tokens associated with Bernoulli random variables that can, in principle, be sampled at once by a single hash (although in practice, this number may be a small constant).
This leads to an efﬁcient sampling scheme to estimate self-attention which relies on speciﬁc modiﬁcations of LSH (based on feasibility of deployment on GPU architectures).
We evaluate our proposed algorithm on the GLUE benchmark with standard 512 sequence length and our method achieves comparable or even slightly better performance than a standard pretrained Transformer.
To evaluate whether our method can indeed handle longer sequences, we conduct experiments on long sequence (4096) language model pretraining and achieve consistent results as standard self-attention, while observing sizable inference speed-ups and memory savings.
I NTRODUCTION
The Transformer model (Vaswani et al., 2017) is incredibly effective across a diverse set of natural language processing (NLP) applications including machine translation (Vaswani et al., 2017), language inference (Devlin et al., 2018) and paraphrasing (Raffel et al., 2019).
Transformer-based models such as BERT (Devlin et al., 2018) are pretrained in an unsupervised manner and later ﬁnetuned on different downstream tasks, often providing state-of-the-art performance on standard benchmarks.
While such models have strong empirical performance, their high computational and memory requirements remain quite high.
Consequently, in the NLP setting, most current models have certain constraints on the sequence length, e.g., BERT and other transformer-based language models (Yang et al., 2019; Liu et al., 2019) limit the sentence length to be at most 512 .
The Multi-Head Self-Attention is central to Transformer based models and provides a ﬂexible global receptive ﬁeld to exchange information among input tokens.
While self-attention provides immense beneﬁts, it is also a key bottleneck in training with long sequences.
In particular, the output of self-attention is a combination of all tokens where coefﬁcients are determined by the similarities among tokens.
While this is empirically beneﬁcial, it involves a sizable resource footprint.
For sequence length n , this leads to a O ( n 2 ) complexity in both time and memory to compute pairwise similarities among all input tokens.
This quadratic cost is a roadblock in attaining potential beneﬁts that may be realizable in various applications by capturing long term context dependencies.
As we will discuss in more detail later, the foregoing issue is a major thrust of several recent and ongoing efforts focused on mitigating the sizable resource requirements of such models.
Our work is inspired by ideas of importance sampling via hashing-based sampling techniques (Spring & Shrivastava, 2017; Charikar & Siminelakis, 2017).
We proposed a Bernoulli based sampling to approximate self-attention, scaling linearly with the input sequence length.
We achieve this by viewing self-attention as a sum of individual tokens associated with Bernoulli random variables whose success probability is determined by the similarities among tokens.
In principle, we can sample all Bernoulli random variables at once with a single hash (although in practice, this number may be a small constant to lower the approximation variance).
This leads to an efﬁcient sampling scheme to estimate self-attention which relies on speciﬁc modiﬁcations of hashing-based importance sampling (based on feasibility of deployment on GPU architectures).
The resulting strategy (You
Only Sample Almost Once, YOSO-Attention) is far more amenable to an efﬁcient and backpropagation friendly implementation, and has a favorable empirical performance proﬁle on natural language modeling tasks.
We evaluate our proposed algorithm on the GLUE benchmark (Wang et al., 2019) with 512 sequence length as well as on long sequence language model pretraining where we see promising results with speed-ups and memory savings.
B ACKGROUND : S ELF -A TTENTION
Self-Attention.
Self-attention is a scaled dot-product attention mechanism to capture token dependencies in the input sequence, which can be deﬁned as,
[Equation]
where Q , K , V ∈ R n × d are embedding matrices from the input sequence, called queries, key and values respectively.
Here, n is the input sequence length, d is the embedding dimension of each token, W Q , W K ,
W V ∈ R
d × d h are learned parameter matrices, d h is the dimension of hidden embedding, and D P is a n × n diagonal matrix which normalizes each row of the exp ( P ) matrix such that the row entries sum up to 1 .
For simplicity, we overload the notations for Q , K , V to denote QW Q , KW K , V W V in our presentation.
Multi-Head Self-Attention.
Multi-Head self- attention in Transformers runs through the scaled dot-product attention multiple times and the attention outputs are concatenated to help the model capture information from multiple representation subspaces Vaswani et al.
Multi-Head Selfattention can be formally written as,
[Equation]
where h is the number of heads, A i , i = 1 , . .
, h are heads with different parameter matrices.
Self-Attention Bottleneck.
A key bottleneck in self-attention is computing the softmax matrix, softmax ( P ) , which requires the calculation of all pairwise input token similarities.
To reduce this cost, we seek to approximate the softmax matrix by viewing self-attention for each query as an expectation of a softmax distribution and computing the approximated self-attention with an efﬁcient sampling mechanism.
In the following sections, we will ﬁrst review LSH-based importance sampling and then propose Bernoulli sampling with LSH to estimate self-attention efﬁciently.
I MPORTANCE S AMPLING VIA L OCALITY S ENSITIVE H ASHING
Importance sampling (Press et al., 2007) helps approximate properties of a target distribution by a weighted average of random draws from another distribution.
It is known (Press et al., 2007) that importance sampling can be directly used for the softmax distribution by drawing samples from a uniform distribution – which avoids sampling from the softmax distribution directly which is harder.
But this leads to a high variance estimate since the softmax distribution is usually concentrated in a small region.
When using this idea for softmax matrix approximation for self-attention in particular , the variance tends to grow with the input sequence length.
Before proceeding, we will summarize an interesting importance sampling method for low variance estimators, speciﬁcally, importance sampling via LSH from (Charikar & Siminelakis, 2017; Spring & Shrivastava, 2017).
LSH-based Importance Sampling .
Consider the case when the angular distance between a key and a query is small.
In this case, the similarity (between the key and the query) as well as the softmax probability will be large.
When viewed through the lens of a nearest neighbor retrieval, the above property coincides with a large collision probability of high similarity key-query pairs, assuming that the neighbor retrieval is implemented via LSH.
Motivated by the link between softmax probability p and LSH collision probability q , Spring & Shrivastava (2017) and Charikar & Siminelakis (2017) suggest using LSH as an efﬁcient sampler for low variance softmax estimators.
(a) Spring & Shrivastava (2017) propose approximating softmax by sampling a set, S , a collection of neighboring keys for each query formed by the union of colliding keys using m hash tables.
The estimator is computed using | S | − 1 (cid:80) i ∈
S p ( q , k i ) q ( q , k i ) v i where q is a query vector,
k i , v i are key and value vectors in the sampling set S , and p ( · , · ) and q ( · , · ) are softmax probability and collision probability of given pairs.
The procedure is equivalent to performing importance sampling without replacement, which involves a dependency among the samples.
Deduplication (avoiding double counting) requires memory to store keys in each hash table and runtime to deduplicate keys for each query.
If the size of hash buckets is skewed, the GPU memory needs depend on the size of the hash bucket and the runtime depends on the size of S .
(b) Charikar & Siminelakis (2017) proposed a Hash based Estimator to simulate a proposal distribution for importance sampling via LSH, which can be easily applied in the context of softmax.
For each hash table, a key is uniformly selected from the bucket that the query is hashed to, for simulating a draw from a proposal distribution.
The estimate is computed as m − 1 (cid:80) mi
=1 p ( q , k i ) | H i ( q ) | q ( q , k i ) v i where | H i ( q ) | denotes the size of hash bucket in the i -th hash table which q is hashed to.
This simulates m samples drawn with replacement from the proposal distribution.
However, the probability of one key being sampled depends not only on (a) the angular distance to the query but also (b) the number of keys within the hash bucket, leading to a sampling dependency among all keys.
Further, using it for self-attention causes a dependence between the sparsity in the softmax matrix and the number of hashes used.
Speciﬁcally, the number of tokens that each query can attend to is bounded by the number of hashes: the procedure samples at most one key for each hash table and so, it adds one additional nonzero to the softmax matrix, at most.
Remark 1.
While LSH-based importance sampling exploits the agreement between high probability p ( · , · ) and high collision probability q ( · , · ) , the alignment is not perfect.
Samples from proposal distribution must be reweighted to compensate for the difference.
Further, for different queries, the likelihood ratios between softmax distribution and proposal distribution w.r.t.
a single key are different.
Therefore, the reweighing has to be done during querying.
Although maintaining hash tables for storing keys is not a major problem in general, the high memory cost for hash tables and computation time for reweighing would inﬂuence efﬁciency when applied to self-attention.
YOSO-A TTENTION
We start from LSH-based importance sampling and seek to address some of the aforementioned issues when it is deployed for approximating self-attention.
Instead of using LSH to simulate sampling from a proposal distribution over tokens, we view attention as a sum of tokens associated with
Bernoulli random variables.
This modiﬁcation relates better to LSH and less with LSH-based importance sampling – the probability of one query colliding with a key is not based on other keys.
This strategy helps avoid the sampling dependency issue in LSH-based importance sampling and offers an opportunity to develop a strategy more amenable to GPUs.
Remark 2.
We assume that the input keys and queries of self-attention are unit length – to unify dot-product similarity in self-attention and cosine similarity in LSH.
This is simple using Neyshabur & Srebro (2015): a temperature variable τ is used to bound the squared (cid:96) 2 norm of all queries and keys and to reconstruct new unit length keys and queries while preserving their pairwise similarities.
We can work with the softmax matrix in angular distance metric and derive our algorithm.
Self-Attention via Bernoulli Sampling .
We aim to approximate self-attention, which uses a softmax matrix to capture the context dependency among tokens via their pairwise similarities.
Assuming that we can represent this context dependency directly using collision probability q ( · , · ) , then the challenges discussed in importance sampling can be resolved.
The coincidence of softmax probability p ( · , · ) and LSH collision probability q ( · , · ) makes q ( · , · ) a sensible starting point for approximating self-attention.
Speciﬁcally, to model dependency based on similarity, the collision probability aligns well with the exponential function in softmax in the domain of interest [ − 1 , 1] in Figure 1: both functions have positive zeroth, ﬁrst and second order derivatives.
Note that (a) positive zeroth order derivative indicates that the dependency is positive, (b) positive ﬁrst order derivative ensures that the dependency based on similarity is monotonic, and (c) positive second order derivative means that low similarity corresponds to almost no dependency.
This leads us to hypothesize that a collision-based selfattention may be as effective as softmax-based self-attention.
It can be formulated as,
[Equation]
where B i ( q , k i ) is a Bernoulli random variable where the success probability is given by the collision probability of q with the keys k i .
Hence, it can be determined by the similarity between q , k i .
In a single hash, each B i ( q , k i ) generates a realization to determine whether the corresponding token will be part of attention output or not.
Conceptually, when sampling from softmax distribution, only one token is sampled as the attention output.
In contrast, Bernoulli sampling determines whether each individual token is a part of the attention output.
In principle, to determine the context dependency among tokens, you only need to sample once (YOSO) using a single hash to generate realizations of all Bernoulli random variables, B i ( q , k i ) , i = 1 , . . . , n .
Speciﬁcally, when keys are hashed to a hash table using a single hash, the realization of B i ( q , k i ) for each query q will beif q collides with k i , otherwise it will be 0 .
To our knowledge, using LSH collision probability to replace softmax dependencies for self-attention has not been studied before.
YOSO-Attention .
By replacing softmax dependency with Bernoulli random variables and using
LSH as an efﬁcient sampler to estimate the success probability, we achieve an efﬁcient self-attention (YOSO-Attention) to approximate softmax-based self-attention.
[Equation]
where B ( Q , K ) is the Bernoulli sampling matrix using m hashes.
[Equation]
where f k , k = 1 , . .
, m are hash functions.
Normalizing Attention .
In standard self-attention, each row of the softmax matrix is normalized so that the dependencies sum up to 1 .
In the above, we have discussed how the pairwise query-key dependencies can be estimated using Bernoulli sampling.
We now present how to normalize the dependency in our method as standard self-attention.
We can ﬁrst estimate the dependencies and then normalize them using the sum of estimated dependencies estimated by B ( Q , K ) 1 where 1 is a
Value Key Code vk3 vk3 vk1 vk
vk0 vk3 vk0 v
[Figure]
[-2 q 
-]
vector of all entries being 1 .
B ( Q , K ) 1 can be computed by Eq. 4 by plugging 1 into V .
To make the estimation of self-attention more efﬁcient, we turn to adopt a (cid:96) 2 normalization to the attention output, similar as Levy et al.
(2015) to use (cid:96) 2 normalization for word embedding.
Thus, attention outputs are invariant of the scaling, B ( Q , K ) 1 , under (cid:96) 2 normalization.
Therefore, we have,
[Equation]
Empirically, we show the (cid:96) 2 normalization does not affect the performance of our method as expected, which can be seen in Fig.
LSH-based Bernoulli Sampling .
Now we discuss how to implement the procedure of using
Bernoulli sampling to approximate self-attention.
While a standard LSH procedure can be used, maintaining hash tables to store keys is inefﬁcient on a GPU – the GPU memory size required for hash table cannot be predetermined and the workload might be skewed due to skewed bucket sizes.
To tackle this issue, we propose LSH-based Bernoulli Sampling by only saving the summation of values corresponding to hashed keys instead of storing a collection of hashed keys. The overview of our algorithm is shown in Figure 2.
To compute Y = B ( Q , K ) V , the procedure proceeds as follows.
For each k ∈ [1 , . . .
, m ] , we sample a hash function f k and create a hash table H k ∈ R 2 τ × d representing 2 τ d -dimensional buckets.
For each key K j, : , we add the value V j, : to the bucket whose index is
hash code f k
K j, : ) , denoted as H kf k
( K j, : ) ,
[Equation]
Note that the size of H k is O (2 τ d ) and is independent of which bucket keys are hashed.
With all keys processed for k ∈ [1 , . . .
, m ] , for each query
Q i, : , we maintain an output vector Y i, : initialized to 0 .
Then, we allocate the bucket in H k using f k ( Q i, : ) for k ∈ [1 , . . .
, m ] and add all corresponding results in buckets to the output vector
Y i, : as
[Equation]
Therefore, each ﬁnal output Y i, : can be computed as,
[Equation]
Remark 3.
The memory and time complexity of this algorithm are O ( m 2 τ d ) and O ( nmd ) respectively,
In addition, both time and memory are independent of the size of hash buckets.
Further, We can improve the memory complexity to O ( m 2 τ ) by reusing hash table and processing a few dimensions each time without increasing time complexity.
The constant τ is small as it controls the decay rate of attention weight with respect to the angular distance between query and key, and it can be chosen to be a function of log 2 ( n ) .
In our experiments, τ is set to log 2 ( n ) .
Speed-up .
While not essential, we ﬁnd that a fast random projection for computing the LSH hash code will be beneﬁcial, since this step takes a large portion of the overall runtime.
As suggested by Andoni et al.
, we use the approximated random projection to reduce time complexity to
O ( nmτ log 2 ( d )) , allowing fast computation of hash codes.
Backpropagation through YOSO-Attention .
For training, we also need to show backward propagation steps for YOSO-Attention.
Here, we discuss this last component of YOSO-Attention which enables end-to-end and efﬁcient training.
For backpropagation, the gradient of the loss L w.r.t.
V can be estimated similar to equation 4,
[Equation]
Time Memory Forward
Backward
Forward
Backward
[Equation]
[Table]
[Equation]
[Table] -Attention
[Equation]
The gradients of L w.r.t.
Q , K are similar, so we only provide the expression for Q ,
[Equation]
where (cid:11) , (cid:12) are element-wise division and multiplication.
The problem with the true gradient is that it goes to inﬁnity as the alignment score between the query and the key approaches 1 , which might lead to divergence.
To avoid this numerical issue, we use a lower bound of the actual derivative of the collision probability, [[( ∇ YOSO L ) V T ] (cid:12) τ 2 (1 − arccos( QK T ) π ) τ ] K , see Fig. 1, which can be efﬁciently estimated via a variation of LSH-based Bernoulli Sampling.
Speciﬁcally, note that the approximation can be decomposed into sum of d LSH-based Bernoulli Sampling,
[Equation]
Therefore, following LSH-based Bernoulli Sampling, the memory complexity is O ( m 2 τ d 2 ) , and time complexity is O ( nmd 2 ) .
The d 2 term can be eliminated by repeatedly using the same hash tables d 2 times without increasing runtime, which improves the memory complexity to O ( m 2 τ ) .
The overall complexity of our method and comparison to standard self-attention are summarized in Table 1.
Further, to address the quadratic dependence on d , in the Appendix, we will discuss a scheme to estimate the same quantity but is linear in d .
R ELATED W ORKS
There are a number of efforts describing ways to reduce the quadratic cost of self attention w.r.t.
input sequence length.
Among these works, Linformer (Wang et al., 2020) suggests that low rank attention might be sufﬁcient and adds linear projections (on the sequence) to ﬁxed size keys and values.
There are also other low rank approximation ideas (Katharopoulos et al., 2020), (Choromanski et al., 2020) using separable functions on queries and keys to replace softmax self-attention.
By assuming the self-attention rank to be independent of input sequence length, these methods can achieve O ( n ) time and memory complexity.
Another direction is to exploit the sparsity of softmax matrix and focus on certain sparsity patterns by only computing softmax dependencies within those patterns, including Sparse Transformer (Child et al., 2019), Longformer (Beltagy et al., 2020), and
Big Bird (Zaheer et al., 2020) and Reformer (Kitaev et al., 2020).
Note that, instead of using LSH as a tool to approximate nearest neighbor search to dynamically determine the sparsity pattern in Reformer, our YOSO-attention takes advantage of the connection of query-key similarity to the LSH collision probability to model the dependency among tokens.
E XPERIMENTS
In this section, we provide the empirical results for the proposed approach.
To evaluate our proposed method, we follow the BERT language model pretraining procedure (Devlin et al., 2018) and evaluate the performance of our method in both intrinsic tasks and multiple downstream tasks in
GLUE benchmark as well as runtime and memory relative to standard self attention.
Previously, we assumed that queries and keys are unit length and described the construction to make it work.
In the experiments, we found that simply applying a (cid:96) 2 normalization on queries and keys and using a temperature τ as a hyperparameter does not degrade the performance of model and yet is more efﬁcient to compute, so we use the simpler version in the experiments.
BERT Pretraining .
Following Devlin et al.
(2018), the model is pretrained on BookCorpus (Zhu et al., 2015) and English Wikipedia.
To evaluate the capacity of model capturing the sentence level information, instead of using Next-Sentence-Prediction (NSP) as sentence level loss in the original
BERT, we adapt the Sentence-Ordering-Prediction (SOP) proposed in ALBERT (Lan et al., 2019) as a more difﬁcult task compared to NSP.
All model are trained with Mask-Language-Modeling (MLM) and SOP objectives.
We used the same hyperparameters for pretraining as Devlin et al.
However, due to the computational resources limit, all models are trained for 500 K steps.
The batch size is set so that around 2 17 tokens are processed per step.
(batch size of 256 for sequence length 512 , and batch size of 32 for sequence length 4096 ).
Each model is trained on an AWS
P3dn.24xlarge instance.
Number of Hashes during Pretraining .
Since the approximation variance decreases as the number of hashes increases, for BERT-base pretraining with YOSO-Attention, we use four different hash settings: 16 hashes, 32 hashes, 64 hashes, and expectation of collision to simulate inﬁnite hashes to demonstrate our method.
We plot MLM validation perplexity and SOP validation loss curves oflength model pretrained with softmax self-attention and YOSO-Attention in the right plots of Figure 3.
The curves of our method using expectation match and slightly exceed softmax self-attention, indicating our method is indeed as capable as self-attention.
It is expected that as the number of hashes increase, the performance of our method will approach the curve using expectation as the approximation become more accurate.
For both MLM and SOP, we conﬁrm that our method is as effective as softmax self-attention.
Number of Hashes during Validation .
YOSO-Attention is a stochastic model.
To make the inference deterministic, as in dropout (Srivastava et al., 2014), we take the expectation as our output.
However, directly computing expectation involves a O ( n 2 ) cost, so we experiment with the effect of different hash settings in validation and simulate expectation as the number of hashes increases.
We plot the MLM perplexity and SOP loss of the same pretrained models using different number of hashes on validation in the center plots of Figure 3.
We observe that as the number of hash increases, the MLM perplexity and SOP loss generally decreases for all pretraining hash settings, which is expected.
Longer Sequences .
To examine whether our method can scale linearly with sequence length, we continue to pretrain BERT-base models using the corresponding 500 K step checkpoints forlength model, and add additional positional embedding as suggested by Longformer.
We observe that compare to 512 sequence length, the small performance gap between YOSO-Attention and softmax self-attention does not increase as suggested in the left plots of Figure 3, providing evidence that the number of hashes can be chosen independent of sequence length.
GLUE Tasks .
In addition to intrinsic tasks, we examined the effectiveness of our method on diverse downstream tasks compare to standard attention and show our method can compete with standard attention even after ﬁnetuning.
We ﬁnetuned all pretrained BERT-base model on MRPC Dolan &
Brockett (2005), RTE Giampiccolo et al.
(2007), SST-2 Socher et al.
QNLI Rajpurkar et al.
QQP Chen et al.
(2018), and MNLI Williams et al.
(2018) tasks in the GLUE benchmarks and report their corresponding Dev metrics.
For large datasets including QNLI, QQP, and MMNL, due to computational resource limits, we cannot do hyperparameter search, so we use a batch size of 32 and optimizer AdamW with learning rate 3e-5 to update our model and ﬁnetune our models for 4 epoches.
For MRPC, RTE, and SST-2, we follow BERT ﬁnetuning to do a hyperparameter search with candidate batch size [ 8, 16, 32 ] and learning rate
[ 2e-5, 3e-5, 4e-5, 5e-5 ] and select the best Dev set result.
Results are listed in Table 2.
We observed that the performance of our method on downstream tasks is comparable with standard self-attention, and even has slightly better results in some hash settings.
Further, the downstream performance of our method generally increases as more hashes are used during pretraining, providing an adjustable trade off between efﬁciency and accuracy. [Table]
[Table]
Comparison to Other Baselines .
Apart from comparing YOSO to standard selfattention, we also evaluated its competitiveness with other efﬁcient attention methods.
To keep the ﬁnancial costs of running these experiments reasonable, instead of training all methods from scratch, we used RoBERTa-base’s pretrained weights as the starting point and trained each model using a batch size 512 and learning rate 5e-5 on BookCorpus (Zhu et al., 2015) and English Wikipedia for 95 K steps.
Then, we ﬁnetuned the models on SST-2, QQP and MNLI. These results are shown in Table 3. We observed that our performance is competitive with other baselines while the memory consumption of YOSO is much less ( 2 . 6 × , 1 . 9 × , 2 . 1 × memory savings compared to Reformer, Longformer, and Linformer respectively, see Table 4).
This has potential ramiﬁcations for training such models with more moderate hardware resources which are much less expensive
[Figure]
[Figure]
[Figure]
[Figure]
Estimation Error .
To assess the effectiveness of our algorithm, using
Q , K from the trained model, we generate attention matrices using our algorithm with different number of hashes and compare it against standard self-attention.
In Figure 5, visually, we see that our method produces similar attention patterns as standard self-attention.
The estimation of attention matrix becomes more accurate as the number of hashes increases.
Further, each output of
YOSO-Attention is a weighted sum of random variables as shown in equation 4; so one may suspect that as the sequence length increases, the variance of YOSO-Attention output might potentially increase.
We did not observe this behavior which may be partly due to the hyperparameter τ
= O (log( n )) that controls the decay rate of LSH collision probability as the similarity changes.
We can also ask whether or not the estimation error of YOSO-Attention for a ﬁxed number of hashes increases as the sequence length increases.
We use Q , K , V generated by the pretrained model and estimate the error between N-YOSO ( Q , K , V ) and E [ N-YOSO ( Q , K , V )] .
As the left plot of Fig.
4 suggests, the relative error of our method stays almost constant as the sequence length increases from 128 to 4096 .
This indicates that using sampling to estimate attention weight based on YOSO-Attention can scale up with sequence length and preserve the same estimation quality without increasing the number of hashes.
Runtime and Memory .
We measure the runtime of our method as sequence length increases.
To show the trend more precisely, we measured the runtime per token as shown in Figure 4 (right).
There is a slight increase in runtime per token as the sequence length increases, but note that the xaxis of the plot is log scale, so the increment is small.
When the sequence length increases by 32 × , the runtime per token only increases by 30% , which is explained by our choice of hyperparameter τ
= O (log( n )) .
Aside from the plot, we report the training and testing efﬁciency of our method as well as three other efﬁcient attention methods against standard self-attention.
The results were measured using Q , K , V of a speciﬁed sequence length generated by a trained model and fed into a BERT-base Multi-Head Attention module multiple times.
The experiments were performed on a single NVIDIA 2080TI.
From Table 4, we can see that while for a standard 512 length sequence, our method has a similar runtime as self-attention, as the sequence length increases, the speed-up and memory savings become signiﬁcant.
While our method offers similar runtime savings as other efﬁcient attention methods, the memory consumption for training (i.e., Backward-Cache) of our method is much lower than all other methods in almost all settings.
[Table]
-Peak: Similar
-Peak
C ONCLUSION
We presented a transformer-based model, YOSO-Attention, that scales linearly in the number of input tokens.
This allows the model to be applicable to a wide range of long document NLP tasks.
Via a randomized sampling based scheme, our model approximates self-attention as a sum of individual tokens associated with Bernoulli random variables that can be sampled at once by a single hash, in principle.
With speciﬁc modiﬁcations of LSH, YOSO-Attention can be efﬁciently deployed within a deep learning framework and various aspects of this idea and our implementation, we expect, will ﬁnd use in other novel settings and applications (e.g., in vision).
A A PPENDIX
In Appendix we provide some details of our method that are left out in the main text.
Backpropogation Derivation When using expectation of LSH collision as attention weights, the attention of one query q to keys k i and associated values v i for all i ∈ { n } is deﬁned as [Equation] 13) then we want to compute the gradient of loss w.r t. q , which we denoted as ∇ q L , with the gradient of loss w.r.t. y denoted, ∇ y L , given.
We start by computing the p -th entry of ∇ q L :
[Equation]
Then use ddx
[Equation]
1 −
[Equation]
and plug it into the Eq.
[Equation]
After swapping the order of two summations, Eq. 13 becomes ∂L
[Equation]
Note that only k ip is different for different entries of ∇ q L , so we can write it as
[Equation]
Equation 11 is the matrix form of above
[Equation]
Note that π (cid:112) 1 − ( QK T ) 2 approaches to 0 as alignment score between the query and the key approaches 1 , so we use the fact that 12 (1 − arccos( x ) π
[Equation]
1 π √ 1 − x 2 for x ∈ [ − 1 , 1] and deﬁne a lower bound to replace the actual gradient
[Equation] Approximating Random Projection in LSH . In the main text, we discussed how to estimate selfattention using Bernoulli sampling via LSH.
The ﬁrst step of using LSH is computing hash code using random projection.
To compute hash codes for a vector x , we proceed as follows.
[Equation]
where R ∈ R ( mτ ) × d , R ij ∼ N (0 , 1) , then the output vector are partition to m τ -dimensional binary hash code.
The time complexity for random project is O ( nmτd ) .
To efﬁciently approximate random projection, we follow the construction used in Andoni et al.
The output of mτ -dimensional vector is divided to mτd d -dimensional vectors, then hash codes are estimated by
[Equation]
(21) where D ji are diagonal matrices with entries uniformly sampled from {− 1 , +1 } , and H is
Hadamard matrix.
This approximation reduce time complexity to O ( nmτ log 2 ( d )) .
Alternative Procedure for Approximating Backpropagation .
In the main text, we provided a procedure as shown in Eq. 4, which use LSH-based Bernoulli Sampling d times as subroutine.
The complexity of this procedure is linear w.r.t.
sequence length n , which is desirable but the runtime can be large if d is relatively large.
Therefore, we provide second procedure, which is linear with respect to d .
The gradient of L w.r.t.
the i -th row of Q is written as
[Equation]
Note that if B ( Q , K ) i,j is zero then the corresponding summation term does not need to be computed.
The alternative procedure counts the number of success in m samples at each entry
B ( Q , K ) i,j and only computes the summation term when B ( Q , K ) i,j is non-zero, and thus the runtime is O ( nnz ( S ( A, B ))
( m + d )) (counting number of success + computing nonzero terms).
In the worst case, nnz ( B ( Q , K )) = n 2 , it would be as expensive as dense matrix multiplications in complexity and even worst in practice due to large memory latency resulted from indirect memory access.
However, in practice, B ( Q , K ) is generally sparse if τ is set properly.
Further, the ﬁrst procedure guarantees a linear complexity scaling of our method for extremely long sequences.
As an improvement, we can dynamically select one from these two method based on runtime, than the time complexity is O (min( nmd 2 , nnz ( B ( Q , K ))
( m + d ))) .