Original: S PARSE A TTENTION WITH L EARNING - TO -H ASH
Revised: S PARSE A TTENTION WITH L EARNING - TO -H ASH

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Transformer has become ubiquitous in sequence modeling tasks.
Revised: Transformer has become ubiquitous in sequence modeling tasks.

Original: 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.
Revised: 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.

Original: To tackle this problem, recent work developed dynamic attention sparsiﬁcation techniques based on Approximate Nearest Neighbor (ANN) methods, where similar queries and keys are allocated to the same hash bucket with high probability.
Revised: To tackle this problem, recent work developed dynamic attention sparsiﬁcation techniques based on Approximate Nearest Neighbor (ANN) methods, where similar queries and keys are allocated to the same hash bucket with high probability.

Original: 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 justiﬁed.
Revised: 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 justiﬁed.

Original: Besides, some of the ANN methods such as Locality-Sensitive Hashing (LSH) are randomized and cannot fully utilize the available real data distributions.
Revised: Besides, some of the ANN methods such as Locality-Sensitive Hashing (LSH) are randomized and cannot fully utilize the available real data distributions.

Original: 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.
Revised: 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.

Original: 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.
Revised: Another advantage of LHA is that it does not impose extra constraints for queries and keys, which makes it applicable to the wide range of pre-trained Transformer models.

Original: 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 classiﬁcation, retrieval, etc.) show the superior performance of LHA over other strong Transformer variants.
Revised: 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 classiﬁcation, retrieval, etc.) show the superior performance of LHA over other strong Transformer variants.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: The Transformer architecture (Vaswani et al., 2017) has been successfully applied to various tasks, including natural language processing (Vaswani et al., 2017;
Revised: The Transformer architecture (Vaswani et al., 2017) has been successfully applied to various tasks, including natural language processing (Vaswani et al., 2017;

Original: Devlin et al., 2018; Dai et al., 2019;
Revised: Devlin et al., 2018; Dai et al., 2019;

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: But self-attention has a quadratic time and space complexity with respect to the sequence length and hence does not scale efﬁciently to long sequences as a result.
Revised: But self-attention has a quadratic time and space complexity with respect to the sequence length and hence does not scale efﬁciently to long sequences as a result.

Original: To address this inefﬁciency 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.
Revised: To address this inefﬁciency 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.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: Besides, the hash functions in LSH are randomized and data-agnostic, which cannot fully utilize the rich information in real-world data distributions.
Revised: Besides, the hash functions in LSH are randomized and data-agnostic, which cannot fully utilize the rich information in real-world data distributions.

Original: In this paper, we address the above limitations of existing ANN-based methods for attention sparsiﬁcation.
Revised: In this paper, we address the above limitations of existing ANN-based methods for attention sparsiﬁcation.

Original: Firstly, we analyze two imbalance issues in LSH-produced sparse attention patterns, i.e., unbalanced hash bucket sizes and unbalanced query-key ratios.
Revised: Firstly, we analyze two imbalance issues in LSH-produced sparse attention patterns, i.e., unbalanced hash bucket sizes and unbalanced query-key ratios.

Original: 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.
Revised: 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.

Original: Thirdly, we propose a novel solution, namely Learning-to-Hash Attention (LHA), for dynamic attention sparsiﬁcation with enhanced model expressiveness.
Revised: Thirdly, we propose a novel solution, namely Learning-to-Hash Attention (LHA), for dynamic attention sparsiﬁcation with enhanced model expressiveness.

Original: LHA directly optimizes our newly deﬁned attention utility metric in an end-to-end manner via separate learnable hash functions for queries and keys, respectively.
Revised: LHA directly optimizes our newly deﬁned attention utility metric in an end-to-end manner via separate learnable hash functions for queries and keys, respectively.

Original: 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 efﬁciently approximate the attention utilities.
Revised: 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 efﬁciently approximate the attention utilities.

Original: Similar to other sparse attention models (Kitaev et al., 2020;
Revised: Similar to other sparse attention models (Kitaev et al., 2020;

Original: Roy et al., 2021), LHA reduces the overall complexity of self-attention from O ( N 2 ) to O ( N 1 . 5 ) for sequence length N .
Revised: Roy et al., 2021), LHA reduces the overall complexity of self-attention from O ( N 2 ) to O ( N 1 . 5 ) for sequence length N .

Original: 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.
Revised: 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.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: 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.
Revised: 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.

Original: Location-based sparse attention methods aim to improve the computational efﬁciency by using pre-speciﬁed global or local sparsiﬁcation patterns over token locations.
Revised: Location-based sparse attention methods aim to improve the computational efﬁciency by using pre-speciﬁed global or local sparsiﬁcation patterns over token locations.

Original: Liu et al.
Revised: Liu et al.

Original: (2018) proposed to alternate coarse attention layers and local attention layers.
Revised: (2018) proposed to alternate coarse attention layers and local attention layers.

Original: Child et al.
Revised: Child et al.

Original: (2019) used a strided sparse attention pattern in image generation.
Revised: (2019) used a strided sparse attention pattern in image generation.

Original: Sukhbaatar et al.
Revised: Sukhbaatar et al.

Original: (2019) imposed sparsity based on the predicted temporal window size for each token.
Revised: (2019) imposed sparsity based on the predicted temporal window size for each token.

Original: Other methods (Zhang et al., 2021; Beltagy et al., 2020; Ainslie et al., 2020; Zaheer et al., 2020) used a pre-speciﬁed subset of locations in the input as the global memory, and only allow non-local attentions from this subset to all the other tokens.
Revised: Other methods (Zhang et al., 2021; Beltagy et al., 2020; Ainslie et al., 2020; Zaheer et al., 2020) used a pre-speciﬁed subset of locations in the input as the global memory, and only allow non-local attentions from this subset to all the other tokens.

Original: Location-based sparse attention cannot leverage more ﬂexible content-based interactions among arbitrary positions, as a limitation.
Revised: Location-based sparse attention cannot leverage more ﬂexible content-based interactions among arbitrary positions, as a limitation.

Original: Content-based sparse attention allows more ﬂexible sparse patterns than location-based ones.
Revised: Content-based sparse attention allows more ﬂexible sparse patterns than location-based ones.

Original: Malaviya et al.
Revised: Malaviya et al.

Original: (2018) used sparsemax to obtain a sparse attention matrix, while (Correia et al., 2019) used entmax .
Revised: (2018) used sparsemax to obtain a sparse attention matrix, while (Correia et al., 2019) used entmax .

Original: These methods require to compute the full attention matrix before sparsiﬁcation and hence cannot reduce the quadratic computation complexity.
Revised: These methods require to compute the full attention matrix before sparsiﬁcation and hence cannot reduce the quadratic computation complexity.

Original: 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).
Revised: 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).

Original: Those methods usually apply an ANN module as the shared hash function to both queries and keys.
Revised: Those methods usually apply an ANN module as the shared hash function to both queries and keys.

Original: Vyas et al.
Revised: Vyas et al.

Original: (2020) and Zhou et al.
Revised: (2020) and Zhou et al.

Original: (2020) sparsiﬁed the attention maps by eliminating redundant queries.
Revised: (2020) sparsiﬁed the attention maps by eliminating redundant queries.

Original: Tay et al.
Revised: Tay et al.

Original: (2020b) designed a differentiable sorting algorithm of internal representations to enable efﬁcient quasi-global local attention.
Revised: (2020b) designed a differentiable sorting algorithm of internal representations to enable efﬁcient quasi-global local attention.

Original: Contemporary to our work,
Revised: Contemporary to our work,

Original: SparseFinder (Treviso et al., 2021) learns sparse attention patterns that approximate entmax attention, but its bucketing strategies are still based on ANN approaches.
Revised: SparseFinder (Treviso et al., 2021) learns sparse attention patterns that approximate entmax attention, but its bucketing strategies are still based on ANN approaches.

Original: In contrast, LHA directly predicts a bucketing strategy (i.e., learnable hash functions) that maximizes the attention utility.
Revised: In contrast, LHA directly predicts a bucketing strategy (i.e., learnable hash functions) that maximizes the attention utility.

Original: 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.
Revised: 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.

Original: Wang et al.
Revised: Wang et al.

Original: (2020b) applied a low-rank decomposition to the attention matrix.
Revised: (2020b) applied a low-rank decomposition to the attention matrix.

Original: Xiong et al.
Revised: Xiong et al.

