Scatterbrain: Unifying Sparse and Low-rank Attention Approximation

By Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, Christopher R'e. 2021

175 citations recorded by the source corpus.

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.

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): “[8] Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré.”
Efficient Transformers: A Survey (2020, background): “Scatter Brain is a new work (Chen et al., 2021) attempts to unify sparse (ﬁxed pattern) attention with low-rank attention.”
A Cookbook of Self-Supervised Learning (2023, methodology): “One approach is with low-rank and/or sparse approximation mechanisms [Kitaev et al., 2020, Choromanski et al., 2020, Wang et al., 2020a, Chen et al., 2021a, Zaheer et al., 2020].”
Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time (2023, methodology): “…Search techniques such as Locality-sensitive hashing (Gionis et al., 1999) and Graph-based indexing (Malkov et al., 2014) for efficient deep neural network training or inference (Zhang et al., 2018; Chen et al., 2019; 2020a; Kitaev et al., 2020; Chen et al., 2021b;a; Liu et al., 2022).”
FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (2023, unknown): “To reduce the computational requirement of attention on such long context, there have been numerous methods proposed to approximate attention [2, 3, 4, 8, 9, 14, 19, 20].”
Pixelated Butterfly: Simple and Efficient Sparse training for Neural Network Models (2021, methodology): “For example, in Transformers, let da and dm be the density of the attention and the MLP layers.”
SqueezeLLM: Dense-and-Sparse Quantization (2023, unknown): “Given this insight, the appropriate strategy is to minimize the memory size even if it may add overhead to arithmetic operations .”
Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time (2023, background): “[19] Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré.”