Sparse Attention with Learning to Hash
Authors: Zhiqing Sun
Yiming Yang
Shinjae Yoo
Venue: ICLR 2022 Poster (2022)
Keywords: Sparse Attention
Transformer
Learning-to-Hash
Natural Language Processing

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 (Learning-to-Hash Attention), which directly learns separate parameterized hash functions for queries and keys, respectively. 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. 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.

Published as a conference paper at ICLR 2022
SPARSE ATTENTION WITH LEARNING-TO-HAN
Zhiqing Sun, Yiming Yang
Shinjae Yoo
Language Technologies Institute, Carnegie Mellon University (zhiqings, yiming) J 8cs. cmu. edu
Brookhaven National Laboratory s jyoo@on1..go gov
ABSTRACT
Transformer has become ubiquitous in sequence modeling tasks. As key com- ponent of Transformer, self-attentie does not scale to long sequences due 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 (A (ANN) methods, where similar queries and keys are allocated 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 the ANN methods such Locality-Senssitee cannot fully utilize the available real data distributions. To overcome these issues, this paper proposes
Hashing (LSH) are randomized and
strategy for sparse attention, namely LHA (Learning- to-Hash Attention), which directly learns separate parameterized hash functions for queries and keys, respectively. 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. Our experiments on evaluation of the WikiTText- 103 dataset for language modeling, the GLUE benehmark for natural language understanding, and the (text/image classification, retrieval, etc.) show the superior performance of LHA over other strong Transformer variants.
Anegge benchmark for multiple tasks
INTRODUCTION
The Transformer architecture (Vaswani et al., 2017) has been successfully applied to various tasks, including natural language processing (Vaswani et al. 2017 Devlin et al. al., 2018; DDa et al., 2019
Liu et al. 2019: Yang et al. (019)) 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 long sequences
result.
To address this inefficiency problem, one the solutions is to approximate the full attention matrix with 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 (A (ANN) approaches, including Locality Sensitive Hashing (LSH) (Kitaev et al. 20200) 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 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, 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 existing ANN-based methods for attention sparsi- fication. 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 quantify how well the sparse patterns approximate the full attention, and we show that deN-derived sparse patterns are substantially inferior to their counterparts. Thirdly, propose novel solution, namely Learning-to-HHar Attention (LHA), for dynamic attention

---

Published as a conference paper at ICLR 2022
sparsification with enhanced model expressiveness. LHA directly optimizes 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 other sparse attention models (Kitaev et al., 2020 Roy et al., 2.21). LHA reduces the overall complexity of self-attention from O(N²) to (NL.
for sequence length N. Our experiments in a wide range of tasks on the evaluation benchmarks for language modeling, natural language understanding, and Boonn-aage--aae Arena show that LHA achieves better performance compared strong transformer baselines.
2 RELATED WORK
Related work can be roughly divided into three categories, i.e., location-based sparse attention, content-based sparse attention, and dense approximation of attention, as outlined below.
Location- based sparse attention methods aim to improve the computational efficiency by using pre-specified global or local sparsification patterns over token locations. [Luu et al. (2018) proposed to alternate coarse attention layers and local attention layers. Child et al. (2019) used strided sparse attention pattern in image generation, Suuhhaatar et al. (2019) imposed sparsity based on the predicted temporal window size for each token. Other methods (Zhang al. 2021: Beltagy et al., 2020 |iinslie et al., 2020; Zaheer et al. (020) used a pre-specified subset of locations in the input the global memory, and only allow non-local attentions from this subset to all the other tokens. Location-basee sparse attention cannot leverage more flexible content-base based interactions among arbitrary positions, as limitation.
Content- based sparse attention allows more flexible sparse patterns than location-based ones. Malaviya et al. (2018) used sparsemax to obtain sparse attention matrix, while (Correia 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 al. 2021 Kitaey al. 2020 Daras 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. TTy et al. (2020b) designed differentiable sorting algorithm of internal representations to enable efficient mnasi-global local attention. Contemporary to our work, SparseFinder (Treviso et al. 202.) 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 few pairs. [Wang et al. (020b) 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. Chhromaaskii et al. (2020) and Peng al. (202I) further proposed an unbiased approximation of softmax with random Fourier features.
Our work in this paper is directly related 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 learn the hash functions for attention sparsification.
RE-EXAMINATION OF CONTENT-BASED SPARSE PATTERNS
3.1 PRELIMINARY
The self-attention mechanism (Vaswani et al.) 2017) can be formulated as the weighted sum of the value vectors V NNdx where the weights are calculated using query vectors Q E N and
2

---

Published as a conference paper at ICLR 2022
key vectors K E R Nxx as:
QK
Attention(Q, A- softmax
vdh
where A denotes the matrix of normalized attention weights, dy the dimension of hidden represen- tations, and N is the sequence length. self-attention has 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 clustering algorithm to queries } and keys {K₃}. 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 hash buckets is defined as:
Sparse-Attention (Qi, K, V)
Σ V,
j:ho (Q)))x()))())
where hg, hQ Rd + [B] are the hash functions for keys and queries, and A. x Ai is the re-normalized attention weights such that Vi, (1shh(Q)))hx((k) Aij 1. In general (Kitaev et al. 2020 Roy et al. 202I). calculating the hash function and performing local attention for each query have the time complexity of ( and O(N/B). respectively. Thus, the overall complexity of self-attention can be reduced from ((N2) to O(N- B+ /B/ ((V15) when B Rs N/B VN.
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(Q.)= Q::0; / MB + MR 11Q:13
G(K,) = K5; M3 MR
12: 0
(3)
where Mq maxq. 11Q:112 and Mr maxk, |1ارl such that [|F(Q.) G(K) const
Qi Kj- Routing Transformer (Roy et al., 2021) and Cluster-former (Wang et al. 2020a) use mini-batch spherical k-means to partition tokens into different clusters.
3.2 BUCKET IMBALANCE ISSUES
Previous content-based sparse attention models take it for granted that the ANN-derived sparse pattern can effectively approximate the full attention. However, 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 effcetively and efficiently:
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 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 LSH a typical ANN method. We apply LSH 10 attention heads in the 3rd layer of Transformer pre-trained on language modeling and obtain the results shown in Figure I (up). We can see that the imbalance issue not only exists in the query-key ratios, but also the bucket sizes. To go step further, we apply LSH to all 16 x 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
Notice that we only consider the setting of single-round hashing in our paper, but our analysis and the proposed LHHA method can be generalized to the multi-round hashing setting.
We use the same I.SH technique in (Kitaev al. 2020), except that do impose constraints 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.
'The detailed experimental setting can he found in the appendix.
3

---

Published as a conference paper at ICLR 2022
Statistics queries in 4x10=400 LSH buckets for 10 attention heads
1000
Queries
Bucket size
Keys
500
2
Overies (SMYFFF) (SMIRF)
Bucket son
1500
1000
500
attention heads
Figure 1: We show the unbalanced bucket sizes and unbalanced query-key ratios with statistics of 10 attention heads in the 3rd layer of a pre-trained Transformer. For each head, we assign the queries and keys of first 1024 tokens in the WikiText-100 validation data into 4 LSH buckets. The buckets are sorted for each attention head according to total bucket sizes (i.e., #query + #key).
sizes 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-ssarsifiee 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 (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 further magnified by the asymmetric transformation.
3.3 PROPOSED METRIC: ATTENTION UTILITY
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. (220),, which considers a practical case (Kitaev al. 2020 Roy et al. 2021 where all clusters strictly contain the same number of queries and keys:
N
vb [B]. H(Q hQ().) = b}1 KK, hg (g(())()))))))))))))))() b]
B
We denote with C the set of all possible assignments in B balanced non-overlapping clusters: = CC,.C2..... ,,r. where T' is the number of possible assignments. We can then define the attention utility AU of each assignment C₁ as:
AU(C.) = Σ
Aij where KY),Q
C. ((Q;, K2),K ni (K₂)}-
(5)
The attention utility represents the aggregation of the sparse attention weights 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, iee,
arg max AU(C,)
Σ
Aij
(6)
C,EC
(Q.,K))C.
is NP-hard.
This theorem is corollary of the NP-hardness of Attention Biclustering in (Daras et al. 2020), and motivates to develop learning- approach optimize the attention utility. Notice that similar
4