Original: (2021) approximated the attention matrix with Nyström approximation.
Revised: (2021) approximated the attention matrix with Nyström approximation.

Original: Katharopoulos et al.
Revised: Katharopoulos et al.

Original: (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.
Revised: (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.

Original: Choromanski et al.
Revised: Choromanski et al.

Original: and Peng et al.
Revised: and Peng et al.

Original: (2021) further proposed an unbiased approximation of softmax with random Fourier features.
Revised: (2021) further proposed an unbiased approximation of softmax with random Fourier features.

Original: Our work in this paper is directly related to the second category, i.e., content-based attention sparsiﬁcation.
Revised: Our work in this paper is directly related to the second category, i.e., content-based attention sparsiﬁcation.

Original: 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 sparsiﬁcation.
Revised: 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 sparsiﬁcation.

Original:  R E - EXAMINATION OF C ONTENT - BASED S PARSE P ATTERNS 3. 
Revised:  R E - EXAMINATION OF C ONTENT - BASED S PARSE P ATTERNS 3.

Original: P RELIMINARY
Revised: P RELIMINARY

Original: The self-attention mechanism (Vaswani et al., 2017) can be formulated as the weighted sum of the value vectors V ∈ R N × d h where the weights are calculated using query vectors Q ∈ R N × d h and key vectors K ∈ R N × d h as:
Revised: The self-attention mechanism (Vaswani et al., 2017) can be formulated as the weighted sum of the value vectors V ∈ R N × d h where the weights are calculated using query vectors Q ∈ R N × d h and key vectors K ∈ R N × d h as:

Original: [Equation]
Revised: [Equation]

Original: where A denotes the matrix of normalized attention weights, d h is the dimension of hidden representations, and N is the sequence length.
Revised: where A denotes the matrix of normalized attention weights, d h is the dimension of hidden representations, and N is the sequence length.

Original: self-attention has a quadratic time and space complexity with respect to the sequence length and hence does not scale efﬁciently to long sequences.
Revised: self-attention has a quadratic time and space complexity with respect to the sequence length and hence does not scale efﬁciently to long sequences.

Original: Content-based sparse attention methods usually apply randomized hash functions or a clustering algorithm to queries { Q i } and keys { K j } , and hope that similar queries and keys are hashed or clustered into the same bucket.
Revised: Content-based sparse attention methods usually apply randomized hash functions or a clustering algorithm to queries { Q i } and keys { K j } , and hope that similar queries and keys are hashed or clustered into the same bucket.

Original: The queries can thus only attend to the keys if both are in the same bucket.
Revised: The queries can thus only attend to the keys if both are in the same bucket.

Original: Formally, a content-based sparse attention strategy with B hash buckets is deﬁned as:
Revised: Formally, a content-based sparse attention strategy with B hash buckets is deﬁned as:

Original: [Equation]
Revised: [Equation]

Original: where h K , h Q : R d h (cid:55)→ [ B ] are the hash functions for keys and queries, and ¯ A ij ∝ A ij is the re-normalized attention weights such that ∀ i, (cid:80) j : h Q ( Q i )
Revised: where h K , h Q : R d h (cid:55)→ [ B ] are the hash functions for keys and queries, and ¯ A ij ∝ A ij is the re-normalized attention weights such that ∀ i, (cid:80) j : h Q ( Q i )

Original: =
Revised: =

Original: h K ( K j ) ¯
Revised: h K ( K j ) ¯

Original: A ij = 1 .
Revised: A ij = 1 .

Original: 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.
Revised: 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.

Original: Thus, the overall complexity 1 of self-attention can be reduced from O ( N 2 ) to O ( N · B + N 2 /B )
Revised: Thus, the overall complexity 1 of self-attention can be reduced from O ( N 2 ) to O ( N · B + N 2 /B )

Original: ≈ O ( N 1 . 5 ) when B ≈ N/B ≈ √ N .
Revised: ≈ O ( N 1 . 5 ) when B ≈ N/B ≈ √ N .

Original: Since the hash functions are not differentiable, Approximate Nearest Neighbor (ANN) methods are used to derive an effective content-based hash function.
Revised: Since the hash functions are not differentiable, Approximate Nearest Neighbor (ANN) methods are used to derive an effective content-based hash function.

Original: Reformer (Kitaev et al., 2020) applies
Revised: Reformer (Kitaev et al., 2020) applies

Original: Locality Sensitive Hashing (LSH) to the tied queries and keys, where several hyper-planes are randomly generated to divide tokens into different buckets.
Revised: Locality Sensitive Hashing (LSH) to the tied queries and keys, where several hyper-planes are randomly generated to divide tokens into different buckets.

Original: SMYRF (Daras et al., 2020) improves
Revised: SMYRF (Daras et al., 2020) improves

Original: Reformer by introducing asymmetric transformation to queries and keys, i.e.,
Revised: Reformer by introducing asymmetric transformation to queries and keys, i.e.,

Revised: F

Original: [Equation]
Revised: [Equation]

Original: where M Q = max Q i || Q
Revised: where M Q = max Q i || Q

Original: i || 

Revised: i ||

Original: and M K = max K j || K j || 2 , such that || F ( Q i ) − G ( K j )
Revised: and M K = max K j || K j || 2 , such that || F ( Q i ) − G ( K j )

Original: || 

Revised: ||

Original: = const −
Revised: = const −

Original: Q i · K j .
Revised: Q i · K j .

Original: 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.
Revised: 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.

Original: 3. B UCKET I MBALANCE I SSUES
Revised: 3. B UCKET I MBALANCE I SSUES

Original: Previous content-based sparse attention models take it for granted that the ANN-derived sparse pattern can effectively approximate the full attention.
Revised: Previous content-based sparse attention models take it for granted that the ANN-derived sparse pattern can effectively approximate the full attention.

Original: However, it is only veriﬁed via empirical evaluation on the down-stream tasks yet, which cannot reﬂect the true attention map approximation ability.
Revised: However, it is only veriﬁed via empirical evaluation on the down-stream tasks yet, which cannot reﬂect the true attention map approximation ability.

Original: Notice that there are two necessary conditions for sparse attention to work effectively and efﬁciently:
Revised: Notice that there are two necessary conditions for sparse attention to work effectively and efﬁciently:

Original: 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.
Revised: 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.

Original: The bucket sizes should be nearly equal in order to effectively reduce the overall complexity.
Revised: The bucket sizes should be nearly equal in order to effectively reduce the overall complexity.

Original: We ﬁrst analyze how badly the two conditions would be violated by LSH 2 , a typical ANN method.
Revised: We ﬁrst analyze how badly the two conditions would be violated by LSH 2 , a typical ANN method.

Original: We apply LSH to 10 attention heads in the 3 rd layer of a Transformer 3 pre-trained on language modeling and obtain the results shown in Figure 1 (up).
Revised: We apply LSH to 10 attention heads in the 3 rd layer of a Transformer 3 pre-trained on language modeling and obtain the results shown in Figure 1 (up).

Original: We can see that the imbalance issue not only exists in the query-key ratios, but also in the bucket sizes.
Revised: We can see that the imbalance issue not only exists in the query-key ratios, but also in the bucket sizes.

Original: To go a step further, we apply LSH to all 16 × 10 = 160 attention heads in the pre-trained Transformer and ﬁnd 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.
Revised: To go a step further, we apply LSH to all 16 × 10 = 160 attention heads in the pre-trained Transformer and ﬁnd 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.

Original: Around 35.9% buckets have the bucket size imbalance problem, where the bucket sizes are twice greater than half smaller than the expected bucket size of 512 .
Revised: Around 35.9% buckets have the bucket size imbalance problem, where the bucket sizes are twice greater than half smaller than the expected bucket size of 512 .

Original: Clearly, neither of the two aforementioned conditions are well satisﬁed in this LSH-sparsiﬁed Transformer model.
Revised: Clearly, neither of the two aforementioned conditions are well satisﬁed in this LSH-sparsiﬁed Transformer model.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: The new analysis results are shown in Figure 1 (down).
Revised: The new analysis results are shown in Figure 1 (down).

Original: We can see that the asymmetric transformation would only exacerbate the imbalance problem, with respect to both query-key ratios and bucket sizes.
Revised: We can see that the asymmetric transformation would only exacerbate the imbalance problem, with respect to both query-key ratios and bucket sizes.

Original: 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 magniﬁed by the asymmetric transformation.
Revised: 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 magniﬁed by the asymmetric transformation.

Original: 3. P ROPOSED M ETRIC : A TTENTION U TILITY
Revised: 3. P ROPOSED M ETRIC : A TTENTION U TILITY

Original: The above analysis shows the imbalance problem in ANN-derived hashing strategies.
Revised: The above analysis shows the imbalance problem in ANN-derived hashing strategies.

Original: To further quantify the approximation quality of sparse attention patterns, we utilize the concept of Attention
Revised: To further quantify the approximation quality of sparse attention patterns, we utilize the concept of Attention

Original: 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:
Revised: 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:

Original: [Equation]
Revised: [Equation]

Original: We denote with C the set of all possible assignments in B balanced non-overlapping clusters:
Revised: We denote with C the set of all possible assignments in B balanced non-overlapping clusters:

Original: C =
Revised: C =

Original: {C 1 , C 2 , . .
Revised: {C 1 , C 2 , . .

Original: , C T } , where T is the number of possible assignments.
Revised: , C T } , where T is the number of possible assignments.

Original: We can then deﬁne the attention utility AU of each assignment C t as:[Equation] C [Equation]
Revised: We can then deﬁne the attention utility AU of each assignment C t as:

Original: =

Original: [Equation]

Original: Q

Original: [Equation]

Original: C

Original: [Equation]
Revised: [Equation]

Original: The attention utility represents the aggregation of the sparse attention weights in an assignment.
Revised: The attention utility represents the aggregation of the sparse attention weights in an assignment.

Original: As we should keep as much sparse attention weights as possible to better approximate a full attention map, attention utility quantiﬁes how well a sparse attention approximates the full one.
Revised: As we should keep as much sparse attention weights as possible to better approximate a full attention map, attention utility quantiﬁes how well a sparse attention approximates the full one.

Original: In fact, we can show that it is computationally intractable to ﬁnd the optimal attention utility:
Revised: In fact, we can show that it is computationally intractable to ﬁnd the optimal attention utility:

Original: Theorem 1.
Revised: Theorem 1.

Original: Finding the assignment that achieves the optimal attention utility, i.e.,
Revised: Finding the assignment that achieves the optimal attention utility, i.e.,

Original: [Equation]
Revised: [Equation]

Original: is NP-hard.
Revised: is NP-hard.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Developing a new metric that can take arbitrary attention masks into account is left for future work.
Revised: Developing a new metric that can take arbitrary attention masks into account is left for future work.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: We also calculate the empirical upper bound by aggregating the attention from each query to its top keys w.r.t dot-products.
Revised: We also calculate the empirical upper bound by aggregating the attention from each query to its top keys w.r.t dot-products.

Original: We present a box-plot of averaged attention utilities for all 16 Transformer layer in Figure 2.
Revised: We present a box-plot of averaged attention utilities for all 16 Transformer layer in Figure 2.

Original: We can see that the LSH methods obtain signiﬁcantly 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.
Revised: We can see that the LSH methods obtain signiﬁcantly 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.

Original: 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).
Revised: 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).

