PaperlayerView exact revision on arXiv

Source-linked AI summary

Scatterbrain: Unifying Sparse and Low-rank Attention Approximation

Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, Christopher Ré

arXiv:2110.15343v1cs.LG

TL;DR

Efficient Transformers need attention approximations that balance accuracy and resource use across tasks. Scatterbrain unifies sparse and low-rank approximations using LSH and kernel features, achieving lower error, large memory savings, and improved benchmark performance. Its scope does not address application-specific privacy, fairness, or training-data issues.

  • Problem

    Efficient attention methods must balance computational and memory efficiency with model quality across diverse tasks, but sparse and low-rank approaches have different strengths.

  • Method

    Scatterbrain combines locality-sensitive hashing for sparse attention with kernel approximation for low-rank attention, inspired by Robust PCA.

  • Results

    Scatterbrain achieves lower approximation error than efficient baselines, reduces attention memory by up to 98% in pre-trained models, and improves end-to-end language-modeling and long-range-task performance.

  • Takeaways & Limitations

    Sparse-plus-low-rank attention improves the efficiency-accuracy trade-off across pre-trained image models, vision Transformers, language modeling, and long-range benchmarks.

  • Takeaways & Limitations

    The work does not address application-specific privacy and fairness issues or problems such as unfathomable training data.

Abstract

from arXiv · show

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.

1 Introduction

Efficient attention methods typically exploit either sparsity or low rank, but their accuracy varies across tasks and attention regimes. Scatterbrain unifies both structures to improve approximation accuracy and efficiency.

  • Long-sequence Transformers face quadratic attention costs, motivating approximations that reduce computational and memory requirements.
  • Sparse and low-rank approximations have different strengths, causing accuracy degradation when applied across diverse tasks.
  • Sparse approximations excel at low-entropy attention, whereas low-rank approximations excel at high-entropy attention; their combination can outperform either alone.
  • Scatterbrain combines locality-sensitive hashing for sparse attention with kernel approximation for the low-rank component, without materializing the full attention matrix.
  • 2.1× lower error than efficient baselines was achieved on real benchmarks, while drop-in use reduced attention memory by up to 98% in pre-trained T2T-ViT and BigGAN.
  • End-to-end training improved WikiText-103 perplexity by up to 1 point and average accuracy on five long-range tasks by up to 5 points.

2 Problem Setting and Related Work

The paper formulates attention approximation as jointly minimizing approximation error and computational or memory cost. Related efficient Transformers primarily use sparse or low-rank approximations, motivating a unified approach.

  • Problem Formulation: Given Q, K, and V, the problem is to approximate softmax(QK^⊤)V while reducing the quadratic complexity of attention computation.
  • Problem Formulation: The approximation procedure is evaluated using two objectives: approximation error and computation or memory cost.
  • Sparse, Low-rank Approximation for Attention Matrices: Sparse methods such as Reformer use LSH to identify large attention entries, while low-rank methods such as Performer use kernelization to avoid explicit quadratic computation.
  • Sparse, Low-rank Approximation for Attention Matrices: Attention structure varies across practical settings, so either sparse or low-rank approximation alone may not robustly approximate a wide range of attention matrices.
  • Sparse + Low-rank and Robust PCA: Robust PCA provides a polynomial-time route to approximately optimal sparse-plus-low-rank matrix decomposition, motivating Scatterbrain's design.

3 Characterization of Sparse + Low-rank Approx. to Attention Matrices

Sparse and low-rank attention approximations are complementary, with their relative accuracy governed by softmax temperature. A clustered-input model and theoretical characterization explain when each structure, or their combination, is effective.

  • 3.1 Motivating Observations: Low-rank and Sparse Structures of Attention Matrices: Sparse approximation is strongest at low softmax temperature, low-rank approximation at high temperature, and robust-PCA combination achieves lower error than either alone.
  • 3.1 Motivating Observations: Low-rank and Sparse Structures of Attention Matrices: Sparse-plus-low-rank approximation can require O(n) parameters for ϵ Frobenius error, whereas sparse or low-rank alone may require O(n^2).
  • 3.2 A Generative Model of How Sparse + Low-rank Structure Can Arise: The generative model samples clustered query vectors around Gaussian cluster centers and defines attention from M_β = exp(β · QQ^⊤).
  • 3.2 A Generative Model of How Sparse + Low-rank Structure Can Arise: In the middle temperature regime, the sparse component captures intra-cluster attention while the low-rank component approximates inter-cluster attention.
  • 3.2 A Generative Model of How Sparse + Low-rank Structure Can Arise: At high temperature, a low-rank approximation with n^1+o(1) parameters can achieve bounded Frobenius error.
  • 3.2 A Generative Model of How Sparse + Low-rank Structure Can Arise: At middle temperature, sparse-plus-low-rank achieves bounded error with n^1+o(1) parameters, while sufficiently low-rank approximations incur Ω(n) error.
  • 3.2 A Generative Model of How Sparse + Low-rank Structure Can Arise: At low temperature, sufficiently low-rank approximations incur error at least Ω(e^β(1−∆^2)).

