Scatterbrain: Unifying Sparse and Low-rank Attention
Authors: Beidi Chen
Tri Dao
Eric Winsor
Zhao Song
Atri Rudra
Christopher Ré
Venue: NeurIPS 2021 Poster (2021)
Keywords: Sparsity
Low-rank
Locality Sensitive Hashing
Kernel Approximation
Efficient Transformers

Abstract: Recent advances in efficient Transformers have exploited either the sparsity or low-rank properties of attention matrices to reduce the computational and memory bottlenecks of modeling long sequences. However, it is still challenging to balance the trade-off between model quality and efficiency to perform a one-size-fits-all approximation for different tasks. To better understand this trade-off, we observe that sparse and low-rank approximations excel in different regimes, determined by the softmax temperature in attention, and sparse + low-rank can outperform each individually. Inspired by the classical robust-PCA algorithm for sparse and low-rank decomposition, we propose Scatterbrain, a novel way to unify sparse (via locality sensitive hashing) and low-rank (via kernel feature map) attention for accurate and efficient approximation. The estimation is unbiased with provably low error. We empirically show that Scatterbrain can achieve $2.1 \times$ lower error than baselines when serving as a drop-in replacement in BigGAN image generation and pre-trained T2T-ViT. On a pre-trained T2T Vision transformer, even without fine-tuning, Scatterbrain can reduce $98\%$ of attention memory at the cost of only $1\%$ drop in accuracy. We demonstrate Scatterbrain for end-to-end training with up to $4$ points better perplexity and 5 points better average accuracy than sparse or low-rank efficient transformers on language modeling and long-range-arena tasks.

Scatterbrain: Unifying Sparse and Low-rank Attention Approximation
Beidi Chen" Tri Dao"t, Eric Winsor
Zhao Song
Atri Rudra
Christopher Ré
Depparment of Computer Science, Stanford University
Adobe Research
Department of Computer Science and Engineering, University at Buffalo, SUNY (beidic, trid, winsor, chrismre) estareeeessssfooooooooooooooo ord. edu, zsongdadobe com, atri Gbuf falo.allo..ooooo.ooo.oooool.ooae edu
Abstract
Recent advances in efficient Transformers have exploited either the sparsity or low-rank properties of attention matrices to reduce the computational and memory bottlenecks of modeling long sequences. However, is still challenging to balance the trade-off between model quality and efficiency to perform a one-size-fits approximation for different tasks. To better understand this trade-off, we observe that sparse and low-rank approximations excel in different regimes, determined by the softmax temperature attention, and sparse + low-rank can outperform each individually. Inspired by the classical robust-PCA algorithm for sparse and low-rank decomposition, we propose Scatterbrain, a novel way to unify sparse (via locality sensitive hashing) and low-rank (via kernel feature map) attention for accurate and efficient approximation. The estimation is unbiased with provably low error. We empirically show that Scatterbrain than baselines when serving as a drop-in replacement in BigGAN image generation and pre-trained ViT-VIT. On a pre-trained T2T Vision transformer, even without fine-tuning, Scatterbrain can reduce 98% of attention memory at the cost only 1% drop in accuracy. We demonstrate Scatterbrain for end-to-end training with up 4 points better perplexity and 5 points better average accuracy than sparse or low-rank efficient transformers language modeling and Long-range-arent tasks.
achieve 2.1 x lower error
Introduction
Transformer models 64 have been adapted in wide variety of applications, including natural language processing 26 [1], image processing [10448), and speech recogmition [33 Training large Transformers requires extensive computational and memory resources, especially when modeling long sequences, mainly due to the quadratic complexity (w.r.t. sequence length) in attention layers. Recent leverage attention approximation the bottleneck by approximating the attention matrices. However, it is challenging to find a robust approximation method that balances the sereeinncy-cccuraee trade-off wide variety of tasks 58 59
advances in efficient transformers 36 17 35 22
overcome
We categorize most of the existing approaches for efficient attention matrix computation into two major groups: exploiting either the sparsity, e.g., Reformer 3T SMYRF 22 or low-rank properties
of the attention matrices, e.g., Linformer [66]. Linear Transformer
35 and Performer [17]. However, these techniques usually have different strengths and focus on the performance of specific tasks, so their approximations still cause accuracy degradation on many other tasks. For instance, according to recent benchmark paper [58] and our experiments, low-rank-based attention might be less effective on hierarchically structured data language modeling tasks, while sparse-base based variants do not perform well on classification tasks.
* Equal contribution. Order determined by coin flip.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).

---