Original: In the next section, we propose an alternative choice of the hash functions to tackle this problem.
Revised: In the next section, we propose an alternative choice of the hash functions to tackle this problem.

Original:  L EARNING - TO -H ASH A TTENTION (LHA)
Revised:  L EARNING - TO -H ASH A TTENTION (LHA)

Original: We propose a novel Learning-to-Hash Attention model in this section.
Revised: We propose a novel Learning-to-Hash Attention model in this section.

Original: The key idea is to learn separate parameterized hash functions for queries and keys, respectively, thus the sparse pattern in
Revised: The key idea is to learn separate parameterized hash functions for queries and keys, respectively, thus the sparse pattern in

Original: 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.
Revised: 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.

Original: Besides, we empirically ﬁnd 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.
Revised: Besides, we empirically ﬁnd 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.

Original: 4. L EARNING - TO -H ASH FOR S PARSE A TTENTION
Revised: 4. L EARNING - TO -H ASH FOR S PARSE A TTENTION

Original: We ﬁrst remind the readers of our deﬁnitions of the sparse attention:
Revised: We ﬁrst remind the readers of our deﬁnitions of the sparse attention:

Original: [Equation]
Revised: [Equation]

Original: where h K , h Q : R d h (cid:55)→ [ B ] are the hash functions for keys and queries, and ¯ A ij ∝ A ij is the renormalized attention weights such that ∀ i, (cid:80) j : h Q ( Q i )
Revised: where h K , h Q : R d h (cid:55)→ [ B ] are the hash functions for keys and queries, and ¯ A ij ∝ A ij is the renormalized attention weights such that ∀ i, (cid:80) j : h Q ( Q i )

Original: =
Revised: =

Original: h K ( K j ) ¯
Revised: h K ( K j ) ¯

Original: A ij = 1 .
Revised: A ij = 1 .

Original: Inspired by the Learning-toHash methods (Wang et al., 2017), we implement the learnable hash functions
Revised: Inspired by the Learning-toHash methods (Wang et al., 2017), we implement the learnable hash functions

Original: h K , h Q : R d h (cid:55)→ [ B ] by deﬁning parameterized functions H Q , H K : R d h (cid:55)→ R B , such that:
Revised: h K , h Q : R d h (cid:55)→ [ B ] by deﬁning parameterized functions H Q , H K : R d h (cid:55)→ R B , such that:

Original: [Equation]
Revised: [Equation]

Original: where H Q and H K can be any parameterized functions such as MLPs.
Revised: where H Q and H K can be any parameterized functions such as MLPs.

Original: Notice that our formulation is a generalization of the non-learnable sparse attention mechanisms in several previous content-based sparse attention models.
Revised: Notice that our formulation is a generalization of the non-learnable sparse attention mechanisms in several previous content-based sparse attention models.

Original: To reproduce the symmetric LSH-based sparse attention scheme in Reformer (Kitaev et al., 2020) or the mini-batch k-means scheme in
Revised: To reproduce the symmetric LSH-based sparse attention scheme in Reformer (Kitaev et al., 2020) or the mini-batch k-means scheme in

Original: Routing Transformer (Roy et al., 2021), we can set:
Revised: Routing Transformer (Roy et al., 2021), we can set:

Original: [Equation]
Revised: [Equation]

Original:  where R ∈ R d h × B 2 is a random matrix, [ · ; · ] denotes the concatenation of two vectors, and µ ∈ R d h × B are the exponentially moving averaged cluster centroids shared by keys and queries.
Revised: [Equation] where R ∈ R d h × B 2 is a random matrix, [ · ; · ] denotes the concatenation of two vectors, and µ ∈ R d h × B are the exponentially moving averaged cluster centroids shared by keys and queries.

Original: The asymmetric transformation in SMYRF (Daras et al., 2020) can also be reproduced by:
Revised: The asymmetric transformation in SMYRF (Daras et al., 2020) can also be reproduced by:

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where a is a random vector, b is a random scalar, r is a scalar parameter, 1 ( · ) :
Revised: where a is a random vector, b is a random scalar, r is a scalar parameter, 1 ( · ) :

Original: [ B ] (cid:55)→ { 0 , 1 } B is the one-hot operation, and F ( · ) and G ( · ) are deﬁned as in Equation 3.
Revised: [ B ] (cid:55)→ { 0 , 1 } B is the one-hot operation, and F ( · ) and G ( · ) are deﬁned as in Equation 3.

Original: 4. O PTIMIZATION O BJECTIVE FOR L EARNING - TO -H ASH
Revised: 4. O PTIMIZATION O BJECTIVE FOR L EARNING - TO -H ASH

Original: Previous ANN-derived sparse attention models construct their hash functions by either randomization or k-mean clustering.
Revised: Previous ANN-derived sparse attention models construct their hash functions by either randomization or k-mean clustering.

Original: In contrast, our hash functions are fully learnable.
Revised: In contrast, our hash functions are fully learnable.

Original: 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.
Revised: 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.

Original: To resolve this issue, we guide the training of our learnable hash functions directly by our proposed attention utility (See Equation 5).
Revised: To resolve this issue, we guide the training of our learnable hash functions directly by our proposed attention utility (See Equation 5).

Original: Notice that the attention utility of each query Q i can be written as (cid:80) j : h Q ( Q i )
Revised: Notice that the attention utility of each query Q i can be written as (cid:80) j : h Q ( Q i )

Original: =
Revised: =

