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

Abstract:
We propose FMMformers, a class of efficient and flexible transformers inspired by the celebrated fast multipole method (FMM) for accelerating interacting particle simulation. FMM decomposes particle-particle interaction into near-field and far-field components and then performs direct and coarse-grained computation, respectively. Similarly, FMMformers decompose the attention into near-field and far-field attention, modeling the near-field attention by a banded matrix and the far-field attention by a low-rank matrix. Computing the attention matrix for FMMformers requires linear complexity in computational time and memory footprint with respect to the sequence length. In contrast, standard transformers suffer from quadratic complexity. We analyze and validate the advantage of FMMformers over the standard transformer on the Long Range Arena and language modeling benchmarks. FMMformers can even outperform the standard transformer in terms of accuracy by a significant margin. For instance, FMMformers achieve an average classification accuracy of 60.74% over the five Long Range Arena tasks, which is significantly better than the standard transformer’s average accuracy of 58.70%.

heading: 1 Introduction
text: 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].

heading: 1.1 Self-attention
text: 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.

heading: 1.2 Contribution
text: 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.

heading: 1.3 Organization
text: 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.

heading: 2 Fast Multipole Method and Self-attention Mechanism
text: 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.

heading: 2.1 Fast multipole method vs. sparse and low-rank matrix approximation
text: 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).

heading: It is called diagonal-plus-semi-separable if
text: 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.

heading: 2.2 Sparse and low-rank patterns in attention maps
text: In this section, we explore the sparse and low-rank structure of the attention matrix A. In particular, we consider the attention matrix A ∈ R256×256 obtained from the standard transformer trained for
WikiText-103 language modeling; see Sec. 4.3 for the experimental details. We randomly select 1000 different attention matrices, and we exclude a banded matrix D with bandwidth 5, 10, and 20 from each of such matrices. Then, we perform singular value decomposition (SVD) to compute the rank of each matrix A−D, and we threshold the small singular values with a magnitude of 10−6. Figure 3 (top row) plots two randomly selected self-attention matrices and the distribution of the rank of the matrix A−D. It is clear that matrix A has only a few large singular values and all other singular values are very small. Moreover, matrix A −D is of low rank, and the rank becomes smaller in general when the bandwidth of D increases, which is consistent with the assumptions in Sec. 2.1, motivating FMMformers.

heading: 3 FMMformer: Practical Near-field and Far-field Attention
text: In this section, we present practical algorithms for implementing the proposed FMMformer defined by (2). In particular, we present fast algorithms for computing the near-field attention DV and the far-field attention LV .

heading: 3.1 Banded matrix modeling of near-field attention
text: We model the near-field attention with the following banded matrix
D = softmax ( bandk (QK>√ D )) , (4)
where the operator bandk(∗) represents taking only the banded part of the matrix ∗ with a bandwidth k (k N ). In practice, there is no need to calculate the matrix product QK>. Instead, we only need to calculate the products of the vectors that correspond to the nonzero entries of the banded matrix bandk(QK>/ √ D). Note that for long sequences, both the time and memory complexity of computing (4) are O(kN).

heading: 3.2 Low-rank matrix modeling of far-field attention
text: We consider practical and efficient low-rank matrix modeling of the far-field attention LV in (1). In principle, any existing off-the-shelf low-rank attention can be integrated into FMMformer to model the far-field attention. In particular, we model the far-field attention by leveraging the kernel trick used in [26, 13, 47], which is flexible in selecting different kernels to modulate the rank of the far-field attention.