4 Scatterbrain: Unifying Sparse and Low-rank Attention

Scatterbrain combines sparse and low-rank attention to address the accuracy-efficiency limits of Robust PCA and standalone sparse or low-rank approximations. Its correction scheme is efficient, exact on selected large entries, and unbiased with lower variance than low-rank baselines.

  • 4 Scatterbrain: Unifying Sparse and Low-rank Attention: Scatterbrain is more accurate than Reformer and Performer while remaining similarly efficient, with low MSE for both small and large attention entries.The qualitative comparison also includes Robust PCA, which is accurate but slow.
  • 4.1 Challenges of Designing an Accurate and Efficient Sparse + Low-rank Approximation: Robust PCA is accurate but too slow and memory-intensive, while simply adding sparse and low-rank approximations double-counts entries on the sparse support.These limitations motivate an approximation that combines both structures without materializing the full attention matrix.
  • 4.2 Scatterbrain: Algorithm Intuition and Description: The sparse-plus-low-rank estimator is exact on the selected support, while outside it follows the low-rank approximation for entries likely to be small.This targets large entries with sparse attention and small entries with low-rank estimation.
  • 4.2 Scatterbrain: Algorithm Intuition and Description: Scatterbrain constructs a low-rank kernel-feature approximation, uses locality-sensitive hashing to select sparse locations, and corrects those locations before combining both outputs.The output is computed as eQ(eK^T V) + SV, avoiding materialization of the full low-rank matrix.
  • 4.3 Scatterbrain: Analysis: Scatterbrain retains unbiasedness while achieving strictly lower variance than Performer.Its variance is reduced because large entries are likely selected by LSH, while small entries have low low-rank-estimation variance.

5 Experiments

Experiments show that Scatterbrain combines sparse and low-rank attention effectively, achieving low approximation error, strong downstream performance, and substantial efficiency gains across vision, language modeling, and long-range tasks.

  • Approximation accuracy: Scatterbrain’s approximation error is within 20% of Robust PCA and up to 2.1× lower than SMYRF and Performer.The comparison uses pre-trained attention matrices and evaluates sparse, low-rank, and sparse+low-rank approximations.
  • End-to-end training: Up to 4 points better performance across diverse sequence-modeling and long-range tasks distinguishes Scatterbrain from Reformer and Performer.On WikiText-103 it achieves 1 point better perplexity than Reformer, while on Long-Range Arena it achieves up to 5 points higher average accuracy than both baselines.
  • Error and task analysis: Under equal parameter budgets, Scatterbrain has the smallest approximation error in most language-modeling layers and consistently outperforms baselines on Long-Range Arena tasks.The LRA comparison uses 12.5% of full computation for each approximation.
  • Scalability: Up to 3× faster and 12× more memory efficient than full attention at sequence length 4096, Scatterbrain remains competitive with SMYRF and Performer.The results attribute this efficiency to combining sparse and low-rank components while retaining their computational advantages.

6 Discussion

Scatterbrain combines sparse and low-rank attention approximation to improve efficiency while preserving approximation quality, but its sparse component remains less hardware-friendly than dense low-rank computation.

  • Limitations: Scatterbrain is not yet as hardware-friendly on GPUs and TPUs as its dense low-rank component because sparse computation is less optimized.The paper notes that more efficient sparse GPU kernels are being developed.
  • Discussion and future work: Scatterbrain unifies sparse and low-rank attention approximation with theoretical guarantees and empirical validation across pretrained and end-to-end Transformer settings.The paper applies it to BigGAN, vision transformers, and vanilla Transformer training.
  • Background and related work: Robust PCA decomposes a matrix into sparse and low-rank components, but conventional optimization can be computationally expensive and may require full matrix materialization.The attention setting motivates efficient alternatives based on LSH and low-rank decomposition without materializing the full attention matrix.
  • Discussion and future work: Scatterbrain provides an unbiased attention-matrix estimate and supports direct replacement of full attention without retraining.This enables evaluation through approximation error and drop-in deployment in pretrained models.
  • Discussion and future work: Longformer and Long-short Transformer can be viewed as special cases of Scatterbrain with local sparse attention and restricted low-rank components.The paper distinguishes Longformer’s global-token mechanism from a true low-rank approximation because of the nonlinearity between its attention steps.