Original: h K ( K j )
Revised: h K ( K j )

Original: A ij .
Revised: A ij .

Original: We ﬁrst calculate the possible attention utilities ψ ( i ) that query Q
Revised: We ﬁrst calculate the possible attention utilities ψ ( i ) that query Q

Original: i can obtain when allocated to all B hash buckets:
Revised: i can obtain when allocated to all B hash buckets:

Original: [Equation]
Revised: [Equation]

Original: Ideally, our hash function would allocate Q i to the bucket that maximize the attention utility, such that Q i can attend to the most salient keys:
Revised: Ideally, our hash function would allocate Q i to the bucket that maximize the attention utility, such that Q i can attend to the most salient keys:

Original: [Equation]
Revised: [Equation]

Original: To achieve this goal, since ψ ( i ) is naturally a normalized categorical distribution, we can simply convert H Q ( Q i ) into a distribution and use the KL divergence between the predicted distribution softmax( H Q ( Q i )) and the desired distribution ψ ( i ) as the optimization objective:
Revised: To achieve this goal, since ψ ( i ) is naturally a normalized categorical distribution, we can simply convert H Q ( Q i ) into a distribution and use the KL divergence between the predicted distribution softmax( H Q ( Q i )) and the desired distribution ψ ( i ) as the optimization objective:

Original: [Equation]
Revised: [Equation]

Original: On the other hand, the possible attention utilities ψ (cid:48)
Revised: On the other hand, the possible attention utilities ψ (cid:48)

Original: ( j ) that a key K i obtains when allocated to all B buckets can be written as:
Revised: ( j ) that a key K i obtains when allocated to all B buckets can be written as:

Original: [Equation]
Revised: [Equation]

Original: While ψ (cid:48) ( j ) is no longer a normalized distribution, we can normalize it and similarly deﬁne the optimization objective L K j also as a KL divergence.
Revised: While ψ (cid:48) ( j ) is no longer a normalized distribution, we can normalize it and similarly deﬁne the optimization objective L K j also as a KL divergence.

Original: During the training stage, our ﬁnal optimization objective is a convex combination of the task-speciﬁc objective and the learning-to-hash objectives (i.e., {L Q i } and {L K j } ) for all queries and keys.
Revised: During the training stage, our ﬁnal optimization objective is a convex combination of the task-speciﬁc objective and the learning-to-hash objectives (i.e., {L Q i } and {L K j } ) for all queries and keys.

Original: Figure 3 in the appendix illustrates the joint training diagram of LHA.
Revised: Figure 3 in the appendix illustrates the joint training diagram of LHA.

Original: 4. A PPROXIMATE A TTENTION U TILITIES
Revised: 4. A PPROXIMATE A TTENTION U TILITIES

Original: The remaining problem is how to efﬁciently compute ψ ( i ) and ψ (cid:48)
Revised: The remaining problem is how to efﬁciently compute ψ ( i ) and ψ (cid:48)

Original: ( j ) .
Revised: ( j ) .

Original: A naive way requires calculating the dot products between all query-key pairs with O ( N 2 ) complexity.
Revised: A naive way requires calculating the dot products between all query-key pairs with O ( N 2 ) complexity.

Original: 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.
Revised: 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.

Original: Let us deﬁne φ : R d h (cid:55)→ R 2 D as the Positive Randomized Features (PRFs) (Choromanski et al., 2020) such that
Revised: Let us deﬁne φ : R d h (cid:55)→ R 2 D as the Positive Randomized Features (PRFs) (Choromanski et al., 2020) such that

Original: [Equation]
Revised: [Equation]

Original: We can approximate ψ ( i ) with the following formula:
Revised: We can approximate ψ ( i ) with the following formula:

Original: [Equation]
Revised: [Equation]

Original: for each bucket b ∈ [ B ] .
Revised: for each bucket b ∈ [ B ] .

Original: Since h K ( K j ) does not change for the queries, we can simply pre-compute (cid:80) j : h K ( K j )=
Revised: Since h K ( K j ) does not change for the queries, we can simply pre-compute (cid:80) j : h K ( K j )=

Original: b φ ( K j ) once to save computation, which reduces the complexity of computing { ψ ( i ) } for all queries from O ( N 2 ) to O ( N · B ) ≈
Revised: b φ ( K j ) once to save computation, which reduces the complexity of computing { ψ ( i ) } for all queries from O ( N 2 ) to O ( N · B ) ≈

Original: O
Revised: O

Original: ( N 1 . 5 ) .
Revised: ( N 1 . 5 ) .

Original: The attention utilities for keys ψ (cid:48) ( j ) can be efﬁciently approximated in a similar way 4 :
Revised: The attention utilities for keys ψ (cid:48) ( j ) can be efﬁciently approximated in a similar way 4 :

Original: [Equation]
Revised: [Equation]

Original: 4. I MPLEMENTATION D ETAILS
Revised: 4. I MPLEMENTATION D ETAILS

Original: Practical content-based sparse attention models require that each hash bucket has the same size, which is crucial in terms of computational efﬁciency on modern hardwares.
Revised: Practical content-based sparse attention models require that each hash bucket has the same size, which is crucial in terms of computational efﬁciency on modern hardwares.

Original: Therefore, we follow Roy et al.
Revised: Therefore, we follow Roy et al.

Original: (2021) and sort the tokens with regard to normalized hash scores, i.e., softmax ( H Q ( Q i )) and softmax ( H K ( K i )) , in each bucket.
Revised: (2021) and sort the tokens with regard to normalized hash scores, i.e., softmax ( H Q ( Q i )) and softmax ( H K ( K i )) , in each bucket.

Original: The hash bucket membership is then determined by the top- k threshold, where k = NB is the bucket size.
Revised: The hash bucket membership is then determined by the top- k threshold, where k = NB is the bucket size.

Original: 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.
Revised: 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.

Original: A pseudo-code implementation for LHA can be found in the appendix.
Revised: A pseudo-code implementation for LHA can be found in the appendix.

Original: In causally masked attentions, the queries cannot attend to keys behind, as is usually the case in language modeling.
Revised: In causally masked attentions, the queries cannot attend to keys behind, as is usually the case in language modeling.

Original: When we use separate hash functions { h Q ( Q i ) } and { h K ( K j } ) for queries and keys, respectively, it is possible that for some query Q i , there exists no such key K j in the same bucket that j ≤ i .
Revised: When we use separate hash functions { h Q ( Q i ) } and { h K ( K j } ) for queries and keys, respectively, it is possible that for some query Q i , there exists no such key K j in the same bucket that j ≤ i .

Original: This would cause serious numerical instability in our implementation.
Revised: This would cause serious numerical instability in our implementation.

Original: To tackle this problem, inspired by Kitaev et al.
Revised: To tackle this problem, inspired by Kitaev et al.

Original: (2020) and Roy et al.
Revised: (2020) and Roy et al.

Original: (2021), we tie the key hashes with query hashes in the case of causal attentions by constructing a joint hash function:
Revised: (2021), we tie the key hashes with query hashes in the case of causal attentions by constructing a joint hash function:

Original: [Equation]
Revised: [Equation]

Original: We ﬁnd that this strategy empirically works better than using other tricks to ﬁx this “no-attentiontarget” problem.
Revised: We ﬁnd that this strategy empirically works better than using other tricks to ﬁx this “no-attentiontarget” problem.

Original: 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.
Revised: 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.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: 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.
Revised: 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.

Original: Due to the space limitations, the detailed hyper-parameter settings are presented in the appendix.
Revised: Due to the space limitations, the detailed hyper-parameter settings are presented in the appendix.

Original: 5. L ANGUAGE M ODELING
Revised: 5. L ANGUAGE M ODELING

Original: Wikitext-103 (Merity et al., 2016) is a large-scale dataset for testing long term dependencies in word-level language models.
Revised: Wikitext-103 (Merity et al., 2016) is a large-scale dataset for testing long term dependencies in word-level language models.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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
Revised: 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

Original: [Table]
Revised: [Table]