heading: 3.2.1 Low-rank attention via kernelization
text: Suppose we model the far-field attention using a rank r matrix L ∈ RN×N , which can be written as the sum of r rank one matrices, i.e.,
L = a1b > 1 + a2b > 2 + · · ·+ arb>r , (5)
where a1,a2, · · · ,ar; b1, b2, · · · , br ∈ RN . Note that LV = (a1b > 1 + a2b > 2 + · · ·+ arb>r )V = a1(b>1 V ) + a2(b>2 V ) + · · ·+ ar(b>r V ), (6)
which indicates that we can compute LV with O(rN) time complexity. Also, we only need to store the vectors u1, · · · ,ur; v1, · · · ,vr ∈ RN , resulting in linear complexity in memory footprint. We borrow the idea of kernelization from the linear transformer [26] for practical implementation of (6). In particular, the authors in [26] generalize the softmax function in (1) to a general kernel function k(qi,kj), i.e.,
v̂i = ∑N j=1 exp(qi,kj)vj∑N j=1 exp(qi,kj)︸ ︷︷ ︸
self-attention
3 =⇒ v̂i = ∑N j=1 k(qi,kj)vj∑N j=1 k(qi,kj)︸ ︷︷ ︸
generalized self-attention
. (7)
If k(qi,kj) = φ(qi)>φ(kj) for a certain feature map φ(·), then we have
v̂i = ∑N j=1 k(qi,kj)vj∑N j=1 k(qi,kj) = ∑N j=1 φ(qi) >φ(kj)vj∑N j=1 φ(qi) >φ(kj) = φ(qi) >∑N j=1 φ(kj)v > j φ(qi)> ∑N j=1 φ(kj) , (8)
3Here, exp(qi,kj) := exp(q>i kj/ √ D).
Note that (8) can be regarded as a rank one approximation of self-attention. We can rewrite (8) into the following compact form
V̂ = φ(Q)(φ(K)>V )
φ(Q)φ(K)> . (9)
To generalize (8) to the rank r approximation, we select a set of linearly independent feature maps {φl(·)}rl=1. Together with the sparse banded matrix modeling of the near-field attention, we propose the following efficient attention model for the FMMformer
V̂ = DV + r∑ l=1 φl(Q)(φl(K) >V ) φl(Q)φl(K)> . (10)
It is evident that both computational time and memory complexity are linear in computing (10). Our design is flexible to selecting feature maps and the sparse banded matrix, which the users can customize. Moreover, causal masking can be implemented easily by truncating the sum from 1 to i in (8) together with masking out the corresponding part of the banded matrix D.
Proposition 1. Let φl(x) ∈ RN (l = 1, 2, · · · , r and r N ) for x ∈ Rn. If {φl(x)}rl=1 are linearly independent at x, then the following matrix L(x) ∈ RN×N has rank r,
L(x) := φ1(x)φ1(x) > + φ2(x)φ2(x) > + · · ·+ φr(x)φr(x)>. (11)
Feature map selection. The feature map selection is crucial for the success of far-field attention modeling. In this work, we adopt the existing successful feature map φ1(x) := elu(x)+1 used in the linear transformer [26] together with φ2(x) := elu(−x) + 1, which is a straightforward modification of φ1(x). Moreover, we consider the third feature map φ3(x) := tanh(x). It is easy to check that φ1(x), φ2(x), and φ3(x) are linearly independent for almost all x. We leave how to design a set of feature maps to optimize the far-field attention modeling as future work.

heading: 3.3 Blending of near-field and far-field attention
text: Based on our experiments, adding a learnable weight in front of each attention component benefits training and generalization. As such, we propose the following scheme to blend the near-field attention and far-field attention
V̂ := (w1D + w2L)V , (12)
where w1 and w2 are two learnable weights, and we enforce their positivity via a sigmoid map.

heading: 4 Experimental Results
text: In this section, we numerically verify the efficiency of FMMformers and empirically analyze the effects of near-field and far-field attention on various benchmarks, including synthetic sequence copy (Sec. 4.1), Long Range Arena (LRA) (Sec. 4.2), and language modeling (Sec. 4.3). We aim to show that: (i) FMMformers are efficient in both computational time and memory footprint. (ii) Multiple kernels benefit learning of the far-field attention. (iii) Blending near-field attention with far-field attention can boost the performance of linear transformers. Throughout this section, we compare FMMformers with linear transformers (linear, r = 1 in (11)), standard softmax transformers (softmax), and softmax transformers that use a banded attention matrix of bandwidth k (bandk). All experiments are conducted on a server with 4 NVIDIA 3090TI GPUs.

heading: 4.1 Synthetic sequence copy task
text: We first consider a synthetic copy task with various sequence lengths, including 128, 256, and 512. In this task, the model has to duplicate a sequence of symbols. Each training and test sample is a sequence of maximum length 128/256/512 with ten different symbols separated by a dedicated separator symbol. We train all transformers for this task using the same setting as in [26].
Boosting performance of linear transformers with near-field attention. We first compare FMMformers, obtained by blending the linear transformer with a banded attention matrix of bandwidths 10, 20, and 30, respectively. Figure 4 shows that for shorter sequences of length 128, all transformers reach similar loss; the standard softmax transformer converges much faster than the linear transformer while blending the linear transformers with near-field attention can improve training. Moreover, the benefits of near-field attention become more significant as the sequence length increases.
Enhancing far-field attention with multi-kernels. After observing that the linear transformer performs poorly as the sequence length increases, we consider augmenting the linear transformer with multiple feature maps; in particular, we consider the three feature maps mentioned above, i.e., φ1(x) = elu(x) + 1, φ2(x) = elu(−x) + 1, and φ3(x) = tanh(x). Figure 5 compares different transformers on different sequence lengths, where rank 2 consists of the feature maps φ1(x) and φ2(x), and rank 3 consists of all three feature maps. These results show that multiple kernels can improve the learning of far-field attention.
Computational and memory complexity. In this part, we compare different transformers in computational time and memory cost. Following [26], we compute the attention and gradient for input sequences with different lengths N ∈ {29, 210, · · · , 216} and measure the peak allocated GPU memory and the required time for each transformer model. We conduct this experiment on an NVIDIA 3090TI with 24GB memory, and we report the time and memory cost per sample in the same way as in [26]. Figure 6 contrasts the time (left) and memory (right) costs of different models.

