Year: 2021

Abstract:
Transformer has become ubiquitous in sequence modeling tasks. As a key component of Transformer, self-attention does not scale to long sequences due to its quadratic time and space complexity with respect to the sequence length. To tackle this problem, recent work developed dynamic attention sparsification techniques based on Approximate Nearest Neighbor (ANN) methods, where similar queries and keys are allocated to the same hash bucket with high probability. However, the effectiveness of those ANN methods relies on the assumption that queries and keys should lie in the same space, which is not well justified. Besides, some of the ANN methods such as Locality-Sensitive Hashing (LSH) are randomized and cannot fully utilize the available real data distributions. To overcome these issues, this paper proposes a new strategy for sparse attention, namely LHA (Learningto-Hash Attention), which directly learns separate parameterized hash functions for queries and keys, respectively. Another advantage of LHA is that it can be used as a plug-and-play replacement for dense attention layers without imposing extra constraints for queries and keys, which makes it applicable to the wide range of pre-trained Transformer models. Our experiments on evaluation of the WikiText-103 dataset for language modeling, the GLUE benchmark for natural language understanding, and the Lang-Range-Arena benchmark for multiple tasks (text/image classification, retrieval, etc.) show the superior performance of LHA over other strong Transformer variants.

heading: 1 INTRODUCTION
text: 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 cannot 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.

heading: 2 RELATED WORK
text: 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 cannot 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 cannot 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.

heading: 3 RE-EXAMINATION OF CONTENT-BASED SPARSE PATTERNS

heading: 3.1 PRELIMINARY
text: 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.

heading: 3.2 BUCKET IMBALANCE ISSUES
text: 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 cannot 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.

heading: 3.3 PROPOSED METRIC: ATTENTION UTILITY
text: The above analysis shows the imbalance problem in ANN-derived hashing strategies. To further quantify the approximation quality of sparse attention patterns, we utilize the concept of Attention Biclustering from (Daras et al., 2020), which considers a practical case (Kitaev et al., 2020; Roy et al., 2021) where all clusters strictly contain the same number of queries and keys:
∀b ∈ [B], |{Qi | hQ(Qi) = b}| = |{Kj | hK(Kj) = b}| = N
B . (4)
We denote with C the set of all possible assignments in B balanced non-overlapping clusters: C = {C1, C2, . . . , CT }, where T is the number of possible assignments. We can then define the attention utility AU of each assignment Ct as:
AU(Ct) = ∑
i,j:(Qi,Kj)∈Ct
Aij , where Ct = {(Qi,Kj)|hCtQ (Qi) = h Ct K (Kj)}. (5)
The attention utility represents the aggregation of the sparse attention weights in an assignment. As we should keep as much sparse attention weights as possible to better approximate a full attention map, attention utility quantifies how well a sparse attention approximates the full one. In fact, we can show that it is computationally intractable to find the optimal attention utility: Theorem 1. Finding the assignment that achieves the optimal attention utility, i.e.,
arg max Ct∈C
AU(Ct) = ∑
(Qi,Kj)∈Ct
Aij (6)
is NP-hard.
This theorem is a corollary of the NP-hardness of Attention Biclustering in (Daras et al., 2020), and it motivates us to develop a learning-based approach to optimize the attention utility. Notice that similar
to other sparse attention models (Kitaev et al., 2020; Roy et al., 2021), the proposed attention utility metric currently does not take causal masks into consideration. Developing a new metric that can take arbitrary attention masks into account is left for future work.
To measure the attention utility of ANN-derived sparse attention patterns, we follow (Kitaev et al., 2020; Daras et al., 2020) and adaptively set hash boundaries to create balanced clusters. The analyzed models include LSH, LSH with asymmetric transformation (i.e., SMYRF), and our proposed Learningto-Hash (L2H) method which will be introduced in the next section. We also calculate the empirical upper bound by aggregating the attention from each query to its top keys w.r.t dot-products. We present a box-plot of averaged attention utilities for all 16 Transformer layer in Figure 2. We can see that the LSH methods obtain significantly less attention utilities than the upper bound or L2H. As a conclusion, the ANN-derived sparse patterns are arguably sub-optimal when used as hash functions for sparse attention. We believe this also partially explains why most of these approaches need to impose extra constraints (e.g. queries and keys are tied or normalized). In the next section, we propose an alternative choice of the hash functions to tackle this problem.