Original: [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original:  sinusoidal absolute positional encoding (Vaswani et al., 2017).
Revised: -layer sinusoidal absolute positional encoding (Vaswani et al., 2017).

Original: The LSH Transformer is implemented in our code base by simply replacing the learnable hash functions by random projections.
Revised: The LSH Transformer is implemented in our code base by simply replacing the learnable hash functions by random projections.

Original: A sequence length of 1536 is used in both training and evaluation stages for our models.
Revised: A sequence length of 1536 is used in both training and evaluation stages for our models.

Original: 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.
Revised: 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.

Original: We also consider a baseline where the hash functions generate content-independent random numbers.
Revised: We also consider a baseline where the hash functions generate content-independent random numbers.

Original: The experimental results are presented in Table 1.
Revised: The experimental results are presented in Table 1.

Original: 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.
Revised: 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.

Original: Surprisingly, we ﬁnd that local Transformer is a very strong baseline for language modeling that outperforms a pure LHA Transformer.
Revised: Surprisingly, we ﬁnd that local Transformer is a very strong baseline for language modeling that outperforms a pure LHA Transformer.

Original: We can also see that a two-layer MLP works better than a linear function as the learnable hash functions in LHA.
Revised: We can also see that a two-layer MLP works better than a linear function as the learnable hash functions in LHA.

Original: The best performance is achieved when we set half of the heads to local attention and the other half to non-local
Revised: The best performance is achieved when we set half of the heads to local attention and the other half to non-local

Original: LHA for each layer.
Revised: LHA for each layer.

Original: Finally, we can see that LHA consistently underperforms the untrainable LSH counterparts under four different settings.
Revised: Finally, we can see that LHA consistently underperforms the untrainable LSH counterparts under four different settings.

Original: For the test set evaluation, we use the best setting from the ablation study.
Revised: For the test set evaluation, we use the best setting from the ablation study.

Original: 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.
Revised: 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.

Original: 5. N ATURAL L ANGUAGE U NDERSTANDING
Revised: 5. N ATURAL L ANGUAGE U NDERSTANDING

Original: 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
Revised: 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

Original: [Table]
Revised: [Table]

Original: GLUE (General Language Understanding Evaluation) dataset (Wang et al., 2018).
Revised: GLUE (General Language Understanding Evaluation) dataset (Wang et al., 2018).

Original: Following the common practice, the maximum sequence length is set to 128.
Revised: Following the common practice, the maximum sequence length is set to 128.

Original: To show the effectiveness of LHA, we choose two competitive baselines in literature: SMYRF (Daras et al., 2020) and Fast Clustered
Revised: To show the effectiveness of LHA, we choose two competitive baselines in literature: SMYRF (Daras et al., 2020) and Fast Clustered

Original: Attention (FCA) (Vyas et al., 2020).
Revised: Attention (FCA) (Vyas et al., 2020).

Original: We also produce an LSH-based sparse attention baseline.
Revised: We also produce an LSH-based sparse attention baseline.

Original: We use a pure LHA model similar to variant (c) in Section 5.1.
Revised: We use a pure LHA model similar to variant (c) in Section 5.1.

Original: We summarize the performance per task in Table 2.
Revised: We summarize the performance per task in Table 2.

Original: We report accuracy for all tasks except STS-B, where we report Pearson correlation.
Revised: We report accuracy for all tasks except STS-B, where we report Pearson correlation.

Original: 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.
Revised: From the table, we can see that LHA performs as good as full attention for all the GLUE tasks [Table] and that LHA outperforms all other methods in the average GLUE score, and has a smaller computational costs in its sparse attention part.

Original: 5. L ONG -R ANGE -A RENA B ENCHMARK
Revised: 5. L ONG -R ANGE -A RENA B ENCHMARK

Original: 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.
Revised: 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.

Original: We follow the apples-to-apples setting 5 of LRA benchmark and compare our method against other efﬁcient attention variants.
Revised: We follow the apples-to-apples setting 5 of LRA benchmark and compare our method against other efﬁcient attention variants.

Original: We use an LHA/local hybrid variant similar to variant (d) in Section 5.1.
Revised: We use an LHA/local hybrid variant similar to variant (d) in Section 5.1.

Original: We consider the tasks of ListOps (Nangia & Bowman, 2018)
Revised: We consider the tasks of ListOps (Nangia & Bowman, 2018)

Original: (LO), byte-level IMDb reviews text classiﬁcation (Maas et al., 2011) (IMDb), byte-level document retrieval on ACL Anthology Network (AAN)
Revised: (LO), byte-level IMDb reviews text classiﬁcation (Maas et al., 2011) (IMDb), byte-level document retrieval on ACL Anthology Network (AAN)

Original: (Radev et al., 2013), and CIFAR10 (Krizhevsky et al., 2009) image classiﬁcation on sequences of pixels (Image).
Revised: (Radev et al., 2013), and CIFAR10 (Krizhevsky et al., 2009) image classiﬁcation on sequences of pixels (Image).

Original: The results are shown in Table 3 and the brief descriptions of the compared baselines can be found in Section 2.
Revised: The results are shown in Table 3 and the brief descriptions of the compared baselines can be found in Section 2.

Original: From the table, we can see that LHA achieves consistent improvements over previous efﬁcient attention models.
Revised: From the table, we can see that LHA achieves consistent improvements over previous efﬁcient attention models.

Original:  C ONCLUSION & F UTURE W ORK
Revised:  C ONCLUSION & F UTURE W ORK

Original: In this paper, we address the limitations of ANN-based sparse attention methods and propose the
Revised: In this paper, we address the limitations of ANN-based sparse attention methods and propose the

Original: Learning-to-Hash Attention (LHA) as our new solution.
Revised: Learning-to-Hash Attention (LHA) as our new solution.

Original: Speciﬁcally, LHA leverages separate learnable hash functions for queries and keys, respectively, and utilizes kernelized techniques for efﬁcient approximation of attention utilities.
Revised: Speciﬁcally, LHA leverages separate learnable hash functions for queries and keys, respectively, and utilizes kernelized techniques for efﬁcient approximation of attention utilities.

Original: The experiments on language modeling, natural language understanding, text classiﬁcation, and image classiﬁcation demonstrated the effectiveness of LHA.
Revised: The experiments on language modeling, natural language understanding, text classiﬁcation, and image classiﬁcation demonstrated the effectiveness of LHA.

Original: 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).
Revised: 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).

Original: E THICS S TATEMENT
Revised: A CKNOWLEDGEMENT

Revised: This research used Perlmutter supercomputer of the National Energy Research Scientiﬁc Computing

Revised: Center, a DOE Ofﬁce of Science User Facility supported by the Ofﬁce of Science of the U.S.

Revised: Department of Energy under Contract

Original: {+No. DE-AC02-05CH
+}
Revised: {+No. DE-AC02-05CH
+}

Revised: using NERSC award NERSC

Revised: DDR-ERCAP0022110.

Revised: E THICS S TATEMENT

Original: This work proposed a more effective sparsiﬁcation pattern for the attention mechanisms.
Revised: This work proposed a more effective sparsiﬁcation pattern for the attention mechanisms.

Original: We believe our proposed model as well as the code to be released can most beneﬁt the ﬁeld of language processing, with the potential to beneﬁt other ﬁelds involving long sequence modeling.
Revised: We believe our proposed model as well as the code to be released can most beneﬁt the ﬁeld of language processing, with the potential to beneﬁt other ﬁelds involving long sequence modeling.

Original: However, it is know that Transformers are also vulnerable to the adversarial attacks (Michel et al., 2019).
Revised: However, it is know that Transformers are also vulnerable to the adversarial attacks (Michel et al., 2019).

Original: 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.
Revised: 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.

Original: R EPRODUCIBILITY S TATEMENT
Revised: R EPRODUCIBILITY S TATEMENT

Original: • We will put a link to an anonymous repository  for our source code as the OpenReview comment.
Revised: • We provide a GitHub repository 6 for our source code.

Original: The hyper-parameters are described in details in the appendix.
Revised: The hyper-parameters are described in details in the appendix.

Original: We also provide a pseudo-code implementaion of our model in the appendix.
Revised: We also provide a pseudo-code implementaion of our model in the appendix.

Original: • All the datasets used in the experiments and the corresponding pre-processing scripts can be found online, including language modeling 6 , GLUE benchmark 7 , Long-Range-Arena benchmark 8 , and time series forecasting 9 .
Revised: • All the datasets used in the experiments and the corresponding pre-processing scripts can be found online, including language modeling 7 , GLUE benchmark 8 , Long-Range-Arena benchmark 9 , and time series forecasting 10 .

Original: The pre-trained model we used (i
Revised: The pre-trained model we used (i

Original: e., RoBERTa) is also publically available 10 .
Revised: e., RoBERTa) is also publically available 11 .

