SMYRF: Efficient Attention using Asymmetric Clustering

Abstract

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.

arXiv comment30 pages, 10 figures

References in corpus (20)

• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Distilling the Knowledge in a Neural Network
• Show, Attend and Tell: Neural Image Caption Generation with Visual Attention
• GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium
• Large Scale GAN Training for High Fidelity Natural Image Synthesis
• Language Models are Few-Shot Learners
• Self-Attention Generative Adversarial Networks
• Longformer: The Long-Document Transformer
• Quantized Neural Networks: Training Neural Networks with Low Precision Weights and Activations
• Activation Functions: Comparison of trends in Practice and Research for Deep Learning
• Training Deep Nets with Sublinear Memory Cost
• DeepFakes: a New Threat to Face Recognition? Assessment and Detection
• Generating Long Sequences with Sparse Transformers
• Reformer: The Efficient Transformer
• Practical and Optimal LSH for Angular Distance
• The Reversible Residual Network: Backpropagation Without Storing Activations
• Sparse Sinkhorn Attention
• Probabilistic Analysis of Linear Programming Decoding
• Music Transformer
• Large Memory Layers with Product Keys