SMYRF: Efficient Attention using Asymmetric Clustering

By Giannis Daras, Nikita Kitaev, Augustus Odena, A. Dimakis. 2020

52 citations recorded by the source corpus.

We propose a novel type of balanced clustering algorithm to approximate attention. Attention complexity is reduced from $O(N^2)$ to $O(N \log N)$, where $N$ is the sequence length. Our algorithm, SMYRF, uses Locality Sensitive Hashing (LSH) in a novel way by defining new Asymmetric transformations and an adaptive scheme that produces balanced clusters. The biggest advantage of SMYRF is that it can be used as a drop-in replacement for dense attention layers without any retraining. On the contrary, prior fast attention methods impose constraints (e.g. queries and keys share the same vector representations) and require re-training from scratch. We apply our method to pre-trained state-of-the-art Natural Language Processing and Computer Vision models and we report significant memory and speed benefits. Notably, SMYRF-BERT outperforms (slightly) BERT on GLUE, while using $50\%$ less memory. We also show that SMYRF can be used interchangeably with dense attention before and after training. Finally, we use SMYRF to train GANs with attention in high resolutions. Using a single TPU, we were able to scale attention to 128x128=16k and 256x256=65k tokens on BigGAN on CelebA-HQ.

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

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

Citations in the corpus

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022, methodology): “Linformer Attention [88], Smyrf [20], and LongShortFormer (LSFormer) [98].”
Hungry Hungry Hippos: Towards Language Modeling with State Space Models (2022, background): “A number of recent works have also explored how to address the shortcomings of attention by approximating the attention computation [9, 16, 35, 39, 59, 63].”
Scatterbrain: Unifying Sparse and Low-rank Attention Approximation (2021, methodology): “, by The sparse matrix S typically has support determined randomly [16] LSH [22, 36], or by clustering [53].”
Visual Attention Methods in Deep Learning: An In-Depth Survey (2022, background): “[91] proposed symmetric Locality Sensitive Hashing (LSH) clustering in a novel way to reduce the size of attention maps, therefore, developing efficient models.”