Scatterbrain: Unifying Sparse and Low-rank Attention

1 Introduction . Transformer models [ 64 ] have been adapted in a wide variety of applications , including natural language processing [ 26 , 7 , 51 ] , image processing [ 10 , 48 ] , and speech recognition [ 43 ] . 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 advances in efficient transformers [ 36 , 17 , 35 , 66 , 22 ] leverage attention approximation to overcome the bottleneck by approximating the attention matrices . However , it is challenging to find a robust approximation method that balances the efficiency-accuracy trade-off on a wide variety of tasks [ 58 , 59 ] . We categorize most of the existing approaches for efficient attention matrix computation into two major groups : exploiting either the sparsity , e.g. , Reformer [ 36 ] , 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 a recent benchmark paper [ 58 ] and our experiments , low-rank-based attention might be less effective on hierarchically structured data or language modeling tasks , while sparse-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 ) . 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 1 left ) . 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 a 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 it is NP-hard to find the optimal combination of sparse and low-rank approximations , Robust PCA [ 9 ] is a polynomial-time solution with tight approximation error . We observe that Robust PCA achieves lower approximation error than sparse or low-rank alone on attention matrices . The difference is most pronounced for “ mid-range ” entropy , where we observe that up to 95 % error reduction is possible . The connection between Robust PCA and attention matrix estimation provides an opportunity to realize a more robust approximation . Specifically , given an attention matrix , one could adaptively perform sparse+low-rank approximation to obtain a low error . However , it 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 or 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 to address the above challenges . Specifically : • In Section 3 , we observe that sparse and low-rank approximation are 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 Locality Sensitive Hashing ( LSH ) to 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 5 , we empirically validate our theory and the proposed method , showing that Scatterbrain 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.1⇥ lower error compared to other efficient baselines on real benchmarks . This leads to a direct application of Scatterbrain as a drop-in replacement to pre-trained full attention , thus reducing up to 98 % of the memory required for attention computations in pre-trained T2T-ViT and BigGAN while maintaining similar quality . Last we show that its superior accuracy and efficiency can improve the efficiency-accuracy trade-offs of Transformer end-to-end training . On the WikiText-103 language modeling task , Scatterbrain achieves up to 1 point better perplexity compared to Reformer and Performer . On 5 benchmark long-range tasks , Scatterbrain improves the average accuracy by up to 5 points.2 2 Problem Setting and Related Work . We first define the approximation problem we aim to solve in this paper . Then we discuss the applications 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 2 Rn⇥d to compute softmax ( QK > ) V . 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 E ⇥ kf ( Q , K ) softmax ( QK > ) k2F ⇤ , and the computation/memory cost C ( f ( · ) ) . 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 , 54 , 17 , 35 , 66 ] . Generally , we can divide most of the techniques into two categories – sparse and low-rank approximations . Reformer [ 36 ] is a 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 O ( n2 ) to O ( nlogn ) . Performer [ 17 ] is an example of the low-rank variant , which uses kernelization to avoid explicit O ( n2d ) computation . One problem of either the sparse or low-rank approximation is that the structure of the attention matrices varies in practice , and it 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 [ 9 ] presents a polynomial algorithm to find the approximately optimal or good combinations of sparse and low-rank approximation of the matrices . The sparse + low-rank matrix structure has been well studied in statistics and signal processing since the late 2000s [ 9 ] . This structure naturally generalizes low-rank [ 33 , 63 ] , and sparse [ 61 ] matrices . Scatterbrain is built on a line of work , e.g. , Bigbird [ 71 ] , Longformer [ 5 ] with the theme of combining multiple types of attention and another one in the optimal transport setting [ 37 ] . 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 we 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 , as they are 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 are complementary , and their ideal combination ( via Robust PCA ) can outperform both ( Section 3.1 ) . 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 . 1 ( left ) , 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 ) . 2Scatterbrain code is available at https : //github.com/HazyResearch/scatterbrain 2 . An ideal combination of sparse and low-rank ( orange line in Fig . 1 left ) , obtained with robust PCA , can achieve lower error than both . Similar observations on other benchmarks and details are presented in Appendix B . 3.2 A Generative Model of How Sparse + Low-rank Structure Can Arise . Sparse + low-rank parameterization is more expressive than either sparse or low-rank alone . Indeed , in the Appendix , we construct a family of attention matrices to show the separation between the approximation capability of sparse + low-rank vs. sparse or low-rank alone : for an n⇥n attention matrix , sparse or low-rank alone requires a O ( n2 ) parameters to get ✏ approximation error in Frobenius norm , while sparse + low-rank only requires O ( n ) parameters . Moreover , we argue here that sparse + low-rank is a natural candidate to approximate generic attention matrices . We describe a generative model of how the sparse + low-rank structure in attention matrices could arise when the elements of 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 3.1 . Generative process of clustered elements in input sequence We describe here a generative model of an input sequence to attention , parameterized by the inverse temperature 2R and the intra-cluster distance 2R . Process 1 . Let Q 2 Rn⇥d , where d ⌦ ( log3/2 ( n ) ) , with every row of Q generated randomly as follows : 1 . For C =⌦ ( n ) , sample C number of cluster centers c1 , ... , cC 2Rd independently from N ( 0 , Id/ p d ) . 2 . For each cluster around ci , sample ni = O ( 1 ) number of elements around ci , of the form zij = ci+rij for j =1 , ... , ni where rij ⇠N ( 0 , Id / p d ) . Assume that the total number of elements is n=n1+···+nC and O ( 1/log 1/4n ) . Let Q be the matrix whose rows are the vectors zij where i=1 , ... , C and j =1 , ... , ni . Let A=QQ > and let the attention matrix be M =exp ( ·A ) . We visualize this generative process in Fig . 2 . Softmax temperature and approx . error We characterize when to use sparse , low-rank , or sparse + low-rank to approximate the attention matrices in Process 1 , depending on the inverse temperature . The intuition here is that the inverse temperature corresponds to the strength of interaction between the clusters . If is large , intra-cluster interaction dominates the attention matrix , the softmax distribution is peaked , and so we only need a sparse matrix to approximate the attention . If is small , then the inter-cluster attention is similar to intra-cluster attention , the softmax distribution is diffuse , and we can approximate it with a low-rank matrix . In the middle regime of , we need the sparse part to cover the intra-cluster attention and the low-rank part to approximate the inter-cluster attention . We formalize this intuition in Theorem 1 ( in bounds below we think of ✏ as a constant ) . All the proofs are in Appendix D. Theorem 1 . Let M , be the attention matrix in Process 1 . Fix ✏2 ( 0,1 ) . Let R2Rn⇥n be a matrix . Consider low-rank , sparse , and sparse + low-rank approximations to M . 1 . High temperature : Assume =o ( logn/logd ) . ( a ) Low-rank : There exists R with no ( 1 ) rank ( and hence n1+o ( 1 ) parameters ) such that kM RkF ✏n . ( b ) Sparse : If R has sparsity o ( n2 ) , then kM RkF ⌦ ( n ) . 2 . Mid temperature : Assume ( 1 2 ) logn O ( logn ) . ( a ) Sparse + low-rank : There exists a sparse + low-rank R with n1+o ( 1 ) parameters with kM RkF ✏n . ( b ) Low-rank : If R is such that n rank ( R ) =⌦ ( n ) , then kM RkF ⌦ ( n ) . ( c ) Sparse : If R has sparsity o ( n2 ) , then kM RkF ⌦ ( n ) . 3 . Low temperature : Assume =⌦ ( logn ) . ( a ) Low-rank : If R is such that n rank ( R ) =⌦ ( n ) , then kM RkF ⌦ ( e ( 1 2 ) ) . ( b ) Sparse : There exists R with sparsity O ( n ) such that kM RkF ✏·e ( 1 2 )

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.