Long-range Sequence Modeling with Predictable Sparse Attention (2022)

Abstract: 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 sub-modules are: (1) A fast Fourier transform based hidden state cross module, which captures and pools L^2 semantic combinations in \mathcal{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 \mathcal{O}(L^2) to \mathcal{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.

## introduction models based on the transformer architecture @xcite have been firmly established as state of the art approaches across a range of domains like language @xcite @xcite , and vision @xcite . 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 token's representations are the weighted average of all tokens using the attention scores. despite its effectiveness, self-attention mechanism'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. @xcite 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 transformer's variants for long-range sequence modeling tasks. @xcite categorize the researches of efficient transformers: (a) fixed patterns or combination of patterns @xcite , in which the field to be attended is pre-defined by fixed pattern. (b) learnable patterns @xcite , in which tokens are sorted or clustered in a data-driven fashion. (c) memory @xcite , in which spacial tokens with global view are introduced to compress the input sequence. (d) low-rank methods @xcite , which adopt low-rank approximations of the self-attention matrix. (e) kernels @xcite , which view the attention mechanism through kernelization. (f) recurrence @xcite , 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-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 (nonzero) 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 @xcite '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 (chunked 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: ## self-attention with pooled hidden state cross 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 selfattention with pooled hidden state cross. ## predictable sparse attention 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. ## experiments we conduct experiments to study the performance of the proposed approach on long sequence modeling tasks. ## conclusion 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 longrange 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.