heading: 4 LEARNING-TO-HASH ATTENTION (LHA)
text: We propose a novel Learning-to-Hash Attention model in this section. The key idea is to learn separate parameterized hash functions for queries and keys, respectively, thus the sparse pattern in LHA can be no longer limited to distance-based hash functions such as LSH or online k-means and adapt to the mismatch of query and key distributions. Besides, we empirically find that LHA can be used as a plug-and-play replacement for dense attention layers, which makes LHA applicable to the wide range of pre-trained Transformer models.

heading: 4.1 LEARNING-TO-HASH FOR SPARSE ATTENTION
text: We first remind the readers of our definitions of the sparse attention: Sparse-Attention(Qi,K, V ) = ∑
j:hQ(Qi)=hK(Kj)
ĀijVj (7)
where hK , hQ : Rdh 7→ [B] are the hash functions for keys and queries, and Āij ∝ Aij is the renormalized attention weights such that ∀i, ∑ j:hQ(Qi)=hK(Kj)
Āij = 1. Inspired by the Learning-toHash methods (Wang et al., 2017), we implement the learnable hash functions hK , hQ : Rdh 7→ [B] by defining parameterized functions HQ, HK : Rdh 7→ RB , such that:
hQ(Qi) = arg max b∈{1,...,B} [HQ(Qi)]b , hK(Kj) = arg max b∈{1,...,B} [HK(Kj)]b , (8)
where HQ and HK can be any parameterized functions such as MLPs.
Notice that our formulation is a generalization of the non-learnable sparse attention mechanisms in several previous content-based sparse attention models. To reproduce the symmetric LSH-based sparse attention scheme in Reformer (Kitaev et al., 2020) or the mini-batch k-means scheme in Routing Transformer (Roy et al., 2021), we can set:
HQ(x) = HK(x) = x[R;−R] or HQ(x) = HK(x) = xµ (9)
where R ∈ Rdh×B2 is a random matrix, [·; ·] denotes the concatenation of two vectors, and µ ∈ Rdh×B are the exponentially moving averaged cluster centroids shared by keys and queries. The asymmetric transformation in SMYRF (Daras et al., 2020) can also be reproduced by:
HQ(x) = 1
(⌊ F (x) · a + b
r
⌋) , HK(x) = 1 (⌊ G(x) · a + b
r
⌋) (10)
where a is a random vector, b is a random scalar, r is a scalar parameter, 1(·) : [B] 7→ {0, 1}B is the one-hot operation, and F (·) and G(·) are defined as in Equation 3.

heading: 4.2 OPTIMIZATION OBJECTIVE FOR LEARNING-TO-HASH
text: Previous ANN-derived sparse attention models construct their hash functions by either randomization or k-mean clustering. In contrast, our hash functions are fully learnable. However, as we use the arg max operation to get the hashing buckets, the parameterized projections in our hash functions cannot be trained in an end-to-end manner. To resolve this issue, we guide the training of our learnable hash functions directly by our proposed attention utility (See Equation 5). Notice that the attention utility of each query Qi can be written as ∑ j:hQ(Qi)=hK(Kj)
Aij . We first calculate the possible attention utilities ψ(i) that query Qi can obtain when allocated to all B hash buckets:
ψ(i) =  ∑ j:hK(Kj)=0 Aij , · · · , ∑ j:hK(Kj)=B−1 Aij  (11) Ideally, our hash function would allocate Qi to the bucket that maximize the attention utility, such that Qi can attend to the most salient keys:
hQ(Qi) = arg max b∈[Nb]
ψ (i) b = arg max
b∈[B] ∑ hK(Kj)=b Aij (12)
To achieve this goal, since ψ(i) is naturally a normalized categorical distribution, we can simply convert HQ(Qi) into a distribution and use the KL divergence between the predicted distribution softmax(HQ(Qi)) and the desired distribution ψ(i) as the optimization objective:
LQi = KL(ψ(i)‖softmax(HQ(Qi))) (13)
On the other hand, the possible attention utilities ψ′(j) that a key Ki obtains when allocated to all B buckets can be written as:
ψ′(j) =  ∑ i:hQ(Qi)=0 Aij , · · · , ∑ i:hQ(Qj)=B−1 Aij  (14) While ψ′(j) is no longer a normalized distribution, we can normalize it and similarly define the optimization objective LKj also as a KL divergence. During the training stage, our final optimization objective is a convex combination of the task-specific objective and the learning-to-hash objectives (i.e., {LQi} and {LKj}) for all queries and keys. Figure 3 in the appendix illustrates the joint training diagram of LHA.