---

Published as a conference paper at ICLR 2022
Boxplot averaged 16x10=
attestion LEtEtt
Upper hound
Learning to-klach
Locality Sensitive
Hesting ILSP
atymirettic
varsformation
02
15
Figure 2: We show the attention utility evaluation of four different sparse attention patterns. For 1,H and LSHs, we assign the queries and keys of first 1024 tokens in the WikiText- 103 validation data into 4 hash buckets. The causal mask is not used when calculating the attention weights.
to other sparse attention models (Kitaev et al. 2020 Roy et al. 2021p, 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 2020 Daras 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 Learning- to-Hasi (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- not-products. We present box-plot of averaged attention utilities for all 16 Transformer layer in Figure 2 We that the LSH methods obtain significantly less attention utilities than the upper bound or L.HH. As conclusion, the ANN-derived sparse patterns arguably sub-optimal when used 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, propose an alternative choice of the hash functions to tackle this problem.
4 LEARNING-TO-HASS ATTENTION (LHA)
We propose a novel Learning-to to--otoHashhhh 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-hased hash functions such as LSH or online k-means and adapt the mismatch of query and key distributions. Besides, we empirically find that LHA can be used as plug-and-pla replacement for dense attention layers, which makes LHA applicable to the wide range of pre-trained Transformer models.
4.1 LEARNING-TO- -HASH FOR SPARSE ATTENTION
We first remind the readers of our definitions of the sparse attention:
Sparse-Attention/ K,
Σ
A
(hxx()))=x((K))
where hx ho normalized attention weights such that Vi, ==)x Au = 1. Inspired by the Learning-to- Hash methods (Wang et al. 2017), we implement the learnable hash functions lK, lig Rdd + [B] by defining parameterized functions Hq, Hx IRdh BB such that:
4 + [B] are the hash functions for keys and queries, and A:G x A,j is the re-
hQ(Q.) = arg max [Q))): hx(KK) = arg max [k((K))) bE(1.....B)
(8)
be be(
where HQ and HK can be any parameterized functions such as MLPs.
Notice that our formulation is 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 al., 2020) or the mini-batch k-means scheme in
Routing Transformer (Roy et al. 2021 we can set:
Ho(x) Hx (x) x: R: R|
Ho(x) Hx(x) z.
5

---

Published as a conference paper at ICLR 2022
where R Rd is random matrix, [ denotes the concatenation of two vectors, and μ Rd4x x are the exponentially moving averaged cluster centroids shared by keys and queries. The asymmetric transformation in SMYRF (Daras al. 2020) can also be reproduced by:
Ho(x) = (1
F F(x) ( (a)))aab))))
(L G(z) (( + b 1)
(x =1
(10)
where a random vector, is random scalar, r is scalar parameter, 1()) [B] + {0. 1).11 R the one-hot operation, and F-) and G(.) are defined as in Equation
4.2 OPTIMIZATION OBJECTIVE FOR LEARNING-TO-HASS
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 arg max operation to get the hashing buckets, the parameterized projections in our hash functions cannot be trained in end-to-end manner. To resolve this issue, 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 Q, can be written as (i.hh(Q.))hhx(k Aij. We first calculate the possible attention utilities (6( that query Qi can obtain when allocated to all B hash buckets:
the
16(00
Σ
A.j.
Σ
Agg
(11)
j:h (K;)=0
Ideally, our hash function would allocate Qi to the bucket that maximize the attention utility, such that Qi attend to the salient keys:
ha (Qi) = arg max vb arg max
Σ
Aij
(12)
bE[N.]
BE(B]
(g(
To achieve this goal, since (t) is naturally normalized categorical distribution, we can simply convert HQ (Qx) into distribution and use the KL divergence between the predicted distribution softmax( Hq (Qi)) and the desired distribution 2b()) the optimization objective:
Lq. K (L (Haimaax(HoQQ)))
(13)
On the other hand, the possible attention utilities zb" that a key K₁ obtains when allocated to all B buckets can be written as:
366
Σ
Agj
Σ
Asg
(14)
1::g(Q;)-B-1 8-1
While vbO is no longer a normalized distribution, we can normalize it and similarly define the optimization objective LK, also as a KI. divergence.
During the training stage, our final optimization objective is a convex combination of the task-specific objective and the learning-to-han ho-hash objectives (i.e., {Lo.} } and (Cx, }) for all queries and keys. Figure 3 in the appendix illustrates the joint training diagram of LHA.
4.3 APPROXIMATE ATTENTION UTILITIES
The remaining problem how to efficiently compute () and 2/, ,, naive way requires calculating
the dot products between all query-key pairs with O(N") complexity. Inspired by the recent advances
in kernelized attention (Katharopoulos et al. 2020 Choromanski et al. 2020 Peng et al.) (02T), we
use random Fourier features (Rahimi et al.
2007 Choromanski et al. 2020) to approximate 45(0) in
an unbiased manner. Let define Φ I R as the Positive Randomized Features (PRFs)
(Choromanski et al. 2020) such that
x xTy
E(Φ(x) Φ(y) = exp
(15)
Vdh

---

Published as a conference paper at ICLR 2022
We approximate 2b(1) with the following formula:
Qi Kj vdn
Ag x
Σ exp
x E ((Q))
Σ o(K))
(16)
(kg
for each bucket E [B]. Since hg (Ky) does not change for the queries, we can simply pre-compute (k) (K,) once to save computation, which reduces the complexity of computing {ıb(0)
for all queries from ((N2) to O(N B) = O(N1-)). The attention utilities for keys vb"o) can be efficiently approximated in similar way
((Q)) (ihol (QQ)) Σ5 o(K,)
(5
Σ
Agj R
E
(((Kg)
(17)
4.4 IMPLEMENTATION DETAILS
Practical content- content-based sparse attention models require that each hash bucket has the same size. which is crucial in terms computational efficiency on modern hardwares. Therefore, harrwfres.Therrffre,, we follow Roy et al. (2021) and sort the tokens with regard to normalized hash scores, i.e., (Hammaxxx@((Q)) softmax softmax(H&<))) (K.)). in each bucket. The hash bucket membership is then determined by the top-k threshold, where T7. the bucket size. Since such 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 v2x to increase the recall of queries and keys. pseudo-code implementation for LHA can found the appendix.
and
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 (Q+) and {hx( (K₃}) 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 S 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 joint hash function:
hg (K,) (Q;) arg max (Ho(Q:), (K,)))
(18)
bc(N.]
We find that this strategy empirically works better than using other tricks to fix this "no-attention- target" problem. Our solution different from Routing Transformer (Roy al. (202T) or Reformer (Kitaev et al. 2020), which impose an extra constraint that queries and keys are tied.
5 EXPERIMENTS
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-Range- Arena. Due to the space limitations, the detailed hyper-parameter settings are presented in the appendix.
5.1 LANGUAGE MODELING
Wikitext 103 Merity et 2016) is 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 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-XI. (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 LHA), we use the
*llass refer Choromanski
2020 for how PRFs stabilize attention renormalization.

---

Published as a conference paper at ICLR 2022
Table Ablation studies the WikiText- 103 validation data in the base setting. Lower perplexity (PPL) is better. All the models have a total of 16 attention layers and 10 heads. Non-Local (NL. i.c.,
LSH LHA) layers when present are always added at the top of the model. Attention size denotes either local window size or hash bucket size. f denotes that the results are taken from Transformer-XL. (Dai et al. 2019).
label
NL Heads
NL Layers
Hash Func.
Att. Size
#Param
Valid PPL
Test PPL
0
768
151M
22.89
Local Transformer
0
384
151M
23.82
0
640
151M
23.09f
24.0f
10
16
Rand. Linear
384
153M
24.64
16
Rand. Linear
384
153M
23.51
LSH Transformer
10
8
Rand. Linear
384
153M
23.76
5
8
Rand. Linear
384
153M
23.53
(a)
10
16
Rand.
384
153M
26.05
(h)
10
16
Linear
384
153M
24.22
LHA Transformer
(c)
16
384
153M
24.00
(d)
16
MLP
384
153M
22.79
23.2
(c)
10
8
384
153M
23.00
(f)
8
MLP
384
153M
22.84
Table 2: Performance of methods per task on the GLUE benchmark, where C' and # denote the
number of queries/keys per cluster and hashing rounds, respectively. t denotes that the results taken from their original papers.
C CoLA
MNLI
MRPC
QNLI
QQP
RTE
SST-2
STS-B
AVG
RoBeRTa-basc
60.9
87.6
88.7
92.7
91.6
68.5
94.5
90.0
84.3
SMYRF
58.9
82.3
89.5
89.3
93.1
87.8
81.4
SMYRF
32
58.8
85.0
87.7
91.1
89.7
68.6
93.2
89.7
83.0
FCAt
59.8
79.4
43.6
74.6
89.4
49.8
94.4
78.9
71.2
i-FCAt
60.1
88.0
87.3
93.0
91.5
70.4
94.7
90.0
84.4
LSH
11
61.5
86.3
88.7
91.1
90.2
93.3
83.5
LSH
22
61.5
86.8
89.7
91.9
90.8
68.4
93.9
88.7
84.0
LSH
45
61.8
87.5
88.7
92.7
91.2
68.1
94.0
89.4
84.2
LHA
11
61.2
85.7
89.1
91.3
90.5
68.0
93.7
89.0
83.6
LHA
22
62.0
86.8
88.0
92.4
91.1
68.6
94.8
89.6
84.2
LHA
45
62.3
87.7
89.1
92.8
91.4
68.5
94.6
89.5
84.5
per-layer sinusoidal absolute positional encoding (Vaswani et al.) (017). The LSH Transformer is implemented in our code base by simply replacing the learnable hash functions random projections. sequence length of 1536 is used both training and evaluation stages for our models. In ablation study, vary 1) the type of non-local attention heads, 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- montent-independen random numbers. The experimental results are presented in Table 1
From the table, we can see that a smaller attention size (384 compared 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 very strong baseline for language modeling that outperforms pure LHA Transformer. We can also see that two-layer MI.P works better than 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 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 evaluation, the ablation study. We can see that the proposed LHA model outperforms Transformer-A XL, which further validate that LHA can help modeling long-term dependency for language modeling tasks.
the best setting from
5.2 NATURAL LANGUAGE UNDERSTANDING
To show the ability our model to approximate arbitrarily complicated attention distributions, we evaluate proposed method the approximation of RoBERTa model (Liu al. 2019) the
8

---

Published as a conference paper at ICLR 2022
Table 3: Accuracy L.,, IMDb, AAN, and Image in Long Range Arena benchmark. Best and second best model per task is shown in boldface and underlined. Throughput is evaluated on IMDb and relative to the vanilla transformer's. and & denote being statistically significantly better (p 0.05 in one-tail proportion test) than vanilla Transformer and the second best model. & denotes that the throughput comparison are run on single NVIDIA Tesla V100 32GB GPU, while previous results (Tay et al. 2020c) are reported on x 4 TPU V3 chips. V, Q and denotes low-rank/kernelized attention, content-based sparse attention, and location-basee sparse attention, respectively.
Accuracy (t)
Throughput (1)
Model
LO
IMDb
AAN
Image
Avg.
1K
2K
3K
4K
Vaswani et al.
2017
Transformer
36.4
64.3
$7.5
42.4
50.2
1.0
1.0
1.0
1.0
Wang et al.
2020b
Linformer
35.7
53.9
52.3
38.6
45.1
1.2
1.9
3.7
5.5
Kitaev et
2020
Reformer
37.3
56.1
53.4
38.1
46.2
0.5
0.4
0.7
0.8
Child et
2019
Sparse Trans."
17.1
63.6
59.6
44.2
46.1
Tay ct al.
2020b
Sinkhorn Trans.
33.7
61.2
53.8
41.2
47.5
1.1
1.6
2.9
3.8
Tay et
2020a
Synthesizer
37.0
61.7
54.7
41.6
48.8
1.1
1.2
2.9
1.4
Zaheer
2020
BigBird
36.0
64.0
59.3
40.8
50.0
0.9
0.8
1.2
Katharopoulos et
2020
Linear Trans.
16.1
65.9
53.1
42.3
44.4
1.1
1.9
3.7
5.6
Choromanski
2020
Performer
18.0
65.4
53.8
42.8
45.0
1.2
1.9
3.8
5.7
Peng et al. 202
36.8
66.0
56.1
1.7
3.4
5.3
Ours LHA
37.9
66.8*
62.8*
45.2"
53.1
1.04
1.14
1.34
1.54
Following the common practice, the maximum sequence length is set to 128. To show the effectiveness of LHA, we choose two competitive baselines literature: SMYRF (Daras et al., (020) and Fast Clustered Attention (FCA) (Vyas al. 2020). We also produce an LSH-based sparse attention baseline. We use pure LHA model similar variant (c) in Section 5.1
GLUE (General Language Understanding Evaluation) dataset (Wang al.
2018
We summarize the performance 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 its sparse attention part.
5.3 LONG-RANGE-ARER BENCHMARK
Long-Range-AAree (LRA) benchmark Tay 2020c) on evaluating model quality under long-context scenarios for Transformer variants. We follow the apples-to-apples setting of LRA benchmark and compare method against other efficient attention variants. We use an LHA/local hybrid variant similar to variant (d) in Section 5.1
recently proposed benchmark focused
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 20133, and CIFARIO (Krizhevsky 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 Section 2 From the table, we can see that LHA achieves consistent improvements over previous efficient attention models.
CONCLUSION & FUTURE WORK
In this paper, we address the limitations of ANN- sparse attention methods and propose the Searning-toHHan 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 LHA model on much larger language modeling datasets, such PC-19 (Rae et al. 2019)
www.ggithub..commgoggg
appl setting

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGEMENT
This research used Perlmutter supercomputer of the National Energy Research Scientific Computing Center, a DOE Office of Science User Facility supported by the Office of Science of the U.S. Department of Energy under Contract No. DE-AC02-05CHI123 using NERSC award NERSC DDR-ERCAP002210.
ETHICS STATEMENT
This work proposed 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. (019).. Since one of the advantage of our model is re-use a pre-trained Transformer model, the effective adversarial perturbations on the pre-trained model could also be effective in our model.
REPRODUCIBILITY STATEMENT
We provide a GitHub repository 6 for our source code. The Hlyperpaarameter are described in details in the appendix. We also provide pseudo-code implementaion of model in the appendix.
All the datasets used in the experiments and the corresponding pre-processing scripts can be found online, including language modeling GLUE benchmark Long-Range-Aie Arenga benchmark and time series forecasting The pre-trained model we used (i.c., RoBERTa) is also publically available
REFERENCES
Joshua Ainslie, Santiago Ontanon, Chris Alberti, Vaclav Cvicek, Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai, Qifan Wang, and Li Yang. Ete: Encoding long and structured inputs transformers. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 268 284, 2020.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint 2rXiv22044.05550 2020.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclau- rin, and Skye Wanderman-Milne Jax: composable transformations of python+ numpy programs. Version 0.1, 55, 2018
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European Conference Computer Vision, pp. 213-229. Springer, 2020.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. In International Conference on Learning Representations, 2020.
10

---

Published as a conference paper at ICLR 2022
Gonçalo M Correia, Vlad Niculae, and André FT Martins. Adaptively sparse transformers. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference Natural Language Processing (EMNLP-IJCNLP), pp. 2174-2184, 14, 2019.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoe Le, and Ruslan Salakhutdi- Transformer-x:: Attentive language models beyond fixed-length context. arXiv preprint arXiv: 2019.
Giannis Daras, Nikita Kitaev, Augustus Odena, and Alexandros G Dimakis. Smyrf: Efficient attention using asymmetric clustering. arXiv preprint arXiv:2010.05311. 2020.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:/810.000 /81::088000 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x 16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 1929, 2020.
Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee. Flax: A neural network library and ecosystem for jax. Version 0.3, 3, 2020.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pp. 5156-5165. PMLR, 2020.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.0445/ 2020.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Peter Liu, Mohammad Saleh, Etienne Pot, Ben Goodrich, Ryan Sepassi, Lukasz Kaiser, and Noam Shazeer. Generating wikipedia by summarizing long sequences. arXiv preprint 1rXiv:1801.1011 2018.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint
2019.
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv
preprint wrXiv:1608.039% 2016.
Iya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv: 2017.
Andrew Maas, Raymond E Daly, Peter T Pham, Dan Huang, Andrew Ng, and Christopher Potts.
Learning word vectors for sentiment analysis. In Proceedings of the 49th annual meeting of the association for computational linguistics: Human language technologies, pp. 142-150, 2011.
Chaitanya Malaviya, Pedro Ferreira, and André FT Martins. Sparse and constrained attention for neural machine translation. arXiv preprint rrXiv:180500824: 2018.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint arXiv: arXiv: 1609.07843, 2016.
Paul Michel, Xian Li, Graham Neubig, and Juan Miguel Pino. On evaluation adversarial perturba- tions for sequencceto-ssequeece models. arXiv preprint arXiv: 2019.
Nikita Nangia and Samuel R Bowman. Listops: A diagnostic dataset for latent tree learning. arXiv preprint
2018.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. arXiv preprint arXiv:2103.02143. 2021.

---

Published as a conference paper at ICLR 2022
Dragomir R Radev, Pradeep Muthukrishnan, Vahed Qazvinian, and Amjad Abu-Jbara. The acl anthology network corpus. Language Resources and Evaluation, 47(4))919-)) 99-444 2013.
Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. arXiv preprini arXiv: 0955011111055..... 2019.
Ali Rahimi, Benjamin Recht, et al. Random features for large-scale kernel machines. NIPS, volume 3, pp. 5. Citeseer, 2007.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53-68, 2021.
Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin. Adaptive attention span in transformers. arXiv preprint arXiv:1905.07799 2019.
Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng. Synthesizer: Rethinking self-attention in transformer models. arXiv preprint arXiv:2005.0077 2020a.
Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da-Cheng Juan. Sparse sinkhorn attention. In International Conference Machine Learning, pp. 9438-9447, PMLR, 2020b.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. arXiv preprint M.Xi::2011.04000 2020c.
Marcos Treviso, António Góis, Patrick Fernandes, Erick Fonseca, and André FT Martins. Predicting attention sparsity in transformers. arXiv preprint arXiv:2109.12/88. 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones. Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. arXiv preprint arXiv:1706.0377 2017.
Apoorv Vyas, Angelos Katharopoulos, and François Fleuret. Fast transformers with clustered attention. Advances in Neural Information Processing Systems, 33, 2020.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv: 2018.
Jingdong Wang, Ting Zhang, Nicu Sebe, Heng Tao Shen, et al. A survey on learning to hash. IEEE transactions on pattern analysis and machine intelligence, 40(4):769-799) 2017.
Shuohang Wang, Luowei Zhou, Zhe Gan, Yen-Chun Chen, Yuwei Fang. Siqi Sun, Yu Cheng, and Jingjing Liu. Cluster-former: Clustering based sparse transformer for long-range dependency encoding. arXiv preprint arXiv:2009.06097, 2020a.
Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.. 04768, 2020b.
Thomas Wolf, Julien Chaumond, Lysandre Debut, Victor Sanh, Clement Delangue, Anthony Moi, Pierric Cistac, Morgan Funtowic7, Joe Davison, Sam Shleifer, et al. Transformers: State-of-the--n natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 38-45 2020.
Haixu Wu, Jiehui Xu, Jianmin Wang, and Mingsheng Long. Autoformer: Decomposition transformers with auto-correlation for long-term series forecasting, arXiv preprint arXiv:2106.1300 2021.
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nystr'," omformer: A nystr)," om-hased algorithm for approximating self-attention. arXiv preprint arXiv: 102. 2021.
Zhilin Yang, Zihang Dai, Yiming Yang. Jaime Carbonell, Ruslan Salakhutdinov, and Quoc V
Le. Xlnet: Generalized autoregressive pretraining for language understanding. arXiv preprint arXiv: /906 08237, 2019.
12

---

Published as a conference paper at ICLR 2022
Andy B Yoo, Morris A Jette, and Mark Grondona. Slurm: Simple linux utility for resource agement. In Workshop on job scheduling strategies for parallel processing, pp. 40.60. Springer, 2003.
Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences, arXiv preprint arXiv:2007.10662 2020.
Hang Zhang, Yeyun Gong, Yelong Shen, Weisheng Li, Jiancheng Lv, Nan Duan, and Weizhu Chen. Poolingformer: Long document modeling with pooling attention. arXiv preprint arXiv:2105.04371, 2021.
Haoyi Zhou, Shanghang Zhang. Jieqi Peng, Shuai Zhang, Jianxin Li, Hui Xiong, and Wancai Zhang. Informer: Beyond efficient transformer for long sequence time-series forecasting. arXiv preprint arXiv:2012.07436, 2020.
13

---

Published as a conference paper at ICLR 2022
Generate attention weights
Optimized towards down-stream tasks
Leanable Functions
Optimized towards attention urilities
Queries
Keys
Genenate spanse puttems
Figure 3: The joint training diagram of LHA, where queries and keys generate attention weights train the learnable hash functions, while the hash functions generate the sparse attention patterns, on which queries and keys are trained towards down-strent tasks.
A EXPERIMENTAL SETTINGS
A.1 RE-EXAMINATION OF CONTENT-BASED SPARSE ATTENTION
For analysis, pre-trained 16-layer and 10-head Transformer the WikiText-103 language modeling benchmark (Merity et al. 2016). A per-layer sinusoidal absolute positional encoding (Vaswani et al. 2017) is injected the queries and keys before self-attention
For Figure I 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 we follow the same experimental setting.
A.2 LEARNING-TO-HALL
During the training stage, our final optimization objective C is a convex combination of the task- specific objective Lrtas (e.g., cross-entropy loss in language modeling) and learning-to-hash objec- tives for all keys and queries, i.e., {Cq,} and { {LK,) (See equation 13
f= (taskk N H
A
L ΣΣΣ
(
L.
L.H + Ckj
(19)
where A, N, L, H denote the loss coefficient, the sequence length, the number of attention layers, and the number attention heads per layer, respectively. Empirically, we set A = 0.05 and found work well across different tasks.
A.3 LANGUAGE MODELING
We the base setting of Transformer-Xi XL (Dai et al.
2019) in experiments, which consists of 16 Transformer layers. For each layer, the hidden size is set to 410, and the number of attention heads set to 10. The dimension of feed-forward layer is set 2100. All codes are implemented based
Flax (Heek et al. 2020) in JAX (Bradbury et al. (0118) The dropout ratio is set to 0.2. The batch size set to 32. We use Adam (Loshchilov & Hutter 2017 the optimizer, and set (8,, (31,32) to (0.9, 0.999). The peak learning rate is set to 3.5e-4. The model is trained for 20k steps with 2k- warm-up up stage with cosine learning rate decay (Loshchilov & Hutter 2016).
A.4 NATURAL. LANGUAGE UNDERSTANDING
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 set to 768, and the number of attention heads is set to 12. The dimension taken from the Transformers Wolf et al. 2020) library. We use AdamW (Loshchilov & Hutter 2017 as the optimizer, and set (51, 32) to (0.9, 0.98). To fine-tune the pre-trained models, we search the optimization hyper-parameters in search space including different batch sizes (16/32/48), learning rates ((1-5) le-5), and the number of epochs (3-5). 10% warm-up schedule and linear learning rate decay (Devlin et al. 2018) is applied.
feed-forward layer is set to 3072. The RoBERTa pre-trained checkpoint is
14

---

Published as a conference paper at ICLR 2022
Table 4: Univariate long sequence time series forecasting results two datasets (lower is better).
denotes that the results are taken from (Zhou et al. 2020
Models
Ours LHA
Informer:
Log Tranransf
Reformers
DeepARt
Propheti
Metrie
MSE
MAF
MSE
MAE
MSE
MAR
MSE
MAE
MSR
MAR
MSR
MAE
0.060 0.191
0.098 0.247
0.103 0.259
0.222
0.389
0.107 0.280
0.115
0.275
ETTm ElTh:
48
0.088 0.236
0.158 0.319
0.161 0.322
0.284
0.445
0.162 0.327
0.168 0.330
168
0.158 0.328
0.183 0.346
0.187 0.355
1.522
1.191
0.239
0.422
1.224
0.763
336
0.219
0.399
0.222 0.387
0.230 0.398
18800
1.124
0.445 0.552
1.549 1.820
720
0.265
0.437
0.269 0.435
0.273 0,463
2.112 1,436
0.707 0.658
2.735
3.253
24
0.016 0.102
0.030 0.137
0.065
0.202
0.095
0.228
0.091
0.243
0.120
0.290
0.021 0.112
0.069
0.203
0.078
0.220
0.249 0.390
0.219
0.362
0.133 0.305
0.054 0.180
0.194 0.372
0.199 0.386
0,920 0.767
0.364
0.496
0.194 0.396
288
0.121 0.277
0.401 0.554
0.411 0,572
1.108 1.245
0.948
0.795
0.452 0.574
672
0.403 0.572
0.512 0.644
0.598 0.702
1.793
1.528
2.437
1.352
2.747 1.174
A.5 LONG-RANGE-AREES BENCHMARK
We follow the apples to apples setting 2 of LRA benchmark and compare our method against other efficient attention variants. Specifically, feed-forward layer size for the IMDb text classification task et al.) 201 1). a 4-layer Transformer with 256 hidden size and 1024 feed-forward layer size for the ListOps task (Radvv al. 2013).
4-layer Transformer with 256 hidden size and 1024
4-layer Transformer with 128 hidden size and 512
layer size for the document retrieval task(Nangia & Bowman 2018), and 4-layer Transformer with 128 hidden size and 64 feed-forward layer size for the CIFAR TO image classification task Krizhevsky et al., 2009). The models are trained
for 20k steps with a stestep warm-up stage with a cosine learning rate decay (Loshchilov & Hutter
2016).
A.6 TIME SERIES FORECASTING
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, and the layer of decoder is set 2. The head number of multi-head attention is chosen from {8, 16). and the dimension of multi-heat attention's 's output is set as 512.
A.7 COMPUTING INFRASTRUCTURE
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 Slurm (Yoo et al. 2003) system. The evaluation of the inference throughput is performed on a stand-alone machine with NVIDIA Tesla V100 32GB GPU.
B ADDITIONAL EXPERIMENTS ON TIME SERIES FORECASTING
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. ETTh; is dataset for lehour-level, while ETTm; is 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 metries. We use an LHA/local hybrid variant similar to variant (d) 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 that LHA significantly improve the performance of ofate-of-the-art for most settings, while slightly underperforms Informer (Zhou et al.) 2020) in two settings of ETTh; dataset. This verifies the effectiveness of LHA on time series data.
www. - g thub. http:p/www..oooooeeeeeoeeoooaooa
ao-aapless paplss-- setting
15

---

Published as a conference paper at ICLR 2022
Table 5: Multivariate long sequence time series forecasting results datasets (lower is better).
denotes that the results are taken from (Zhou et al. 2020
Models
Ours
LHA
Informert
LogTranst
Reformert
LSTnet
Metric
MSE
MAE
MSE
MAE
MSE
MAE
MSE
MAE
MSE
MAE
0.475
0.469
0.577
0.549
0.686
0.604
0.991
0.754
1.293
0.901
ETTh
0.569
0.544
0.685
0.625
0.766
0.757
1.313
0,906
1.456
0.960
168
0.883
0.706
0.931
0.752
1.002
0.846
1.824
1.138
1.997
1.214
336
0.970
0.793
1.128
0.873
1.362
0.952
2.117
1.280
2.655
1.369
720
1.021
0.805
1.215
0.896
1.397
1.291
2,415
1.520
2.143
1.380
24
0.056
0.191
0.323
0.369
0.419
0.412
0.724
0.607
1.968
1.170
0.088
0.236
0.494
0.503
0.507
0.583
1.098
0.777
1.999
1.215
0.158
0.328
0.678
0.614
0.768
0.792
1.433
0,945
2.762
1.542
0.219
0.399
1.056
0.786
1.462
1.320
1.820
1.094
1.257
2.076
672
0.265
0.437
1.192
0.926
1.669
1.461
2.187
1.232
1.917
2.941
In multivariate forecasting setting. all the seven features are the prediction targets. The multivariate evaluation results can be found in Ta We can see that LHA significantly improve the performance of state-of-the-ar for all settings, which further demonstrates the effectiveness of the proposed I.HA method.
C PROOF
We follow the notations (Daras et al. 2020). Let C the set of all possible assignments B balanced non-overlapping clusters. A specific assignment is denoted by C, and there are T possible such assignments:
c² = (c C ,...Cr
(20)
c = (Qi.... Q2 K Ki...., K& a GsUU, vie (1....,.) L}
cP (c1, -. CB)
Cx Gay = Ø Vea, Cy E CP
(21)
We have the following lemma, which is referred the max-mass problem: Lemma 1.
DDaras el al. 2020) The optimization problem:
max EQi
Qu K
(22)
CREER
(Q.,K)ECP
NP-hard.
The main idea of the proof is to show that solving polynomially the above problem would mean that we could also solve polynomial time the 3-DM, which is known 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
Proof of Theorem Let (Qi) and (K,} denote the query and key sets that consider for comput- ing the attention utility. We can first construct the new query set Q' and key set K"' such that
eock,
Q K2-
where
E EK
Vah
4e Vdn )
This can be achieved by applying SVD the matrix of set Q", such that
Next, we construct another query
Q
Q¹'
Σド Q K3
The problem of finding the optimal attention utility, i.e..
arg max wa..maxXuCC AU(C))
Σ
Ag
C,EC
(Q.,
16

---

Published as a conference paper at ICLR 2022
is thus equivalent the problem in Lemma with Q" and X'' the query and key sets, which is proven to be NP-hard.
D ALGORITHM
We present a detailed pseudo-code implementation for LHA in Algorithm |
Algorithm 1 Single-layer Single-hear head Learning-to tonee--t-HHaee Hash Attention (LHA)
1: Queries, Keys and Values: Q, K, V RNdd
2: Query/Key Hash Functions: Ho, H
3: Attention Utility Loss Weight: À
4: if causal mask then
5: Qscore f softmax(Hq) (Q) (K))))
B x N x dr
6:
Kscore + Q scoore
7:
else
8:
Qscore +
softmax (Q))
D B N x de B x N x de
9:
Kscore softmax H (K))
10: W + N/B
D bucket size x N log(NN B X N log( N) W da
11: Qidz +- top- { (V2W)//// (Qscore)
12: K Kids <- top- (2W) (WW((K) (Ksoore)
Q' + gather( Q Quaz)
14: K' + gather( K, Kida)
B X W X da
V' + gather(V Kidw)
W de,
16: A + Q'(K")*
D B X W X W X dj.
if causal mask then
18: A + mask(A)
D B W X W B W W
A + softmax (A).
20: V + A-V"
D B X W X W X dn.
merge Kide, V*)
W d
22: 2 AUq(Q. K)
D B X N, Equation
16
23: S AUx (Q, K)
B N, Equation T7
24: LAU A K KT[1bllsscnee + A KT(bb']] secor))
25: return X, CAU
E A COMPLETE DERIVATION OF THE LHA OBJECTIVE
Let Q, K, E Nxd, denote the query, key, and value vectors of the attention mechanism, where N is the sequence length and d, is the hidden size. Let Φ Rd + R²D be Positive Randomized Features (PRFs) (Choromanski et al. 2020) such that:
xTy y Vdh
E(Φ(x) x(y)] = cxp
We can approximate the aggregated attention utility (Eq.
5 by:
Ki-Ky
Σ
Ag
Σ
exp
o E
((Q.)-
Σ
Φ(K,)
Von
J.hk
(=b
(k)
17

---

Published as a conference paper at ICLR 2022
divergence tsaming to hash kays layer
Lyy
32
oper
1500
Kye
(40%
400
0.2
100
035
Ugpe
Lyy
010
Layer
0.1
033
180
X59
100g
960
650
6131
Negatve eмerоу querys
negeeve erenpyy of keys byer
Lo
-14
.15
22
23
428
3030
200
300
3.00
800
1300
3000
3E1
400
300
800
Figure 4: The training dynamics of the KL divergence and negative entropy in LHA for queries and keys.
Σ =j exp
Q.KK vdh
3 = Σ
Ay
Σj= exp
(e.Kg
(bo(Q.)=b
vin
E
Φ(Qi) Ejihk Φ(K,)
E
(((()))) E/ H(K,)
((Qi). -1
X
((Qi) Σx ((K,)
2(Qi) j: ((Q)) En Φ(K,)
x E
Σ
k(K))
where X denotes that the attention approximation converges under the positive random feature condition (Katharopoulos et al. 2020). Next, the Learning-tor HasHash objectives for each query and each key can be defined as:
Lo. K (T (H
(23)
CK.
(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 saminngtt--hasee objectives for all keys and queries in each attention head for each layer:
A == (1 Lrask N. I.
L. I. ( Ckg
(25)
H
18

---

Published as a conference paper at ICLR 2022
during interence
during training (ms)
softmar attention
LIAA with
175
LHA
LNA with attertion
125
25
10.0
7.
24
5-
00
soquenic longth (x10F)
soquence longth (x10f)
Figure 5: The training and inference latency of LHA and softmax attention. The latency is measured single NVIDIA Tesla A100 GPU with batch size of 1.
smyrt
140
Ish
smyrt
12h
num of attention hnads
12h
num. of ettertion heads
optimal
optimal
100
150
100
60
20
500
1000
1500
2000
10
bucket size
query key ratio
Figure 6: We show the histogram of hash bucket statistics for all 16 10 160 attention heads in the pre-trained 16-layer Transformer. For each head, assign the queries and keys of first 1024 tokens in the WikiText-103 validation data into 4 LSH buckets. (left) The histogram the bucket sizes buckets. (right) The histogram of the query-key ratios for all 160 x 640 buckets. The dashed red lines denote the optimal case where each hash bucket has the same size and same number of queries and keys.
160 x = 640
F TRAINING DYNAMICS
We plot the training dynamics of the KL divergence (i.e., LQ and Lx) and the negative entropy of the bucket-wise attention utilities (i.e., negative entropy of vp for queries and v/ for keys) Figure KL divergence is the optimization objective of LHA, while negative entropy measures the diversity of bucket-wise attention utility. The model we use a pure LHA model with 16 layers and 10 LHA heads, trained 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 KI. divergence first increase (due to the increase of the diversity of bucket-wise attention utility), and then decrease (due to learning-to- -ooming-to-hash optimization).
G ADDITIONAL ANALYSIS ON TRAINING & INFERENCE EFFICIENCY
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 to-hash losses (i.e., the KL divergence terms). We study a single-layer attention, which has 10 attention heads Fillustaeess and the hidden size of 410. Figure illustrates the latency in both training and inference stage for
19

---

Published as a conference paper at ICLR 2022
softmax attention and LHA with different attention size. We see that LHA achieve more significant speedup when facing longer sequences.
H ADDITIONAL ANALYSIS ON THE BUCKET IMBALANCE ISSUES
To further analyze how the bucket imbalance issues are alleviated by the proposed LHA, we plot the histograms of bucket sizes and query-ker key ratios for all 16 10 160 attention heads in pre-trained Transformer. The results are shown in Figure 6 We can see that learning-to-hass attention (12h) has more buckets close to the optimal bucket size (i.e., 512) and has more balanced query-key ratios for
hash buckets.
Notice that 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 token sorting strategy (Roy et al. 2021) to maintain the same bucket sizes. When calculating the query-key ratios, we do count the buckets which have no queries or keys.
20


Reviews:
[{"review_id": "c9Lk_AOYqa3", "reviewer": "Reviewer_fhAP", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "### Strengths:\n* The attention sparsification is important to reduce the complexity of Transformer when applied to long sequence.\n\n* The imbalance of bucket size and query-key ratios are studied, which is closely related to efficiently reducing the complexity and performance on down-stream tasks.\n\n* The attention utility is proposed as a metric to measure how well the attention weights are preserved.\n\n* The attention utility is used to train the learnable hash functions since the argmax operation makes it impossible to train the hash functions from the down-stream tasks.\n\n* An approximation of the attention utilities is proposed with random Fourier features to reduce the complexity.\n\n### Weaknesses:\n\n* What is the statistics of the bucket size and query-key ratios for the proposed LHA? Is it significantly better than LSH?\n\n* It is not clear how to formulate the final training objective and balance the loss terms.\n\n* It is not clear how to apply LHA as a plug-and-play replacement for dense attention layers in pre-trained Transformer models. Since the LHA method introduces the learnable hash functions h_k and h_Q, the hash functions should be trained on the target dataset.\n\n* The implementation is inconsistent with the analyze. By using the token sorting method in Roy et al. 2021, the validation of attention bi-clustering is no longer guaranteed. However, the attention utility is meaningful only if the  attention bi-clustering is guaranteed. \n\n\n### Minor:\n\n* The notation of key and query is confused. Q_i should be the query, and a key should use K_i (before (14)). \n\n* To make the paper self-contained, it is necessary to include a brief introduction to the compared baselines in Table 3. Adding a column to show what type of attention sparsification would help the reader to understand the comparison.\n\n* Below(14): is not longer --> is no longer\n\n* Below (16): h_K(K_j) does change --> does not change", "summary_of_the_paper": "This paper proposes a learning to hash attention (LHA) to learn sparse attention for Transformer.\nThe 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. \nExperiments on several applications validate the effectiveness of the proposed LHA.", "summary_of_the_review": "This paper propose to learn separate hash functions for keys and queries with the guidance of attention utilization, which is the metric proposed in this paper. However, the implementation is inconsistent with the analyze and important method details are missing. At this point, I would tend to vote this paper slightly below the bar. If the authors can address my concerns, I would be happy to increase my score.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "ZIZvo9aTE9h", "reviewer": "Reviewer_wDdo", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The authors figure the bucket imbalance problem in ANN-derived content-based sparse attention and analyze its weakness on the imbalance problem. Attention Biclusteringis introduced to find the optimal attention utility. The learning to hash-based attention model is proposed to improve the effectiveness of sparse attention. Experiments on different applications support their claims.\n\nConcerns:\n1. It is well-known that LSH is data-independent hashing, which is formulated by random projection with some pre-defined metrics. Notably, using learning to hash is better than the LSH and its variants, in most cases, for representation learning. Therefore, the authors need to clarify why replacing LSH with learning to hash models is useful in sparse attention.\n2. Such a simple replacement could fully support your work in the present form. The authors should give full reasons for your contributions.\n3. The authors fail to convince the reviewer what are the connections between Theorem 1 and your proposed method. To me, you can directly claim there are some limitations on LSH-based methods, and implementing learning to hash can improve its representation capabilities on feature learning.\n4. Efficiency is one of the proposed methods, and more analysis on the efficiency and efficacy is necessary.\n", "summary_of_the_paper": "This work studies the content-based sparse attention in transformer and how to improve the self-attention part, i.e., efficiency and effectiveness.", "summary_of_the_review": "The authors try to improve some limitaions on the content-based sparse attention when using LSH-produced sparse attention patterns. A Learning-to-Hash Attention is formulated to enhance its model expressiveness. Experiments show its usefulness of the proposed methods. However, there are questionable points that should be clarified.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "9HloATeo8sA", "reviewer": "Reviewer_zEkG", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n\nI really enjoyed reading this paper. It starts by departing itself from related works by investigating the critical issue of not having balanced buckets, showing that this characteristic is a common pitfall of previous popular works. Therefore, the foundation is clear and has practical implications.\n\nMoreover, quantifying the degree to which an approximate self-attention can imitate the full attention is very informative and might guide future work on fast transformers. The authors acknowledge that computing this number is an NP-hard problem, but they go further and try to optimize this metric during training. \n\nThe idea of having a learnable function prior to bucketing is not entirely novel (see a contemporary idea by Treviso et al., 2021), but the formulation is concise and clear, enabling generalizations to previous approaches.\n\nThe experiments span a large set of tasks in NLP and CV with increasingly large input lengths. LM experiments show that the proposed method (trainable LSH) outperforms Reformer (untrainable LSH) on multiple settings. Results on the GLUE benchmark show that the proposed method achieves results on par with the original RoBERTa model, which is a good sanity check. The experiments on Long Range Arena give the final flavor to the paper, with the proposed method outperforming related works on all considered tasks while being cheaper (higher throughput) to train than a standard transformer.\n\n\nCons:\n\nHowever, I was also able to spot some concerns in this paper. Concretely:\n\n- The core of the proposed method relies on optimizing the attention utility, which relies on the Performer's softmax kernel to regularize the training regime. The paper opts to use a KL between the softmax kernel and LHA to that end. It would be interesting to see how this KL term evolves during training.\n\n- The choice of 10 heads from the 3rd layer in section 3.2 looks arbitrary. In fact, it is widely known that different heads at different layers focus on distinct tokens. Thus, what is the motivation for that choice? Moreover, from which pre-trained model were the results extracted? This is not clear in the paper.\n\n- Figure 1 shows the bucket size of queries and keys for each attention head using LSH attention. Since LSH attention was originally designed to use tied queries and keys (i.e., queries = keys), how can queries and keys have different bucket sizes? Also, it is unclear if queries and keys were tied for all experiments using LSH attention in the rest of the paper.\n\n- Which setting regarding the number of attention heads and layers was used for experiments on the GLUE benchmark and LRA? That is, how many layers and heads were set to use LHA?\n\n- In section 3.2, it is said that \"queries should attend to enough keys to get a good approximation of the full-attention.\". Is there a precise notion of what is \"enough\" in this context? That is, for recovering the true softmax distribution, queries should attend to ALL keys. In contrast, if we aim to recover an entmax distribution, queries can attend to only a subset of keys (see the sparse-consistency property in Treviso et al., 2021). So, what would change in LHA if we want to approximate entmax instead of softmax?\n\n\nMinor comments: \n\n- While Equation 16 is easy to follow, there is a big jump in Equation 17. \n- In section 4.3, \"Since h_K(K_j) does change for the queries\". Perhaps \"Since h_K(K_j) does NOT change for the queries\"?\n\n\nRefs:\n\nTreviso, M., Góis, A., Fernandes, P., Fonseca, E., & Martins, A. F. (2021). Predicting Attention Sparsity in Transformers. arXiv preprint arXiv:2109.12188.\n\n\nUpdate: \n\nI thank the authors for their responses and efforts to improve the paper. The newly added experiments alongside the discussion addressed my concerns. I believe this paper provides a new angle to study the efficiency of bucket-based methods, proposing a method that overcomes issues found in previous approaches. Therefore, the paper presents a step forward for this field, and I recommend it to be accepted.", "summary_of_the_paper": "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.\n", "summary_of_the_review": "Overall, I think this paper is well-written, has a clear and practical motivation, and provides an elegant solution to the quadratic bottleneck issue in transformers. Most of my concerns are about the lack of correctness in some parts of the paper, such as choosing hyperparameters or evaluating LSH attention. However, the authors can easily address these concerns in the rebuttal period. Therefore, I am learning towards acceptance. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper adds to the literature of efficient sparse attention for long-range transformer architectures: a learned hash function is proposed by building successfully upon contributions from previous work. A similar idea appears in contemporary work, but with clear and complementary differences.

The reviewers are convinced of the importance of attention complexity, bucket imbalance issues, and agree that learning-to-hash is is a promising solution. The authors have clarified almost all outstanding concerns, in some cases adding valuable new results (e.g. timing experiments.)

I echo the reviewers' concern and stress to the authors to clarify the precise meaning of "plug-and-play", as it may be misinterpreted (e.g., no fine-tuning needed? or just no change to model but still fine-tuning is needed.)

Some of the cited papers are accepted at conferences, please update your bib file with the correct information for credit attribution.
Decision: Accept (Poster)