Fast Transformers with Clustered Attention

Apoorv Vyas1,2 Angelos Katharopoulos1,2 François Fleuret1,2
1 Idiap Research Institute 2 École Polytechnique Fédérale de Lausanne

Paper

Colab

Code

Docs

Blog
Transformers have been proven a successful model for a variety of tasks in sequence modeling. However, computing the attention matrix, which is their key component, has quadratic complexity with respect to the sequence length, thus making them prohibitively expensive for large sequences. To address this, we propose clustered attention, which instead of computing the attention for every query, groups queries into clusters and computes attention just for the centroids. To further improve this approximation, we use the computed clusters to identify the keys with the highest attention per query and compute the exact key/query dot products. This results in a model with linear complexity with respect to the sequence length for a fixed number of clusters. We evaluate our approach on two automatic speech recognition datasets and show that our model consistently outperforms vanilla transformers for a given computational budget. Finally, we demonstrate that our model can approximate arbitrarily complex attention distributions with a minimal number of clusters by approximating a pretrained BERT model on GLUE and SQuAD benchmarks with only 25 clusters and no loss in performance.
Acknowledgements
SNSF logo Angelos Katharopoulos and Apoorv Vyas are supported by SNSF for their research.