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