Long-range Sequence Modeling with Predictable Sparse Attention

Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 234 - 243 May 22-27, 2022 c⃝2022 Association for Computational Linguistics Long-range Sequence Modeling with Predictable Sparse Attention Yimeng Zhuang, Jing Zhang, Mei Tu Samsung Research China - Beijing (SRC-B) {ym.zhuang, jing97.zhang, mei.tu}@samsung.com 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 redun- dant. Therefore, in this paper, we design an effi- cient Transformer architecture, named Fourier Sparse Attention for Transformer (FSAT), for fast long-range sequence modeling. We pro- vide 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 L2 semantic combinations in O(L log L) time complexity. (2) A sparse attention ma- trix estimation module, which predicts domi- nant 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 in- dex of dominant elements. The overall com- plexity about the sequence length is reduced from O(L2) to O(L log L). Extensive exper- iments (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 computa- tional costs, and achieves new state-of-the-art results on the Long Range Arena benchmark. 1 Introduction 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 Trans- former architecture perceiving long-range context heavily relies on the multi-head self-attention mech- anism, 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 ex- tend Transformer for very long sequences, such as document-level text tasks, high-resolution im- ages, 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 mod- eling tasks. Tay et al. (2020c) categorize the re- searches of efficient Transformers: (a) Fixed pat- terns 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 …

Content truncated to 500 words