Original: A E XPERIMENTAL S ETTINGS A. R E - EXAMINATION OF C ONTENT - BASED S PARSE A TTENTION
Revised: A E XPERIMENTAL S ETTINGS A. R E - EXAMINATION OF C ONTENT - BASED S PARSE A TTENTION

Original: For analysis, we pre-trained a 16-layer and 10-head Transformer on the WikiText-103 language modeling benchmark (Merity et al., 2016).
Revised: For analysis, we pre-trained a 16-layer and 10-head Transformer on the WikiText-103 language modeling benchmark (Merity et al., 2016).

Original: A per-layer sinusoidal absolute positional encoding (Vaswani et al., 2017) is injected to the queries and keys before self-attention.
Revised: A per-layer sinusoidal absolute positional encoding (Vaswani et al., 2017) is injected to the queries and keys before self-attention.

Original: For Figure 1, we analyzed the queries and keys for the ﬁrst 1024 tokens in the WikiText-103 validation data, and set the number of LSH buckets B = 4 .
Revised: For Figure 1, we analyzed the queries and keys for the ﬁrst 1024 tokens in the WikiText-103 validation data, and set the number of LSH buckets B = 4 .

Original: We apply LSH to all the 10 attention heads in the 3 rd Transformer layer.
Revised: We apply LSH to all the 10 attention heads in the 3 rd Transformer layer.

Original: For Figure 2, we follow the same experimental setting.
Revised: For Figure 2, we follow the same experimental setting.

Original: A. L EARNING - TO -H ASH
Revised: A. L EARNING - TO -H ASH

Original: During the training stage, our ﬁnal optimization objective L is a convex combination of the taskspeciﬁc objective L task (e.g., cross-entropy loss in language modeling) and learning-to-hash objectives for all keys and queries, i.e., {L Q i } and {L K j } (See equation 13).
Revised: During the training stage, our ﬁnal optimization objective L is a convex combination of the taskspeciﬁc objective L task (e.g., cross-entropy loss in language modeling) and learning-to-hash objectives for all keys and queries, i.e., {L Q i } and {L K j } (See equation 13).

Original: [Equation]
Revised: [Equation]

Original: (19) where λ , N , L , H denote the loss coefﬁcient, the sequence length, the number of attention layers, and the number of attention heads per layer, respectively.
Revised: (19) where λ , N , L , H denote the loss coefﬁcient, the sequence length, the number of attention layers, and the number of attention heads per layer, respectively.

Original: Empirically, we set Λ = 0 . 05 and found it work well across different tasks.
Revised: Empirically, we set Λ = 0 . 05 and found it work well across different tasks.

Original: A. L ANGUAGE M ODELING
Revised: A. L ANGUAGE M ODELING

Original: We use the base setting of Transformer-XL (Dai et al., 2019) in our experiments, which consists of 16 Transformer layers.
Revised: We use the base setting of Transformer-XL (Dai et al., 2019) in our experiments, which consists of 16 Transformer layers.

Original: For each layer, the hidden size is set to 410, and the number of attention heads is set to 10.
Revised: For each layer, the hidden size is set to 410, and the number of attention heads is set to 10.

Original: The dimension of feed-forward layer is set to 2100.
Revised: The dimension of feed-forward layer is set to 2100.

Original: All codes are implemented based on
Revised: All codes are implemented based on

Original: Flax (Heek et al., 2020) in JAX (Bradbury et al., 2018).
Revised: Flax (Heek et al., 2020) in JAX (Bradbury et al., 2018).

Original: The dropout ratio is set to 0.2.
Revised: The dropout ratio is set to 0.2.

Original: The batch size is set to 32.
Revised: The batch size is set to 32.

Original: We use AdamW (Loshchilov & Hutter, 2017) as the optimizer, and set ( β 1 , β 2 ) to (0.9, 0.999).
Revised: We use AdamW (Loshchilov & Hutter, 2017) as the optimizer, and set ( β 1 , β 2 ) to (0.9, 0.999).

Original: The peak learning rate is set to 3.5e-4.
Revised: The peak learning rate is set to 3.5e-4.

Original: The model is trained for 20k steps with a 2k-step warm-up stage with a cosine learning rate decay (Loshchilov & Hutter, 2016).
Revised: The model is trained for 20k steps with a 2k-step warm-up stage with a cosine learning rate decay (Loshchilov & Hutter, 2016).

Original: A. N ATURAL L ANGUAGE U NDERSTANDING
Revised: A. N ATURAL L ANGUAGE U NDERSTANDING

Original: We use the base setting of RoBERTa (Liu et al., 2019) in our experiments, which consists of
Revised: We use the base setting of RoBERTa (Liu et al., 2019) in our experiments, which consists of

Original: Transformer layers.
Revised: Transformer layers.

Original: For each layer, the hidden size is set to 768, and the number of attention heads is set to 12.
Revised: For each layer, the hidden size is set to 768, and the number of attention heads is set to 12.

Original: The dimension of feed-forward layer is set to 3072.
Revised: The dimension of feed-forward layer is set to 3072.

Original: The RoBERTa pre-trained checkpoint is taken from the Transformers (Wolf et al., 2020) library.
Revised: The RoBERTa pre-trained checkpoint is taken from the Transformers (Wolf et al., 2020) library.

Original: We use AdamW (Loshchilov & Hutter, 2017) as the optimizer, and set ( β 1 , β 2 ) to (0.9, 0.98).
Revised: We use AdamW (Loshchilov & Hutter, 2017) as the optimizer, and set ( β 1 , β 2 ) to (0.9, 0.98).

Original: To ﬁne-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).
Revised: To ﬁne-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).

Original: A 10% warm-up schedule and linear learning rate decay (Devlin et al., 2018) is applied.
Revised: A 10% warm-up schedule and linear learning rate decay (Devlin et al., 2018) is applied.

Original: [Table]
Revised: [Table]

Original: A. L ONG -R ANGE -A RENA B ENCHMARK
Revised: A. L ONG -R ANGE -A RENA B ENCHMARK

Original: We follow the apples-to apples setting 11 of LRA benchmark and compare our method against other efﬁcient attention variants.
Revised: We follow the apples-to apples setting 12 of LRA benchmark and compare our method against other efﬁcient attention variants.

Original: Speciﬁcally, we use a 4-layer Transformer with 256 hidden size andfeed-forward layer size for the IMDb text classiﬁcation 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 classiﬁcation task(Krizhevsky et al., 2009).
Revised: Speciﬁcally, we use a 4-layer Transformer with 256 hidden size andfeed-forward layer size for the IMDb text classiﬁcation 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 classiﬁcation task(Krizhevsky et al., 2009).

Original: The models are trained for 20k steps with a 2k-step warm-up stage with a cosine learning rate decay (Loshchilov & Hutter, 2016).
Revised: The models are trained for 20k steps with a 2k-step warm-up stage with a cosine learning rate decay (Loshchilov & Hutter, 2016).

Original: A. T IME S ERIES F ORECASTING
Revised: A. T IME S ERIES F ORECASTING

Original: We follow the experimental setup of Informer (Zhou et al., 2020).
Revised: We follow the experimental setup of Informer (Zhou et al., 2020).

Original: Speciﬁcally, 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.
Revised: Speciﬁcally, 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.

Original: The layer of encoder is chosen from { 6 , 4 , 3 , 2 } and the layer of decoder is set as 2.
Revised: The layer of encoder is chosen from { 6 , 4 , 3 , 2 } and the layer of decoder is set as 2.

Original: 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.
Revised: 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.

Original: A. C OMPUTING I NFRASTRUCTURE
Revised: A. C OMPUTING I NFRASTRUCTURE

Original: 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.
Revised: 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.

Original: The evaluation of the inference throughput is performed on a stand-alone machine with 1 NVIDIA Tesla V100 32GB GPU.
Revised: The evaluation of the inference throughput is performed on a stand-alone machine with 1 NVIDIA Tesla V100 32GB GPU.

Original: B A DDITIONAL E XPERIMENTS ON T IME S ERIES F ORECASTING
Revised: B A DDITIONAL E XPERIMENTS ON T IME S ERIES F ORECASTING

Original: We also evaluate our model on time series forecasting tasks.
Revised: We also evaluate our model on time series forecasting tasks.

