FMMformer: Efficient and Flexible Transformer via Decomposed Near-field and Far-field Attention

1 Introduction . Transformers [ 58 ] have achieved state-of-the-art performance in sequence processing tasks , including machine translation and language modeling [ 58 , 2 , 15 , 4 , 61 , 16 , 9 ] . Also , transformers can effectively transfer knowledge from a pre-trained model to tasks with limited supervision [ 43 , 44 , 16 , 64 , 34 ] . Transformers rely on the attention mechanism and particularly self-attention as a fundamental building block for their modeling [ 5 , 58 , 27 ] . 1.1 Self-attention . The self-attention mechanism is used to learn long-range dependencies while enabling parallel processing of the input sequence . For a given input sequence X : = [ x1 , x2 , · · · , xN ] > ∈ RN×Dx of ∗Equal contribution and Co-first author †Please correspond to : wangbaonj @ gmail.com or chenlong @ math.uci.edu 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . N feature vectors that have been encoded in a Dx-dimensional vector space , self-attention transforms X into an output sequence V̂ in the following two steps : Step 1 . Project the input sequence X into three matrices via the following linear transformations Q = XW > Q ; K = XW > K ; V = XW > V , where WQ , WK ∈ RD×Dx , and WV ∈ RDv×Dx are the weight matrices . We denote Q : = [ q1 , · · · , qN ] > , K : = [ k1 , · · · , kN ] > , and V : = [ v1 , · · · , vN ] > , where the vectors qi , ki , vi for i = 1 , · · · , N are the query , key , and value vectors , respectively . Step 2 . For each query vector qi for i = 1 , · · · , N , we compute the output vector v̂i as follows v̂i = N∑ j=1 softmax ( q > i kj√ D ) vj , ⇐⇒ V̂ = ( QK > √ D ) V : = AV , ( 1 ) where the softmax function is applied to each row of the matrix ( QK > ) / √ D. For long sequences , the computational time and memory footprint of transformers are dominated by ( 1 ) . It is evident that the memory cost is O ( N2 ) to store the attention matrix A . Also , the computational complexities of computing the matrix-matrix products QK > and AV are bothO ( N2 ) . These limitations impede the application of transformers to many important settings that involve very long sequences [ 33 , 25 , 39 ] . When applying self-attention for long sequence modeling , we have to limit the context window to a reasonable size to make it computationally feasible , limiting the effectiveness of learning long-range dependencies . Efficient transformer models have been proposed , including leveraging sparse and low-rank attention . Many of the existing efficient transformers gain computational and memory efficiency at the cost of significant accuracy degradation . 1.2 Contribution . Leveraging the idea of the fast multipole method ( FMM ) [ 19 ] , we propose a class of efficient , flexible , and expressive transformers , namely FMMformers . At the core of FMMformers is to replace the self-attention V̂ = AV in ( 1 ) with the following matrix-matrix product V̂ : = ( D +L ) V , ( 2 ) where D is a banded matrix with bandwidth k N and L is a low-rank matrix of rank r N . In practice , we normalize matrix D + L such that the sum of each row is 1 ; for the sake of presentation , we ignore this normalization step below . Both DV and LV can be computed with linear computational and memory complexity ; they model the near-field and far-field attention , respectively . FMMformers are flexible in designing the sparse banded matrix and the low-rank matrix for modeling near-field and far-field attention . In particular , we can control the bandwidth of the banded matrix D and the rank of the low-rank matrix L for expressivity and efficiency tradeoff . In addition to the efficiency and flexibility , FMMformers gain significant accuracy improvement over linear transformers and can even outperform the standard transformer in terms of accuracy . We illustrate the idea of FMMformers in Figure 1 : Instead of modeling the full attention by a dense unstructured matrix , we employ a sparse banded matrix to model the near-field attention and several rank one matrices to model the far-field attention . 1.3 Organization . We structure this paper as follows : In Sec . 2 , we briefly review the celebrated FMM and establish the connection between FMM and self-attention . In Sec . 3 , we present a practical implementation of FMMformer that leverages existing techniques for low-rank approximation of the self-attention mechanism . We validate and empirically analyze the efficiency and accuracy of FMMformers in Sec . 4 . We discuss related works in Sec . 5 . The paper ends up with concluding remarks . Technical proofs and more experimental details are provided in the Appendix . 2 Fast Multipole Method and Self-attention Mechanism . In this section , we review FMM and present an algebraic interpretation of FMM , see Sec . 2.1 . Then , in Sec . 2.2 , we explore the structure of the attention matrix , showing that FMM can be used to accelerate the self-attention mechanism . 2.1 Fast multipole method vs. sparse and low-rank matrix approximation . FMM is a numerical method that was originally developed to speed up the calculation of longrange forces in the n-body problem [ 19 ] and has been regarded as one of the top 10 algorithms in scientific computing in the 20th century [ 14 ] . The key idea is that the far-field interaction can be well-approximated by separable low-rank matrices , while the near-field interaction can be calculated directly . We use the following simple example to illustrate mathematical reasoning . Without ambiguity , we reuse notations in the previous section and assume : ( A1 ) A ( i , j ) = g ( |qi − kj | ) depends on the distance of two vectors qi and kj , where A ( i , j ) is the ( i , j ) -th entry of the matrix A ∈ RN×N . ( A2 ) The function g ( s ) is smooth for s 6= 0 . ( A3 ) The function g satisfies g ( st ) = g ( s ) g ( t ) . One noticeable example in the physical application is the gravitational potential g ( |qi − kj | ) = 1/|qi − kj | , for which the key vectors { kj } are the location of source particles and the query vectors { qi } are the location of the target points . Assumption ( A3 ) is not essential , which is presented here for the convenience of proof and can be replaced by other separable forms , e.g. , g ( st ) = g ( s ) + g ( t ) . The near-field and far-field are defined through the distance |qi − kj | . We now explain the low-rank approximation based on the well-separated condition . For illustrative purpose , we assume the index set { 1 , 2 , . . . , N } is partitioned into two groups { T1 , T2 } . Definition 1 . Group T1 is called well-separated from T2 if there exists a vector k∗ and a number δ ∈ ( 0 , 1 ) such that |kj − k∗| ≤ δ|qi − k∗| ∀i ∈ T1 , j ∈ T2 . The vector k∗ is a representative vector of { kj , j ∈ T2 } , e.g. , the center of vectors in T2 . For any qi , i ∈ T1 , it is far away from { kj , j ∈ T2 } and the far-field interaction A ( i , j ) , i ∈ T1 , j ∈ T2 can be approximated by g ( |qi − k∗| ) , i.e . each row of A ( T1 , T2 ) , the submatrix of A with the row index set T1 and the column index set T2 , is constant . For example , when calculating the gravitation of a galaxy from the Earth , we can simply treat the galaxy as one single point , although the galaxy may contain hundreds of millions of stars . By including p terms of the Taylor series , the approximation can be more accurate by using rank p instead of rank 1 matrix approximation . Lemma 1 . Let { T1 , T2 } be two well-separated index sets . Assume ( A1 ) - ( A3 ) hold . For any ε > 0 , the sub-matrix A ( T1 , T2 ) can be uniformly approximated by a rank p matrix to a tolerance ε > 0 in the sense that : there exists rank p matrices U ∈ R|T1|×p , V ∈ R|T2|×p , with p ≥ C| logδ | for some positive constant C , such that |A ( i , j ) − ( UV > ) ( i , j ) | ≤ , ∀i ∈ T1 , j ∈ T2 . The applicability of the analytic kernel function g was limited to partial differential equations or integral equations where Green ’ s function satisfying ( A1 ) - ( A3 ) . In the application of machine learning , it is hard to verify ( A1 ) - ( A3 ) . Instead , we use the definition of diagonal-plus-semi-separable matrices from the book [ 6 , Definition 1.10 ] . We use MATLAB/Numpy notation tril ( K , p ) to denote the lower triangular matrix with zeros above the pth-subdiagonal of K and similar notation triu ( K , p ) for the upper triangular part . Definition 2 . [ 6 , Definition 1.10 ] A matrix A ∈ RN×N is called ( p , q ) -semi-separable if there exist matrices U , V ∈ RN×p and W , Z ∈ RN×q such that A = triu ( UV > , 0 ) + tril ( WZ > , 1 ) . It is called diagonal-plus-semi-separable if . A = D + triu ( UV > , 1 ) + tril ( WZ > , 1 ) . with some diagonal matrix D. Definition 2 can be naturally extended to include a banded matrix D and sum of several low-rank matrices . Moreover , One can verify the semi-separable property of matrix K by checking the decay of singular values of the matrix . As often used in low-rank approximation methods , the numerical rank or ε-rank of a matrix K , for a tolerance ε , is the number of singular values of K that are greater than ε‖K‖2 . The low rank approximation relies on the well separateness of two subsets . Based on a hierarchical partition of the index set , aH-matrix [ 21 ] can be constructed ; see Figure 2 for an illustration . Further compression leads toH2-matrix [ 20 , 22 ] and the hierarchically semi-separable ( HHS ) matrix [ 11 , 62 ] . Other variants include hierarchically block-separable ( HBS ) [ 36 ] , and hierarchically off-diagonal low-rank ( HODLR ) [ 3 ] matrices , etc . In our application , we write the decomposition as A = D + r∑ l=1 φl ( Q ) φ > l ( K ) . ( 3 ) In the query and key spaces , the vectors qi and kj may not be well-separated . Then nonlinear feature maps φl ( · ) , l = 1 , · · · , r to higher dimensions , which are trainable , can be used to make the mapped datasets more separable . In ( 3 ) , each kernel function φl : RN → RN is a vector function of length N . We can mimic the hierarchical decomposition used inH-matrix to use low rank kernel approximation φl ( Q ( 1 : N/2 , : ) φ > l ( K ( N/2 + 1 : N , : ) with halved length . Such approximation can be recursively applied to get a multilevel decomposition in the low rank approximation component . Our numerical results show that a simple one level near-field and far-field decomposition is good enough .

Summary: A fast and memory-efficient transformer architecture based on a sparse attention mechanism inspired by the fast multipole method (FFM). It decomposes attention matrices into the sum of a banded matrix (capturing near-range interactions) and a low rank matrix (capturing long-range interactions). This low rank matrix is itself expressed as a sum of rank 1 matrices, each characterized by a given feature map. Experiments on the LRA benchmark and WikiText-103 show promising results in terms of test accuracy/perplexity Vs training/inference speedups.