heading: 4.3 APPROXIMATE ATTENTION UTILITIES
text: The remaining problem is how to efficiently computeψ(i) andψ′(j). A naive way requires calculating the dot products between all query-key pairs withO(N2) complexity. Inspired by the recent advances in kernelized attention (Katharopoulos et al., 2020; Choromanski et al., 2020; Peng et al., 2021), we use random Fourier features (Rahimi et al., 2007; Choromanski et al., 2020) to approximate ψ(i) in an unbiased manner. Let us define φ : Rdh 7→ R2D as the Positive Randomized Features (PRFs) (Choromanski et al., 2020) such that
E[φ(x) · φ(y)] = exp ( xTy√ dh ) (15)
We can approximate ψ(i) with the following formula:
ψ (i) b = ∑ j:hK(Kj)=b Aij ∝ ∑ j:hK(Kj)=b exp ( Qi ·Kj√ dh ) ∝ E φ(Qi) · ∑ j:hK(Kj)=b φ(Kj)  (16) for each bucket b ∈ [B]. Since hK(Kj) does not change for the queries, we can simply pre-compute∑
j:hK(Kj)=b φ(Kj) once to save computation, which reduces the complexity of computing {ψ(i)} for all queries from O(N2) to O(N · B) ≈ O(N1.5). The attention utilities for keys ψ′(j) can be efficiently approximated in a similar way4:
ψ ′(i) b = ∑ i:hQ(Qi)=b Aij ∝∼ E  ∑ j:hQ(Qi)=b
φ(Qi) φ(Qi) · ∑N j=1 φ(Kj)
 · φ(Kj)  (17)

heading: 4.4 IMPLEMENTATION DETAILS
text: Practical content-based sparse attention models require that each hash bucket has the same size, which is crucial in terms of computational efficiency on modern hardwares.Therefore, we follow Roy et al. (2021) and sort the tokens with regard to normalized hash scores, i.e., softmax(HQ(Qi)) and softmax(HK(Ki)), in each bucket. The hash bucket membership is then determined by the top-k threshold, where k = NB is the bucket size. Since such a bucketing strategy no longer guarantees the validation of attention bi-clustering, as each query or each key can be assigned to zero or more than one clusters, we further empirically enlarge the hash bucket sizes by √ 2× to increase the recall of queries and keys. A pseudo-code implementation for LHA can be found in the appendix.
In causally masked attentions, the queries cannot attend to keys behind, as is usually the case in language modeling. When we use separate hash functions {hQ(Qi)} and {hK(Kj}) for queries and keys, respectively, it is possible that for some query Qi, there exists no such key Kj in the same bucket that j ≤ i. This would cause serious numerical instability in our implementation. To tackle this problem, inspired by Kitaev et al. (2020) and Roy et al. (2021), we tie the key hashes with query hashes in the case of causal attentions by constructing a joint hash function:
hK(Ki) = hQ(Qi) = arg max b∈[Nb] [HQ(Qi)b +HK(Ki)b] (18)
We find that this strategy empirically works better than using other tricks to fix this “no-attentiontarget” problem. Our solution is different from Routing Transformer (Roy et al., 2021) or Reformer (Kitaev et al., 2020), which impose an extra constraint that queries and keys are tied.

heading: 5 EXPERIMENTS
text: In this section, we conduct experiments to verify the effectiveness of our approach on several benchmark datasets covering language modeling, natural language understanding, and Long-RangeArena. Due to the space limitations, the detailed hyper-parameter settings are presented in the appendix.

