Efficient Content-Based Sparse Attention with Routing Transformers (2003.05997v5)

Abstract: 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 LLMing 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.

Efficient Content-Based Sparse Attention with Routing Transformers

The paper "Efficient Content-Based Sparse Attention with Routing Transformers" addresses the computational inefficiency inherent in self-attention mechanisms used in sequence modeling. Self-attention, renowned for its applicability across various sequence modeling tasks, suffers from quadratic complexity concerning sequence length, making it computationally expensive for long sequences. This work introduces the Routing Transformer, a model that mitigates these limitations by combining the flexibility of content-based sparse attention with the efficiency of local sparse attention patterns.

Main Contributions

The Routing Transformer proposes a novel method to achieve sparse attention by dynamically learning to attend only to relevant content. This involves:

1. Sparse Routing Module: The model uses an online $k$-means clustering mechanism to implement a sparse routing module, reducing attention complexity from $O(n^2d)$ to $O(n^{1.5}d)$, where $n$ and $d$ denote the sequence length and hidden dimension, respectively.
2. Local and Content-Based Attention: By integrating local attention strategies with content-based clustering, the model achieves a balance between computational efficiency and the ability to capture long-range dependencies in the data.

Experimental Results

The Routing Transformer exhibits superior performance when compared to existing sparse attention models. Notably:

• On the Wikitext-103 dataset, the model achieves a perplexity of 15.8, outperforming previous models with a perplexity of 18.3.
• In image generation on the ImageNet-64 dataset, it reduces bits per dimension to 3.43 compared to 3.44 of its counterparts.
• The model also achieves a test perplexity of 33.2 on the PG-19 dataset, setting a new benchmark.

These results demonstrate that the Routing Transformer not only enhances computational efficiency but also maintains or improves upon the predictive performance of prior models.

Implications and Future Directions

The proposed methodology presents substantial implications for the fields of natural language processing and generative tasks:

• Efficiency and Scalability: The reduction in complexity allows for more feasible deployment of attention mechanisms on lengthy sequences, which is beneficial for tasks involving large-scale text and image data.
• Potential Applications: The underlying clustering-based routing mechanism has broader applicability, potentially impacting domains requiring efficient processing of sparse, high-dimensional data, such as 3D point clouds and social network analysis.

Future research may focus on further optimizing the clustering mechanism to enhance balance and efficiency. Exploring the integration of the Routing Transformer with other architectures or adapting it to different modalities could further expand its utility in AI applications.

In conclusion, the Routing Transformer offers a significant advancement in the efficient processing of long sequences by leveraging dynamic, content-sensitive sparse attention, setting the stage for further innovation in scalable AI models.