A.3 Locality Sensitive Hashing for Efficient Neural Network Training

The paper situates LSH within efficient neural-network training and structured matrix approximation, while identifying the difficulty of finding suitable structured approximations for attention matrices.

  • Locality Sensitive Hashing: LSH accelerates similarity search by using indexing structures that trade search quality for sub-linear search time.LSH has also been applied to estimation problems and efficient neural-network training.
  • Locality Sensitive Hashing: Efficient neural-network training can be framed as similarity search, with weights treated as search data and inputs as queries.Prior methods use this view to perform sparse forward and backward computations via maximum inner-product search.
  • Structured matrices: Sparse-plus-low-rank matrices offer asymptotically fast matrix-vector multiplication and subquadratic parameter storage.The paper lists sparse, low-rank, and fast-transform matrices as examples of structured matrices.
  • Structured matrices: Fixed structured transforms are difficult to use for attention because efficient algorithms for finding the closest structured matrix to a given attention matrix are lacking.This limits direct approximation of attention with more general structured matrix classes.
  • Research goals: The paper studies where sparse and low-rank approximations excel and whether combining them yields better approximation.It reports negatively correlated errors, entropy-dependent regimes, and potential gains from sparse-plus-low-rank structure.
  • Experimental setup: Approximation errors are measured with the Frobenius norm on attention matrices collected from IMDb, WikiText103, and BigGAN models.The study uses attention matrices from vanilla Transformers and a pretrained BigGAN, with additional collection procedures described for WikiText103 and BigGAN.

B.2 Observation 1: Sparse and low-rank approximation errors are negatively correlated

Across IMDb, WikiText-103, and BigGAN-ImageNet, sparse and low-rank approximation errors are statistically significantly negatively correlated, indicating complementary strengths.

  • Correlation analysis: Sparse and low-rank approximation errors are negatively correlated across IMDb, WikiText-103, and BigGAN-ImageNet.The reported correlations use Spearman’s rank, Pearson, and Kendall’s Tau coefficients.
  • Correlation analysis: The experiments fix the parameter budget K and compare ideal top-K row entries for sparse approximation with top-K eigenvalues for low-rank approximation.Three standard statistical correlation tests are applied to the resulting errors.
  • Interpretation: The negative correlation suggests that a property of the softmax attention distribution determines when sparse or low-rank approximation performs better.The next observation tests this relationship using attention-distribution entropy.

B.3 Observation 2: Sparse approximation error is lower when softmax entropy is low and low-rank approximation error is lower error when entropy is high

Approximation quality depends on attention entropy: sparse methods suit peaked, low-entropy distributions, while low-rank methods suit diffuse, high-entropy distributions. Their combination can reduce error across entropy regimes.

  • Observation 2: Figure 7 relates softmax-distribution entropy and approximation error across IMDb, WikiText103, and BigGAN-ImageNet.Its left plots compare sparse and low-rank errors, while its right plots relate error to entropy.
  • Observation 2: For high-entropy attention distributions, low-rank matrices approximate attention well; for low-entropy distributions, sparse matrices are better suited.Entropy is computed row-wise and averaged, with the attention matrices padded to 1024×1024.
  • Observation 3: Sparse and low-rank approximation errors are negatively correlated, so their strengths can complement one another across attention matrices.This complementarity motivates combining the two approximation types.
  • Observation 3: Sparse-plus-low-rank approximation can achieve lower error than either method alone across the entropy range when the mix ratio is chosen appropriately.The ideal combination is described using robust PCA.
  • Observation 3: Combining sparse and low-rank attention requires an efficient decomposition, accurate estimation without full-matrix materialization, and sufficient expressive power.The paper identifies these as central computational, accuracy, and representation challenges.

C Scatterbrain Algorithm and Implementation Details