heading: 5.1 LANGUAGE MODELING
text: Wikitext-103 (Merity et al., 2016) is a large-scale dataset for testing long term dependencies in word-level language models. It contains 103M training tokens from 28K articles, with an average length of 3.6K tokens per article, which allows testing the ability of long-term dependency modeling. We use this dataset as a probe dataset to perform various ablations to tease apart the effect of various hyper-parameter choices on the model performance.
We follow the base setting of the state-of-the-art Transformer-XL (Dai et al., 2019) model, which contains 16 Transformer layers with 10 heads per layer. For local attentions, we use the relative positional encoding (Dai et al., 2019), while for non-local attentions (i.e., LSH or LHA), we use the
4Please refer to (Choromanski et al., 2020) for how PRFs stabilize attention renormalization.
per-layer sinusoidal absolute positional encoding (Vaswani et al., 2017). The LSH Transformer is implemented in our code base by simply replacing the learnable hash functions by random projections. A sequence length of 1536 is used in both training and evaluation stages for our models. In ablation study, we vary 1) the type of non-local attention heads, 2) the type of learnable hash function in LHA (i.e., linear or two-layer MLP), 3) the number of non-local attention heads, and 4) the number of non-local attention layers. We also consider a baseline where the hash functions generate content-independent random numbers. The experimental results are presented in Table 1.
From the table, we can see that a smaller attention size (384 compared to 768 in local Transformer) would decrease the model performance by 0.9, while adding LSH or LHA heads can improve the performance by 0.3 and 1.0, respectively. Surprisingly, we find that local Transformer is a very strong baseline for language modeling that outperforms a pure LHA Transformer. We can also see that a two-layer MLP works better than a linear function as the learnable hash functions in LHA. The best performance is achieved when we set half of the heads to local attention and the other half to non-local LHA for each layer. Finally, we can see that LHA consistently underperforms the untrainable LSH counterparts under four different settings. For the test set evaluation, we use the best setting from the ablation study. We can see that the proposed LHA model outperforms Transformer-XL, which further validate that LHA can help modeling long-term dependency for language modeling tasks.

heading: 5.2 NATURAL LANGUAGE UNDERSTANDING
text: To show the ability of our model to approximate arbitrarily complicated attention distributions, we evaluate our proposed method on the approximation of RoBERTa model (Liu et al., 2019) on the

heading: Model LO IMDb AAN Image Avg. 1K 2K 3K 4K
text: GLUE (General Language Understanding Evaluation) dataset (Wang et al., 2018). Following the common practice, the maximum sequence length is set to 128. To show the effectiveness of LHA, we choose two competitive baselines in literature: SMYRF (Daras et al., 2020) and Fast Clustered Attention (FCA) (Vyas et al., 2020). We also produce an LSH-based sparse attention baseline. We use a pure LHA model similar to variant (c) in Section 5.1.
We summarize the performance per task in Table 2. We report accuracy for all tasks except STS-B, where we report Pearson correlation. From the table, we can see that LHA performs as good as full attention for all the GLUE tasks, and that LHA outperforms all other methods in the average GLUE score, and has a smaller computational costs in its sparse attention part.

heading: 5.3 LONG-RANGE-ARENA BENCHMARK
text: Long-Range-Arena (LRA) benchmark (Tay et al., 2020c) is a recently proposed benchmark focused on evaluating model quality under long-context scenarios for Transformer variants. We follow the apples-to-apples setting5 of LRA benchmark and compare our method against other efficient attention variants. We use an LHA/local hybrid variant similar to variant (d) in Section 5.1.
We consider the tasks of ListOps (Nangia & Bowman, 2018) (LO), byte-level IMDb reviews text classification (Maas et al., 2011) (IMDb), byte-level document retrieval on ACL Anthology Network (AAN) (Radev et al., 2013), and CIFAR10 (Krizhevsky et al., 2009) image classification on sequences of pixels (Image). The results are shown in Table 3 and the brief descriptions of the compared baselines can be found in Section 2. From the table, we can see that LHA achieves consistent improvements over previous efficient attention models.

