You Only Sample (Almost) Once: Linear Cost Self-Attention Via Bernoulli Sampling

1 INTRODUCTION . 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 finetuned 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 flexible global receptive field to exchange information among input tokens . While self-attention provides immense benefits , 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 coefficients are determined by the similarities among tokens . While this is empirically beneficial , it involves a sizable resource footprint . For sequence length n , this leads to a O ( n2 ) complexity in both time and memory to compute pairwise similarities among all input tokens . This quadratic cost is a roadblock in attaining potential benefits 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 efficient sampling scheme to estimate self-attention which relies on specific modifications 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 efficient and backpropagation friendly implementation , and has a favorable empirical performance profile 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 . 2 BACKGROUND : SELF-ATTENTION . Self-Attention . Self-attention is a scaled dot-product attention mechanism to capture token dependencies in the input sequence , which can be defined as , A ( Q , K , V ) = softmax ( QWQ ) ( KWK ) T√dh︸ ︷︷ ︸ P V WV = DP exp ( P ) V WV ( 1 ) where Q , K , V ∈ Rn×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 , WQ , WK , WV ∈ Rd×dh are learned parameter matrices , dh is the dimension of hidden embedding , and DP 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 QWQ , KWK , V WV 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 . ( 2017 ) . Multi-Head Selfattention can be formally written as , MultiHead ( Q , K , V ) = Concat ( A1 ( Q , K , V ) , · · · , Ah ( Q , K , V ) ) W ( 2 ) where h is the number of heads , Ai , 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 efficient sampling mechanism . In the following sections , we will first review LSH-based importance sampling and then propose Bernoulli sampling with LSH to estimate self-attention efficiently . 3 IMPORTANCE SAMPLING VIA LOCALITY SENSITIVE HASHING . 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 , specifically , 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 efficient 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 ∑ i∈S p ( q , ki ) q ( q , ki ) vi where q is a query vector , ki , vi 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 ∑m i=1 p ( q , ki ) |Hi ( q ) | q ( q , ki ) vi where |Hi ( 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 . Specifically , 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 influence efficiency when applied to self-attention . 4 YOSO-ATTENTION . 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 modification 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 ` 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 . Specifically , 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 , first and second order derivatives . Note that ( a ) positive zeroth order derivative indicates that the dependency is positive , ( b ) positive first 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 self-attention may be as effective as softmax-based self-attention . It can be formulated as , n∑ i=1 Bi ( q , ki ) vi ( 3 ) where Bi ( q , ki ) is a Bernoulli random variable where the success probability is given by the collision probability of q with the keys ki . Hence , it can be determined by the similarity between q , ki . In a single hash , each Bi ( q , ki ) 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 , Bi ( q , ki ) , i = 1 , . . . , n. Specifically , when keys are hashed to a hash table using a single hash , the realization of Bi ( q , ki ) for each query q will be 1 if q collides with ki , 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 efficient sampler to estimate the success probability , we achieve an efficient self-attention ( YOSO-Attention ) to approximate softmax-based self-attention . YOSO ( Q , K , V ) = B ( Q , K ) V ; E [ YOSO ( Q , K , V ) ] = ( 1− arccos ( QK T ) π ) τ V ( 4 ) where B ( Q , K ) is the Bernoulli sampling matrix using m hashes . B ( Q , K ) i , j = 1 m m∑ k=1 1fk ( Qi , : ) =fk ( Kj , : ) where fk , k = 1 , . . . , m are hash functions . ( 5 ) 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 first estimate the dependencies and then normalize them using the sum of estimated dependencies estimated by B ( Q , K ) 1 where 1 is a 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 efficient , we turn to adopt a ` 2 normalization to the attention output , similar as Levy et al . ( 2015 ) to use ` 2 normalization for word embedding . Thus , attention outputs are invariant of the scaling , B ( Q , K ) 1 , under ` 2 normalization . Therefore , we have , N-YOSO ( Q , K , V ) = ` 2 ( B ( Q , K ) V ) ( 6 ) Empirically , we show the ` 2 normalization does not affect the performance of our method as expected , which can be seen in Figure 3 . 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 inefficient on a GPU – the GPU memory size required for hash table can not 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 fk and create a hash table Hk ∈ R2τ×d representing 2τ d-dimensional buckets . For each key Kj , : , we add the value Vj , : to the bucket whose index is hash code fk ( Kj , : ) , denoted as Hkfk ( Kj , : ) , Hkfk ( Kj , : ) ←H k fk ( Kj , : ) + Vj , : ( 7 ) Note that the size of Hk is O ( 2τd ) and is independent of which bucket keys are hashed . With all keys processed for k ∈ [ 1 , . . . , m ] , for each query Qi , : , we maintain an output vector Yi , : initialized to 0 . Then , we allocate the bucket in Hk using fk ( Qi , : ) for k ∈ [ 1 , . . . , m ] and add all corresponding results in buckets to the output vector Yi , : as Yi , : ← Yi , : +Hkfk ( Qi , : ) , : ( 8 ) Therefore , each final output Yi , : can be computed as , Yi , : = m∑ k=1 n∑ j=1 1fk ( Qi , : ) =fk ( Kj , : ) Vj , : = n∑ j=1 B ( Q , K ) i , jVj , : ( 9 ) Remark 3 . The memory and time complexity of this algorithm are O ( m2τ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 ( m2τ ) 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 log2 ( n ) . In our experiments , τ is set to log2 ( n ) . Speed-up . While not essential , we find that a fast random projection for computing the LSH hash code will be beneficial , since this step takes a large portion of the overall runtime . As suggested by Andoni et al . ( 2015 ) , we use the approximated random projection to reduce time complexity to O ( nmτ log2 ( 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 efficient training . For backpropagation , the gradient of the loss L w.r.t . V can be estimated similar to equation 4 , ∇V L = ( ( 1− arccos ( QKT ) π ) τ ) T ( ∇YOSOL ) ≈ B ( K , Q ) ( ∇YOSOL ) ( 10 ) The gradients of L w.r.t . Q , K are similar , so we only provide the expression for Q , ∇QL = ( ( ∇YOSOL ) V T ) ( τ ( 1− arccos ( QK T ) π ) τ−1 ) ( π √ 1− ( QKT ) 2 ) ) K ( 11 ) where , are element-wise division and multiplication . The problem with the true gradient is that it goes to infinity 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 , [ [ ( ∇YOSOL ) V T ] τ2 ( 1 − arccos ( QKT ) π ) τ ] K , see Figure 1 , which can be efficiently estimated via a variation of LSH-based Bernoulli Sampling . Specifically , note that the approximation can be decomposed into sum of d LSH-based Bernoulli Sampling , ( ∇̂QL ) i , : = d∑ l=1 ( ∇YOSOL ) i , l n∑ j=1 B ( Q , K ) i , j ( Vj , l τ 2 Kj , : ) ( 12 ) Therefore , following LSH-based Bernoulli Sampling , the memory complexity is O ( m2τd2 ) , and time complexity is O ( nmd2 ) . The d2 term can be eliminated by repeatedly using the same hash tables d2 times without increasing runtime , which improves the memory complexity to O ( m2τ ) . The overall complexity of our method and comparison to standard self-attention is 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 .

Summary: The paper proposes to replace the weighted average of the values in standard self-attention with the average of values sampled in a way that the expectation is close to the result of self-attention. In particular, the authors associate with each query-key pair, a Bernoulli random variable with expected value close to the exponential of the dot-product. Sampling these variables and averaging the values per query is formulated in an efficient way using locality-sensitive hashing.