Scatterbrain approximates attention by combining sparse and low-rank components, using locality-sensitive hashing and kernel feature maps with normalization. Theoretical constructions show that this combination can achieve parameter-efficiency and approximation accuracy unavailable to either component alone.

  • Algorithm: Scatterbrain combines a kernel-feature low-rank approximation with an LSH-based sparse approximation and normalizes their summed outputs.The algorithm computes low-rank and sparse outputs separately, then applies diagonal normalization to their sum.
  • Expressiveness: Sparse + low-rank approximation can require O(n) parameters where sparse or low-rank alone requires O(n^2) for comparable Frobenius error.This separation is demonstrated for constructed attention-matrix families.
  • Expressiveness: Sparse + low-rank uses the sparse component for large diagonal entries and the low-rank component for the remaining entries.The construction exploits matrices whose diagonal is difficult for low-rank approximation while their off-diagonal structure is difficult for sparse approximation.
  • Expressiveness: For one matrix family, sparse + low-rank estimation uses O(γ^-1n^1.5 log n) parameters with Frobenius error at most γ√n.The corresponding sparse or low-rank estimators have error at least Ω(√n) under the stated rank or sparsity conditions.
  • Estimator construction: The estimator matches the attention matrix on selected sparse locations while approximating the remaining structure through its low-rank component.This construction underlies the unbiasedness and variance analysis developed later.

D.3 Scatterbrain: Analysis

Scatterbrain combines sparse and low-rank estimators by selecting some attention entries with LSH and using the low-rank estimator elsewhere. Under stated assumptions, it remains unbiased and has lower variance than the Performer estimator, while experiments evaluate it across language, classification, and image settings.

  • Estimator analysis: Scatterbrain is unbiased because it combines two estimators that are each unbiased for the unnormalized attention entry.On selected indices it uses the exact sparse-side value, and elsewhere it uses the low-rank estimator.
  • Estimator analysis: Scatterbrain variance is (1 − p) times Performer variance and is therefore strictly lower when p is positive.The selection probability p is determined by the LSH collision bound under the theorem’s unit-sphere and proximity assumptions.
  • Evaluation settings: The evaluation covers WikiText103, Copy, Long Range Arena tasks, GLUE, BigGAN, and pre-trained vision transformers.The listed settings include autoregressive language modeling, long-range classification, image generation, and image classification.
  • Implementation: Inference allocates relatively more memory to the low-rank component when observed attention entropy is high.The implementation describes this allocation rule as an inference-time heuristic based on an observed subset of attention matrices.

E.3.1 Memory Budget

Scatterbrain maintains an advantage over sparse and low-rank baselines across memory budgets and can combine Performer with Reformer to improve the efficiency–accuracy trade-off. The framework also supports alternative sparse components, though local attention remains weak on Copy.

  • Memory-budget ablation: Scatterbrain outperforms sparse and low-rank baselines across a sweep of WikiText-103 parameter budgets.The ablation varies the memory budget and reports validation perplexity and speed-up relative to full attention.
  • Memory-budget ablation: Combining Performer with Reformer achieves the same accuracy using half the memory and runs faster than Reformer in the reported setting.The comparison is presented as exploiting the sparse + low-rank structure of attention.
  • Alternative components: Scatterbrain can replace Reformer with local or block-sparse attention while retaining the same combination framework.With local or block-sparse attention, the sparse support is fixed rather than input-adaptive.
  • Alternative components: Local attention is reasonably competitive on WikiText-103 but performs poorly on Copy because Copy requires non-local attention lookups.This result applies to the additional local-attention experiments.

E.4 Analysis

Across analyzed tasks, Scatterbrain generally has the smallest approximation error, while sparse or low-rank baselines are weaker in particular regimes. Additional experiments report strong performance when replacing attention in BERT and pre-trained vision models.

  • Approximation-error analysis: Scatterbrain has the smallest approximation error in most language-modeling cases, while sparse attention has the largest error.The language-modeling pattern is reported to match the end-to-end results.
  • Approximation-error analysis: Scatterbrain also has the smallest approximation error in the reported classification analysis, while Performer is worst on listed tasks.The passage connects this pattern to the task-dependent behavior of kernel or low-rank approximations.
  • GLUE fine-tuning: Replacing BERT attention with Scatterbrain outperforms the other tested baselines on most downstream GLUE tasks at half the full-attention parameter budget.The experiment fine-tunes BERT on nine downstream tasks.
  • Further discussion: The framework can extend to alternatives such as Routing Transformer or Nystromformer and may support choosing approximation methods using attention entropy.These extensions and selection opportunities are presented as future applications of the framework’s observations.
Loading 2110.15343v1…