Original: We use the ETT (Electricity Transformer
Revised: We use the ETT (Electricity Transformer

Original: Temperature) dataset (Zhou et al., 2020), which contains 2-year data from two separated counties in
Revised: Temperature) dataset (Zhou et al., 2020), which contains 2-year data from two separated counties in

Original: China.
Revised: China.

Original: ETTh 1 is a dataset for 1-hour-level, while ETTm 1 is a dataset for 15-minute-level.
Revised: ETTh 1 is a dataset for 1-hour-level, while ETTm 1 is a dataset for 15-minute-level.

Original: Each data point consists of the target value “oil temperature” and 6 power load features.
Revised: Each data point consists of the target value “oil temperature” and 6 power load features.

Original: The train/val/test is 12/4/4 months.
Revised: The train/val/test is 12/4/4 months.

Original: The Mean Squared Error (MSE) metric and Mean Average Error (MAE) metric are used as the evaluation metrics.
Revised: The Mean Squared Error (MSE) metric and Mean Average Error (MAE) metric are used as the evaluation metrics.

Original: We use an LHA/local hybrid variant similar to variant (d) in Section 5.1.
Revised: We use an LHA/local hybrid variant similar to variant (d) in Section 5.1.

Original: In univariate forecasting setting, all the seven features are used as input and “oil temperature” is the prediction target.
Revised: In univariate forecasting setting, all the seven features are used as input and “oil temperature” is the prediction target.

Original: The univariate evaluation results can be found in Table 4.
Revised: The univariate evaluation results can be found in Table 4.

Original: We can see that
Revised: We can see that

Original: LHA signiﬁcantly improve the performance of state-of-the-art for most settings, while slightly underperforms Informer (Zhou et al., 2020) in two settings of ETTh 1 dataset.
Revised: LHA signiﬁcantly improve the performance of state-of-the-art for most settings, while slightly underperforms Informer (Zhou et al., 2020) in two settings of ETTh 1 dataset.

Original: This veriﬁes the effectiveness of LHA on time series data.
Revised: This veriﬁes the effectiveness of LHA on time series data.

Revised: [Table]

Original: [Table]
Revised: [Table]

Original: In multivariate forecasting setting, all the seven features are the prediction targets.
Revised: In multivariate forecasting setting, all the seven features are the prediction targets.

Original: The multivariate evaluation results can be found in Table 5.
Revised: The multivariate evaluation results can be found in Table 5.

Original: We can see that LHA signiﬁcantly improve the performance of state-of-the-art for all settings, which further demonstrates the effectiveness of the proposed LHA method.
Revised: We can see that LHA signiﬁcantly improve the performance of state-of-the-art for all settings, which further demonstrates the effectiveness of the proposed LHA method.

Original: C P ROOF
Revised: C P ROOF

Original: We follow the notations in (Daras et al., 2020).
Revised: We follow the notations in (Daras et al., 2020).

Original: Let C B the set of all possible assignments in B balanced non-overlapping clusters.
Revised: Let C B the set of all possible assignments in B balanced non-overlapping clusters.

Original: A speciﬁc assignment is denoted by C Bt and there are T possible such assignments:
Revised: A speciﬁc assignment is denoted by C Bt and there are T possible such assignments:

Original: [Equation]
Revised: [Equation]

Original: We have the following lemma, which is referred as the max-mass problem:
Revised: We have the following lemma, which is referred as the max-mass problem:

Original: Lemma 1.
Revised: Lemma 1.

Original: (Daras et al., 2020)
Revised: (Daras et al., 2020)

Original: The optimization problem:
Revised: The optimization problem:

Original: [Equation]
Revised: [Equation]

Original: is NP-hard.
Revised: is NP-hard.

Original: 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.
Revised: 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.

Original: Please refer to (Daras et al., 2020) for the detailed constructive proofs.
Revised: Please refer to (Daras et al., 2020) for the detailed constructive proofs.

Original: Next, we show how to use this lemma to prove Theorem 1.
Revised: Next, we show how to use this lemma to prove Theorem 1.

Original: Proof of Theorem 1.
Revised: Proof of Theorem 1.

Original: Let { Q i } and { K j } denote the query and key sets that we consider for computing the attention utility.
Revised: Let { Q i } and { K j } denote the query and key sets that we consider for computing the attention utility.

Original: We can ﬁrst construct the new query set Q (cid:48) and key set K (cid:48) such that
Revised: We can ﬁrst construct the new query set Q (cid:48) and key set K (cid:48) such that

Original: [Equation]
Revised: [Equation]

Original: d
Revised: d

Original: [Equation]
Revised: [Equation]

Original:  where Q (cid:48) i ∈ Q (cid:48) , K (cid:48)
Revised: , where Q (cid:48) i ∈ Q (cid:48) , K (cid:48)

Original: j ∈ K (cid:48)
Revised: j ∈ K (cid:48)

Original: This can be achieved by applying SVD to the matrix of (cid:16) e Qi · Kj √ d h (cid:17) ij .
Revised: This can be achieved by applying SVD to the matrix of (cid:16) e Qi · Kj √ d h (cid:17) ij .

Original: Next, we construct another query set Q (cid:48)(cid:48) , such that
Revised: Next, we construct another query set Q (cid:48)(cid:48) , such that

Original: [Equation]
Revised: [Equation]

Original: The problem of ﬁnding the optimal attention utility, i.e.,
Revised: The problem of ﬁnding the optimal attention utility, i.e.,

Original: [Equation]
Revised: [Equation]

Original: is thus equivalent to the problem in Lemma 1, with Q (cid:48)(cid:48) and K (cid:48) as the query and key sets, which is proven to be NP-hard.
Revised: is thus equivalent to the problem in Lemma 1, with Q (cid:48)(cid:48) and K (cid:48) as the query and key sets, which is proven to be NP-hard.

Original: D A LGORITHM
Revised: D A LGORITHM

Original: We present a detailed pseudo-code implementation for LHA in Algorithm 1.
Revised: We present a detailed pseudo-code implementation for LHA in Algorithm 1.

Original: 1: Queries, Keys and Values: Q, K, V ∈ R N × d h 2: Query/Key Hash Functions:  3: Attention Utility Loss Weight: λ 4: if causal mask then 5:
Revised: 1: Queries, Keys and Values: Q, K, V ∈ R N × d h 2: Query/Key Hash Functions: H Q , H K 3: Attention Utility Loss Weight: λ 4: if causal mask then 5:

Original: Q score ← softmax( H Q ( Q ) + H K ( K ))
Revised: Q score ← softmax( H Q ( Q ) + H K ( K ))

Original: (cid:46) B × N × d h 6:
Revised: (cid:46) B × N × d h 6:

Original: K score ← Q score 7: else 8:
Revised: K score ← Q score 7: else 8:

Original: Q score ← softmax( H Q ( Q ))
Revised: Q score ← softmax( H Q ( Q ))

Original: (cid:46) B × N × d h 9:
Revised: (cid:46) B × N × d h 9:

Original: K score ← softmax( H K ( K ))
Revised: K score ← softmax( H K ( K ))

Original: (cid:46) B × N × d h 10: W ← N/B (cid:46) bucket size 11: Q idx ← top- {√ 2W } ( Q score ) (cid:46) B
Revised: (cid:46) B × N × d h 10: W ← N/B (cid:46) bucket size 11: Q idx ← top- {√ 2W } ( Q score ) (cid:46) B

Original: × N log( N ) 12: K idx ← top- {√ 2W } ( K score ) (cid:46) B
Revised: × N log( N ) 12: K idx ← top- {√ 2W } ( K score ) (cid:46) B

Original: × N log( N ) 13: Q (cid:48)
Revised: × N log( N ) 13: Q (cid:48)

Original: ← gather( Q, Q idx ) (cid:46) B
Revised: ← gather( Q, Q idx ) (cid:46) B

Original: × W × d h 14: K (cid:48) ← gather( K, K idx ) (cid:46) B
Revised: × W × d h 14: K (cid:48) ← gather( K, K idx ) (cid:46) B

Original: × W ×
Revised: × W ×

Original: d h 15: V (cid:48) ← gather( V, K idx ) (cid:46) B
Revised: d h 15: V (cid:48) ← gather( V, K idx ) (cid:46) B

Original: × W × d h 16:
Revised: × W × d h 16:

Original: A ← Q (cid:48)
Revised: A ← Q (cid:48)

Original: ( K (cid:48) )
Revised: ( K (cid:48) )

Original: T (cid:46) B
Revised: T (cid:46) B

