Self-attention mechanism has been shown to be an effective approach for capturing global context dependencies in sequence modeling, but it suffers from quadratic complexity in time and memory usage. Due to the sparsity of the attention matrix, much computation is redundant. Therefore, in this paper, we design an efficient Transformer architecture, named Fourier Sparse Attention for Transformer (FSAT), for fast long-range sequence modeling. We provide a brand-new perspective for constructing sparse attention matrix, i.e. making the sparse attention matrix predictable. Two core submodules are: (1) A fast Fourier transform based hidden state cross module, which captures and pools L 2 semantic combinations in O ( L log L ) time complexity. (2) A sparse attention matrix estimation module, which predicts dominant elements of an attention matrix based on the output of the previous hidden state cross module. By reparameterization and gradient truncation, FSAT successfully learned the index of dominant elements. The overall complexity about the sequence length is reduced from O ( L 2 ) to O ( L log L ) . Extensive experiments (natural language, vision, and math) show that FSAT remarkably outperforms the standard multi-head attention and its variants in various long-sequence tasks with low computational costs, and achieves new state-of-the-art results on the Long Range Arena benchmark. Models based on the Transformer architecture (Vaswani et al., 2017) have been firmly established as state of the art approaches across a range of domains like language (Brown et al., 2020; Clark et al., 2020; Devlin et al., 2018), and vision (Carion et al., 2020; Dosovitskiy et al., 2020). The Transformer architecture perceiving long-range context heavily relies on the multi-head self-attention mechanism, in which the relevance of every token pairs is computed to decide the attention scores and to-ken's representations are the weighted average of all tokens using the attention scores. Despite its effectiveness, self-attention mech-anism's quadratic time and memory complexity about the sequence length is an obstacle to extend Transformer for very long sequences, such as document-level text tasks, high-resolution images, videos, etc. Shen et al. (2021, 2018) elaborate the issue of high computational complexity. For instance, more than 68GB GPU memory and 1.6T multiply-accumulation operations are required for a 64  64  32 3D feature volume. Great efforts have been made to develop Trans-former's variants for long-range sequence modeling tasks. Tay et al. (2020c) categorize the researches of efficient Transformers: (a) Fixed patterns or combination of patterns (Beltagy et al., 2020; Zaheer et al., 2020), in which the field to be attended is pre-defined by fixed pattern. (b) Learnable patterns (Kitaev et al., 2020; Tay et al., 2020a), in which tokens are sorted or clustered in a data-driven fashion. (c) Memory (Ma et al., 2021; Lee et al., 2019), in which spacial tokens with global view are introduced to compress the input sequence. (d) Low-rank methods (Tay et al., 2021; Wang et al., 2020), which adopt low-rank approximations of the self-attention matrix. (e) Kernels (Katharopoulos et al., 2020; Choromanski et al., 2020a,b), which view the attention mechanism through kernelization. (f) Recurrence (Rae et al., 2019), which connects multiple segments via recurrence structure. Despite their variety, approximating the quadratic-cost attention matrix by the sparsity of attention matrix is the common idea. In this paper, we propose predictable sparse attention, and name it as Fourier Sparse Attention for Transformer (FSAT) due to fast Fourier transform is a key operation in our method. FSAT is a brand-new perspective of efficient Transformer, i.e. learning the sparse structure of an attention ma-234 trix in end-to-end fashion. Specifically, we firstly compute the semantic relevance of token pairs and then use it to predict the indices of dominant (non-zero) elements of an attention matrix, and finally attention scores are filled according to the predicted sparse structure. In this process, two problems have to be solved: (1) Efficiently capturing semantic relevance of L 2 token pairs where L is the length of an input sequence. (2) Learning discrete indices with gradient descent algorithm. To this end, we propose pooled hidden state cross to efficiently calculate and compress semantic relevance in O ( L log L ) time complexity. For end-to-end training, we get continuous and meaningful gradients for learning discrete indices by reparameterization and gradient truncation. Consequently, FSAT is out of the scope of Tay et al. (2020c)'s taxonomy. It's worth noting that predictable sparse attention is different from the methods of learnable patterns. Although these methods use learnable algorithms to sort or cluster tokens, they still exploit fixed patterns (chun-ked patterns). Instead, FSAT directly predicts the sparse structure of an attention matrix without any pre-defined pattern. In order to fit the predicted sparse attention matrix, the key and value vectors in self-attention mechanism are projected from pooled hidden state cross vectors, which can be viewed as 2-order features of the tokens. As an extra benefit, model's expressiveness may increase. Therefore, unlike some efficient Transformer variants which approximate the quadratic-cost attention matrix at the expense of accuracy, FSAT not only reduces computational complexity but also improves model accuracy in some tasks. On Long Range Arena benchmark, FSAT outperforms the Transformer and several recent efficient self-attention methods by a large margin. To summarize, our contributions are as follows:  We propose Fourier Sparse Attention for Transformer (FSAT) to extend Transformer for long sequences. The overall complexity about the sequence length is reduced from O ( L 2 ) to O ( L log L ) . We introduce the pooled hidden state cross to implement FSAT. Empirically, extensive experiments (natural language, vision, and math) demonstrate the advantages of our proposed methods, and new state-of-the-art results are achieved on the Long Range Arena benchmark. Tay et al. (2020c) have provided a comprehensive overview of existing efficient Transformers. Some promising models are compared with our method in the experiments. Big bird (Zaheer et al., 2020) uses random, sliding window and global attention to build hybrid attention pattern. Performer (Choromanski et al., 2020a,b) utilizes orthogonal random features to approximate softmax-attention kernels with linear complexity. Linformer (Wang et al., 2020) achieves linear complexity by adopting random projections based on the JL lemma to compress the attention length to a fix length. Longformer (Beltagy et al., 2020) combines local windowed attention with task-motivated global attention for long documents. Reformer proposed in Kitaev et al. (2020) clusters similar tokens by locality-sensitive-hashing, and dot-product attention is performed inside clusters. Feature crosses, which synthesize crossing combinations of features, is a widely used technique for extending features' predictive ability in machine learning. For example, Takahashi et al. (2018) demonstrate their gender identification system leveraging synergy of both texts and images by feature cross technique. Yu et al. (2018) and Seo et al. (2016) utilize crossed feature to design a trilinear attention function. Chen et al. (2021) explore how to search the best feature crosses by submodular optimization. More researches involving feature crosses focus on feature selection (Zadeh et al., 2017; Wei et al., 2015; Hoque et al., 2014; Nie et al., 2010; Guyon and Elisseeff, 2003; Kwak and Choi, 2002; Rogati and Yang, 2002; Weston et al., 2000). Recently, Fourier transform in Transformer has garnered interest. Choromanski et al. (2020a,b) propose Performer by approximating softmax attention-kernels via orthogonal random Fourier features. Tamkin et al. (2020) propose BERT + Prism model using spectral filters in the activations of neurons for producing multi-scale representations, and got positive experimental results at 235 1 ( ) i f x W 1 W 2 i x ij c k c Sum-pooling (via FFT) L DW 1 W 2 L j x 2 ( ) j f x i x  j x  Figure 1: The process of feature mapping, hidden state cross, and sum-pooling along antidiagonals, corresponding to Formula More radically, Lee-Thorp et al. (2021) reform Transformer by replacing the entire self-attention sub-layer with discrete Fourier transforms along sequence dimension and hidden dimension respectively. In this section, we start by explaining the motivation of introducing pooled hidden state cross, then introduce how to compute pooled hidden state cross, and finally discuss the way to equip self-attention with pooled hidden state cross. Three desiderata motivate our use of pooled hidden state cross: (1) Long-range semantic dependency and relevance can be captured by hidden state crosses, since the combinations of every token pair have been included. Capturing semantic relevance is also the basis of the predictable sparse attention proposed in the next section. (2) Hidden state cross is a way to extract 2-order token features, intuitively it may generate more expressive feature representations. (3) Crossing and pooling hidden states are conducted depth-wisely, so that they can be efficiently implemented via fast Fourier transform. Inspired by the feature cross technique, we propose the concept of hidden state cross. Briefly speaking, feature cross technique (a.k.a feature combination) synthesizes new feature xy by multiplying feature x and feature y . We extend it to the level of hidden state of deep learning models. Specifically, given the hidden states of a token sequence x 0 ,    , x L  1 , we define c ij = f 1 ( x i )  f 2 ( x j ) (1) as the crossed hidden state vector of the i -th and the j -th token, where f (  ) is a parameterized nonlinear feature mapping function, subscripts indicate containing different parameters, and  is the Hadamard product. We expect that the semantic combination of two tokens can be learned and encoded in the crossed hidden state vector. Problems arise when computing the hidden state crosses of L 2 token pairs. Firstly, the computational complexity is O ( L 2 ) about the sequence length L , which is computationally prohibitive for long sequences. Secondly, the output should be L 2 vectors which is too large to be attended in the Transformer model. To alleviate the problems, crossed hidden states are sum-pooled in this paper. Figure 1 illustrates the computation. The pooled hidden state cross c k represents the sum of crossed hidden states along the k -th antidiagonal. Therefore, the output vectors are compressed from L 2 vectors, i.e. { c ij } i,j  [0 ,L  1] , to 2 L  1 vectors, i.e. { c k } k  [0 , 2 L  2] . Formula 2 can be efficiently implemented by Fast Fourier Transform (FFT). Specifically, hidden states are firstly non-linearly converted by the feature mapping functions, and then 1D discrete Fourier transform is applied along the sequence dimension to transform the mapped hidden states into frequency domain, crossing and pooling are then conducted via multiplication in frequency domain, finally by applying inverse 1D discrete Fourier transform the pooled hidden state crosses are transformed back from frequency domain. By the Hermitian property, the imaginary part of the output is zero. Thus, we can safely only keep the real part of the output and avoid involving complex numbers into the model. Formally, C 0 =  ( F  1 ( F ( f 1 ( X ))  F ( f 2 ( X )))) (3) in which, X  RL  D denotes the matrix consisting of the L D -dimensional hidden states. are 1D Fourier transform and inverse 1D Fourier transform respectively. means keeping the real part of complex numbers. C 0  R (2 L  1)  D is the output matrix of pooled hidden state crosses. The computational complexity about the sequence length is reduced from O ( L 2 ) to O ( L log L ) by FFT. Formula 3 has the output matrix of shape (2 L  1)  D , which doubles the sequence length. To keep the computational complexity of attention not increasing, the length is needed to be reduced. As is shown in Figure 1, sum-pooling along antidiagonal produces a symmetric token combination (cross) about a central token. Specifically, in even-numbered antidiagonals, the central token is the k 2 -th token. For instance, the token combinations of the 10th antidiagonal include token 5-5, 4-6, 3-7, and so on. In odd-numbered antidiagonals, the symmetric center is between  k 2  -th token and  k 2  -th token. For instance, the 11th antidiagonal includes the token combinations of token 5-6, 4-7, 3-8, etc. Therefore, we can reduce the length by merging the consecutive even-numbered and odd-numbered antidiagonals so that token combinations of near symmetric centers are together. In the previous example, token combinations of the 10th antidiagonal and the 11th antidiagonal are summed up. Besides, token combinations of two same tokens are subtracted, e.g. token combination 5-5. Formally, C = LN ( C 1 + C 2  f 1 ( X )  f 2 ( X )) (4) where C 1  RL  D (padding a row of zero to align its the length to C 2 ) and C 2  RL  D are the odd-numbered and even-numbered rows of matrix C 0 respectively, LN denotes layer normalization which ensures stable training, C  RL  D is the output. In this section, we revise the multi-head self-attention to utilize our pooled hidden state cross. The output of an attention layer is calculated as follows. where d is the dimension of a single head, superscript h denotes the h -th head, (cid:81) denotes the concatenation operation of H heads along the last dimension,  is a row-wise scoring function (e.g. softmax), W o  R Hd  D is the output projection matrix. It's worth noting that the key, and value in attention mechanism are revised. Due to the sparsity of attention matrix, most elements of the attention matrix are close to zero, for the sake of simplicity, we call those elements, which are much greater than zero, dominant elements. Base on the pooled hidden state cross, we propose predictable sparse attention, which predicts dominant elements of attention matrix, to avoid computing the full attention matrix. In this section, we describe the predictable sparse attention by a weighted directed sparse graph, in which the vertexes are the L query/key vectors of the input sequence, its directed edges represent that the head vertex attends to the tail vertex in the attention mechanism, and two weights (i.e. attention score and confidence) are assigned to each edge. Figure 2 illustrates a sub-graph with a single key vertex and its in-neighbors. The attention matrix is the adjacency matrix of the graph. For the multihead attention, each head has a graph computed independently. The i -th output vector of the pro-237 posed predictable sparse attention is defined as A ( X , C ) i = ( H (cid:89) h =1 (( q hi K h TN hi  d )  s hi  N hi ) V hN hi ) W o (7) where row vector q h , and matrices K h , V h  RL  d are respectively the query, key, and value projected from X or C following Formula 6, N hi represents the out-neighbors set of the i -th vertex in the h -th head's directed graph, when N hi is written at subscript, it means only extracting the ma-trix's rows corresponding to the vertexes in N hi , and s hi  N hi represents the confidence vector consisting of the confidence scores of the edges pointing from i -th vertex to the vertexes in N hi in the h -th head's directed graph. The challenge of the predictable sparse attention is to find out which elements in the attention matrix are dominant under the condition of not computing the full attention matrix. We utilize pooled hidden state cross, because semantic combination vectors contain the information of the relevance of token pairs. We introduce attention confidence to help the model learning the sparse structure of an attention matrix. Specifically, we define the confidence of the i -th query vector q i attending to the j -th key vector k j as s i  j =  ( i |  I j ,  2 ) (8) in which,  denotes the probability density function of Gaussian distribution,  I j is the index of the dominant query vector, which attends to the key vector k j with a dominant attention score (i.e. edge  I j  j corresponds to a dominant element in the attention matrix),  2 is a hyper-parameter representing the variance. We have this definition because of the observation that the query vectors far away from the dominant query vector have decreasing probabilities of attending to the key vector. The key to make a sparse attention matrix predictable is how to back-propagate gradients through the predicted discrete indices. In this paper, discrete indices are reparameterized. A dominant index matrix I  RL  M is predicted based on the pooled hidden state cross C : I =  ( CWI + b I )  L max (9) where WI  RD  M and b I  RM are the learnable weight and bias respectively,  (  ) is the sigmoid function, L max is the maximum sequence length that the model supports. Since there may be multiple dominant query vectors for a key vector, the hyper-parameter M presumes the maximum number of dominant query vectors for a single key vector. Given the sparse graph described by an index matrix I  NL  M , whose value in the j -th row m -th column I jm indicates that there is a directed edge pointing from the I jm -th query vector to the j -th key vector, the confidence score of each edge can be calculated as follows. where  I jm is the m -th predicted dominant index of the j -th key vector predicted by Formula 9. Therefore, applying the chain rule, the gradient of confidence scores from a loss function can continue to be propagated through Formula 9 and Formula 10 to matrix C . The index matrix I  NL  M decides which edges are considered and which edges are ignored in the sparse graph. In this paper, two types of index matrix are adopted, a predicted index matrix I p =  I  and a random index matrix I r  U [0 ,N  1] . The process for learning sparse attention matrix can be viewed as a process for searching right indices, it is a process of exploring new knowledge and exploiting existing knowledge. Therefore, in training, the sparse graph is decided by the union of I p (ex-ploitation) and I r (exploration), and, in inference, only I p is used. When back-propagation, the gradient of attention confidence is truncated into the range (  , 0 ] for stable convergence. Because, for gradient descent algorithm, positive gradients will decrease the confidence values on edges, it means that the gradients prevent the model from considering these edges in sparse attention mechanism, and tune the model's parameters to change its predicted dominant-indices. But due to the discreteness of indices, changing the predicted dominant-indices 238 Models ListOps Text Retrieval Image Pathfinder Avg. slightly larger or smaller does not ensure moving closer to the correct dominant indices. On the contrary, negative gradients indicate hitting the correct dominant-indices, and the model should be tuned using these gradients. In terms of computational complexity, the proposed predictable sparse attention has lower computational cost in time and memory usage. Specifically, the computational complexity for pooled hidden state cross includes the feature mapping O ( LD 2 ) , and fast Fourier transform O ( LD log L ) . The computational complexity involved in Formula 7 includes computing attention probability O ( LMD ) , computing attention confidence O ( LDM ) , and matrix multiplications about the value matrix and projection matrices O ( LMD + LD 2 ) . The overall computational complexity is O ( LD 2 + LD log L + LMD ) . Since M is always small (e.g. M = 4 ), for long sequences, this complexity is much smaller than O ( L 2 D + LD 2 ) which is the complexity of standard multi-head attention. In memory usage, sparse attention has no need to store the full attention matrix, thus the memory complexity is reduced from O ( L 2 H + LD ) to O ( LMH + LD ) . We conduct experiments to study the performance of the proposed approach on long sequence modeling tasks. As the primary goal, we evaluate the proposed Fourier Sparse Attention for Transformer (FSAT) on multiple tasks requiring long-context perception. We test our models on the Long Range Arena (LRA) benchmark (Tay et al., 2020b), since it is specifically designed for evaluating the performance of efficient Transformers on various long sequence tasks, and there are quite a number of baseline models evaluated on this benchmark. The LRA benchmark includes five tasks of different kinds and modalities (natural language, vision, and math) in order to simulate meaningful real-world tasks under the long-context scenario. ListOps This task requires models to compute the output value of mathematical expression with a hierarchical structure and operators. The sequence lengths are up to 2K. Text A byte-level text classification task to probe the model's reasoning ability with compositional, unsegmented characters. Character 239 Steps per second  Peak Memory Usage  Model 1K 2K 3K 4K 1K 2K 3K 4K Transformer 1.0 1.0 1.0 1.0 1.00 1.00 1.00 1.00 Local Attention 1.1 1.7 3.2 5.3 0.49 0.29 0.19 0.14 Linformer 1.2 1.9 3.7 5.5 0.44 0.21 0.18 0.1 Reformer 0.5 0.4 0.7 0.8 0.56 0.37 0.28 0.24 Sinkhorn Trans. sequences are truncated or padded to a fixed maximum length of 4K in this task. Retrieval A byte-level document retrieval task tests model's ability to compress long sequences into representations suitable for similarity-based matching. Image An image classification task evaluates a model's performance of perceiving 2D spatial relations between input pixels. Images are flattened to sequences of length 1K pixels. Pathfinder A binary image classification task tests if a model can capture long-range spatial dependencies by judging if two points are connected by a path consisting of dashes in an image with distractor paths. 32  32 images are flattened to sequences of length 1K pixels. We compare our model with a number of promising models, including vanilla Trans-former(Vaswani et al., 2017), a local attention baseline, Sparse Transformer(Child et al., 2019), Longformer(Beltagy et al., 2020), Linformer(Wang et al., 2020), Reformer(Kitaev et al., 2020), Sinkhorn Transformer(Tay et al., 2020a), Synthe-sizer(Tay et al., 2021), Big Bird(Zaheer et al., 2020), Linear Transformer(Katharopoulos et al., 2020), Performer(Choromanski et al., 2020a,b), and more recent models, such as FNet(Lee-Thorp et al., 2021), Nsytr  o mformer(Xiong et al., 2021), and Luan(Ma et al., 2021). We run our experiments on the LRA benchmark with the configurations based on Tay et al. (2020b) open source codebase. Specifically, we follow the original data preprocessing, data split, and keep roughly equivalent model parameters for a fair comparison with the baselines reported in Tay et al. (2020b). An exception is that we reproduce the experiments of the Retrieval task for a longer training of 30K steps because models are not fully converged in 5K training steps. Ma et al. (2021); Lee-Thorp et al. (2021); Xiong et al. (2021) also pointed the same issue. We also re-run the vanilla Transformer using our Pytorch implementation. For the proposed FSAT, the default value of the hyper-parameter M is 4, and the variance  2 is empirically set to L max which is the maximum sequence length of each task. To ensure roughly equivalent model parameters, we reduce the dimension of FFN layer from 4 times of the hidden size to 2 times for FSAT to offset the increased parameters of non-linear feature mapping functions (Formula 1). In our code, sparse matrix multiplications involved in FSAT model are implemented via scatter operations at batch-level for better efficiency. Median results of 5 runs are reported in the tables. Table 1 summarizes the performance of a number of models on the LRA benchmark. As we can see, the proposed model clearly outperforms all previously published approaches, and achieves new state-of-the-art performance on four of the five datasets, and a 5 . 8% absolute improvement over average performance, which validates the effectiveness of the proposed FSAT model. It is noteworthy that we separately report the average accuracy 240 S p ee d / m s Length /# tokens FSATTrans 70 60 50 40 30 20 10 0 6k 5k 4k 3k 2k 1k 0 Length /# tokens FSATTrans M e mm o r y . without the Text task, this is because we find that convolution layers have a significant impact on this task. In the experiment, we adopt a depth-wise separable convolution layer with a kernel size of 5 as the feature mapping function when computing hidden crosses, the accuracy significantly increases from 65 . 95% to 80 . 24% . Therefore, considering the particularity of the Text task, we report its result separately. We suspect that the small data size may be the reason of its particularity. The time and memory efficiency of our model and competing approaches are summarized in Table 2. Compared with the vanilla Transformer, our FSAT significantly reduces the computational complexity with faster training speed and lower memory usage, which demonstrates that directly predicting the sparse structure of attention matrix is an effective way for building efficient Transformer architecture. The limitation of FSAT is that extra operations (e.g. gathering, slicing) of linear complexity are involved, so that FSAT can not bring parallel superiority into full play. Even so, among all compared Transformer variants, FSAT achieves promising results in time and memory efficiency. The trend of computation complexity increasing is shown in Figure 3. This is in line with expectations, FSAT has a linear rate of increase, its advantage is especially obvious for sequences longer than 4K tokens. This demonstrates the potential of the proposed predictable sparse attention for the tasks with much longer sequences, e.g. 3D feature volume. An ablation study is conducted to verify the necessity of our proposed model components. In Table 3, we report FSAT models with the different number of predicted dominant indices. The results show ListOps Retrieval Image Avg. that for each key vector in the attention mechanism about 4 predicted dominant query vectors are enough for the model to produce high accuracy. We also remove the sparse attention module, and test the architecture of only integrating the pooled hidden state cross into the attention mechanism, corresponding to Formula 5. We call this architecture Fourier Attention for Transformer (FAT). It can be seen from the results of FAT, better results can be obtained with the pooled hidden state cross in some tasks, which supports our hypothesis that the 2-order token feature may generate more expressive feature representations. It is noteworthy that without the gradient truncation, or only using random indices I r or I p , the performance significantly drops. Besides, the depth-wise convolution based feature mapping performs worse than a fully-connected structure base feature mapping. In this paper, we proposed a new efficient Transformer architecture, FSAT, which directly predicts the sparse structure of attention matrix to avoid computing the quadratic-cost full self-attention. The proposed approach has the advantages in long-range sequence modeling tasks meanwhile reaching a balance between time, memory, and accuracy. We showed the effectiveness of the proposed method in modeling long sequence context using the Long Range Arena benchmark. Experimental results showed that state-of-the-art performance is achieved by our proposed method.
abstain abstain method objective abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain objective abstain method abstain abstain objective abstain abstain abstain abstain abstain abstain abstain abstain abstain objective objective method objective abstain method abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain other abstain abstain abstain abstain abstain abstain abstain abstain abstain method method abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain method abstain abstain other method abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain other abstain method abstain method abstain abstain abstain abstain abstain abstain abstain abstain abstain method abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain other abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain result abstain abstain abstain result abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain result abstain abstain objective abstain objective objective