heading: 6 CONCLUSION & FUTURE WORK
text: In this paper, we address the limitations of ANN-based sparse attention methods and propose the Learning-to-Hash Attention (LHA) as our new solution. Specifically, LHA leverages separate learnable hash functions for queries and keys, respectively, and utilizes kernelized techniques for efficient approximation of attention utilities. The experiments on language modeling, natural language understanding, text classification, and image classification demonstrated the effectiveness of LHA. For future work, we would like to validate the effectiveness of LHA model on much larger language modeling datasets, such as PG-19 (Rae et al., 2019).
5www.github.com/google-research/long-range-arena#apples-to-apples-setting

heading: ETHICS STATEMENT
text: This work proposed a more effective sparsification pattern for the attention mechanisms. We believe our proposed model as well as the code to be released can most benefit the field of language processing, with the potential to benefit other fields involving long sequence modeling. However, it is know that Transformers are also vulnerable to the adversarial attacks (Michel et al., 2019). Since one of the advantage of our model is to re-use a pre-trained Transformer model, the effective adversarial perturbations on the pre-trained model could also be effective in our model.

heading: REPRODUCIBILITY STATEMENT
text: • We will put a link to an anonymous repository for our source code as the OpenReview comment. The hyper-parameters are described in details in the appendix. We also provide a pseudo-code implementaion of our model in the appendix.
• All the datasets used in the experiments and the corresponding pre-processing scripts can be found online, including language modeling6, GLUE benchmark7, Long-Range-Arena benchmark8, and time series forecasting9. The pre-trained model we used (i.e., RoBERTa) is also publically available10.

heading: Generate attention weights
text: Generate sparse patterns
Optimized towards down-stream tasks
Optimized towards attention utilities
Figure 3: The joint training diagram of LHA, where queries and keys generate attention weights to train the learnable hash functions, while the hash functions generate the sparse attention patterns, on which queries and keys are trained towards down-stream tasks.

heading: A EXPERIMENTAL SETTINGS

heading: A.1 RE-EXAMINATION OF CONTENT-BASED SPARSE ATTENTION
text: For analysis, we pre-trained a 16-layer and 10-head Transformer on the WikiText-103 language modeling benchmark (Merity et al., 2016). A per-layer sinusoidal absolute positional encoding (Vaswani et al., 2017) is injected to the queries and keys before self-attention.
For Figure 1, we analyzed the queries and keys for the first 1024 tokens in the WikiText-103 validation data, and set the number of LSH buckets B = 4. We apply LSH to all the 10 attention heads in the 3rd Transformer layer. For Figure 2, we follow the same experimental setting.

heading: A.2 LEARNING-TO-HASH
text: During the training stage, our final optimization objective L is a convex combination of the taskspecific objective Ltask (e.g., cross-entropy loss in language modeling) and learning-to-hash objectives for all keys and queries, i.e., {LQi} and {LKj} (See equation 13).
L = (1− λ)Ltask + λ
N · L ·H N∑ i=1 L∑ h=1 H∑ h=1 (LL,HQi + L L,H Kj ) (19)
where λ, N , L, H denote the loss coefficient, the sequence length, the number of attention layers, and the number of attention heads per layer, respectively. Empirically, we set Λ = 0.05 and found it work well across different tasks.

heading: A.3 LANGUAGE MODELING
text: We use the base setting of Transformer-XL (Dai et al., 2019) in our experiments, which consists of 16 Transformer layers. For each layer, the hidden size is set to 410, and the number of attention heads is set to 10. The dimension of feed-forward layer is set to 2100. All codes are implemented based on Flax (Heek et al., 2020) in JAX (Bradbury et al., 2018). The dropout ratio is set to 0.2. The batch size is set to 32. We use AdamW (Loshchilov & Hutter, 2017) as the optimizer, and set (β1, β2) to (0.9, 0.999). The peak learning rate is set to 3.5e-4. The model is trained for 20k steps with a 2k-step warm-up stage with a cosine learning rate decay (Loshchilov & Hutter, 2016).

