Efficient Content-Based Sparse Attention with Routing Transformers

By Aurko Roy, Mohammad Taghi Saffar, Ashish Vaswani, David Grangier. 2020

773 citations recorded by the source corpus.

Self-attention has recently been adopted for a wide range of sequence modeling problems. Despite its effectiveness, self-attention suffers from quadratic computation and memory requirements with respect to sequence length. Successful approaches to reduce this complexity focused on attending to local sliding windows or a small set of locations independent of content. Our work proposes to learn dynamic sparse attention patterns that avoid allocating computation and memory to attend to content unrelated to the query of interest. This work builds upon two lines of research: It combines the modeling flexibility of prior work on content-based sparse attention with the efficiency gains from approaches based on local, temporal sparse attention. Our model, the Routing Transformer, endows self-attention with a sparse routing module based on online k-means while reducing the overall complexity of attention to O(n1.5d) from O(n2d) for sequence length n and hidden dimension d. We show that our model outperforms comparable sparse attention models on language modeling on Wikitext-103 (15.8 vs 18.3 perplexity), as well as on image generation on ImageNet-64 (3.43 vs 3.44 bits/dim) while using fewer self-attention layers. Additionally, we set a new state-of-the-art on the newly released PG-19 data-set, obtaining a test perplexity of 33.2 with a 22 layer Routing Transformer model trained on sequences of length 8192. We open-source the code for Routing Transformer in Tensorflow.1

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Model Architectures, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Model Architectures

In Efficiency & Compression

Citations in the corpus

PaLM: Scaling Language Modeling with Pathways (2022, background): “Sparsity in the sequence length is an area that allows for training efficiently with extremely long sequences (Zaheer et al., 2020; Tay et al., 2020; Choromanski et al., 2020; Kitaev et al., 2020; Roy et al., 2020; Child et al., 2019; Gale et al., 2019).”
Rethinking Attention with Performers (2020, methodology): “(Roy et al., 2020) which applies k-means clustering to learn dynamic sparse attention regions, or (Kitaev et al.”
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022, methodology): “These methods range from sparse-approximation [49, 71] to low-rank approximation [11, 48, 81], and their combinations [3, 8, 89].”
Self-attention Does Not Need $O(n^2)$ Memory (2021, methodology): “So the space complexity of transformers is a particular concern, c.f. Kitaev et al. (2020); Roy et al. (2021); Zaheer et al. (2020).”
Efficiently Scaling Transformer Inference (2022, neutral): “Several approaches (Gupta and Agrawal, 2020) to improve the inference efficiency of Transformer models focus on model architecture improvements, for example efficient attention layers (Roy et al., 2020; Choromanski et al., 2020; Kitaev et al., 2020; Sukhbaatar et al., 2019; Child et al., 2019), dist”
SCROLLS: Standardized CompaRison Over Long Language Sequences (2022, methodology): “The language modeling community typically uses perplexity to measure how well models predict the next token, a practice that has been adopted by several works on efficient transformer architectures (Roy et al., 2021; Choromanski et al., 2020; Tay et al., 2020a; Peng et al., 2021).”
Variational Diffusion Models (2021, methodology): “The specific diffusion processes used by Song and Ermon [2019] and Sohl-Dickstein et al. [2015] can be seen as special cases of the proposed model. Song and Ermon [2019] use αt = 1, called variance-exploding diffusion processes by Song et al.”
Perceiver: General Perception with Iterative Attention (2021, methodology): “…non-dense attention patterns (Ho et al., 2019; Wang et al., 2020a; Belt-agy et al., 2020; Child et al., 2019; Correia et al., 2019; Ye et al., 2019; Roy et al., 2020), approximating or otherwise simplifying the matrix multiplications used in QKV attention (Choromanski et al., 2021; Peng et al.,…”
Improved Denoising Diffusion Probabilistic Models (2021, methodology): “We did not attempt to calculate NLL using DDIM, since Song et al. (2020a) does not present NLL results or a simple way of estimating likelihood under DDIM.”
Sparse is Enough in Scaling Transformers (2021, background): “Sparse attention-based approaches have made the attention layer more efficient, especially for long sequences, by incorporating additional combinatorial mechanisms, as in [40], or selecting a subset of tokens this layer attends to [31, 5, 19, 37, 15, 4] or other approaches [12].”
Generating Images with Sparse Representations (2021, background): “…Transformer variants (Child et al., 2019; Kitaev et al., 2020; Choromanski et al., 2021; Dhariwal et al., 2020; Zaheer et al., 2021; Roy et al., 2020) make long-sequence training more practical, care is still required to optimize model hyperparameters against available memory when…”
Deformable DETR: Deformable Transformers for End-to-End Object Detection (2020, background): “A similar idea is proposed by Roy et al. (2020), where k-means ﬁnds out the most related keys.”
BASE Layers: Simplifying Training of Large, Sparse Models (2021, background): “An orthogonal strand of work has improved the efﬁciency of transformer attention mechanisms, often by making them sparse (Child et al., 2019; Correia et al., 2019; Roy et al., 2020).”
Long Range Arena: A Benchmark for Efficient Transformers (2020, background): “Generative Modeling / Language Modeling This generative modeling task requires predicting the next character, word, or pixel and is a staple in xformer evaluations (Roy et al., 2020; Kitaev et al., 2020).”
Efficient Transformers: A Survey (2020, methodology): “For example, the k -means clustering in Routing Transformer (Roy et al., 2020) can also be interpreted as a form of global model memory approach, since one can view the centroids as parameterized model memory.”
Pre-Trained Models: Past, Present and Future (2021, neutral): “Compared to predefined chunks, some works (Roy et al., 2021; Kitaev et al., 2020) find that using learnable parameters to assign tokens into chunks results in better performance.”
AudioLM: A Language Modeling Approach to Audio Generation (2022, background): “While several works have explored efficient alternatives to self-attention [39], [40], [41], another solution to this scaling problem is to work with mappings of the natural signals to a compact, discrete representation space.”
ETC: Encoding Long and Structured Inputs in Transformers (2020, methodology): “The Routing Transformer (Roy et al., 2020) learns dynamic sparse attention patterns using online k -means, reducing complexity to O ( n 1 .”