You Only Sample (Almost) Once: Linear Cost Self-Attention Via Bernoulli Sampling
Authors: Zhanpeng Zeng
Yunyang Xiong
Sathya N. Ravi
Shailesh Acharya
Glenn Fung
Vikas Singh
Venue:  (2021)
Keywords: self-attention
efficient
linear complexity
language model
transformer
BERT

Abstract: Transformer-based models have come to dominate the landscape in a wide range of natural language processing (NLP) applications. The heart of the transformer model is the self-attention mechanism, which captures the interactions of token pairs in the input sequences and consequently, depends quadratically on the input sequence length. It is known that training such models on longer sequences is quite expensive, and often, prohibitively so. We show that a Bernoulli sampling attention mechanism based on Locality Sensitive Hashing (LSH), decreases the quadratic complexity to linear. We bypass the quadratic cost by considering self-attention as a sum of individual tokens associated with Bernoulli random variables that can, in principle, be sampled at once by a single hash (although in practice, this number may be a small constant). This leads to an efficient sampling scheme to estimate self-attention which relies on specific modifications of LSH (based on feasibility of deployment on GPU architectures). We evaluate our proposed algorithm on the GLUE benchmark with standard 512 sequence length and our method achieves comparable or even slightly better performance than a standard pretrained Transformer. To evaluate whether our method can indeed handle longer sequences, we conduct experiments on long sequence (4096) language model pretraining and achieve consistent results as standard self-attention, while observing sizable inference speed-ups and memory savings.

Under review as a conference paper at ICLR 2021
YOU ONLY SAMPLE (ALMOST) ONCE: LINEAR COST SELF-ATTENTION VIA BERNOULLI SAMPLING
Anonymous authors
Paper under double-blind review
ABSTRACT
Transformer-basee models have come to dominate the landscape in a wide range of natural language processing (NLP) applications. The heart of the transformer model is the self-attention mechanism, which captures the interactions of token pairs in the input sequences and consequently, depends quadratically on the input sequence length. It is known that training such models longer sequences is quite expensive, and often, prohibitively so. We show that a Bernoulli sampling attention mechanism based on Locality Sensitive Hashing (LSH), decreases the quadratic complexity to linear. We bypass the quadratic cost by considering self- attention sum of individual tokens associated with Bernoulli random variables that can, in principle, be sampled at once by a single hash (although in practice, this number may be small constant). This leads an efficient sampling scheme to estimate self-attention which relies on specific modifications of LSH (based on feasibility of deployment on GPU architectures). We evaluate our proposed algorithm on the GLUE benchmark with standard 512 sequence length and our method achieves comparable or even slightly better performance than standard pretrained Transformer. To evaluate whether our method can indeed handle longer sequences, we conduct experiments on long sequence (4096) language model pre- training and achieve consistent results as standard self-attention, while observing sizable inference speed-ups and memory savings.
INTRODUCTION
The Transformer model (Vaswani et al. 2017) incredibly effective across diverse set of natu- ral language processing (NLP) applications including machine translation (Vaswani et al., (017)) language inference (Devlin et al. 2018)) and paraphrasing (Raffel et al. 2019). Transformer-base models such BERT (Devlin al. 2018) are pretrained in unsupervised manner and later finetuned on different downstream tasks, often providing state-of-the-an performance on standard benchmarks. While such models have strong empirical performance, their high computational and memory requirements remain quite high. Consequently, in the NLP setting, most current models have certain constraints the sequence length, e.g., BERT and other transformer-based language models (Yang et al., 2019 Liu et al., 2019) limit the sentence length to be at most 512.
The Multi-Head Self-Attentio is central Transformer based models and provides flexible global receptive field to exchange information among input tokens. While self-attention provides immense benefits, is also a key bottleneck in training with long sequences. In particular, the output of self-attention is combination of all tokens where coefficients are determined by the similarities among tokens. While this empirically beneficial, it involves a sizable resource footprint. For sequence length n. this leads to O(n²) complexity in both time and memory to compute pairwise similarities among all input tokens. This quadratic cost is a roadblock in attaining potential benefits that may be realizable in various applications by capturing long term context dependencies. As will discuss in more detail later, the foregoing issue is a major thrust several recent and ongoing efforts focused on mitigating the sizable resource requirements of such models.
Our work is inspired by ideas of importance sampling via bashing-based sampling techniques (Spring & Shrivastava, 2017 Charikar & Siminclakis 2017). We proposed Bemoulli based sam- pling to approximate self attentonn, scaling linearly with the input sequence length. We achieve this by viewing self-artention sum of individual tokens associated with Bernoulli random variables

---