heading: A.4 NATURAL LANGUAGE UNDERSTANDING
text: We use the base setting of RoBERTa (Liu et al., 2019) in our experiments, which consists of 12 Transformer layers. For each layer, the hidden size is set to 768, and the number of attention heads is set to 12. The dimension of feed-forward layer is set to 3072. The RoBERTa pre-trained checkpoint is taken from the Transformers (Wolf et al., 2020) library. We use AdamW (Loshchilov & Hutter, 2017) as the optimizer, and set (β1, β2) to (0.9, 0.98). To fine-tune the pre-trained models, we search the optimization hyper-parameters in a search space including different batch sizes (16/32/48), learning rates ((1-5) * 1e-5), and the number of epochs (3-5). A 10% warm-up schedule and linear learning rate decay (Devlin et al., 2018) is applied.

heading: A.5 LONG-RANGE-ARENA BENCHMARK
text: We follow the apples-to apples setting11 of LRA benchmark and compare our method against other efficient attention variants. Specifically, we use a 4-layer Transformer with 256 hidden size and 1024 feed-forward layer size for the IMDb text classification task(Maas et al., 2011), a 4-layer Transformer with 256 hidden size and 1024 feed-forward layer size for the ListOps task(Radev et al., 2013), a 4-layer Transformer with 128 hidden size and 512 feed-forward layer size for the document retrieval task(Nangia & Bowman, 2018), and a 4-layer Transformer with 128 hidden size and 64 feed-forward layer size for the CIFAR10 image classification task(Krizhevsky et al., 2009). The models are trained for 20k steps with a 2k-step warm-up stage with a cosine learning rate decay (Loshchilov & Hutter, 2016).

heading: A.6 TIME SERIES FORECASTING
text: We follow the experimental setup of Informer (Zhou et al., 2020). Specifically, the input length of recurrent component is chosen from {24, 48, 96, 168, 336, 720} for the ETTh1, and chosen from {24, 48, 96, 192, 288, 672} for the ETTm dataset. The layer of encoder is chosen from {6, 4, 3, 2} and the layer of decoder is set as 2. The head number of multi-head attention is chosen from {8, 16}, and the dimension of multi-head attention’s output is set as 512.

heading: A.7 COMPUTING INFRASTRUCTURE
text: All the model training are conducted on a machine with 4 NVIDIA Ampere A100 40GB GPUs and 64 AMD EPYC 7713 64-Core Processor in a Slurm (Yoo et al., 2003) system. The evaluation of the inference throughput is performed on a stand-alone machine with 1 NVIDIA Tesla V100 32GB GPU.

heading: B ADDITIONAL EXPERIMENTS ON TIME SERIES FORECASTING
text: We also evaluate our model on time series forecasting tasks. We use the ETT (Electricity Transformer Temperature) dataset (Zhou et al., 2020), which contains 2-year data from two separated counties in China. ETTh1 is a dataset for 1-hour-level, while ETTm1 is a dataset for 15-minute-level. Each data point consists of the target value “oil temperature” and 6 power load features. The train/val/test is 12/4/4 months. The Mean Squared Error (MSE) metric and Mean Average Error (MAE) metric are used as the evaluation metrics. We use an LHA/local hybrid variant similar to variant (d) in Section 5.1.
In univariate forecasting setting, all the seven features are used as input and “oil temperature” is the prediction target. The univariate evaluation results can be found in Table 4. We can see that LHA significantly improve the performance of state-of-the-art for most settings, while slightly underperforms Informer (Zhou et al., 2020) in two settings of ETTh1 dataset. This verifies the effectiveness of LHA on time series data.
11www.github.com/google-research/long-range-arena#apples-to-apples-setting
In multivariate forecasting setting, all the seven features are the prediction targets. The multivariate evaluation results can be found in Table 5. We can see that LHA significantly improve the performance of state-of-the-art for all settings, which further demonstrates the effectiveness of the proposed LHA method.

