Scatterbrain: Unifying Sparse and Low-rank Attention Approximation

Abstract

Recent advances in efficient Transformers have exploited either the sparsity or low-rank properties of attention matrices to reduce the computational and memory bottlenecks of modeling long sequences. However, it is still challenging to balance the trade-off between model quality and efficiency to perform a one-size-fits-all approximation for different tasks. To better understand this trade-off, we observe that sparse and low-rank approximations excel in different regimes, determined by the softmax temperature in attention, and sparse + low-rank can outperform each individually. Inspired by the classical robust-PCA algorithm for sparse and low-rank decomposition, we propose Scatterbrain, a novel way to unify sparse (via locality sensitive hashing) and low-rank (via kernel feature map) attention for accurate and efficient approximation. The estimation is unbiased with provably low error. We empirically show that Scatterbrain can achieve 2.1x lower error than baselines when serving as a drop-in replacement in BigGAN image generation and pre-trained T2T-ViT. On a pre-trained T2T Vision transformer, even without fine-tuning, Scatterbrain can reduce 98% of attention memory at the cost of only 1% drop in accuracy. We demonstrate Scatterbrain for end-to-end training with up to 4 points better perplexity and 5 points better average accuracy than sparse or low-rank efficient transformers on language modeling and long-range-arena tasks.

NeurIPS 2021

References in corpus (27)

• An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
• Language Models are Few-Shot Learners
• Longformer: The Long-Document Transformer
• XLNet: Generalized Autoregressive Pretraining for Language Understanding
• Linformer: Self-Attention with Linear Complexity
• A Simpler Approach to Matrix Completion
• Generating Long Sequences with Sparse Transformers
• Pointer Sentinel Mixture Models
• Reformer: The Efficient Transformer
• Pay Less Attention with Lightweight and Dynamic Convolutions
• Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention
• Practical and Optimal LSH for Angular Distance
• Long Range Arena: A Benchmark for Efficient Transformers
• Rethinking Attention with Performers
• Structured Transforms for Small-Footprint Deep Learning
• CondenseNet: An Efficient DenseNet using Learned Group Convolutions
• Hopfield Networks is All You Need
• Long-Short Transformer: Efficient Transformers for Language and Vision
• Luna: Linear Unified Nested Attention
• SLIDE : In Defense of Smart Algorithms over Hardware Acceleration for Large-Scale Deep Learning Systems
• Learning Fast Algorithms for Linear Transforms Using Butterfly Factorizations
• Nyströmformer: A Nyström-Based Algorithm for Approximating Self-Attention
• Learning Space Partitions for Nearest Neighbor Search
• Kaleidoscope: An Efficient, Learnable Representation For All Structured Linear Maps
• Sub-Linear Memory: How to Make Performers SLiM
• A Tale of Two Efficient and Informative Negative Sampling Distributions
• Climbing the WOL: Training for Cheaper Inference