Source-linked AI summary
Efficient Content-Based Sparse Attention with Routing Transformers
Aurko Roy, Mohammad Saffar, Ashish Vaswani, David Grangier
TL;DR
Quadratic attention cost limits Transformer scaling to long sequences, while existing sparse methods trade content flexibility for efficiency. The Routing Transformer learns content-based routing with clustering and combines it with local attention, achieving strong results across language and image generation benchmarks.
Problem
Quadratic time and space complexity makes self-attention difficult to scale to long sequences.
Method
Routing Transformer combines clustering-based content routing with local attention to select sparse, content-relevant contexts without constructing a full attention matrix.
Results
Routing Transformer redefines the state of the art on Wikitext-103, PG-19, and ImageNet-64, while coming close on enwik-8; representative results are 15.8 vs 18.3 perplexity on Wikitext-103.
Takeaways & Limitations
The approach provides an efficient attention mechanism for long-sequence autoregressive generation and may suit naturally sparse domains such as 3D point clouds, social networks, and protein interactions.
Takeaways & Limitations
Sparse attention is more memory-expensive than Transformer-XL-like recurrence and scales to fewer layers for the same problem; Routing Transformer is also roughly 1.7× slower than Local Transformer on PG-19.
Abstract
from arXiv · showhide
Self-attention has recently been adopted for a wide range of sequence modeling problems. Despite its effectiveness, self-attention suffers from quadratic compute 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\left(n^{1.5}d\right)$ from $O\left(n^2d\right)$ 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.
1 Introduction
Self-attention provides direct access to prior context but incurs quadratic cost, making long-sequence modeling difficult. The Routing Transformer combines content-based routing with local attention to improve efficiency while retaining broad access to relevant context.
- Self-attention computes each position’s representation from weighted prior inputs and can directly focus on any earlier context.
- Quadratic time and space complexity in sequence length n makes self-attention difficult to scale to long sequences.Long sequences occur in domains including music, images, speech, video, and document-level translation.
- Prior sparse-attention methods restrict attention to fixed or learned local and strided temporal contexts, limiting patterns independently of content.
- Content-based sparse attention permits arbitrary patterns but may require constructing a full dense attention matrix before sparsification.
- Routing Transformer clusters attention contexts so each query attends to selected elements in its assigned cluster, combining content flexibility with natively sparse computation.The routing strategy is inspired by spherical k-means for Maximum Inner Product Search.
- 15.8 vs 18.3 perplexity on Wikitext-103, 33.2 vs 33.6 on PG-19, and 3.43 vs 3.44 bits/dim on ImageNet-64 are reported with comparable or fewer layers and heads.The experiments cover language modeling and unconditional image generation.
2 Related Work
Prior efficient-attention work mainly uses temporal sparsity, while content-based methods allow richer patterns but often retain quadratic construction costs. Routing attention is motivated by bridging these efficiency and flexibility limitations.
- Attention with Temporal Sparsity: Temporal sparsity reduces attention by limiting fixed local context or temporal resolution, including learned context lengths and cached extended contexts.
- Attention with Content-Based Sparsity: Content-based sparse attention enables richer attention patterns through sparsemax-like formulations, but full attention matrices may still be instantiated before sparsification.
- Attention with Content-Based Sparsity: Routing attention addresses this gap by allowing arbitrary content-based sparsity without instantiating non-zero entries of the full attention matrix.
- Attention with Content-Based Sparsity: LSH-based routing uses fixed random hyperplanes, whereas the Routing Transformer learns space-partitioning centroids with mini-batch spherical k-means.Both approaches approximate Maximum Inner Product Search for dot-product attention.
- Sparse Computation beyond Attention: Related sparse-computation methods include gating, reinforcement-learning-based conditional computation, sparse memory reads, and sparsely gated mixture-of-experts.
3 Self-Attentive Auto-regressive Sequence Modeling
Autoregressive Transformer models factor sequence probabilities into next-symbol predictions and use stacked self-attention modules to refine representations. Each self-attention layer projects inputs into queries, keys, and values, forms attention weights, and aggregates value representations.
- Autoregressive sequence models represent a sequence probability as a product of conditional next-symbol probabilities.
- Transformer architectures model these conditional distributions with learned parameters and have achieved strong results across language, image, and music generation.
- Each Transformer attention module refines representations by taking weighted averages of representations from previous modules.
- A self-attention layer receives n vectors of dimension d and applies three learned linear projections to form queries, keys, and values.
- The key and query matrices determine an n×n attention matrix whose rows are normalized with softmax.
- In autoregressive models, queries attend only to keys from previous time steps, and the next representation is computed as AV.
- The quadratic complexity of attention motivates sparse-attention approaches for long sequences.
4 Efficient Content-Dependent Sparse Attention
The paper defines sparse attention by selecting query-specific key sets and proposes learning these patterns from content rather than positions alone. Routing attention uses shared online k-means clusters to restrict attention, preserving expressive routing while reducing computation and memory.
- Sparse attention: Sparse attention assigns each query a set of key positions it may attend to, forming the sequence-wide sparsity pattern S.Classical causal and local attention define these sets from positions, independently of query and key vectors.
- Content-dependent sparsity: Content-based sparsity sets S = f(x), allowing arbitrary data-informed patterns and potentially greater sparsity for queries with longer-lasting effects.Unlike fixed local or strided patterns, this formulation can accommodate data without a clear ordering over observations.
- Content-dependent sparsity: Correia et al. infer content-based sparsity after instantiating a full attention matrix, whereas the proposed approach is designed to remain natively sparse.The paper motivates a clustering-based formulation to avoid full-matrix construction before sparsification.
- Routing attention with clustering: Routing attention clusters queries and keys with shared online k-means centroids, then considers only query-key pairs assigned to the same cluster.Centroids are shared across sequences and learned online with the remaining model parameters.
- Routing attention with clustering: Unit-normalized queries and keys make nearest-centroid routing relevant to maximum inner product search, while same-cluster assignments preserve large attention weights as non-zero entries.The paper connects dot-product attention importance to MIPS and argues that clustering retains high-weight interactions.
- Routing attention with clustering: O(n^1.5d) complexity replaces O(n^2d) when k = √n, with routing comparisons contributing O(nkd) and within-cluster dot products contributing O(n^2d/k).Balanced clusters make each query compare with roughly n/k keys; top-k assignment adds O(n log n), which is dominated by O(n^1.5d).
5 Experiments
Across text and image generation, Routing Transformer achieves strong results with content-based sparse attention, including new state-of-the-art results on several benchmarks.
- Wikitext-103: 15.8 vs 18.3 perplexity: Routing Transformer improves on Transformer-XL on Wikitext-103 while using fewer self-attention layers and no segment-level recurrence.The reported values are test perplexities for Routing Transformer and Transformer-XL, respectively.
- enwik-8: 0.99 vs 0.98 perplexity: Routing Transformer is competitive on enwik-8 but remains slightly worse than Adaptive Transformer.The reported perplexity is for the enwik-8 benchmark.
- ImageNet-64: 3.425 vs 3.437 bits/dim: Routing Transformer improves on the previous state of the art on ImageNet-64.Scaled-up local attention obtains 3.48 bits/dim, while the Routing Transformer result is reported against Child et al. (2019).
- PG-19: 33.2 vs 33.6 test perplexity: a 22-layer Routing Transformer improves on the 36-layer Compressive Transformer on PG-19 while generating sequences of length 8192.A 24-layer local attention model obtains 39.3 and a 36-layer Transformer-XL obtains 36.3 on the same benchmark.
6 Analysis
The analysis finds that local attention supplies strong local representations, while routing layers add complementary global consistency. Routing improves quality over local attention but incurs a speed trade-off on the tested hardware.
- Local vs Global: Local attention achieves 3.009 versus full attention’s 2.983 bits per dim, while adding two routing layers with four heads reaches 2.986 bits per dim.The result shows that a small number of routing layers nearly closes the local-attention gap to full attention.
- Local vs Global: 2.975 bits per dim is achieved by the best CIFAR-10 model with four routing layers, four routing heads, and an attention window of 512.Increasing the window to 1024 improves performance across settings, while routing-only models perform worse than configurations including local attention.
- Local vs Global: The analysis evaluates local–routing attention divergence with Jensen-Shannon divergence across layers on Wikitext-103, averaging means and standard deviations over 10 runs.The divergence is computed over sequences of length 4096 and is bounded by 0.6931 under the stated logarithm convention.
- Local vs Global: Routing attention selects high-dot-product token pairs through approximate MIPS, complementing local attention’s local consistency with global consistency.The analysis attributes this mechanism to spherical k-means over the global token set.
- Wall-clock time: Routing Transformer step times are slower than local attention on TPUv3, with local attention roughly 1.7× faster on PG-19.The paper attributes this trade-off to limited TPU support for sparse operations and treats wall-clock efficiency as secondary to memory-efficient approximation.
7 Conclusion
The conclusion presents Routing Transformer as content-based sparse attention that avoids full attention matrices while retaining content-similarity-based routing. Experiments report state-of-the-art results across major long-sequence language and image-generation benchmarks.
- 7 Conclusion: Routing Transformer selects sparsity patterns from content similarity without computing a full attention matrix.It is motivated by non-negative matrix factorization and has complexity comparable to local attention models.
- 7 Conclusion: Routing Transformer redefines the state of the art on Wikitext-103, PG-19, and ImageNet-64, while coming close on enwik-8.The experiments cover text and image generation and also identify complementary routing patterns relative to local attention.
- 7 Conclusion: The work contributes an efficient attention mechanism for modeling long sequences and autoregressive generative modeling.The paper also suggests possible applications to naturally sparse domains such as 3D point clouds, social networks, and protein interactions.
A Samples from Routing Transformer
The paper presents samples generated by a Routing Transformer trained on PG-19 sequences of length 8192. Generation uses nucleus sampling with p = 0.8 and temperature 1.0.
- A Samples from Routing Transformer: PG-19 samples are generated with nucleus sampling at p = 0.8 and temperature 1.0.The model was trained using sequence length 8192.
A.1 Sample - I
The passage describes religious and ecclesiastical developments involving the Councils of Nice and Basle, the Swiss Church, Zwingli, and Luther’s influence.
- A union called the Papal Council opposed the pretensions of the Council of Basle during the early history of the Council of Nice.
- The Councils of Nice and Basle arranged for representatives to present proposals concerning their disputes to other councils.
- Zwingli introduced a new conception of Church government that connected civil and ecclesiastical organization and influenced Protestant lands.
- The Reformation encouraged study of Scripture and contributed to a more liberal conception of Christian life.
- Luther’s German translation and theological writings reached many readers and supplied valuable religious and cultural information.
A.2 Sample - II
The passage recounts Louis XVIII’s conduct during political and military disturbances, emphasizing his composure, anxiety, and reliance on ministers and military commanders.
- Louis XVIII initially showed no interest in the proceedings of the day before reacting to a cannon-shot from St. Cloud.
- The King generally displayed little uneasiness, but extreme agitation or reverie could reveal changes in his expression.
- His anxiety concerned his son’s safety and his own weakness and inexperience in affairs of state.
- During the crisis, the King sent for General Bugeaud and discussed the conduct of the forces and the danger threatening France.
A.3 Sample - III
The passage recounts the development and reception of theatrical works, including adaptations inspired by Bulwer’s Pelham and performances involving Mr. Kean and Miss O’Neil.
- Mr. Green wrote The Adventures of Major de la Motte on the principles of Bulwer’s Pelham after an unsuccessful attempt to obtain an introduction for the earlier work.
- The public was satisfied with one of Mr. Green’s dramas, which brought a clever young performer into general notice.
- A compromise transferred the play to the company while Mr. Kean assumed the role of Sir Giles Overreach.
- The narrator found Mr. Kean’s readings intensely enjoyable and repeatedly read favored scenes aloud.
- Mr. Kean and Miss O’Neil created a sensation when they made their debut in The Hunchback.
A.4 Sample - IV
The passage presents descriptions of birds and natural features in Canada and nearby regions, emphasizing distribution, physical characteristics, and habitat observations.
- Northern warblers commonly summer in Canada but may winter in South America, requiring observation across both seasonal ranges.
- Some bird species have highly localized distributions, with only a few breeding haunts recorded in the British Islands.
- The Black Guillemot is described by its black bill, feet, and wings contrasted with ash-grey upper parts.
- Trees near Kingston reached more than fifty feet, while sheltered areas supported abundant low-growing aromatic shrubs.