CATEGORIZATION
Iegui
5
SPARSE V
K
V
Specse Larz-nak
00-
SCATTERBRAIN
Streess-Lorrank
(()((())))) SV
Rigi REA)
Sq de No? (16)
0.3
(Q))
LOWRANK V
Kende Hank
01
MKK
6: R
F2:R (1....))
Estropy
softmax (QK*) V
Figure 1: Leftt regimes that sparse+low-rane approximation is more accurate, based on the entropy of the attention matrices. Right: Scatterbrain Workflow. For the attention layer in Transformers, after computing Query Q. Key K, and Value V matrices, we approximate softmax(QIt VV with two components: (i) sparse SV (ii) low-rank (())
We observe that sparse and low-rank approximations are complementary for many attention matrices in practice, and sparse+low-rank could outperform each individually (Figure Ileft). We empirically categorize the regimes in which sparse or low-rank approximation achieves better error based on the softmax temperature of attention (of which the entropy of softmax distribution can be used as proxy). We expect that sparse methods perform well if the attention depends on a few entries (low entropy softmax). In contrast, low-rank methods do better if the attention depends on a mixture of many components (high entropy softmax). This explains the phenomenon that current sparse and low-rank-based approaches excel on different kinds of tasks. A natural question is whether one could understand and unify the strength of both approaches. While is NP-hard to find the optimal combination of sparse and low-rank approximations, Robust PCA [ is a polynomial-ting with tight approximation error. We observe that Robust PCA achieves lower approximation error than sparse or low-rank alone attention matrices. The difference is most pronounced for "mid-range" entropy, where we observe that up to 95% error reduction is possible.
solution
The connection between Robust PCA and attention matrix estimation provides an opportunity to realize a more robust approximation. Specifically, given attention matrix, one could adaptively perform sparse+low-rank approximation to obtain a low error. However, comes with three challenges: (i) How to decompose the attention matrices into sparse and low-rank components and estimate them efficiently and accurately; Robust PCA is accurate but slow and requires materializing the full attention, while straightforward addition of sparse and low-rank attention will be inaccurate due to double counting. (ii) It is not clear if there is a theoretical guarantee that sparse + low-rank approximation is strictly better than sparse low-rank in some regimes, though we observe the separation empirically. (iii) How does the lower approximation error transfer to end-to-end performance in real tasks.
In this paper, we propose Scatterbrain, an accurate and efficient robust estimation of attention matrices with theoretical guarantees address the above challenges. Specifically:
In Section we observe that sparse and low-rank approximation complementary and demonstrate that sparse + low-rank structure arises naturally when elements in the input sequence form clusters. We theoretically characterize and analyze the regimes where sparse, low-rank, and sparse+low-rank excel, dictated by the softmax temperature of attention.
In Section 4 inspired by the classical Robust PCA algorithm, we propose Scatterbrain, which efficiently combines sparse and low-rank matrices to approximate attention. In particular, we use Lo- cality Sensitive Hashing (LSH) identify large entries of the attention matrix (after softmax) without materializing the full matrix and then leverage kernel approximation to parameterize the low-rank part. We prove that our method has a strictly lower approximation error than the low-rank baseline. In Section accurately approximates the attention matrix, is memory efficient for long sequences, and works well across different tasks. First, we show that its approximation accuracy is close to our oracle Robust PCA and achieves 2 lower error compared to other efficient baselines on real benchmarks. This leads to direct application of Scatterbrain a drop-in replacement to pre-trained full attention, thus reducing up to 98% of the memory required for attention computations pre-trained TiT-ViT and BigGAN while maintaining similar quality. Last we show that its superior accuracy and
5 we empirically validate our theory and the proposed method, showing that Scatterbrain

---

efficiency can improve the efficiency-acceeee the WikiText- 103 language modeling task. Scatterbrain achieves up point better perplexity compared to Reformer and Performer. On 5 benchmark long-range tasks, Scatterbrain improves the average accuracy by up 5o points
trade-off's of Transformer end-to-end training. On
2 Problem Setting and Related Work
We first define the approximation problem we aim to solve in this paper. Then we discuss the applica- tions of sparse and low-rank techniques in efficient Transformers and introduce robust PCA algorithm.
Problem Formulation: In the attention matrix approximation problem, we are given three matrices, query, key, and value, Q, K, V mn compute voftmax(QK") VV. We seek to reduce the quadratic complexity of softmax(QK) (applied row-wise) with low approximation error. More precisely, for an approximation procedure f. we minimize two objectives, the approximation error Il% and the computation//eeoory cost C()(-)).
E -
Sparse, Low-rank Approximation for Attention Matrices: Recent work exploits the sparsity patterns or finds a low-rank mapping of the original attention matrices to overcome the computational
and memory bottlenecks in Transformers 36 22
35 66).
Generally, can divide most of the
techniques into two categories- sparse and low-rank approximations. Reformer (3 is representative sparse variant that uses LSH [3] to retrieve or detect the locations of the attention matrices with large values and reduce the computation from ((n²) ) to (nlogn). Performer Di is an example of the low-rank variant, which uses kernelization to avoid explicit (n²2d) computation. One problem of either the sparse or low-rank approximation is that the structure of the attention matrices varies in practice, and is challenging to perform robust approximation on a wide range of attention matrices. For example, Wang et al. [66] observes that attentions tend to have more low-rank structures in lower layers and Ramsauer et al. [52] shows that they are sparser in the later stage of the training. Ideally, we want to unify the strength of both techniques, but it is NP-hard to find the best combination of sparse and low-rank approximation.
Sparse + Low-rank and Robust PCA: Fortunately, classical Robust PCA g presents a polynomial algorithm to find the approximately optimal or good combinations of sparse and low-rank approx- imation of the matrices. The sparse low-rank matrix structure has been well studied in statistics and signal processing since the late 2000s [Р This structure naturally generalizes low-rank [33
63 and
sparse [I]] matrices. Scatterbrain is built on a line of work, c.g., Bigbird [71] Longformer [ with the theme of combining multiple types of attention and another one in the optimal transport setting [r7 However, despite the multitude of papers, this sparse + low-rank matrix approximation has not been rigorously studied in the context of attention matrices. We undertake this study and show how can relax the sparse + low-rank approximation from robust PCA, making it efficient while still retaining PCA's accuracy. In fact, our results shed further light on why Bigbird or Longformer work, they special cases of a single principled structure. An extended discussion of related work is in Appendix A
3 Characterization of Sparse + Low-rank Approx. to Attention Matrices
We motivate the use of sparse + low-rank approximation of the attention matrices with the key observation that for many attention matrices, sparse and low-rank approximation complementary, and their ideal combination (via Robust PCA) can outperform both (Section B.T). Furthermore, we argue that the sparse + low-rank structure can arise naturally when elements in the input sequence form clusters, as dictated by the softmax temperature (Section 3.2
3.1 Motivating Observations: Low-rank and Sparse Structures of Attention Matrices
We empirically characterize regimes where sparse and low-rank approximation are well-suited, based on the softmax temperature (for which we use the softmax distribution entropy is a proxy). Specifically, in Fig. [[eft), we present the approximation error of the original attention matrices and the approximation (sparse or low-rank) of matrices sampled from a 4-layer Transformer trained on IMDb reviews classification [58] We make two observations:
1. Sparse and low-rank approximation are complementary: sparse excels when the softmax temperature scale is low (i.e., low entropy), and low-rank excels when the softmax temperature is high (i.e., high entropy).
Scatterbrain code is available at httpp///w/ggggtt...././/.///eeea
3

---

2.
An ideal combination of sparse and low-rank (orange line in Fig. llet), obtained with robust PCA, can achieve lower error than both.
Similar observations on other benchmarks and details are presented in Appendix
3.2 A Generative Model of How Sparse + Low-rank Structure Can Arise
Sparse low-rank parameterization is more pressive than either sparse or low-rank alone. In- deed, in the Appendix, we construct a family of attention matrices to show the separation be- tween the approximation capability of sparse low-rank vs. sparse or low-rank alone: for an nnn attention matrix, sparse or low-rank alone requires a (n²) parameters to get approxi- mation error in Frobenius norm, while sparse low-rank only requires (n) parameters.
Figure 2: Visualization of the generative process, for three different values of the intra-cluster distance (small, medium, and large). The vectors from the input sequence (rows of Q) form clusters that lie approximately on the unit sphere. Different colors represent different clusters.
Moreover, we argue here that sparse + low-rank a natural candidate to approximate generic at- tention matrices. We describe generative model of how the sparse + low-rank structure in attention matrices could arise when the elements the input sequence form clusters. Under this process, we characterize how the softmax temperature dictates when we would need sparse, low-rank, or sparse low-rank matrices to approximate the attention matrix. This result corroborates the observation in Section
Generative process of clustered elements in input sequence We describe here generative model of an input sequence attention, parameterized by the inverse temperature 8 IR and the intra-cluster distance A E R.
Process 1. Let Q Bnxd, where d 2 Dilog 3/2 (n))), with every row of Q generated randomly as follows:
1. For (oCC))(()) sample C number cluster centers CI,..., EIR independently from (O.ITa/ Vd). For each cluster around Gs sample n; = =(1) number of elements around Ci of the form Zij- = C Grij j 1,, where Tij ~> ~N (0,,La// Assume that the total number of elements is n==
and A O1L/lg" n/...
Let Q be the matrix whose rows are the vectors Ziy where i= I and j= I,...,n,, Let A A= QQ and let the attention matrix be M3 M
We visualize this generative process in Fig
Softmax temperature and approx. error We characterize when to use sparse, low-rank, or sparse + low-rank to approximate the attention matrices in Process [ depending on the inverse temperature B. The intuition here that the inverse temperature corresponds to the strength of interaction between clusters. If 3 is large, intra-cluster interaction dominates the attention matrix, the softmax distribution peaked, and so we only need a sparse matrix to approximate the attention. If small, then the inter-cluster attention is similar to intra-cluster attention, the softmax distribution is diffuse, and we can approximate it with low-rank matrix. In the middle regime of 3, we need the sparse part to cover the intra-cluster attention and the low-rank part approximate the inter-cluster attention.
We formalize this intuition in Theorem I bounds below we think of constant), All the proofs are in Appendix[] D
Let Mg, be the attention matrix in Process 7 Fix E (0.1). Ler R Rx be a matrix. Consider low-rank, sparse, and sparse + low-rank approximations to Mg.
Theorem 1.
1. High temperature: Assume 3= ollogn//ood)
(a) Low-rank: There exists R with nl(1) rank (and hence n too()) parameters) such that [|
S
(b) Sparse: If R has sparsity o(n²), then || R|r 2(nn)
2. Mid temperature: Assume A logn < $ Olognn).

---

(a) Sparse low-rank: There exists a sparse low-rank R with n ++(1)) parameters with 1Ma
R| F en.
(b) Low-rank: I is such that n narkank(( (1)) then (1())))))0(m)
(c) Sparse: If R has sparsity o(n²), then ||M 2
Low temperature: Assume 3= (loogn)..
(a) Low-rank: If is such that n (R) then 1 20
(b) Sparse: There exists R with spansity (n) such that M3 Se ee
4 Scatterbrain: Unifying Sparse and Low-rank Attention
We present Scatterbrain, and show that approximates at- tention accurately and efficiently. Section 4.1 describes the challenges designing an accurate and efficient approx- imation, and how obvious baselines such as Robust PCA or a simple combination of sparse attention and low-rank attention fail to meet both criteria. Section 4 demonstratees how Scatterbrain address the challenges (Fig. contains
Accuracy
High
Robest PCA
Scoterfcaki (pucse rosit)
Performr (low-aak)
Referer (spaase
schematic of Scatterbrain). In Section 433 we show that Scatterbrain unbiased with provably lower variance than low-rank baselines such as Performer.
Efficiency
Low
High
Fig. shows qualitative comparison between different methods of approximating the attention matrix: Robust PCA is accurate but slow, sparse (e.g., Reformer), and low-rank (e.g., Performer) attention are fast and memory- efficient but may not be very accurate, while Scatterbrain is more accurate than its sparse and low-rank counterparts
Figure 3: Qualitative comparison of ap- prox. accuracy and efficiency, among Ro- bust PCA, sparse (Reformer) and low- rank (Performer) attention, and Scatter- brain. Scatterbrain more accurate while being efficient.
while remaining just efficient.
More details about the efficient implementation of Scatterbrain are in Appendix
4.1 Challenges of Designing Accurate and Efficient Sparse Low-rank Approximation
We seek a sparse+ low-rank approximation of the attention matrix Fl A that is both accurate and efficient. The natural theoretical baseline of Robust PCA is too slow and requires too much memory, while the most straightforward way of combining sparse attention and low-rank attention fails due to double counting on the support of the sparse attention.
If the goal is accuracy, Robust PCA is the most studied algorithm to find a sparse + low-rank approximation to a given matrix. It relaxes the NP- harhard problem of finding the best sparse + low- rank approximation into convex optimization problem, with the nuclear norm and & constraints. Even though it can be solved in polynomial time, it is orders of magnitude too slow to be used in each iteration of a training loop. Moreover, requires materializing the attention matrix, which defeats the main purpose of reducing compute and memory requirements.
On the other hand, one efficient way to get sparse + low-rank approximation of an attention matrix is to simply add the entries of sparse approximation S (say, from Reformer) and a low-rank approximation QK for Q. e Ruxm (say, from Performer). The sparse matrix S typically has support determined randomly I6 by LSH 3d 22]. or by clustering (54]. On the support of S, which likely includes the locations of the large entries of the attention matrix A, the entries of S match those of A. One can multiply (S+ QK =V QQ( V) efficiently because S sparse, and grouping Q(K" V) reduces the matrix multiplication complexity when m & n, from O(n²m) to ((nd). The approximation S + Q matches QKT outside supp(S), hence it could be accurate there if QKT is accurate. However, S QKT will not be accurate on the support of S due to the contributions from both S and from QKT Adjusting QKT to discount the contribution from is difficult, especially if we want to avoid materializing QKT for efficiency.
Forr simplicity of discussion, we consider the unnormalized attention matrix A ==exp(QKT] omitting the usual scaling of Vd and the softmax normalization constant.
5

---

4.2 Scatterbrain: Algorithm Intuition and Description
The simple insight behind our method is that on the support of the sparse matrix S, instead of trying to match the entries of the attention matrix A, we can set the entries of S to discount the contribution from the low-rank part QK' This way, the approximation S+ ++QKT T will match A exactly on the support of S, and will match QK' outside supp(S), which means will still accurate there if QK is accurate. We do not need to materialize the full matrix QKT as need subset of its entries is required, hence our approximation will be compute and memory efficient.
Scatterbrain thus proceeds in three steps: we construct a low-rank approximation OR ~ A, and construct sparse matrix S such that S S+QK" matches A on the support of S, then finally multiply SV and Q(K T V) and combine the result. More specifically:
1. Low-rank Approximation. We define a procedure LOWRANK that returns two matrices Q,K E Rax such that QK approximates A. In particular,
a randomized kernel feature map vm exp (Wa- x ||=:1/// /2) with W E RRnd randomly sampled, entry- wise, from the standard normal distribution N (0.1). We apply 6 to each row vector of Q KK matrices, and denote (Q) and K = ()) (row-wise). Note that we do not materialize QKT Sparse Approximation. define a procedure SPARSE that returns sparse matrix S that matches A A-QK" on SS) In particular, using a family of locality sensitive hash functions, compute the hash codes of each query and key vectors in Q.K K matrices (row-wise). Let S be the set of locations (i.j) where qi and kj have the same hash codes (i.e, fall into the same hash bucket). Let S be the sparse matrix whose support is S, and for each (i.j) ES, define
Φ: - where <(x) J.m
Su exp((, k,) (1))-((p ((q.) = ky) q. kyy
where qi, Ki kj are the i-th and j-th rows of Q,K, Q respectively. Note that we do not materialize QKT
3. Scatterbrain Approximation. With Q, K returned from LOWRANK and S from SPARSE, we compute the (unnormalized) attention output with
Ō= Ō=(²K (()KK+SVVVVV- = VV
The precise algorithm, including the normalization step, as well as the causal/unidirectinnll variant, is described in Appendix C We also note Scatterbrain's flexibility: it can use different kinds of low-rank and sparse approximation as its sub-components. The combination of Reformer and Performer is simply one instance of Scatterbrain. Instead of using Reformer as sparse component, we could use local attention [ or random block-sparse attention [16]. Instead of using Performer 17 as
low-rank component, we could also use Linear attention 335 or global tokens as in BigBird
The Scatterbrain method would work exactly the same way. As long as the low-rank component is unbiased (e.g.. Performer), its combination with any sparse component in Scatterbrain would yield an unbiased estimator of the attention matrix as shown below.
10-
Reformer (sparse)
Performer (low-rankk
Sontterbeain
(spane---oo---p low-mank))
MISE
4.3 Scatterbrain: Analysis
Our method combines a low-rank approximation QKT (which has rank m < n) with a sparse approximation S. We argue that it is accurate (lower approximation error than baselines) and efficient (scaling the same as sparse or low-rank alone). The main insight of the analysis is that our approximation is exact for entries on the support of S (picked by LSH), which are likely to be large. For entries not in the support of S (likely to be small), our approximation matches the low-rank part (Performer) QKT, which is unbiased and has low variance for these entries. As result, Scatterbrain retains the unbiasedness of Performer 17 but with strictly lower variance.
Figure 4: Per-entry MSE for differ- ent approximations, across a range of magnitude of q k.Scatterbrain has low MSE for both small and large entries, thus outperforming its sparse (Reformer) and low-rank (Performer) counterparts.

---

We compare Scatterbrain to its low-rank baseline (Performer) and sparse baseline (Reformer). Performer also based on the kernel approximation o, and simply uses QK' T to approximate the attention matrix A. Reformer uses LSH to identify large entries A, then compute sparse matrix S such that Si exp((g) k2) for i)
Accuracy: Because of the way S is defined in Eq. D gK + matches A= exp(QKT) exactly on locations (i.j) ES, which are locations with likely large values. This addresses a weakness of low-rank methods (e.g., Performer) where the low-rank estimate is not accurate for locations with large values. We analyze the expectation and variance per entry of our estimator below (proof in Appendix
Theorem 2. Define (qq,k) = exp(a' kt). Gpte as Performer's estimator and Gabe as Scatterbrain estimator. Denote S IR as the unit sphere. Suppose q,k Sd-1 are such that ||g-kl <T. Then: (q,1:)]= (1-p)- Varlapte (q,bi) Var Varlopfe (g.k)]
Bl (q, = (q,kk, Varla where p= p=exp( (100------00 Ind 1---ndd--((nndd))
(3)
Hence Scatterbrain is unbiased. similar to Performer [177] but with strictly lower variance. The variance is small if exp(g' k)) is small (since Var(G k)) will be small), or if exp(1 k) is large (since the probability of not being selected by LSH, 1 p. will be small). In Fig. 4 we plot the per-entry MSE of different methods from Theorem when approximating the unnormalized softmax attention exp(QKT) ) Scatterbrain can approximate well both small entries (similar to the low-rank baseline, Performer), as well as large entries (similar to the sparse baseline, Reformer). Thus Scatterbrain has much lower MSE than Performer for large entries, and lower MSE than Reformer for small entries.
Efficiency: In Eq. 2 the computation SV is efficient because S is sparse, and [(K T because of the way we associate matrix multiplication (scaling as (nond) instead of (n² dn), which is much bigger if m << n)).
V) is efficient
We validate these two properties of our approach in Section
5 Experiments
We validate three claims that suggest Scatterbrain provides an accurate and efficient approximation to attention matrices, allowing to outperform its sparse and low-rank baselines on benchmark datasets.
In Section 5.1 we evaluate the approximation error and testing accuracy of different approximation methods on pre-trained models such as BigGAN and Vision Transformer. We show that the approxi- mation by Scatterbrain is close to the Robust PCA oracle and up to 2.1 lower approximation error than other efficient baselines.
In Section 5.2 we validate that when trained end-to-end, Scatterbrain outperforms baselines (sparse or low-rank attention) on a wide variety of benchmark tasks, including language modeling, classification, and the Long-range Arena (LRA) benchmarks. Scatterbrain achieves up to points higher average accuracy on the LRA benchmark compared to Performer and Reformer.
In Section [.3 we demonstrate the scalability of Scatterbrain, showing that has comparable memory and time usage with simpler baselines (sparse or low-rank alone) across range of input sequence lengths (Section 5.3)) while requiring up 12x smaller memory than full attention.
All details (hyperparameters, data splits, etc.), along with additional experiments, are in Appendix
5.1 Scatterbrain's 's Approximation Accuracy
We evaluate Scatterbrain's 's approximation ac- euracy in three steps: (1) compare it with of Robust PCA (sparse+low-randl
Table 1: Top-1 Accuracy of pre-trained T2T Vision Transformer on ImageNet with different attention re- placements. Error represents the average normalized approximation error to full attention.
our theoreti- cal foundation and oracle (2) compare it with SMYRI [22], Performer 117. which are pop- ular variants of sparse and low-rank approx- imation to attention respectively and a naive baseline that directly adds SMYRF and Per- former, (3) evaluate the inference accuracy when replacing full attention with Scatterbrain
Atfention
Top-1 1 Acc
Error (avg)
Full Attention
81.7%
SMYRF
79.8%
11.4%
Performer
80.1%
7.5%
Buseline SMYRF + Performer
79.7%
12.6%
Scatterbrain
80.7%
5.3%
'SMYRFia is a variant of Reformer that does not require the key and query to be the same, which is necessary for experiments in this section.
7

---

Seatteebrat RoasttPCA PCA
Stzef
100
Soger
Stgat
547
Prrkomm
Peitouer
Scatorrhmiit
Sonteclesin
Seactueleam
Apgros Unoe
Pill
Full
04
20
-0
04-
13
300
600
1.
Xlenory
Figure 5: First: approximation comparison between Scatterbrain and its "lowerbound" Robust PCA. Second: comparison of eπтоr entropy among SMYRF, Performer and Scatterbrain, three represen- tatives of sparse, low-rank and sparse+low-ras approximations. Third and forth: Inception score (higher is better) and FID score (lower is better) of different attention variants for pretrained BigGAN.
approximation. Scatterbrain achieves erгor within 20% of the oracle robust PCA, and up to 2.1x lower error than SMYRF and Performer. When serving as drop-in replacement for full attention, even without training, Scatterbrain can reduce the attention memory of Vision Transformer by 98% the cost of only 0.8% drop of accuracy.
Setup: We use the attention matrices from pre- trained BigGAN and T2T- BigGAN is a state-of- the-art model in Image Generation for ImageNet. BigGAN has a single attention layer at resolution 64 64 (4096 queries). T2T-ViT has 14 attention layers. Scatterbrain sets the ratio between SMYRF and Performer based on the entropy of an observed subset of attention matrices in different layers. We allocate more memory to the low-rank component compared to the sparse part if the entropy high.
Scatterbrain and Robust PCA: We first show that Scatterbrain approximates pre-trained attention matrices 10"x x faster while its approximation error is within 20% on average. We also provide an example visualization on 100 attention matrices from the BigGAN generation process in Figure
Scatterbrain vs. SMYRF and Performer: We show that Scatterbrain approximates pre-trained dense attention matrices with very low error compared to sparse (Reformer) or low-rank (Performer). Measuring Frobenius approx. error on the BigGAN image generation task, Scatterbrain achieves 2x lower error compared to Performer.
Drop-in replacement for full attention: We show that accurate approximation directly leads to efficient Inference. We replace BigGAN's dense attention with a Scatterbrain layer without other modifications. In (right two), we show Inception and FID scores for Scatterbrain and other baselines under different memory budgets. Similarly, we use TiT-ViT 77), which a token-to-token vision Transformer pre-traie pre-trained of Scatterbrain for each layer and the end-to-end testing accuracy after replacing full attention with Scatterbrain and other baselines. Notably, Scatterbrain achieves 80. 7% Top-1 accuracy, which is only 1% drop from the original 81.7% by full attention reducing up to 98% of the memory usage.
ImageNet [25] In Table D show the average approximation error
5.2 End-to-end Training Performance
Scatterbrain's 's accurate approximation of attention matrices allows it outperform other efficient Transformer methods on benchmark tasks. Across a range of diverse tasks, both commonly used autoregressive tasks (sequence modeling) and benchmark long-range classification tasks (Long-Range Arena), Scatterbrain outperforms Performer (low-rank baseline) and Reformer (sparse baseline) by up to 4 points.
5.2.1 Auto-regressive Tasks
On the standard language modeling task of Wikitext- Scatterbrain obtains point better perplexity than Reformer (sparse baseline), coming within 1.5 points of full attention.
Settings: We compare the performance of Scatterbrain against Reformer and Performer on one popular synthetic task, Copy, and one large language modeling task: WikiText103 46). Reformer is a representative approxiapppooiitiiooon based variant and Performer is www-raak-approximaaiio variant. The base model vanilla Transformer [6 We observed that generally allocating more memory budget to sparse tends to perform better, so Scatterbrain sets the ratio to 3:1 (sparse: low-rank component) for simplicity. The statistics each dataset and model hyper- papeer-prraetters are Appendix E
We report the best results of each method in perplexity.
8

---

Table The performance of Scatterbrain, Reformer, Performer and Full-Attention on -ange- Arena benchmarks and 2 popular language modeling tasks. We fix the same number of parameters (1/8 of the full) used for approximating the attention matrix for each method.
Attention
Copy (ppl)
WikiText- (ppl)
Attention
ListOps
Text
Retrieval
Image
Pathfinder
Avg
Full Attention
25,258
Full Attention
38.2
63.29
80.85
41.78
73.98
59.62
Reformer
6.8
27.68
Refermer
36.85
58.12
78.36
28.3
67.95
53.9
Performer
66
Performer
35.75
62.36
78.83
39.71
68.6
57.05
Scatterbrain
26.72
Scatterbrain
38.6
64.55
80.22
43.65
69,91
59.38
Results: Table 2 shows the testing perplexity for Scatterbrain and other baselines under the same parameter budget (each approximation is only allowed to compute 1 of the full computation). Scatter- brain achieves comparable perplexity compared to the full attention Transformer model on Copy, and WikiText- 103. Notably, Scatterbrain achieves points lower perplexity on Copy and point lower on WikiText- 103 compared to Reformer, while Performer does not train stably on auto-regressive tasks (loss does not go down).
We also analyze the results by visualizing the error of Reformer (sparse), Performer (low- rank), and Scatterbrain (sparse low-rank) given the same number of parameters when approximating the full attention matrices for each attention layer during training (Appendix F The conclusion is for language modeling tasks, sparse+low-rand
Analysis:
has the smallest approximation error in most of the cases, and sparse has the largest error, which matches with the end-to-end results. It also confirms the observation in the popular benchmark paper [58] that kernel or low-rank based approximations are less effective for hierarchical structured data.
5.2.2 Classification Tasks
On a suite of long-range
benchmark tasks (Long Range Area), Scatterbrain outperforms Reformer (sparse baseline) and Performer (low-rank baseline) by up to points on average.
We compare the performance of Scatterbrain against Reformer and Performer on ListOps, two classifications: byte-level IMDb reviews text classification, image classification sequences of pixels, a text retrieval, and pathfinder tasks. The datasets are obtained from the Long Range Arena (LRA) Benchmark [58], which is recent popular benchmark designed for testing efficient Transformers. Similar to the auto-regressive tasks above, we use Reformer and Performer as baselines. The base model is also vanilla Transformer. We follow the evaluation protocol from 58 the best accuracy of each method.
Settings:
We report
Results: Table 2 shows the individual and average accuracy of each task for Scatterbrain and other baselines under the same parameters budget. Specially, each approximation is only allowed to 12.5% of the full computation. We can see Scatterbrain is very close to full attention even with a large reduction in computation and memory. Further more, it outperforms all the other baselines consistently on every task and achieves more than 5 point average accuracy improvement than sparse- based approximation Reformer and more than 2 point average accuracy improvement than low-rank-baad variant Performer.
Analysis: Similarly, in order to analyze the performance of Re- former, Performer and Scatter- brain. we visualize their approx- imation еттоr given the same number parameters when ap- proximating the full attention matrices for each attention layer during training (Appendix E.. We again find that Scatterbrain has the smallest approximation error, while Performer is the worst on ListOps and Reformer has the largest error on classifica- tion tasks, which matches with
Reformer
Reformet
Perfoctnet
Performer
Scazterbrain
Scatterbrain
50
Altn
Metuory (GIB)
Full
100
1054
1008
10384
1024
4036
10381
Sequence Lengtth
Seqgoence Length
Figure 6: Speed and memory required by different efficien attention methods. Scatterbrain is competitive with SMYRF (sparse baseline) and Performer (low-rank baseline), while up to 3x faster and 12x more memory efficient than full attention for sequence length 4096.

---

the end-to-end results and confirms our observations earlier (sparse and low-rank approximation excel in different regimes).
5.3 Scatterbrain's Efficiency, Scaling with Input Sequence Length
We include ablation studies the scalability of Scatterbrain in Fig. K showing that it is as computation and memory-efficient as simpler baselines such as SMYRF and Performer, while up to 3 faster and 12x more memory efficient than full attention for sequence length 4096. This demonstrates that our combination of sparse and low-rank inherits their efficiency.
We report run times and memory consumption of the sequence lengths ranging from 512 to 32768. We use a batch size of 16 for all rums and conduct experiments V100 GPU. Since the efficiency would be largely conditioned on hardware and implementation details, we perform best-effort fair comparisons. We adapt the Pytorch implementation from prtorch-fastt transfrnnffo fransforers ormers and implement Scatterbrain similarly without any customized cuda kernels.
library for our baselines
6 Discussion
Limitations, As Scatterbrain has sparse attention as a component, it is not yet as hardware friendly (on GPUs and TPUs) the low-rank component, which uses the very optimized dense matrix multiplication. This is the same limitation suffered by other sparse attention methods, but we are excited that more efficient sparse GPU kernels are being developed BT 29
Potential negative societal impacts. Our work seeks understand the role of matrix approximation (and potentially energy savings) in the attention layer, which may improve a wide range of applications, each with their own potential benefits and harms. For example, making it language modeling more compute and memory efficient might facilitate spreading misinformation, and better image and video processing may make automatic surveillance easier. To mitigate these risks, one needs address application-specific issues such as privacy and fairness, going beyond the error metrics we considered. Specially, for language models (LMs). while our work partially addresses the issue of environmental cost of LMs raised in [6, it does not address other issues such as unfathomable training data [6]
Discussion and future work. In this work, we make an observation on the sparse + low-rank structure of the attentions in Transformer models and theoretically characterize the regimes where sparse, low-rank and sparse low-rank excel, based on the softmax temperature of the attention matrices. Motivated by this observation, we present Scatterbrain, a novel way to unify the strengths of both sparse and low-rank methods for accurate and efficient attention approximation with provable guarantees. We empirically verify the effectiveness of Scatterbrain on pretrained BigGAN, vision transformers, as well as end-to to-end training of vanilla transformer. We anticipate that the study of this core approximation problem can prove useful in other contexts, such as generalized attention layers with other non- non-linearity beside softmax, and wide output layer in language modeling or extreme-classfiiatton.
Acknowledgments
We thank Xun Huang, Sarah Hooper, Albert Gu, Ananya Kumar, Sen Wu, Trenton Chang, Megan Leszczynski, and Karan Goel for their helpful discussions and feedback on early drafts of the paper.
We gratefully acknowledge the support of NIH under No. U54EB020405 (Mobilize), NSF under Nos. CCF1763315 (Beyond Sparsity), CCF1 C631563078 (Volume to Velocity), and 1937301 (RTML): ONR under No. N0001411 772266 (Unifying Weak Supervision); ONR N00014-20-1-24888 Understanding and Applying Non-Euclidean Geometry in Machine Learning; N000142012275 (NEPTUNE): the Moore Foundation, NXP, Xilinx. LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC. ARM, Hitachi, BASF, Accenture, Ericsson, Qualcomm, Analog Devices, the Okawa Foundation, American Family Insurance, Google Cloud, Salesforce, Total, the HAI-AWS Cloud Credits for Research program, the Stanford Data Science Initiative (SDSI), and members of the Stanford DAWN project: Facebook, Google, and VMWare. The Mobilize Center is a Biomedical Technology Resource Center, funded by the NIH National Institute of Biomedical Imaging and Bioengineering through Grant P41EB027060. The U.S. Government is authorized to reproduce and distribute reprints for Govermmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or recom- mendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of NIH, ONR, or the U.S. Government. Atri Rudra's 's research is supported by NSF grant
10

---

References
[1] Keivan Alizadeh, Ali Farhadi, and Mohammad Rastegari. Butterfly transform: An efficient FFT based neural architecture design. In The Conference on Computer Vision and Pattern Recognition (CVPR), 2020.
[2] Alexandr Andoni, Piotr Indyk, Thijs Laarhoven, llya Razenshteyn, and Ludwig Schmidt. Practical and optimal ish for angular distance. In C. Cortes, N. D. Lawrence, D. D. Lee, M. Sugiyama, and R. Garnett, editors, Advances in Neural Information Processing Systems (Neur!PS), pages
1225-1233 2015.
[3] Alexandr Andoni, Piotr Indyk, Thijs Laarhoven, Ilya Razenshteyn, and Ludwig Schmidt. Practical and optimal LSH for angular distance. In Proceedings of the 28th International Conference on Neural Information Processing Systems-Volume 1, pages 125--1333,200 2015.
[4] R Artusi, P Verderio, and E Marubini. Bravais- pearsso and spearman correlation coefficients: meaning, test of hypothesis and confidence interval. The International journal of biological markers, (722)118-111 2002.
[5] Iz Beltagy, Matthew Peters, and Arman Cohan, Longformer: The long-document transformer. arXiv preprint ar 2xxv:2004.00/500 2020.
[6] Emily M. Bender, Timnit Gebru, Angelina McMillan- Majoor, and Margaret Mitchell. On the dangers of stochastic parrots: Can language models be too big? In Proceedings of the 2021 ACM Conference on Fairness, Accountability, and Transparency, New York, NY, USA, 2021.
[7] Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005. 2020.
[8] Emmanuel J Candès and Benjamin Recht. Exact matrix completion via convex optimization. Foundations of Computational mathematics, 917717777,
[9] Emmanuel J Candès, Xiaodong Li, Yi Ma, and John Wright. Robust principal component analysis? Journal of the ACM 58) 2011.
[10] Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European Conference on Computer Vision, pages 213-229 Springer, 2020.
[11] Beidi Chen and Anshumali Shrivastava. Densified winner take all (wta) hashing for sparse datasets. In Uncertainty in artificial intelligence, 2018.
[12] Beidi Chen, Anshumali Shrivastava, and Rebecca C Steorts. Unique entity estimation with application to the syrian conflict. The Annals Applied Statistics, 1239) 2018.
[13] Beidi Chen, Yingchen Xu, and Anshumali Shrivastava. Fast and accurate stochastic gradient estimation. 2019.
[14] Beidi Chen, Tharun Medini, James Farwell, Charlie Tai, Anshumali Shrivastava, al. SLIDE: In defense of smart algorithms over hardware acceleration for large-scale deep learning systems. Proceedings of Machine Learning and Systems, 2:291-308 2020.
[15] Beidi Chen, Zichang Liu, Binghui Peng, Zhaozhuo Xu, Jonathan Lingjie Li, Tri Dao, Zhao Song, Anshumali Shrivastava, and Christopher Ré. Mongoose: A learnable ish framework for efficient neural network training. In The International Conference on Learning Representations, 2021.
[16] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint 1:Xiv:9904.10009 2019.
[17] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers, arXiv preprint arXiv:2009.1479.0 2020.
11

---

[18] Shabnam Daghaghi, Tharun Medini, Nicholas Meisburger, Beidi Chen, Mengnan Zhao, and Anshumali Shrivastava. A tale of two efficient and informative negative sampling distributions. In International Conference on Machine Learning, pages 2319-2329. PMLR, 2021.
[19] Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoe V Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. arXiv preprint arXiv: 2019.
[20] Tri Dao, Albert Gu, Matthew Eichhorn, Atri Rudra, and Christopher Ré. Learning fast algorithms for linear transforms using butterfly factorizations. In The International Conference on Machine Learning (ICML), 2019.
[21] Tri Dao, Nimit Sohoni, Albert Gu, Matthew Fichhorn, Amit Blonder, Megan Leszc7ynski, Atri Rudra, and Christopher Ré. Kaleidoscope: An efficient, learnalee representation for all structured linear maps. In The International Conference on Learning Representations (ICLR), 2020.
[22] Giannis Daras, Nikita Kitaev, Augustus Odena, and Alexandros G Dimakis. Smyrf: Efficient attention using asymmetric clustering. arXiv preprini arXiv:2010.053/5, 2020.
[23] Christopher De Sa, Christopher Re, and Kunle Olukotun. Global convergence of stochastic gradient descent for some non-convex matrix problems. In International Conference on Machine Learning, pages
PMLR, 2015.
[24] Christopher De Sa, Albert Gu, Rohan Puttagunta, Christopher Ré, and Atri Rudra. A two-pronged progress in structured dense matrix vector multiplication. In Proceedings of the Twenty-Ninh Annual -CM-SIAM Symposium on Discrete Algorithus,, pages (060-1079) SIAM, 2018.
[25] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pages 248-255, 2009. doi: 10. 10.1
[26] Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 18Xiv:1810.048000 2018.
[27] Yihe Dong, Piotr Indyk, Ilya Razenshteyn, and Tal Wagner. Learning space partitions for nearest neighbor search. In International Conference on Learning Representations (ICLR), 2019.
[28] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghami, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.1 / 1929, 2020.
[29] Trevor Gale. Matei Zaharia, Cliff Young, and Erich Elsen. Sparse GPU kernels for deep learning.
In Supercomputing, 2020.
[30] Aristides Gionis, Piotr Indyk, Rajeev Motwani, et al. Similarity search in high dimensions via hashing. In Vidb, volume 99, pages 518-529, 1999.
[31] Scott Gray, Alec Radford, and Diederik P Kingma. GPU kernels for block-sparse weights. arXiv preprint arXiv: /7 117771..9222 1.0924,, 3, 2017.
[32] Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. In Advances in neural information processing systems (NeurIPS). 2020.
[33] Harold Hotelling. Analysis of a complex of statistical variables into principal components. Journal of educational psychology, 24(6))417, 1933.
[34] Piotr Indyk and Rajeev Motwani. Approximate nearest neighbors: towards removing the curse of dimensionality. In Proceedings of the thirtieth annual ACM symposium on Theory of computing. pages 604-613, 1998.
[35] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rons: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pages 5156-5165. PMLR, 2020.
12

---

[36] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In The International Conference on Machine Learning (ICML), 2020.
[37] Johannes Kliepera, Marten Lienen, and Stephan Günnemann. Scalable optimal transport in high dimensions for graph distances, embedding alignment, and more. In Proceedings of the 38th International Conference on Machine Learning. PMLR, 2021.
[38] Alex Krizhevsky, Geoffrey Hinton, el. Learning multiple layers of features from tiny images. 2009.
[39] Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self-supervised learning of language representations. In The International Conference on Learning Representations (ICLR), 2020.
[40] Valerii Likhosherstov, Krzysztof Choromanski, Jared Davis, Xingyou Song, and Adrian Weller. Sub-linear memory: How to make performers slim, arXiv preprint arXiv:2012. 11346 2020.
[41] Drew Linsley, Junkyung Kim, Vijay Veerabadran, and Thomas Serre. Learning long-range spatial dependencies with horizontal gated-recurrent units. arXiv preprint arXiv:/805.08.0 2018
[42] Zichang Liu, Zhaozhuo Xu, Alan Ji, Jonathan Li, Beidi Chen, and Anshumali Shrivastava. Climbing the wol: Training for cheaper inference. arXiv preprint arXiv:2007.012330, 2020.
[43] Haoneng Luo, Shiliang Zhang, Ming Lei, and Lei Xie. Simplified self-attention for transformer- based end-to-end speech recognition. In 2021 IEEE Spoken Language Technology Workshop (SI.T), pages 75-81. IEFE, 2021.
[44] Xuezhe Ma, Xiang Kong, Sinong Wang. Chunting Zhou, Jonathan May, Hao Ma, and Luke Zettlemoyer. Luna: Linear unified nested attention. arXiv:2106.015400 2021.
[45] Andrew Maas, Raymond E Daly, Peter T Pham, Dan Huang, Andrew Y 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, 2011.
[46] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, arXiv preprint arXiv: 1609 (074077433. 2016.
[47] Nikita Nangia and Samuel R Bowman. Listops: A diagnostic dataset for latent tree learning. arXiv preprint arXiv: NoXiv:/884.. 0804000048.8 2018.
[48] Niki Parmar, Ashish Vaswani. Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. International Conference on Machine Learning, pages 4055-4064. PMLR, 2018.
[49] Dragomir R Radev, Pradeep Muthukrishnan, Vahed Qazvinian, and Amjad Abu-Jbara. The acl anthology network corpus. Language Resources and Evaluation, 47(4):919 19499994,2299 2013.
[50] Jack Rae, Anna Potapenko, Siddhant M Jayakumar, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. In The International Conference on Learning Representations (ICIR), 2020.
[51] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. arXiv preprint arXiv: /910. 10683, 2019.
[52] Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlović, Geir Kjetil Sandve, et al. Hopfield networks is all you need. arXiv preprint arXiv:2008.022 02217, 2020..
[53] Benjamin Recht. A simpler approach to matrix completion. Journal of Machine Learning Research, 12(12), 2011.
[54] 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.
13

---

[55] Anshumali Shrivastava and Ping Li. Asymmetric Ish (alsh) for sublinear time maximum inner product search (mips). Advances in Neural Information Processing Systems (NeurIPS), pages 23211-332,, 2014.
[56] Vikas Sindhwani, Tara N. Sainath, and Sanjiv Kumar. Structured transforms for small-footprint deep learning. In Advances in Neural Information Processing Systems, pages 3088-3096, 2015.
[57] Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin. Adaptive attention span in transformers. In Proceedings the Annual Meeting of the Association for Computational Linguistics, 2019.
[58] 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 arXiv:2011 1.04006, 2020.
[59] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. arXiv preprint arXiv:2009.06733: 2020.
[60] Richard Taylor.
Interpretation of the correlation coefficient: a basic review, Journal of diagnostic medical sonography, 6(1):35-39, 1990.
[61] Reginald P Tewarson and Reginald P Tewarson. Sparse matrices, volume 69. Academic Press New York, 1973.
[62] Anna Thomas, Albert Gu, Tri Dao, Atri Rudra, and Christopher Ré. Learning compressed transforms with low displacement rank. In Advances in neural information processing systems (NeurIPS), pages 9052-9060, 2018.
[63] Madeleine Udell and Alex Townsend. Why are big data matrices approximately low rank? SIAM Journal on Mathematics of Data Science, 1(1):144-166) 2019.
[64] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. arXivv:7066. 03762, 2017.
[65] 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: 1804 07461, 2018.
[66] Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04766 2020.
[67] Felix Wu, Angela Fan, Alexei Baevski, Yann N Dauphin, and Michael Auli Pay less attention with lightweight and dynamic convolutions. In The International Conference on Learning Representations (ICLR), 2019.
[68] Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nystromformer: A Nystrom- algorithm for approximating self-attention. arXiv preprint arXiv:2102.03902 2021.
[69] Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Ruslan Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding. arXiv preprint /906::990600880 2019.
[70] Li Yuan, Yunpeng Chen, Tao Wang. Weihao Yu, Yujun Shi. Francis EH Tay, Jiashi Feng, and Shuicheng Yan. tokenss-to-tokn vit: Training vision transformers from scratch on imagenet. arXiv preprint arXiv:2101.11988 2021.
[71] 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.
[72] Chen Zhu, Wei Ping, Chaowei Xiao, Mohammad Shoeybi, Tom Goldstein, Anima Anandkumar, and Bryan Catanzaro. Long- short transformer: Efficient transformers for language and vision. arXiv preprint axXiv:2107.02192 2021.
14


Reviews:
[{"review_id": "pmGnL2_ziLk", "reviewer": "Reviewer_QK3S", "summary": "The work proposes to combine sparse attention and low-rank attention to form a better approximation to full attention. Theoretical analysis shows sparse attention is particularly good for approximating low-entropy attention while low-rank attention is good for the high-entropy case. By combining the two, the proposed Scatterbrain is able to outperform both low-rank and sparse attention alone on post-training approximation with BigGAN and T2T-ViT as well as end-to-end training on language modeling and long range arena. As a result, Scatterbrain makes the gap between full attention and linear-complexity attention closer.", "questions": "", "limitations": "", "rating": 7, "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": "The idea of combining sparse attention and low-rank attention is intuitively reasonable and well motivated by the theoretical analysis. Although there is still a clear gap compared to the full attention, the empirical results is already encouraging. \n\nOne suggestion is to add some study of the impact of parameter budget under language modeling (as it's more realistic) as well as the speed & memory test (Fig 6). This will provide additional helpful insights, especially for practical applications. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The limitations are clearly discussed in the paper and I don't see any potential negative societal impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "HFhRj_l6tLr", "reviewer": "Reviewer_Gfk7", "summary": "This submission presents an efficient approximation method of attention matrices. This work demonstrates the effectivenss of the proposed method really well in the domain of dense attention, e.g., Transformer or self-attention.\n\nThe authors provided thorough motivational examples, both theoretical and empirical analyses, and discussions. \n\nThis work achieves a new trade-off point that works likely to be better than sparse or low-rank counterparts by using sparse + low-rank decomposition.\nThe proposed method maintains both favorable properties of Performer and Reformer, and extends the working regime of the algorithm by sacrificing a bit of cost over Reformer.\n\nThe developed algorithm is simple and exact in Supp(S). Also, the authors design the experiment very well, and show effectively that the proposed scatterbrain has lower approximation errors. In various tasks, it is shown that the proposed method maintains the efficiency of Reformer at a similar level, while extending the successful working regimes.\n\nThe proposed method is practical in that the pre-trained attention-based models can be sped up by replacing the attention computation as a drop-in replacement that maintains the original quality as much as possible.\n\nI'm happy to be selected as a reviewer of this work. I'm enjoying reading this work, and appreciate the authors submitting high-quality work.", "questions": "", "limitations": "", "rating": 9, "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": "\nPros\n- A new good trade-off between speed and accuracy is achieved\n- Very well-motivated and thoroughly analyzed\n- The proposed algorithm is well-designed such that it is easy to replace the pre-trained components of Transformer.\n\nCons\n- I'm biased; thus, I couldn't find any major cons.\n- Discussion about the rank hyper-parameter and its experiment is not discussed.\n\n\n<Comments>\n\n* This work mainly focuses on the efficiency improvement of the pre-trained attention-based models. However, in practice, the most bottleneck to adopt such large attention-based models is on training. This reviewer is aware of the fine-tuning experiment in the supplementary material, but it would be more impactful to demonstrate full training from scratch. This is likely to amplify the gap with Reformer or Performer (especially against Performer), because gradient efficiency of the scatterbrain is much higher.\n\n* There must be a balance hyper-parameter between rank and sparsity. It acts as another trade-off controller. The discussion and phase transition on how to set the proper balance would be helpful to readers.\n\n\n\n================= After rebuttal =================\n\nThe authors' responses answer this reviewer's questions and comments appropriately.\nThis reviewer keeps supporting the original rating.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors discussed the positive effect of the proposed method on efficiency and energy saving. Also, they warned about spreading fake data due to the advance of the generative models in multiple domains, e.g., NLP, image, and video, and its extension to privacy and fairness issues. They also argued the necessity of addressing those problems individually. \n\nThe discussion is reasonable and agreed upon.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "Fju-1T5pZL", "reviewer": "Reviewer_6aFS", "summary": "This paper proposes ScatterBrain a combination of sparse and low-rank attention and shows that it outperforms both Reformer and Performer on language modeling, Long range arena and vision tasks. ", "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": "Essentially, despite what this paper is trying hard*not* to be, this seems like a simple combination of sparse (Reformer) LSH attention with low rank (Performer) attention. I would rate the novelty quite low especially with the way the methods are combined (just adding Reformer attention with Performer's low rank attention). \n\nI think many prior work (BigBird, ETC, Luna, Long short Transformers etc) have already shown that a combination of multiple sparse patterns or global components are useful. Along this line, I think the narrative of multiple attention being useful is not entirely new either. \n\nIf this paper was pitching itself \"evaluating the effectiveness of a combination of different xformer variants\", I can see how this is useful but the number of examples enumerated here is only a single method. I would encourage the authors to reframe this paper as such and evaluate more components rather than pitching a single combination (Reformer + Performer). Still, it is great to know that this outperforms standalone Performer and Reformer but i think the contribution is not really enough. \n\nFurthermore, I also have some concerns about experiments, namely:\n\n1. The performance of Performer on Copy and Wikitext is horrible. Have the authors tuned the model carefully? I think the kernel choice in Performer is important, it would be good to know if Performer was adequately tuned. \n2. For Long Range arena, Reformer and Performer are only reported. Given that there are other models evaluated in LRA, why did the authors not compare with the other models? For example, did the model outperform BigBird in the LRA paper? \n3. For language modeling, can the authors provide a strong local attention baseline? I think local attention should always at least be reported whenever xformer comes into play, given that it is the \"simple naive baseline\" \n\nMore questions:\n1. The sparse variant here is Reformer, which is form of learnable sparsity. What about fixed sparse patterns, like Sparse Transformer or Longformer/Sparse Transformer?\n\nMore thoughts and suggestions: \n1. Try combining different sparse and low-rank models, and draw insights from that. Are there benefits of different composition of attention types? What about different types of sparse models.\n2.  What about interleaving layers? Does interleaving Reformer and Performer give on the necessary inductive bias as ScatterBrain? \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, i have no concerns. ", "review_text": "", "extra_scores": {"ethical_concerns": "N/A ", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "0.5"}, "extra_text": {}}, {"review_id": "-6aiEsgFkf9", "reviewer": "Reviewer_Hndw", "summary": "Recently, sparse- and low-rank approximations are proposed to reduce the computational cost of the Transformer. However, each approximation works well in different areas. The authors show that the softmax temperature (entropy) determines which approximation is more appropriate. To enjoy the best of both approaches, this paper proposes Scatterbrain, which uses both sparse- (as Reformer) and low-rank (as Performer) attention computation. The key idea is to make sparse matrix S to match (A – QK) instead of A. Experiments on both pre-trained networks and end-to-end training show that Scatterbrain can achieve much smaller approximation error and better performance, respectively.", "questions": "", "limitations": "", "rating": 8, "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": "Overall, the idea and the motivation are clear. The paper is well-written and easy to follow. The selection of baselines is appropriate, and they are properly compared with the proposed algorithm. The analysis of the attention temperature is novel and meaningful. The result is also well visualized (ex: Fig 4, 5). Using robust-PCA as the most accurate approximation also helps readers to get the sense of the ‘golden’ combination.\n\nSuggestions:\n\n(1)\tI think figure 2 needs more explanations. (… what points and color mean, how they are visualized, etc.) In fact, is the figure necessary?\n\n(2)\tIt would be great if a simpler way (just add S + QK, line (198-200)) is also compared in experiments.\n\n(3)\tIn figure 6, the legend should be Reformer -> SMYRF.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors well addressed the limitations and potential negative social impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: In this paper, the authors discussed the effect of sparse and low-rank approximation for the transformer. The proposed combination leads to better performances comparing to the existing fast transformer with each approximation scheme only.  The paper is well-motivated and the theoretical and empirical analyses are reasonable. Therefore, I recommend the accept this submission. 

The paper can be further improved by taking the reviewers' suggestions about the experiments, e.g., full-training-from-the-scratch  (Reviewer Gfk7), comparison on LRA (6aFS), tradeoff between memory vs. accuracy (QK3S), and comparing with more baselines (Hndw). 

Decision: Accept (Poster)