heading: 4.2 Long Range Arena (LRA) Benchmark
text: In this experiment, we evaluate our model on tasks that involve longer sequence lengths in the Long Range Arena benchmark [54]. We show that the FMMformer outperforms the baseline linear transformer and standard softmax transformer [58], justifying the advantage of the FMMformer in capturing long-range dependencies. We provide model and training details in the Appendix.
Datasets and metrics. We consider all five tasks in the LRA benchmark, including Listops [38], byte-level IMDb reviews text classification [35], byte-level document retrieval [42], CIFAR-10 image classification on sequences of pixels [30], and Pathfinder [32]. These tasks involve long sequences of length 2K, 4K, 4K, 1K, and 1K, respectively. We follow the setup/evaluation protocol in [54] and report the test accuracy for individual task and the average result across all tasks.
Results. We summarize our results in Table 1. Like in the copy task, we observe that adding near-field attention modeled by banded attention matrices improves the performance of linear transformers. More interestingly, using bandwidth 5 already yields good results across all LRA tasks while significantly reducing the computational and memory cost of calculating the attention matrix. For example, in the byte-level document retrieval [42] task, a banded matrix with bandwidth 5 only accounts for 0.125% of the corresponding full attention matrix. The FMMformer with 1 kernel (blending a banded matrix of bandwidth 5 with the linear transformer using feature map φ1(x)) outperforms the linear transformer and yields similar or better results than the standard softmax transformer in all tasks. Furthermore, the FMMformer with 2 kernels (blending a banded attention matrix of bandwidth 5 with the linear transformer using feature maps φ1(x) and φ2(x)) further improves the FMMformer with 1 kernel, justifying the need of better low-rank approximation for the far-field attention. Across tasks, the FMMformer obtains the best average accuracy. Also, it is worth noting that tasks in the LRA benchmark cover different data modalities include text and images. Good performance of the FMMformer on these tasks demonstrates that the advantages of our model over the linear and standard transformers are consistent across data modalities.

heading: 4.3 Language Modeling on WikiText-103
text: Experiments on the copy task in Sec. 4.1 illustrate the effect of combining near-field and far-field attention. Results on the LRA benchmark in Sec. 4.2 show the ability of our FMMformer to capture very long-range dependency and extend to different data modalities. Now our goal is to confirm the advantage of the FMMformer on a large-scale application. We consider the word-level language modeling task on WikiText-103 [37].
Datasets and metrics. WikiText-103 consists of articles from Wikipedia and is a dataset with long contextual dependencies. The training set is made up of about 28K articles containing 103M running words; this corresponds to text blocks of about 3600 words. The validation and test sets are composed of 218K and 246K running words, respectively. Each of them contains 60 articles and about 268K words. Our experiment follows the standard setting [37, 47] and split the training data into L-word independent long segments. For evaluation, we use a batch size of 1, and go through the text sequence with a sliding window of size L. We consider only the last position for computing perplexity (PPL) except in the first segment, where all positions are evaluated as in [2, 47].
Results. Table 2 shows the validation and test perplexity of our models versus the linear and standard softmax transformer on WikiText-103. Here, we also compare with the linear attention with the fast weight trick proposed in [47]. Consistent with previous experiments, the FMMformer outperforms the linear transformer with or without fast weight. The standard softmax transformer obtains the best results in this task, but the gap between the FMMformer and the standard transformer is very small when a larger bandwidth is used for near-field attention in the FMMformer. This is justified by the improvement in terms of PPL of the FMMformer with a near-field attention of bandwidth 20 compared to the FMMformer with a near-field attention of bandwidth 5. Also, FMMformer with 2 kernels (φ1(x) and φ2(x)) still improves over FMMformer with 1 kernel (φ1(x)). Consider the linear complexity of computational time and memory advantage of FMMformers, the small performance gap of FMMformers to standard softmax transformers can potentially be overcome by using the multilevel or hierarchical near-field and far-field decomposition.

