Fast Transformers with Clustered Attention

By Apoorv Vyas, Angelos Katharopoulos, Franccois Fleuret. 2020

182 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence

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

Citations in the corpus

Nyströmformer: A Nystöm-based Algorithm for Approximating Self-Attention (2021, methodology): “Landmark points (inducing points (Lee et al. 2019)) can be selected by using K-means clustering (Zhang, Tsang, and Kwok 2008; Vyas, Katharopoulos, and Fleuret 2020).”
FNet: Mixing Tokens with Fourier Transforms (2021, background): “…2018; Beltagy et al., 2020; Ainslie et al., 2020; Zaheer et al., 2020; Wang et al., 2020; Tay et al., 2020b,a; Kitaev et al., 2020; Roy et al., 2021; Vyas et al., 2020; Liu et al., 2018) or via linearization of the attention matrix (Katharopoulos et al., 2020; Choromanski et al., 2021; Peng et…”
A Survey of Transformers (2021, background): “Prototype Clustered Attention[138], Informer[170]”
Separable Self-attention for Mobile Vision Transformers (2022, methodology): “To improve the efficiency of MHA, the second line of research uses similarity measures to group tokens [8, 13, 14].”
Long Range Language Modeling via Gated State Spaces (2022, background): “Over the years, a variety of improvements have been proposed to alleviate this quadratic complexity (Tay et al., 2020; Choromanski et al., 2021; Ramsauer et al., 2021; Wang et al., 2020; Katharopoulos et al., 2020; Peng et al., 2021; Ainslie et al., 2020; Zaheer et al., 2020; Beltagy et al., 2020; D”
Attention mechanism in neural networks: where it comes and where it goes (2022, background): “It is shown that self-attention can lead to improvements to discriminative constituency parser [81], and speech recognition as well [82, 83].”
RomeBERT: Robust Training of Multi-Exit BERT (2021, methodology): “C L ] 24 J a n 2021 computation efﬁciency by means of feature clustering (Vyas et al., 2020; Zheng et al., 2020; Ki-taev et al., 2020), low rank approximation (Choro-manski et al., 2020), attention sparse regularization (Gao et al., 2021), and reversible architecture (Kitaev et al., 2020).”
A Practical Survey on Faster and Lighter Transformers (2021, background): “[122] proposed the improved cluster attention that separately consider the k keys with the highest attention for each cluster.”
Multimodal Co-Attention Transformer for Survival Prediction in Gigapixel Whole Slide Images (2021, background): “…mechanisms have since seen widespread adoption across many different domains outside of neural machine translation such as language model pretraining [14, 55, 56], visual recognition [46, 16], visual question answering [34, 29, 64, 28, 44, 63], graph neural networks [54], and point clouds [31, 18].”
Efficient Transformers: A Survey (2020, background): “A key characteristic of learning patterns is to determine a notion of token relevance and then assign tokens to buckets or clusters (Vyas et al., 2020; Wang et al., 2020b).”
\infty-former: Infinite Memory Transformer (2021, background): “Some propose using sparse attention mechanisms, either with data-dependent patterns (Kitaev et al., 2020; Vyas et al., 2020; Tay et al., 2020a; Roy et al., 2021; Wang et al., 2021) or data-independent patterns (Child et al.”
Learned Token Pruning for Transformers (2021, background): “These can be broadly categorized as follows: (i) efficient architecture design [8, 21, 24, 26, 27, 36, 45, 47, 50, 60]; (ii) knowledge distillation [23, 38, 42–44]; (iii) quantization [2, 4, 13, 15, 25, 40, 58, 59, 61]; and (iv) pruning.”
Visual Attention Methods in Deep Learning: An In-Depth Survey (2022, background): “Clustered Attention : To address the computational inefficiency of transformers, Vyas et al. [168] proposed a clustered attention mechanism that relies on the idea that correlated queries follow the same distribution around Euclidean centers.”