heading: C PROOF
text: We follow the notations in (Daras et al., 2020). Let CB the set of all possible assignments in B balanced non-overlapping clusters. A specific assignment is denoted by CBt and there are T possible such assignments:
CL = {CB1 , CB2 , ...CBT }. (20)
CBt = {c1, c2, ..., cB} :
{ ci = {Qi1, ..., QiN
B ,Ki1, ...,K i N B } ci ⊆ Q ∪K, ∀i ∈ {1, ..., L} cx ∩ cy = ∅ ∀cx, cy ∈ CBt
(21) We have the following lemma, which is referred as the max-mass problem: Lemma 1. (Daras et al., 2020) The optimization problem:
max CBt ∈CB ∑ (Qi,Kj)∈CBt Qi ·Kj (22)
is NP-hard.
The main idea of the proof is to show that solving polynomially the above problem would mean that we could also solve in polynomial time the 3-DM, which is known to be NP-complete. Please refer to (Daras et al., 2020) for the detailed constructive proofs. Next, we show how to use this lemma to prove Theorem 1.
Proof of Theorem 1. Let {Qi} and {Kj} denote the query and key sets that we consider for computing the attention utility. We can first construct the new query set Q′ and key set K′ such that
Q′i ·K ′j = eQi·Kj√ dh , where Q′i ∈ Q′,K ′j ∈ K′
This can be achieved by applying SVD to the matrix of (
eQi·Kj√ dh ) ij . Next, we construct another query
set Q′′, such that Q′′i =
Q′i∑n j=1Q ′ i ·K ′j
The problem of finding the optimal attention utility, i.e.,
arg max Ct∈C
AU(Ct) = ∑
(Qi,Kj)∈Ct
Aij
is thus equivalent to the problem in Lemma 1, with Q′′ and K′ as the query and key sets, which is proven to be NP-hard.

heading: D ALGORITHM
text: We present a detailed pseudo-code implementation for LHA in Algorithm 1.
Algorithm 1 Single-layer Single-head Learning-to-Hash Attention (LHA) 1: Queries, Keys and Values: Q,K, V ∈ RN×dh 2: Query/Key Hash Functions: 3: Attention Utility Loss Weight: λ 4: if causal mask then 5: Qscore ← softmax(HQ(Q) +HK(K)) . B ×N × dh 6: Kscore ← Qscore 7: else 8: Qscore ← softmax(HQ(Q)) . B ×N × dh 9: Kscore ← softmax(HK(K)) . B ×N × dh 10: W ← N/B . bucket size 11: Qidx ← top-{ √ 2W}(Qscore) . B ×N log(N)
12: Kidx ← top-{ √
2W}(Kscore) . B ×N log(N) 13: Q′ ← gather(Q,Qidx) . B ×W × dh 14: K ′ ← gather(K,Kidx) . B ×W × dh 15: V ′ ← gather(V,Kidx) . B ×W × dh 16: A← Q′(K ′)T . B ×W ×W × dh 17: if causal mask then 18: A← mask(A) . B ×W ×W 19: A← softmax(A). . B ×W ×W 20: V ′ ← A · V ′ . B ×W ×W × dh 21: X ← scatter(Kidx, V ′) . B ×W × dh 22: ψ = AUQ(Q,K) . B ×N , Equation 16 23: ψ′ = AUK(Q,K) . B ×N , Equation 17 24: LAU = λ ·KL(ψ‖Qscore) + λ ·KL(ψ′‖Kscore) 25: return X,LAU

heading: E A COMPLETE DERIVATION OF THE LHA OBJECTIVE
text: Let Q,K, V ∈ RN×dh denote the query, key, and value vectors of the attention mechanism, where N is the sequence length and dh is the hidden size. Let φ : Rdh 7→ R2D be Positive Randomized Features (PRFs) (Choromanski et al., 2020) such that:
E[φ(x) · φ(y)] = exp ( xTy√ dh )
We can approximate the aggregated attention utility (Eq. 5) by:
ψ (i) b = ∑ j:hK(Kj)=b Aij ∝ ∑ j:hK(Kj)=b exp ( Qi ·Kj√ dh ) ∝ E φ(Qi) · ∑ j:hK(Kj)=b φ(Kj) 
ψ ′(i) b = ∑ i:hQ(Qi)=b Aij =
∑ j:hQ(Qi)=b exp ( Qi·Kj√ dh ) ∑N
j=1 exp ( Qi·Kj√ dh ) ∝ E [ φ(Qi) · ∑ j:hK(Kj)=b φ(Kj) ]
E [ φ(Qi) · ∑N j=1 φ(Kj) ] ∝∼ E [ φ(Qi) · ∑ j:hK(Kj)=b φ(Kj)
φ(Qi) · ∑N j=1 φ(Kj)
]
∝ E  ∑ j:hQ(Qi)=b
φ(Qi) φ(Qi) · ∑N j=1 φ(Kj)
 · φ(Kj) 
where ∝∼ denotes that the attention approximation converges under the positive random feature condition (Katharopoulos et al., 2020). Next, the Learning-to-Hash objectives for each query and each key can be defined as:
LQi = KL(ψ(i)‖softmax(HQ(Qi))) (23)
LKi = KL(ψ′(i)‖softmax(HK(Ki))) (24)
The final optimization objective L is a convex combination of the task-specific objective Ltask (e.g., cross-entropy loss in language modeling) and learning-to-hash objectives for all keys and queries in each attention head for each layer:
L = (1− λ)Ltask + λ
N · L ·H N∑ i=1 L∑ h=1 H∑ h=1 (LL,HQi + L L,H Kj ) (25)

heading: F TRAINING DYNAMICS
text: We plot the training dynamics of the KL divergence (i.e., LQ and LK) and the negative entropy of the bucket-wise attention utilities (i.e., negative entropy of ψ for queries and ψ′ for keys) in Figure 4. KL divergence is the optimization objective of LHA, while negative entropy measures the diversity of bucket-wise attention utility. The model we use is a pure LHA model with 16 layers and 10 LHA heads, trained on WikiText-103 for 6000 steps. Each LHA head has 8 hash buckets. From the plots, we can see that the values of negative entropy consistently increase during training for both queries and keys, while the values of KL divergence first increase (due to the increase of the diversity of bucket-wise attention utility), and then decrease (due to learning-to-hash optimization).

heading: G ADDITIONAL ANALYSIS ON TRAINING & INFERENCE EFFICIENCY
text: To further analyze the efficiency of the proposed LHA compared to softmax attention, we measure the latency of LHA and softmax attention for both training and inference stage, varying by sequence length. Notice that in the inference stage, the LHA model needs not to calculate the learning-to-hash losses (i.e., the KL divergence terms). We study a single-layer attention, which has 10 attention heads and the hidden size of 410. Figure 5 illustrates the latency in both training and inference stage for
softmax attention and LHA with different attention size. We can see that LHA can achieve more significant speedup when facing longer sequences.

heading: H ADDITIONAL ANALYSIS ON THE BUCKET IMBALANCE ISSUES
text: To further analyze how the bucket imbalance issues are alleviated by the proposed LHA, we plot the histograms of bucket sizes and query-key ratios for all 16× 10 = 160 attention heads in a pre-trained Transformer. The results are shown in Figure 6. We can see that learning-to-hash attention (l2h) has more buckets close to the optimal bucket size (i.e., 512) and has more balanced query-key ratios for hash buckets.
Notice that to compare the statistics of the bucket size and query-key ratios between LSH and LHA, we fine-tune learnable hash functions for the same pre-trained Transformer. Notice that in this experiment, we directly use the highest-ranked bucket as the hash bucket for each query and key, instead of using a token sorting strategy (Roy et al. 2021) to maintain the same bucket sizes. When calculating the query-key ratios, we do not count the buckets which have no queries or keys.
I VISUALIZATION
We visualize the full attention map of an attention head in a Transformer model pre-trained on WikiText-103 dataset in Figure 7, and visualize its LHA counterpart and LSH counterpart in Figure 8 and Figure 9, respectively. We do not apply the causal mask when calculating the attention weights. The attention utility of LHA and LSHare 0.461 and 0.335, respectively. From the visualization, we can see that LHA is able to better approximate the full attention with learnable hash functions.

Review summaries:
This paper proposes a learning to hash attention (LHA) to learn sparse attention for Transformer. The proposed LHA addresses the limitation of ANN-based sparse attention method by separate learnable hash functions for queries and keys and utilizes kernelized techniques for efficient approximation of attention utilities.  Experiments on several applications validate the effectiveness of the proposed LHA.

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.