Under review as a conference paper at ICLR 2021
whose success probability is determined by the similarities among tokens. In principle, sample all Bernoulli random variables at once with a single hash (although in practice, this number may be small constant to lower the approximation variance). This leads to efficient sampling scheme to estimate self-attention which relies on specific modifications of hashing-hased importance GPU architectures). The resulting strategy (You
sampling (based feasibility of deployment Only Sample Almost Once, YOSO-Attention) is far more amenable to an efficient and backpropaga- tion friendly implementation, and has a favorable empirical performance profile on natural language modeling tasks. We evaluate our proposed algorithm on the GLUE benchmark (Wang et al.) 2019) with 512 sequence length well as on long sequence language model pretraining where we see promising results with speed-ups and memory savings.
2 BACKGROUND: SELF-ATTENTION
Self-Attention. Self-attention is a scaled dot- dot-product attention mechanism to capture token depen- dencies in the input sequence, which can be defined as,
(KWk) Vdn
A(Q, KV V) softmax
VWr = Dp (P) VWr
(1)
where K, V E Radd are embedding matrices from the input sequence, called queries, key and values respectively. Here, n is the input sequence length, d is the embedding dimension of each token, Wo. WK. Wy E Rdxdd, are learned parameter matrices, d₁, is the dimension of hidden embedding, and Dp is a such that the row entries sum up to 1. For simplicity, we overload the notations for Q, K, V to denote QWo KWK, VWv in our presentation.
diagonal matrix which normalizes each row of the exp (P) matrix
Multi-Head Self-Attention. Multi-Head self-attention in Transformers runs through the scaled dot-product attention multiple times and the attention outputs are concatenated help the model capture information from multiple representation subspaces Vaswani et al. (2017). Multi-Head Self- attention can be formally written as,
MultiHead (Q, K, V) = Concat (Ar (Q.K. (,,
where h is the number heads, 1, are heads with different parameter matrices.
Self-Attention Bottleneck. A key bottleneck in self-attention is computing the softmax matrix, softmax (P), which requires the calculation of all pairwise input token similarities. To reduce this cost, seek to approximate the softmax matrix by viewing self-aftention for each query an expectation of a softmax distribution and computing the approximated self-attention with an effi- cient sampling mechanism. In the following sections, we will first review LSH-based importance sampling and then propose Bernoulli sampling with LSH to estimate self-attention efficiently.
3 IMPORTANCE SAMPLING VIA LOCALITY SENSITIVE HASHING
Importance sampling (Press et al. 2007) helps approximate properties of a target distribution by weighted average of random draws from another distribution. It is known (Press et al. 2007) that importance sampling can be directly used for the softmax distribution by drawing samples from uniform distribution which avoids sampling from the softmax distribution directly which is harder. But this leads to a high variance estimate since the softmax distribution is usually concentrated in small region. When using this idea for softmax matrix approximation for self-attention in particular, the variance tends to grow with the input sequence length. Before proceeding, will summarize an interesting importance sampling method for low variance estimators, specifically, importance sampling via LSH from (Charikar & Siminelakis 2017 Spring & Shrivastava 2017).
LSH-based Importance Sampling. Consider the case when the angular distance between a key and query is small. In this case, the similarity (between the key and the query) as well as the softmax probability will be large. When viewed through the lens nearest neighbor retrieval, the above property coincides with large collision probability of high similarity key-query pairs, assuming that
2

---

Under review as a conference paper at ICLR 2021
the neighbor retrieval is implemented via LSH. Motivated by the link between softmax probability p and LSH collision probability 9. Spring & Shrivastava (0177) and Charikar & Siminelakis (2017)
suggest using LSH as efficient sampler for low variance softmax estimators.
(a) Spring & Shrivastava (2117) propose approximating softmax by sampling a set, S, a collection of neighboring keys for each query formed by the union of colliding keys using hash tables. The estimator is computed using [S] | Eies plg.ku) olq.k.) on where q is a query vector, k;, v, are key and value vectors in the sampling set S, and p(:,-) and q(::-) are softmax probability and collision probability of given pairs. The procedure is equivalent to performing importance sampling without replacement, which involves a dependency among the samples. Deduplication (avoiding double counting) requires memory to store keys in each hash table and runtime to deduplicate keys for each query. If the size of hash buckets is skewed, the GPU memory needs depend on the size of the hash bucket and the runtime depends on the size of S.
(b) Charikar & Siminelakis (2017) proposed a Hash based Estimator to simulate a proposal distri- bution for importance sampling via LSH, which can be easily applied in the context of softmax. For each hash table, a key is uniformly selected from the bucket that the query is hashed to, for simu- lating draw from a proposal distribution. The estimate computed as m E== p(q.k. (.. v; 9lq.k.)
where [g)( denotes the size of hash bucket in the i-th hash table which is hashed to. This sim- ulates m samples drawn with replacement from the proposal distribution. However, the probability of one key being sampled depends not only on (a) the angular distance to the query but also (b) the number of keys within the hash bucket, leading to sampling dependency among all keys. Further, using it for self-attention causes a dependence between the sparsity in the softmax matrix and the number of hashes used. Specifically, the number of tokens that each query can attend to is bounded by the number of hashes; the procedure samples at most one key for each hash table and so, it adds one additional nonzero to the softmax matrix, at most.
Remark I. While LSH-based importance sampling exploits the agreement between high probability p(:--) and high collision probability q(- -). the alignment is not perfect. Samples from proposal distribution must be reweighted to compensate for the difference. Further, for different queries, the likelihood ratios between softmax distribution and proposal distribution w.r.t. single key different. Therefore, the reweighing has to be done during querying. Although maintaining hash tables for storing keys is not major problem in general, the high memory cost for hash tables and computation time for reweighing would influence efficiency when applied to self-attention.
4 YOSO-ATTENTION
We start from LSH-based importance sampling and seek to address some of the aforementioned issues when it is deployed for approximating self-attention. Instead of using LSH to simulate sam- pling from proposal distribution over tokens, we view attention Bernoulli random variables. This modification relates better to LSH and less with ASSe- im- portance sampling the probability of one query colliding with a key is not based other keys. This strategy helps avoid the sampling dependency issue in LSH-based importance sampling and offers an opportunity develop strategy more amenable GPUs.
sum of tokens associated with
We assume that the input keys and queries of self self-attention are unit length to unify dot-product similarity in self-attention and cosine similarity in LSH. This is simple using Neyshabur & Srebro (2015): a temperature variable J is used to bound the squared &, norm of all queries and keys and to reconstruct new unit length keys and queries while preserving their pairwise similarities. We can work with the softmax matrix in angular distance metric and derive our algorithm.
Remark 2.
Self-Attention via Bernoulli Sampling. We aim to approximate self-attention, which uses a soft- max matrix to capture the context dependency among tokens via their pairwise similarities. As- suming that we can represent this context dependency directly using collision probability g(:s),, then the challenges discussed in importance sampling can be resolved. The coincidence of soft- max probability :):.) and LSH collision probability .(---) makes q(::)) a sensible starting point for approximating self-attention. Specifically, model dependency based on similarity, the col- lision probability aligns well with the exponential function in sofumax in the domain of inter- est [-1, Figure| both functions have positive zeroth, first and second order derivatives.
3

---

Under review as a conference paper at ICLR 2021
Note that (a) positive zeroth order derivative indicates that the depen- dency is positive, (b) positive first der derivative ensures that the depen- dency based similarity is mono- tonic, and (c) positive second order derivative means that low similar-
exporenti
[2] wetooee exponentill hattre collzinn
attention weight.
0.8
collaiion probuba May
kwer bound of
0.6
0.4
-1.0
-0.5
0.0
0.5
3.0
-1.0
-0.5
0.0
0.5
1.0
ity corresponds to almost no depen- dency. This leads us hypothesize that a collision-based self-attention may be as effective as softmax-based self-attention. It can be formulated
alignment score
aligmmen score
Figure
compare attention weights using exp(r(2 1)) with the collision probability of concatenating T hyperplane hashes
Charikar 2002 (1 arccos(z)/nt for = 8. We plot exp(7(x I]] that the range is between 0 and but without
changing the actual attention weights softmax. We also plot the derivative of exponential function and collision probability, as well as a lower hound we will use later during backpropagation. Our method can be viewed as using LSH collision probability to estimate biased approximation of exponential function.
as,
Σ B.(q,k;
(3)
where (g k;) is a Bernoulli ran- dom variable where the success probability is given by the collision probability of q with the keys kis Hence, it can be determined by the similarity between q, k; In a single hash, each (g k;) generates a realization to determine whether the corresponding token will be part of attention output or not. Conceptually, when sampling from softmax distribution, only one token is sampled as the attention output. In contrast, Bernoulli sampling determines whether each individual token is a part of the attention output. In principle, to determine the context dependency among tokens, you only need to sample once (YOSO) using a single hash to generate realizations of all Bernoulli random variables, (q k = 1, ....., n. Specifically, when keys are hashed to hash table using a sin- gle hash, the realization of (2 k;) for each query q will be 1 if q collides with k; otherwise will be O. To our knowledge, using LSH collision probability to replace softmax dependencies for self-attention has not been studied before.
YOSO-Attention. By replacing softmax dependency with Bernoulli random variables and using LSH as an efficient sampler to estimate the success probability, we achieve an efficient self-ateention (YOSO-Attention) approximate softmax-base self-attention.
arrcos (QK)Q
YOSO(Q, K, V) B(Q, K)V; E[[OSO(Q, K, V)
V
(4)
π
where BBQ, is the Bernoulli sampling matrix using m hashes.
1 m
B(Q, K)+,
fA(KK)))((K)))))) where fk, k ===1,,. L.... are hash functions.
(5)
Normalizing Attention. In standard self-attention, each row of the softmax matrix is normalized so that the dependencies sum up 1. the above, we have discussed how the pairwise query-key dependencies can be estimated using Bernoulli sampling. We now present how to normalize the dependency in our method as standard self-attention. We can first estimate the dependencies and then normalize them using the sum of estimated dependencies estimated by B(Q, K)1 where is vector of all entries being B Q, K) K)1 can be computed by Eq. 4 by plugging 1 into V. To make the estimation of self-attention more efficient,
turn adopt l normalization to the attention output, similar as Levy et al. (2015) to use f2 normalization for word embedding. Thus, attention outputs invariant of the scaling, BiQ, K)1, under & normalization. Therefore, we have,
N-YOSO K, V) (B( K)V)
(6)
Vidua
Cosda
Quer
Daigon
X
0
My
Figure 2: Overview of The hash table stores the sum values associated with hashed keys.
4