Original: × W × W × d h 17: if causal mask then 18:
Revised: × W × W × d h 17: if causal mask then 18:

Original: A ← mask( A ) (cid:46) B × W × W 19: A ← softmax( A ) .
Revised: A ← mask( A ) (cid:46) B × W × W 19: A ← softmax( A ) .

Original: (cid:46) B × W × W 20: V (cid:48)
Revised: (cid:46) B × W × W 20: V (cid:48)

Original: ← A · V (cid:48) (cid:46) B × W × W
Revised: ← A · V (cid:48) (cid:46) B × W × W

Original: ×
Revised: ×

Original: d h 21: X ← scatter( K idx , V (cid:48) )
Revised: d h 21: X ← merge( K idx , V (cid:48) )

Original: (cid:46) B
Revised: (cid:46) B

Original: × W ×
Revised: × W ×

Original: d h 22: ψ = AU Q ( Q, K ) (cid:46) B × N , Equation23: ψ (cid:48) =
Revised: d h 22: ψ = AU Q ( Q, K ) (cid:46) B × N , Equation23: ψ (cid:48) =

Original: AU K ( Q, K ) (cid:46) B × N , Equation24:
Revised: AU K ( Q, K ) (cid:46) B × N , Equation24:

Original: L AU = λ · KL ( ψ (cid:107) Q score ) + λ · KL ( ψ (cid:48) (cid:107) K score ) 25: return X, L AU
Revised: L AU = λ · KL ( ψ (cid:107) Q score ) + λ · KL ( ψ (cid:48) (cid:107) K score ) 25: return X, L AU

Original: E A C OMPLETE D ERIVATION OF THE LHA O BJECTIVE
Revised: E A C OMPLETE D ERIVATION OF THE LHA O BJECTIVE

Original: Let Q, K, V ∈ R N × d h denote the query, key, and value vectors of the attention mechanism, where
Revised: Let Q, K, V ∈ R N × d h denote the query, key, and value vectors of the attention mechanism, where

Original: N is the sequence length and d h is the hidden size.
Revised: N is the sequence length and d h is the hidden size.

Original: Let φ : R d h (cid:55)→ R 2 D be Positive Randomized
Revised: Let φ : R d h (cid:55)→ R 2 D be Positive Randomized

Original: Features (PRFs) (Choromanski et al., 2020) such that:
Revised: Features (PRFs) (Choromanski et al., 2020) such that:

Original: [Equation]
Revised: [Equation]

Original: h (cid:19)
Revised: h (cid:19)

Original: We can approximate the aggregated attention utility (Eq. 5) by: ψ ( i ) b
Revised: We can approximate the aggregated attention utility (Eq. 5) by: ψ ( i ) b

Original: = (cid:88)
Revised: = (cid:88)

Original: j : h K ( K j )=
Revised: j : h K ( K j )=

Original: b
Revised: b

Original: A ij ∝ (cid:88)
Revised: A ij ∝ (cid:88)

Original: j : h K ( K j )=
Revised: j : h K ( K j )=

Original: b exp (cid:18) Q i · K j √ d h (cid:19)
Revised: b exp (cid:18) Q i · K j √ d h (cid:19)

Original: ∝ E  
Revised: ∝ E  

Original: φ ( Q i ) · (cid:88)
Revised: φ ( Q i ) · (cid:88)

Original: j : h K ( K j )=
Revised: j : h K ( K j )=

Original: b φ ( K j ) 
Revised: b φ ( K j ) 

Original: 
Revised: 

Original: [Equation]
Revised: [Equation]

Original: where ∝∼ denotes that the attention approximation converges under the positive random feature condition (Katharopoulos et al., 2020).
Revised: where ∝∼ denotes that the attention approximation converges under the positive random feature condition (Katharopoulos et al., 2020).

Original: Next, the Learning-to-Hash objectives for each query and each key can be deﬁned as:
Revised: Next, the Learning-to-Hash objectives for each query and each key can be deﬁned as:

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: The ﬁnal optimization objective L is a convex combination of the task-speciﬁc objective L task (e
Revised: The ﬁnal optimization objective L is a convex combination of the task-speciﬁc objective L task (e

Original: [Equation]
Revised: [Equation]

Original: g., cross-entropy loss in language modeling) and learning-to-hash objectives for all keys and queries in each attention head for each layer:
Revised: g., cross-entropy loss in language modeling) and learning-to-hash objectives for all keys and queries in each attention head for each layer:

Original: [Equation]
Revised: [Equation]

Original: F T RAINING D YNAMICS
Revised: F T RAINING D YNAMICS

Original: We plot the training dynamics of the KL divergence (i.e., L Q and L K ) and the negative entropy of the bucket-wise attention utilities (i.e., negative entropy of ψ for queries and ψ (cid:48) for keys) in Figure 4.
Revised: We plot the training dynamics of the KL divergence (i.e., L Q and L K ) and the negative entropy of the bucket-wise attention utilities (i.e., negative entropy of ψ for queries and ψ (cid:48) for keys) in Figure 4.

Original: KL divergence is the optimization objective of LHA, while negative entropy measures the diversity of bucket-wise attention utility.
Revised: KL divergence is the optimization objective of LHA, while negative entropy measures the diversity of bucket-wise attention utility.

Original: The model we use is a pure LHA model with 16 layers and 10 LHA heads, trained on WikiText-103 for 6000 steps.
Revised: The model we use is a pure LHA model with 16 layers and 10 LHA heads, trained on WikiText-103 for 6000 steps.

Original: Each LHA head has 8 hash buckets.
Revised: Each LHA head has 8 hash buckets.

Original: 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 ﬁrst increase (due to the increase of the diversity of bucket-wise attention utility), and then decrease (due to learning-to-hash optimization).
Revised: 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 ﬁrst increase (due to the increase of the diversity of bucket-wise attention utility), and then decrease (due to learning-to-hash optimization).

Original: G A DDITIONAL A NALYSIS ON T RAINING & I NFERENCE E FFICIENCY
Revised: G A DDITIONAL A NALYSIS ON T RAINING & I NFERENCE E FFICIENCY

Original: To further analyze the efﬁciency 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.
Revised: To further analyze the efﬁciency 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.

Original: Notice that in the inference stage, the LHA model needs not to calculate the learning-to-hash losses (i.e., the KL divergence terms).
Revised: Notice that in the inference stage, the LHA model needs not to calculate the learning-to-hash losses (i.e., the KL divergence terms).

Original: We study a single-layer attention, which has 10 attention heads and the hidden size of 410.
Revised: We study a single-layer attention, which has 10 attention heads and the hidden size of 410.

Original: Figure 5 illustrates the latency in both training and inference stage for softmax attention and LHA with different attention size.
Revised: Figure 5 illustrates the latency in both training and inference stage for softmax attention and LHA with different attention size.

Original: We can see that LHA can achieve more signiﬁcant speedup when facing longer sequences.
Revised: We can see that LHA can achieve more signiﬁcant speedup when facing longer sequences.

Original: H A DDITIONAL A NALYSIS ON THE B UCKET I MBALANCE I SSUES
Revised: H A DDITIONAL A NALYSIS ON THE B UCKET I MBALANCE I SSUES

Original: 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
Revised: 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

Original: Transformer.
Revised: Transformer.

Original: The results are shown in Figure 6.
Revised: The results are shown in Figure 6.

Original: 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.
Revised: 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.

Original: Notice that to compare the statistics of the bucket size and query-key ratios between LSH and LHA, we ﬁne-tune learnable hash functions for the same pre-trained Transformer.
Revised: Notice that to compare the statistics of the bucket size and query-key ratios between LSH and LHA, we ﬁne-tune learnable hash functions for the same pre-trained Transformer.

Original: 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.
Revised: 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.

Original: When calculating the query-key ratios, we do not count the buckets which have no queries or keys.
Revised: When calculating the query-key ratios, we do not count the buckets which have no queries or keys.

Original: I V ISUALIZATION

Original: We visualize the full attention map of an attention head in a Transformer model pre-trained on

Original: WikiText-103 dataset in Figure 7, and visualize its LHA counterpart and LSH counterpart in Figureand Figure 9, respectively.

Original: We do not apply the causal mask when calculating the attention weights.

Original: The attention utility of LHA and LSHare 0.461 and 0.335, respectively.

Original: From the visualization, we can see that LHA is able to better approximate the full attention with learnable hash functions.