Sparse Attention with Learning to Hash

1 INTRODUCTION . The Transformer architecture ( Vaswani et al. , 2017 ) has been successfully applied to various tasks , including natural language processing ( Vaswani et al. , 2017 ; Devlin et al. , 2018 ; Dai et al. , 2019 ; Liu et al. , 2019 ; Yang et al. , 2019 ) , computer vision ( Carion et al. , 2020 ; Dosovitskiy et al. , 2020 ) , and time series forecasting ( Zhou et al. , 2020 ; Wu et al. , 2021 ) . Such a success is mainly due to the self-attention component , which enables each token to directly interact with any other tokens in the entire sequence . But self-attention has a quadratic time and space complexity with respect to the sequence length and hence does not scale efficiently to long sequences as a result . To address this inefficiency problem , one of the solutions is to approximate the full attention matrix with a sparse one , as the softmax operation is dominated by the largest elements . Some recent efforts focus on dynamic learning of sparse attention patterns via Approximate Nearest Neighbors ( ANN ) approaches , including Locality Sensitive Hashing ( LSH ) ( Kitaev et al. , 2020 ; Daras et al. , 2020 ) and mini-batch spherical k-means ( Roy et al. , 2021 ; Wang et al. , 2020a ) . The queries and keys are hashed or clustered into different buckets , hoping that the queries and keys in the same bucket are similar with a high probability . The effectiveness of those ANN approaches rely on the assumption that the ( transformed ) query and key vectors should lie in the same space , which could be sub-optimal in dealing with different sparsity patterns , as analyzed in this paper ( Section 3.2 ) . Besides , the hash functions in LSH are randomized and data-agnostic , which can not fully utilize the rich information in real-world data distributions . In this paper , we address the above limitations of existing ANN-based methods for attention sparsification . Firstly , we analyze two imbalance issues in LSH-produced sparse attention patterns , i.e. , unbalanced hash bucket sizes and unbalanced query-key ratios . Secondly , we design a new metric called attention utility to quantify how well the sparse patterns approximate the full attention , and we show that ANN-derived sparse patterns are substantially inferior to their counterparts . Thirdly , we propose a novel solution , namely Learning-to-Hash Attention ( LHA ) , for dynamic attention sparsification with enhanced model expressiveness . LHA directly optimizes our newly defined attention utility metric in an end-to-end manner via separate learnable hash functions for queries and keys , respectively . As for reducing the computational complexity in the training phase , LHA uses unbiased kernelized attention techniques ( Choromanski et al. , 2020 ; Peng et al. , 2021 ) to efficiently approximate the attention utilities . Similar to other sparse attention models ( Kitaev et al. , 2020 ; Roy et al. , 2021 ) , LHA reduces the overall complexity of self-attention from O ( N2 ) to O ( N1.5 ) for sequence length N . Our experiments in a wide range of tasks on the evaluation benchmarks for language modeling , natural language understanding , and Long-Range-Arena show that LHA achieves better performance compared to strong transformer baselines . 2 RELATED WORK . Related work can be roughly divided into three categories , i.e. , location-based sparse attention , content-based sparse attention , and dense approximation of attention , as outlined below . Location-based sparse attention methods aim to improve the computational efficiency by using pre-specified global or local sparsification patterns over token locations . Liu et al . ( 2018 ) proposed to alternate coarse attention layers and local attention layers . Child et al . ( 2019 ) used a strided sparse attention pattern in image generation . Sukhbaatar et al . ( 2019 ) imposed sparsity based on the predicted temporal window size for each token . Other methods ( Zhang et al. , 2021 ; Beltagy et al. , 2020 ; Ainslie et al. , 2020 ; Zaheer et al. , 2020 ) used a pre-specified subset of locations in the input as the global memory , and only allow non-local attentions from this subset to all the other tokens . Location-based sparse attention can not leverage more flexible content-based interactions among arbitrary positions , as a limitation . Content-based sparse attention allows more flexible sparse patterns than location-based ones . Malaviya et al . ( 2018 ) used sparsemax to obtain a sparse attention matrix , while ( Correia et al. , 2019 ) used entmax . These methods require to compute the full attention matrix before sparsification and hence can not reduce the quadratic computation complexity . Approximate Nearest Neighbor ( ANN ) methods address this limitation by calculating the content-based sparse patterns in advance ( Roy et al. , 2021 ; Kitaev et al. , 2020 ; Daras et al. , 2020 ; Wang et al. , 2020a ) ( which will be discussed more in Section 3.1 ) . Those methods usually apply an ANN module as the shared hash function to both queries and keys . Vyas et al . ( 2020 ) and Zhou et al . ( 2020 ) sparsified the attention maps by eliminating redundant queries . Tay et al . ( 2020b ) designed a differentiable sorting algorithm of internal representations to enable efficient quasi-global local attention . Contemporary to our work , SparseFinder ( Treviso et al. , 2021 ) learns sparse attention patterns that approximate entmax attention , but its bucketing strategies are still based on ANN approaches . In contrast , LHA directly predicts a bucketing strategy ( i.e. , learnable hash functions ) that maximizes the attention utility . Another line of research explored low-rank or kernelized dense approximation of attention matrices , instead of computing the attention scores exactly for only a few pairs . Wang et al . ( 2020b ) applied a low-rank decomposition to the attention matrix . Xiong et al . ( 2021 ) approximated the attention matrix with Nyström approximation . Katharopoulos et al . ( 2020 ) utilized the association property of Key-Query-Value multiplication and reduce the quadratic complexity to linear complexity with kernelized approximation to the softmax operation . Choromanski et al . ( 2020 ) and Peng et al . ( 2021 ) further proposed an unbiased approximation of softmax with random Fourier features . Our work in this paper is directly related to the second category , i.e. , content-based attention sparsification . Specially , we address the limitations of existing ANN-based methods by modeling queries and keys in separate vector spaces and by proposing a novel approach to learn the hash functions for attention sparsification . 3 RE-EXAMINATION OF CONTENT-BASED SPARSE PATTERNS . 3.1 PRELIMINARY . The self-attention mechanism ( Vaswani et al. , 2017 ) can be formulated as the weighted sum of the value vectors V ∈ RN×dh where the weights are calculated using query vectors Q ∈ RN×dh and key vectors K ∈ RN×dh as : Attention ( Q , K , V ) = A · V = softmax ( QKT√ dh ) · V , ( 1 ) where A denotes the matrix of normalized attention weights , dh is the dimension of hidden representations , and N is the sequence length . self-attention has a quadratic time and space complexity with respect to the sequence length and hence does not scale efficiently to long sequences . Content-based sparse attention methods usually apply randomized hash functions or a clustering algorithm to queries { Qi } and keys { Kj } , and hope that similar queries and keys are hashed or clustered into the same bucket . The queries can thus only attend to the keys if both are in the same bucket . Formally , a content-based sparse attention strategy with B hash buckets is defined as : Sparse-Attention ( Qi , K , V ) = ∑ j : hQ ( Qi ) =hK ( Kj ) ĀijVj , ( 2 ) where hK , hQ : Rdh 7→ [ B ] are the hash functions for keys and queries , and Āij ∝ Aij is the re-normalized attention weights such that ∀i , ∑ j : hQ ( Qi ) =hK ( Kj ) Āij = 1 . In general ( Kitaev et al. , 2020 ; Roy et al. , 2021 ) , calculating the hash function and performing local attention for each query have the time complexity of O ( B ) and O ( N/B ) , respectively . Thus , the overall complexity1 of self-attention can be reduced fromO ( N2 ) toO ( N ·B+N2/B ) ≈ O ( N1.5 ) whenB ≈ N/B ≈ √ N . Since the hash functions are not differentiable , Approximate Nearest Neighbor ( ANN ) methods are used to derive an effective content-based hash function . Reformer ( Kitaev et al. , 2020 ) applies Locality Sensitive Hashing ( LSH ) to the tied queries and keys , where several hyper-planes are randomly generated to divide tokens into different buckets . SMYRF ( Daras et al. , 2020 ) improves Reformer by introducing asymmetric transformation to queries and keys , i.e. , F ( Qi ) = [ Qi ; 0 ; √ M2Q +M 2 K − ||Qi||22 ] , G ( Kj ) = [ Kj ; √ M2Q +M 2 K − ||Kj ||22 ; 0 ] , ( 3 ) where MQ = maxQi ||Qi||2 and MK = maxKj ||Kj ||2 , such that ||F ( Qi ) −G ( Kj ) ||22 = const− Qi · Kj . Routing Transformer ( Roy et al. , 2021 ) and Cluster-former ( Wang et al. , 2020a ) use mini-batch spherical k-means to partition tokens into different clusters . 3.2 BUCKET IMBALANCE ISSUES . Previous content-based sparse attention models take it for granted that the ANN-derived sparse pattern can effectively approximate the full attention . However , it is only verified via empirical evaluation on the down-stream tasks yet , which can not reflect the true attention map approximation ability . Notice that there are two necessary conditions for sparse attention to work effectively and efficiently : 1 . The number of queries and the number of keys in each bucket should be reasonably balanced , as queries should attend to enough keys to get a good approximation of the full-attention ; 2 . The bucket sizes should be nearly equal in order to effectively reduce the overall complexity . We first analyze how badly the two conditions would be violated by LSH2 , a typical ANN method . We apply LSH to 10 attention heads in the 3rd layer of a Transformer3 pre-trained on language modeling and obtain the results shown in Figure 1 ( up ) . We can see that the imbalance issue not only exists in the query-key ratios , but also in the bucket sizes . To go a step further , we apply LSH to all 16× 10 = 160 attention heads in the pre-trained Transformer and find that around 61.3 % buckets have the query-key imbalance problem , where the query-key ratios are either greater than 2:1 or smaller than 1:2 . Around 35.9 % buckets have the bucket size imbalance problem , where the bucket 1Notice that we only consider the setting of single-round hashing in our paper , but our analysis and the proposed LHA method can be generalized to the multi-round hashing setting . 2We use the same LSH technique as in ( Kitaev et al. , 2020 ) , except that we do not impose extra constraints to queries and keys . This is because we would like to develop a plug-and-play replacement for dense attention layers without imposing extra constraints for queries and keys 3The detailed experimental setting can be found in the appendix . sizes are twice greater than half smaller than the expected bucket size of 512 . Clearly , neither of the two aforementioned conditions are well satisfied in this LSH-sparsified Transformer model . There can be several possible reasons that cause the imbalance problem : 1 ) the Euclidean distance metric in ANN methods does not monotonically decrease with the dot-product metric used in attention mechanism ; 2 ) the queries and keys are from different distribution and not normalized , and thus restricts the effectiveness of ANN methods . To investigate the root cause , we apply the SMYRF asymmetric transformation ( Equation 3 ) to queries and keys , which creates a monotonic relation between Euclidean distances and dot-products . The new analysis results are shown in Figure 1 ( down ) . We can see that the asymmetric transformation would only exacerbate the imbalance problem , with respect to both query-key ratios and bucket sizes . Therefore , we can conclude that the root cause of the imbalance problem is the mismatch between the query and key distributions , that could be further magnified by the asymmetric transformation .

Summary: This paper introduces a new method based on learnable hash functions to reduce the O(N^2) cost of self-attention in transformers to O(N^1.5). As previously known for bucket-based approaches, this cost is only achieved when buckets are balanced. The paper investigates the effectiveness of related approaches regarding bucket imbalance issues by showing statistics for several attention heads of a pre-trained transformer.  A precise metric is designed to quantify this notion of efficiency ("attention utility"), and later this metric is optimized by learning separate parameterized hash functions for queries and keys. To be able to optimize this metric, the authors use the unbiased approximation to the softmax function proposed by Performer (Choromanski et al., 2020) as a regularization term. Experiments on several NLP and CV tasks show that the proposed method achieves better results than previous fast-transformers while being faster than a standard transformer.