---

Under review as a conference paper at ICLR 2021
Empirically, we show the &g normalization does not affect the performance of method as pected, which can be seen in Figure 3
LSH-based Bernoulli Sampling. Now discuss how to implement the procedure of using Bernoulli sampling to approximate self-attention. While a standard LSH procedure can be used, maintaining hash tables store keys is inefficient hash table cannot be predetermined and the workload might be skewed due to skewed bucket sizes. To tackle this issue, we propose LSH-based Bemoulli Sampling by only saving the summation of values corresponding to hashed keys instead of storing a collection of hashed keys.
GPU the GPU memory size required for
The overview of our algorithm is shown in Figure 2 To compute Y = K)V, the procedure proceeds as follows. For each & E [l...., ml,, we sample hash function f& and create a hash table H E R² xd representing 2 d-dimensional buckets. For each key K Kj we add the value Vj, to the bucket whose index is hash code f& j;;), denoted as H (K(,).))
H (k,,.) + H (..K,.)) + V₁::
(7)
Note that the size of H is (2 and independent of which bucket keys are hashed. With all keys processed for k E [L...., m.) for each query Qi... maintain output vector Y,, ini- tialized to O. Then, we allocate the bucket in H using fu (u,) for k E [I,...,m] and add all corresponding results in buckets to the output vector Y., as
Y, Y., H (c.))
(8)
Therefore, each final output y. can be computed as,
Y
I. k=13=1
V Σ B(Q.K).,VV
Remark 3. The memory and time complexity of this algorithm are O(m2" d) and ((mmd) respec- tively, In addition, both time and memory independent of the size of hash buckets. Further, We can improve the memory complexity to (mm2)) by reusing hash table and processing few dimen- sions each time without increasing time complexity. The constant is small as it controls the decay rate of attention weight with respect to the angular distance between query and key, and it can be chosen be a function of (n). In our experiments, 7 is set (n.).
Speed-up. While not essential, we find that a fast random projection for computing the LSH hash code will be beneficial, since this step takes a large portion of the overall runtime. As suggested by Andoni et al. (2015), we use the approximated random projection to reduce time complexity to O (nmm (d)), allowing fast computation of hash codes.
Backpropagation through YOSO-Attention For training, we also need to show backward propa- gation steps for YOSO-Attention Here, we discuss this last component of YOSO-Attention which enables end-to-end and efficient training.
For backpropagation, the gradient of the loss L w.r.t V can be estimated similar to equation
VVL = ((1 (11 arccos QKT) Y ( (YoooL)) se B(K, (V Yoso
(10)
π
The gradients of w.r.t. Q. are similar, so we only provide the expression for Q,
VQl.= (VVroooo) @ (7(1- arccos(QQ' ) e n (T
K (I1)
π
where 2, O are element- division and multiplication. The problem with the true gradient is that it goes to infimity as the alignment score between the query and the key approaches 1, which might lead to divergence. To avoid this numerical issue, we use lower bound of the actual derivative of the collision probability, VoVoooL)VVT O 3
arcras(QK")) Y"],, see Figure which can be efficiently estimated via a variation of LSH-based Bernoulli Sampling. Specifically, note that the approximation can be decomposed into of LSH-based Bernoulli Sampling.
(VoL), Σ Liosoo i. Σ B(Q, K) K), (V,
(12)
5

---

Under review as a conference paper at ICLR 2021
Therefore, following LSH-based Bernoulli Sampling, the memory complexity is O(m27d and time complexity is O(nmd). The d' term can be climinated by repeatedly using the same hash tables df times without increasing runtime, which improves the memory complexity to O(m27). The overall complexity of our method and comparison to standard self-attention is summarized in Table | Further, to address the quadratic dependence d, in the Appendix, will discuss scheme to estimate the same quantity but is linear in d.
5 RELATED WORKS
There are number of efforts describing ways to reduce the quadratic cost of self attention w.r.. input sequence length. Among these works, Linformer (Wang et al., (020) suggests that low rank attention might be sufficient and adds linear projections (on the sequence) to fixed size keys and values. There are also other low rank approximation ideas (Katharopoulos et al. 2020), (Choro-) manski et al. 2020) using separable functions on queries and keys replace softmax self-attention By assuming the self-attention rank to be independent of input sequence length, these methods can achieve (n) time and memory complexity. Another direction to exploit the sparsity of softmax matrix and focus on certain sparsity patterns by only computing softmax dependencies within those patterns, including Sparse Transformer (Child et al. 2019), Longformer (Beltagy et al. 2020), and Big Bird (Zaheer et al. 2020) and Reformer (Kitaev et al. 2020). Note that, instead of using LSH as a tool to approximate nearest neighbor search to dynamically determine the sparsity pattern Re- former, our YOSO-attention attention takes advantage of the connection of query-key similarity to the LSH collision probability to model the dependency among tokens.
6 EXPERIMENTS
In this section, we provide the empirical results for the proposed approach. To evaluate our pro- posed method, we follow the BERT language model pretraining procedure (Devlin et al. (0018) and evaluate the performance of method in both intrinsic tasks and multiple downstream tasks in GLUE benchmark as well as runtime and memory relative to standard self attention. Previously, we assumed that queries and keys unit length and described the construction make it work. In the experiments, we found that simply applying a f2 normalization on queries and keys and using temperature
hyperparameter does not degrade the performance of model and yet is more efficient to compute, so we use the simpler version in the experiments.
Following Devlin et al. (0018), the model is pretrained on BookCorpus ZZuu et al. 2013) and English Wikipedia. To evaluate the capacity of model capturing the sentence level information, instead of using Next-Sentence- Sentennce--Peeiiition (NSP) as sentence level loss in the original BERT, we adapt the as a more difficult task compared to NSP. All model are trained with Loak-LLangage--oodaee Modeling (MLM) and SOP objectives. We used the same hyperparameters for pretraining as Devlin et al.) (2018). However, due to the computational resources limit, all models are trained for 500K steps. The batch size is set so that around 2 tokens are processed per step. (batch size of 256 for sequence length 512, and batch size of 32 for sequence length 4096).
BERT Pretraining.
Orderring-reedicttonn (SOP) proposed in ALBERT (Lan et al., 2019)
Number of Hashes during Pretraining. Since the estimation variance decreases as the number of hashes increases, to evaluate the trade-off between efficiency and performance in YOSO, we test on four hash settings: 16 hashes, 32 hashes, 64 hashes, and expectation of collision to simulate infinite hashes. We plot MLM validation perplexity and SOP validation loss curves of 512 length model pretrained with softmax self-attention and YOto-Attention in the right plots of Figure 3] The curves of our method using expectation match and slightly exceed softmax self-attention, indicating our method is indeed as capable as self-attention. is expected that the number of hashes increase, the performance of our method will approach the curve using expectation as the approximation
Memory
Forward
Backward
Forward
Backward
Self Attentio
O(n
(nndd)
O(n")
O(n²)
YOSO-Attctton
Opnor log, (d) mmd)
O(ama")
O(mmr + m2
O(m2)) )
Table 1: Time/memory complexity of self-attention and YOSO-attention in forward/backward computation

---

Under review as a conference paper at ICLR 2021
Validation Parplesity
Vo5D-14
Ge Mettent
Validation Perplexity
3.4
Gef Votention
vosD- 16
MUS4 Valldation Perplexity
11
YOSO 32
Sett Amhentton
50-
voss 32
5.4
vosp 32
10
YOS3 B4
Yose 64
Vosp
YO50
5.2
VOSO
8-
1.0
7-
1.8
200K
B00K
400K
500K
256 number hashes
64
312
4.
20K
90K
100K
Se Attencc on
0.180
Solf Ansuition
Self Amunioin
SOP Validation LoGs
3.4
VOSO is
0.365
V050
Velidetion Less
0.4
1050 92
VOSO- 32
VOSO
YO6O 64
0.3-
VOSD- 64
VOSS-E
0.35
VOGG- 64
VOSD- E
0.3
A350
0.145
3.
0.142
0.1
3.0
0.0
number
Figure 3: (a) The left two plots are results on MLM and SOP for 512 sequence length. We report MLM validation perplexity and SOP validation loss for each 2K training steps. (b) The middle two plots are results for MLM and SOP when using different number of hashes validation. Since the runtime of YOSO-Attention linear with respect the number of hashes, these two plot directly reflect the equivalent relation hetween performance vs inference time. (c) The right two plats are results on on ML. and SOP for 4096 sequence length. YOSO-x means the model is pretrained with YOSO-Attention using x hashes with E heing expectation.
become more accurate. For both MLM and SOP, we confirm that our method is as effective as softmax self-attention.
Number of Hashes during Validation. YOsO-Attention a stochastic model. To make the in-
ference deterministic, in dropout (Srivastava et al. 2014). we take the expectation as output. However, directly computing expectation involves a O cost, so we experiment with the effect of different hash settings in validation and simulate expectation as the number of hashes increases. We plot the MLM perplexity and SOP loss of the same pretrained models using different number of hashes on validation in the center plots of Figure 3 We observe that as the number of hash increases, the MLM perplexity and SOP loss generally decreases for all pretraining hash settings.
Pretraining on Longer Sequence. To examine whether our method can scale linearly with se- quence length, we continue to pretrain BERT-base models using the corresponding 500K step check- points for 512 length model, and add additional positional embedding as suggested in |Beltayy et al. 2020 We observe that compared to 512 sequence length, the small performance gap between YOSO-Attention and softmax eeffaatteniion does not increase as suggested in the left plots of Figure providing evidence that the number of hashes can be chosen independent of sequence length.
GLUE Tasks. In addition to in- trinsic tasks, we examined the effec- tiveness of our method on diverse downstream tasks and ask how our method compares with standard at- tention even after finetuning. We finetuned all pretrained BERT-base model on MRPC (Dolan & Brock- et 2005). RTE (Giampiccolo et al.
Method
MRFC
SST-2
QNLI
MNLI-m/mm
Self-Attention
88.3
70.8
91.1
90.3
87.3
82.4/82.4
YOSO 16
87.1
68.6
90.7
88.5
85.3
79.6/79.5
YOSO-32
87.3
71.8
90.9
89.0
86.3
80.5/80.7
YOSO-64
88.1
69.7
91.5
89.5
87.0
81.6/81
YOSO-E
88.1
90.1
82.2/82.9
YOSO-16-E
87.8
69.3
91.5
89.3
86.8
81 1L.0/81.4
YOSO-32-E
87.8
70.4
91.1
90.1
86.8
80.8/8144
YOSO-64-E
88.3
72.2
91.7
90.0
87.2
81. 11.8/82.8
Table 2: Dev set results on GLUE tasks. We report Fl score for MRPC and QQP and accuracy for others. YOSO-x means the same as in Figure YOSO-x-E means that YOSO-x is finetuned on downstream tasks using expectation.
2007) SST-2 (Socher et al. 2013).
QNLI (Rajpurkar et al. 2016). QQP
(Cen et al. 2018). and MNLI
Williams et al. 2018) tasks in the GLUE benchmarks and report their corresponding dev metrics. For large datasets including QNLI, QQP, and MMNL, due to extensive resource needs, we cannot do hyperparameter search, so we used a batch size of 32 and learning rate 3e-5 to update our model and finetune our models for 4 epochs. For MRPC, RTE, and SST-2, we follow BERT finetuning to do hyperparameter search with candidate batch size {8, 16, 32) and learning rate {2e-5, 3e-5, dc-5,5e-5 and select the best dev set result. Results are listed in Table We observed that YOSO's performance on downstream tasks is comparable with standard attention, and even has slightly better
7

---

Under review as a conference paper at ICLR 2021
results in some hash settings. Further, the downstream performance of YOSO generally increases as more hashes are used, providing an adjustable trade-off between efficiency and accuracy.
Longer Sequence Task. To further evaluate YOSO long sequence tasks, extended the po- sitional embeddings of trained YOSO-64 model and used it as an initialization to train a 4096 length YOSO-128 model using a batch size of 64 and learning rate 5e-5 on BookCorpus (Zhu al., 2013). English Wikipedia, one third of the Stories (Trinh & Le 2018), and one third of Realnews (Zellers al. 2019) for 100K steps, similar Longformer pretraining (Beltagy et al.) 2020). Then, we finetuned our model on WikiHop (Welbl et al. 2018). Due to the computational resource limits, we only tested a small set of hyperparameters (batch size 32, learning rate {le-5, 2c-5, 4e-5), number of epochs 10). The dev accuracy is 73.7 for YOSO- 1-8-1128-E,, which is comparable to 73.8 Longformer-512 (see caption in Table 3 without hyperparameter search but slightly worse than 75.0 that Longformer-511 achieves with hyperparameter search.
Comparisons Baselines. Apart from compar- ing YOSO to standard self-attention, we also evalu- ated its competitiveness with other efficient attention methods. To keep the financial costs of these exper- iments reasonable, instead of training all methods from scratch, used RoBERTa-base pretrained weights as the starting point and trained each model using batch size 512 and learning rate 5e-5 Book- Corpus (Zhu et al. 2015) and English Wikipedia for 95K steps. Then, we finetuned the models SST-2, QQP, and MNLI. These results are shown in Table We observed that our performance competitive with other baselines while the memory consumption of YOSO much less (2.6x, 1.9x, 2.1x mem- ory savings compared to Reformer, Longformer, and Linformer respectively, see Backward-Cacks in Ta- ble 4.. This has potential ramifications for training such models with more moderate hardware resources which are much less expensive. Further, no- tice that YOSO is potentially applicable to wider range of applications, especially where the input sequence represents an unordered set of high dimensional points (where spatial locality of the input sequence may hold).
Method
SST-2
QQP
MNLI-m/mm
YOSO.E 32
93.5
87.3/90.5
84.4/84.1
YOSO-64
$7.9/90.9
Reformer-2
92.9
87.8/91.0
85.6/85.5
Longtommer-64
94.9
S8.49144
85.3/85.2
Table 3:
Dev results SST-2, QQP, and MNLL. We report both F1 score and accuracy QQP and accuracy for others. Reformer-x: Re- former using HuggingFace implementation (Wolf et al. 2019) using hashes. Longformer-x: Long- former using HuggingFace implementation (Woll
al. 2019 using sliding window size X.
YOSO X is similar to description in Figure [ We did not re-train Linformer to keep compute costs reasonable; when trained from scratch Linformer Wang al. 2020) reports accuracy 93.4 for SST- and 90.8 for QQP.
0.4
D4
BA
B4
64
64
02
Seiieitetionn
6.0
YOS3-8
YO50-18
86
Y062-12
Figure 4: Attention matrices generated by self-attention and YOSO-Attention with different hash settings using the same input. Notice that the patterns are preserved well.
Estimation Error. To assess the effectiveness of our algorithm, using Q, K from the trained model, we generated attention matrices using our algorithm with different number of hashes and compare against standard self-attention. In Figure 4 visually,
see that method produces similar atten- tion patterns as standard self-attention. The estimation of attention matrix becomes more accurate as the number of hashes increases. Further, each output of YOSO-Attention is a weighted sum of random variables as shown in equation B so one may suspect that as the sequence length increases, output might potentially increase. We did not observe this behavior
the variance of YOSO-Attentic which may be partly due to the hyperparameter 7 ((log(n)) that controls the decay rate of LSH collision probability the similarity changes. We can also ask whether not the estimation error of YOSO-Attention for fixed number of hashes increases as the sequence length increases. We use Q,K, K. V generated by the pretrained model and estimate the error between N-SOOO K, V) and EEEN-NOOOOO YOSO(Q) Q, K, V)). As the left plot of Figure 5 suggests, the relative error of our method stays almost constant as the sequence length increases from 128 to 4096. This indicates that using sampling to estimate attention weight based on YOSO-Attention can scale up with sequence length and preserve the same estimation quality without increasing the number of hashes.
8

---

Under review as a conference paper at ICLR 2021
Runtime and Memory. We mea- sure the runtime of our method as sequence length increases. To show the trend more precisely, we mea- sured the runtime per token as shown in Figure 5 (right). There is a slight increase in runtime per token as the sequence length increases, but note that the r-axis of the plot is log scale, so the increment is small. When the sequence length increases by 32x, the runtime per token only increases by 30%, which is explained by our choice of hyperaarameter r (log(n)). Aside from the plot, we report the training and testing ef- ficiency of our method as well as three other efficient attention meth-
0.
Relative Error
Rundimm per Tokan
G6-
27
24
314 203
3/2
20
354
20
Sequence Length
Sequence Length
Figure 5: (a) The relative error the left plot is defined as
[
[EIN YOSO V. N- Q. Villas j. The relative error is es- [[EN K.
timated by computing BIN- MON-YOSO(Q, K, V)j based on collision probability, then estimating N- N-YOSO(Q,) K, V) multiple times, finally computing the mean of relative error of each estimate as an estimation of the outer expectation. estimated by estimating N- N-YOSO(Q: K, V) multiple times and measuring the total time elapsed and then dividing the total time by number of iterations and sequence length to get runtime per token.
(b) The runtime per token
ods against standard self-attention. The results were measured using Q. K, V of specified se- quence length generated by a trained model and fed into BERT- base Multi-Head Attention module multiple times. The experiments were performed on a single NVIDIA 2080TI. From Table we can see that while for a standard 512 length sequence, our method has similar runtime as self-attention, as the sequence length increases, the speed-up and memory savings become significant. While our method offers similar runtime savings as other efficient attention methods, the memory consumption for training (i.e., Backward-Cache of our method is much lower than all other methods in almost all settings.
Sequence Length
Model
1024
2048
1024
2048
Forward/ / Barkward Time (ms)
Backward Cache 'Training Peak
''lsstigg Peak Memary (MB)
Seli-Allemiie
0.60088
1.872.29
5.978..04
21.2729.73
517 70748
1797 7777601175
6717
659
25897 391872565
YOSO-16
1009
2.793.73
4.07 7,35
527
34.
1277 110
1367
YOSO-32
0.71/1.21
1.4922.8
3:11 /4.88
/.31
18/44/40
36/89/880
71/178/ 161
142/ 355 345
YOSO-64
19971.91 1.91
2.26/3.9
4.77/7884
10.21/ 16.02
19/ 166
39 1440/1131
/ 263
154 / 5559 /549
Y08O-128
19112
4,0916.45
9.36/
19.41/2833
22/
45 45/242/10
483/
178/ 967
Litoomer- 128
0427031
0.367 1.03
1.7272.16
3.33/ 4.76
23733717
60766733
007716771 71
1667217
Linfommer-256
0,66/
1.36
2.16/2288
6.02/6.16
38/55 24
76/111/49
136/233/ 100
239/ //0//1/// /22
Reformas-2
0.73/1.14
1.5272.29
6.3379.19
48/61 73
967 76
1927 243 151
5841 4487
Retormas4
1,22/205
14.1/4.1
18.19
10.49/ 16.36
114/71
/
688 9L1 1572
Reformes
2,17/3.85
4.66
9.04 71.4
18.84/ 30.3
162/ 220/ 139
32411 278
647/878/ /55
1294 1756 /111
Longlormer
0.171.71
1.373.37
76.87
6,197158
35743725
71736735
1427 172799
2217 307 7779
Longformer-11 128
0.9/207
/4.12
18.41
41/48//32
193/
345
Longfoemer 256
1.09/2.74
2.1915.56
4.45 10089
8.47 22.91
51/64/ /46
104/ 129 91
208 /58///88/ 182
350 481 7363
Longformer 512
1.42/3.95
2.96/.65
6.05 17.86
11.54/ 36.14
72/ 100 100//00
148/211/150
299 / /317
538/ / 7667 585
Table 4: The reported values pertain single instance. Time is estimated by averaging total runtime and then dividing by hatch size, while memory measured hy dividing total memory consumption by batch size. Linformer-x: Linformer model where the sequence is projected to a fixed length. The rest of method names are consistent with Table 3 Backward- Cachhe amount of memory required to cache certain tensors for back- ward pass (e.g.. activation and inputs which is required for gradient calculation or is cached to avoid redundant calculation). Since the Backward-Cache will accumulate as model grows deeper, this part is usually the most significant portion of memory consumption. Training-Peak: maximum memory needed during training includ- ing temporary tensors and cached tensors for backward pass. Some tensors are not required for backward pass, and they are only held in memory temporarily. As model grows deeper, the amount of memory for temporary tensors does not grow. Testing- Testing-Peak: Similar to Seaining-Paak, but nothing is cached for backward pass.
7 CONCLUSION
We presented transformer- model, YOSO-Attention, that scales linearly in the number of in- put tokens. This allows the model to be applicable to a wide range of long document NLP tasks. Via a randomized sampling based scheme, our model approximates self-attention tokens associated with Bernoulli random variables that can be sampled at once by a single hash, in principle. With specific modifications of LSH, YOSO-Attention can be efficiently deployed within a deep learning framework and various aspects of this idea and our implementation, we expect, will find use in other novel settings and applications (e.g., in vision).
sum of individual

---

Under review as a conference paper at ICLR 2021
REFERENCES
Alexandr Andoni, Piotr Indyk, Thijs Laarhoven, Ilya Razenshteyn, and Ludwig Schmidt. Practical and optimal Ish for angular distance. In Advances in neural information processing systems, pp. 1225-1233, 2015.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020.
Moses Charikar and Paris Siminelakis. sions. In 2017 IEEE 58th Annual Symposium on Foundations of Computer Science (FOCS), pp. 1032-1043. IEEE, 2017.
-esttddstttttssross for kernel density in high dimen-
Moses S Charikar. Similarity estimation techniques from rounding algorithms. In Proceedings of the thiry-fourth annual ACM symposium on Theory of computing, pp. 380-388, 2002.
Zihan Chen, Hongbo Zhang, Xiaoji Zhang, and Leqi Zhao. Quora question pairs, 2018.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. ArXiv, abs/1904. 10509, 2019.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Davis, Tamás Sarlós, D. Belanger, Lucy J. Colwell, and Adrian Weller. Masked language modeling for pro- teins via linearly scalable long-context transformers. ArXiv, 00420000.0355,, 2020.
Jacob Devlin, Ming- Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
William Dolan and Chris Brockett. Automatically constructing a corpus sentential paraphrases. In Proceedings of the International Workshop on Paraphrasing, 2005.
Danilo Giampiccolo, Bernardo Magnini, Ido Dagan, and Bill Dolan. The third PASCAL recognizin textual entailment challenge. In Proceedings of the ACL-PASCAL workshop on texiual entailment and paraphrasing, pp. 1-9. Association for Computational Linguistics, 2007.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. arXiv preprint arXiv:2006.1622 16236. 2020.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. ArXiv, 15s/20011.04511, 2020.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Sori- cut. Albert: A lite bert for self- learning of language representations. arXiv preprint arXiv:1909. 1942, 2019.
Omer Levy, Yoav Goldberg, and Ido Dagan. Improving distributional similarity with lessons learned from word embeddings. Transactions of the Association for Computational Linguistics, 3:211- 225, 2015, doi: 10.1162/taclaa.000 a00134. URL
/wwww
Q15 1016
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 arXiv: 1907. 11692, 2019.
Behnam Neyshabur and Nathan Srebro. On symmetric and asymmetric Ishs for inner product search. In International Conference on Machine Learning, pp. 1926-1934, 2015.
William H Press, Saul A Teukolsky, William T Vetterling, and Brian P Flannery. Numerical recipes 3rd edition: The art of scientific computing. Cambridge university press, 2007.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint arXiv: 1910. 10683, 2019.
10

---

Under review as a conference paper at ICLR 2021
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pp. 2383 23922299,, 2016.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality sentiment treebank. In Proceedings of EMNLP, pp. 1631-1644 2013.
Ryan Spring and Anshumali Shrivastava. plers and estimators for partition function computation in log-linear models. arXiv preprint arXiv:1703.05166 2017.
unbiased and efficient class of Ish-based sam-
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way prevent neural networks from overfitting. The journal of machine learning research, 15(1):1929-1958, 2014.
Trieu H Trinh and Quoc Le. A simple method for commonsense reasoning. arXiv preprint Ntl::880000287,, 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
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. 2019. In the Proceedings of ICLR.
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. ArXiv, (bss/200.00468,, 2020.
Johannes Welbl, Pontus Stenetorp, and Sebastian Riedel. Constructing datasets for multi-hop read- ing comprehension across documents. Transactions of the Association for Computational Lin- guistics, 6:287-302, 2018.
Adina Williams, Nikita Nangia, and Samuel R. Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of NAACL-HI.T, 2018.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Huggingface's transformers: State-of-the-ar natural language processing. ArXiv, pp. arXiv-1910, 2019.
Zhilin Yang, Zihang Dai, Yiming Yang. Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. XInet: Generalized autoregressive pretraining for language understanding. Advances in neural information processing systems, pp. 5753-5763, 2019.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in Neural Information Processing Systems, 33, 2020.
Rowan Zellers, Ari Holtzman, Hannah Rashkin, Yonatan Bisk, Ali Farhadi, Franziska Roesner, and Yejin Choi. Defending against neural fake news. Advances in Neural Information Processing Systems, pp. 9054-9065, 2019.
Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. Proceedings of the IEEE international conference on computer vision, pp. 19 19-27, 2015.
11

---

Under review as a conference paper at ICLR 2021
A APPENDIX
In Appendix, we provide some details of our method that are left out in the main text.
Backpropogation Derivation When using expectation of LSH collision attention weights, the attention of one query to keys and associated values o² for all E (1.... n) is defined as
arccos( (²' k²)
(13)
then we want to compute the gradient loss w.r.t. q. which we denoted as V L. with the gradient of loss w.r.t. y denoted, V VyL, given. We start by computing the p-th entry of V LLL:
OL.
aL dyj
OL a Jy
arccos (9" kt
(14)
pqp
By dqp
dqp
Then use
d du
arcoos(2) 20)
T(1- srcentt))) --
and plug it into the Eq.
13
zvI-xx
arcoos( ( Ae
T-1
aL
àL
(15)
0q
dy
(q" kijz
p
After swapping the order of two summations, Eq.
13 becomes
arcccsi (q Ak')
AT.
(16)
0qp
TV (q" k2)
p
Note that only k is different for different entries of V so we can write it as
(1 arccos(q "9- =k')
r-1
VGL Σ (VVL)
17
k²
(17)
o TV/ (q k:)
Equation II is the matrix form of above
VoL= (VrosoL) @ (F(1- arccos QK?
i (TVI-(Q K
(18)
π
Note that π #V1 - (QKTY approaches to as alignment score between the query and the key - for x [-1, and define lower
approaches 1, so we use the fact that [1 arccos(x)
bound replace the actual gradient
V = ((V (V)) ①
[111 arccos QKT)) 1)) K
(19)
Approximating Random Projection in LSH. In the main text, we discussed how to estimate self- attention using Bernoulli sampling via LSH. The first step of using LSH is computing hash code using random projection. compute hash codes for a vector æ, we proceed as follows.
F: _d
1)
F(x) sign(Ræ))
(20)
where R x Rij (O. 1), then the output vector partition 7-dimensional binary hash code. The time complexity for random project is Ommrrd). To efficiently approximate random projection, we follow the construction used in Andoni et al. (2015). The output of mr-dimensional vector is divided to mi d'
d-dimensional vectors, then hash codes are estimated by
F
sign( HD HD,
HD x))
(21)
where D are diagonal matrices with entries uniformly sampled from [-1, and H is Hadamard matrix. This approximation reduce time complexity O(nmr (og)
12

---

Under review as a conference paper at ICLR 2021
Alternative Procedure for Approximating Backpropagation. In the main text, provided procedure as shown in Eq. [2] which use LSH-based Bernoulli Sampling times as subroutine. The complexity of this procedure is linear w.r.t sequence length which is desirable but the runtime can be large if d is relatively large. Therefore, we provide second procedure, which is linear with respect d. The gradient of L w.r.t. the i-th row of Q is written
T (VoL), = ((Vroso B(Q K)J K ==1
(22)
Note that if B(Q,K) k. is zero then the corresponding summation term does not need be computed. The alternative procedure counts the number of success in m samples at each entry B(Q, K),, and only computes the summation term when B(Q, K) non-zero, and thus the runtime is O(nnz(S(A))) Bm) + d)} (counting number of success computing nonzero terms). In the worst case, nnz (B(. K)) n², it would be as expensive as dense matrix multiplications in complexity and even worst in practice due large memory latency resulted from indirect memory access. However, in practice, BQQ. K) is generally sparse procedure guarantees linear complexity scaling of our method for extremely long sequences. As an improvement, we can dynamically select one from these two method based on runtime, than the time complexity is O (nmnn nn (m) d))).
set properly. Further, the first
13


Reviews:
[{"review_id": "eC-sAowKE_", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary\n\nThe paper proposes to replace the weighted average of the values in standard self-attention with the average of values sampled in a way that the expectation is close to the result of self-attention. In particular, the authors associate with each query-key pair, a Bernoulli random variable with expected value close to the exponential of the dot-product. Sampling these variables and averaging the values per query is formulated in an efficient way using locality-sensitive hashing.\n\n### Strengths\n\n- Using sampling to approximate self-attention is novel and promising.\n- The LSH formulation where the values are averaged in the bucket is a clever way to avoid the pairwise interactions between queries and keys.\n- Training with hashing and evaluating using the expectation is interesting and provides evidence for the approximation quality of YOSO.\n\n### Weaknesses\n\n1. In the evaluation there is never an explicit comparison with respect to both time and performance. I appreciate that given a large enough sequence length YOSO will always be faster but will it be good enough?\n\n2. One of the most important parts of the methodology, the gradient computation, is the least clearly written. For instance, equation 11 contains $\\nabla_{Attn}$ which is never defined and subsequent equations contain $\\nabla_{YOSO}$ which seems to contradict equation 11. Moreover, how is equation 11 derived? Is it the gradient of the expectation?\n\n3. In table 3, the performance is only measured with respect to inference. Given that the most computationally intensive part of the method is the backward pass, a comparison with respect to wall-clock time per epoch, as well as total training time would be very informative.\n\n### Reasons for recommendation\n\nI find the idea very elegant and interesting however, the experimental section is somewhat lacking. There is no clear evaluation of the trade-off between speed and performance. The MLM task, although significant and demanding, contains sequences of small length, otherwise why not show a graph of performance vs inference-time.", "extra_scores": {}, "extra_text": {}}, {"review_id": "VTl7JKi4-75", "reviewer": "AnonReviewer2", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper tries to improve the efficiency of (multi-head) self-attention by reducing the computational complexity from a quadratic one to a linear one. The authors propose to use Bernoulli sampling to approximate the self-attention's softmax distribution through importance sampling via LSH, which makes a linear cost self-attention possible. \n\nPros:\n1. This paper provides a thorough solution of how to accelerate self-attention via an approximation by Bernoulli sampling and LSH. \n2. The experiments show that the proposed approach could achieve considerable speedup while preserving model performance. \n\nCons:\n1. The authors should give a more direct comparison with an important related work the reformer, which also uses LSH to speedup the computation of attention. \n2. The experiments were mainly conducted on MLM on GLUE, which lacks generalization among tasks. Adding more tasks such as MT or autoregressive/causal LM would make the experimental part more solid and convincing. \n\n\n\n---------\nMinors:\n- Fig 2, 3 and Tab 1 are not cross-refed in the main body\n- Format of citation: seems all of the citations are of this format - authors (year), which is not correct when citations do not act as subjective of objective in the sentence. Please check the format guideline.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "GdNUZpUxH9L", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 2, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper presents a linear-time attention model based on importance sampling and locality sensitive hashing. The idea is to use Bernoulli sampling to approximate the self-attention - which is quadratic in complexity. \n\nHonestly, this is a very crowded space and already many models exist (https://arxiv.org/abs/2009.06732). The authors are aware of these works, cite them and yet there is no comparison. \n\nThis method seems to be rooted in LSH and a very natural question is how does compare to Reformers. There is not even a sparse transformer or local attention baseline in the experiments. This raises questions about whether this paper will even make any impact at all. (comparisons with longformer is done only on speed/memory but not qualitatively). why? \n\nI also find other flaws with the model. If sampling is used, this essentially makes the model stochastic (correct me if Im wrong here). but there are undesirable properties of this such as having non-deterministic inference.\n\nAnother flaw is that method potentially introduces a lot of instability in training. I think the authors could comment a little on this. Transformers are already notoriously difficult to train and I figure that this method would probably make it way harder for practitioners to get the hyperparameters correct. I think playing with the scaling and normalization of the self-attention weights is something non-ideal, and unless the authors can show this is reasonable stable i am not convinced. \n\nI also find it difficult to understand the choices of tasks. It seems like GLUE benchmark is used, yet most of the tasks (like SST) are relatively shorter sequences. I think the authors need to explore datasets that showcase the model's ability on longer sequences. Artificially raising sequence len during pretraining is not really sufficient to be convincing that the model is doing something useful for longer sequences (since the masked out tokens really depend on local context). \n\nMy constructive feedback to the authors to improve the paper is to have reasonable baselines for comparison. The datasets are also not appropriate. I would suggest some actually long-range tasks in order to showcase the model's capabilities. \n\nAt the rate of the number of new models that tackle this problem, I suspect it would be wise to wrap up your sleeves and add actual efficient transformer baselines.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "0WJ9QAJTEUT", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This article presents YOSO, an locality sensitive sampling based attention mechanism for large scale language modeling. \n\n\n\nStrength:\nA new idea of applying  locality sensitive sampling to approximate attention matrix in the transformer\n\n\nWeakness:\n1. Comparison of Complexity: [1] presents the complexity of different efficient transformers. For linformer[2], the time and memory complexity is O(nk). Is there any justification of LSH sampling equipped YOSO with complexity more than O(nm\\tau log(d)+nmd)?\n2.Experiments: YOSO takes linformer as baselines. However, the pre-training experiment part does not provide steps vs ppl of linformer with YOSO in Figure 4. What is the comparison result of YOSO with linformer on iteration wise convergence? Also, linformer demonstrates better accuracy in downstream tasks such as SST-2. Is there any comparison to an explanation that can analyze this difference in performance?\n3.Efficiency: YOSO demonstrates an advantage over linformer and longformer in memory and runtime. However, is there any analysis on why YOSO achieves this superiority with higher complexities? Are there any system-level advantages that YOSO can show?\n\nSome discussions: \nReformer[3] design an attention mechanism that computations are held in the neighbor tokens inside the hash buckets. YOSO also uses hash based sampling to compute attention via neighbor tokens that have high collision probability. On the other hand, linformer introduces a more global view for attention by the low rank projection. Is there any analysis of the local vs global intuition?\n\n[1]Efficient Transformers: A Survey https://arxiv.org/pdf/2009.06732.pdf\n\n[2]Linformer: Self-Attention with Linear Complexity https://arxiv.org/abs/2006.04768\n\n[3] Reformer: The Efficient Transformer https://arxiv.org/abs/2001.04451", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper presents an interesting idea for making self-attention efficient. Several reviewers were not satisfied with the experiments because it did not include runtime and sought after benchmarks.  Rebuttal did a good job of clarifying a few of those with newly added experiments that make the paper stronger. However, the new experiments are in limited settings as well as the real advantage over LSH baselines require more investigation.   This could make a good paper in the future if the experiments are made more rigorous with standard tasks and benchmarks. 
Decision: Reject