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.1× 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. * Equal contribution. Order determined by coin flip.

Revisiting Token Pruning for Object Detection and Instance Segmentation

“…Transformer Acceleration Various methods have been explored for optimizing Transformers' high computational cost, including designing alternative lightweight attention formulations [11,28,31,46,50,54,68], removing unnecessary network modules [17,40,53] approximating attention multiplications with low-rank decompositions [6,12,55], distilling knowledge into a more efficient student network [48,51,69], and extending network quantization techniques for Transformers [1,18,30,49,67]. Furthermore, acceleration techniques specific to ViTs have been proposed [19,34,41,44,47,61,63] by exploiting the redundancy in the input patches to early drop tokens for saving computation.…”

Linearization Weight Compression and In-Situ Hardware-Based Decompression for Attention-Based Neural Machine Translation

“…Another way to reduce the memory requirements and computation complexity is low-rank approximation [38]. In [39], a method which exploits both sparsity and low-rank approximation is proposed, which is also called as Scatterbrain. It is shown that the Scatterbrain can outperform the methods which employ only either sparsity or low-rank approximation, illustrating that the sparsity and low-rank approximation can be exploited synergistically.…”

Model Compression and Hardware Acceleration for Neural Networks: A Comprehensive Survey

“…This combination significantly hinders deployment on devices with constrained computational and memory resources, particularly in real-time applications such as autonomous driving [13] and virtual reality [14], where fulfilling low latency requirements and producing a highquality user experience are crucial. This underscores the pressing need for advancements in model compression techniques such as pruning [15], quantization [16], knowledge distillation [17], and low-rank factorization [18]. Moreover, the rapid adoption of ViTs can be attributed not only to algorithmic innovations and data availability but also to enhancements in processor performance.…”