heading: 5 Related works.
text: Low-rank transformers. Low-rank approximation of the self-attention matrix A has been a popular method in reducing the quadratic computational and memory complexity of transformers to linear. Linearized attention that leverages kernelization tricks can be considered as the rank one approximation of the self-attention matrix [60, 26, 13, 49]; the choice of the feature map function is crucial for the success of linearized attention. Fast weight memories [48] have been used to improve memory capacity of linearized attention [47]. The Nyström method has also been leveraged for developing efficient attention with linear computational complexity [63]. Many other low-rank attention models exist, e.g., [8, 45, 50, 40]. FMMformers employ low-rank attention to model far-field attention; in principle, the merits of existing low-rank attention can be integrated into FMMformers.
Sparse transformers. Attention matrices have been enforced with different sparsity patterns to gain efficiency, including fixed sparsity patterns [41, 39, 7, 1, 65], a combination of different sparsity patterns [12, 24], and data-dependent/learnable sparsity patterns [33, 60, 52, 53, 29, 46, 59]. Informer [66] is another efficient attention scheme using a sparse query. Note that the existing sparsity pattern can be very complicated, while we adopt a sparse banded matrix to model the near-field attention.
Other efficient transformers. Reformer reduces the cost of self-attention to O(N logN) via locality-sensitive hashing [29]. Linformer [60] and Longformer [7] obtain the linear complexity using random projection and local window attention, respectively. Galerkin transformer in [10] uses a Galerkin self-attention for the encoder. See [55] for a review of efficient transformers.
Sparse and low-rank interpretation of FMM. The fast multipole method is introduced by Greengard and Rokhlin [19] for the efficient computation of gravitational/electrostatic potentials and fields. Applications of FMM to machine learning can be found in [18, 31, 57]. FMM can be generalized algebraically for efficient computation of the dense matrix-vector product. Algebraic counterpart of FMM includeH-matrix [21],H2-matrix [20, 22], hierarchically semi-separable (HSS) [11, 62], hierarchically block-separable (HBS) [36], and hierarchically off-diagonal low-rank (HODLR) [3] matrices. The common feature is to compress the off-diagonal sub-matrices by low-rank approximations.

heading: 6 Concluding Remarks
text: In this paper, we proposed FMMformers, a class of efficient and flexible transformers with linear time and memory complexity, inspired by the fast multipole method. In FMMformers, we decompose the full attention into near-field and far-field attention; we model the near-field attention with a sparse banded matrix and model the far-field attention using a low-rank matrix leveraging ideas of the linear transformer [26]. We validate the efficiency of FMMformers on various benchmark tasks, including synthetic sequence copy, LRA benchmark, and WikiText-103 language modeling. Our numerical results show that FMMformers consistently outperform the linear transformer on all benchmarks and outperform the standard softmax transformer on the LRA tasks. In our work, we select linearly independent feature maps to enhance the learning of far-field attention. It is natural to ask how to design a set of feature maps to optimize the performance of FMMformers? Furthermore, we leave the application of FMMformers for improving the vision transformer [17, 56] as future work.

heading: 7 Acknowledgement
text: This material is based on research sponsored by NSF grants DMS-1924935, DMS-1952339 and DMS-2012465, DOE grant DE-SC0021142, and ONR grant N00014-18-1-2527 and the MURI grant N00014-20-1-2787. We thank Professor Jack Xin for helpful discussions.

Review summaries:
The paper proposed the new efficient transformer (FMMformers), which decomposes the attention computation into direct and coarse-grained computations. The coarse-grained computation is handled by the low-rank matrix and the direct computation is handled by a banded matrix. The proposed method shows better performance on the LRA dataset and comparable performance on the language modeling with the original transformer  

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.

The work proposes FMMformer, which borrows ideas from fast multipole method (FMM) to approximate the quadratic Softmax attention with a linear complexity. Specifically, the approximation consists of a near-field Softmax term (i.e. local attention) and a far-field multi-kernel linear attention. Some theoretical analysis is included to support the idea. The proposed model is evaluated on one synthetic copy task, long range arena and wikitext-103 language modeling. On the three tasks, the proposed model has a better quality compared to linear attention, though the gap between FMMformer and the standard transformer is still significant.