PaperlayerView exact revision on arXiv

Source-linked AI summary

Linformer: Self-Attention with Linear Complexity

Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, Hao Ma

arXiv:2006.04768v3cs.LGstat.ML

TL;DR

Transformer self-attention has O(n^2) time and space complexity, making large models costly for long sequences. Linformer approximates self-attention with a low-rank factorization implemented through linear projections, reducing the operation to O(n). It achieves comparable downstream performance while improving inference efficiency.

  • Problem

    Transformer self-attention has O(n^2) time and space complexity with respect to sequence length, creating a major efficiency bottleneck.

  • Method

    Linformer approximates the self-attention matrix as low-rank and decomposes attention into smaller attentions through linear projections.

  • Results

    Linformer achieves O(n) self-attention complexity and has 1.5× faster inference time with a 1.7× larger maximum batch size at n = 512 and k = 128.

  • Takeaways & Limitations

    Linformer provides a more time- and memory-efficient Transformer while maintaining comparable downstream performance to standard Transformer models.

Abstract

from arXiv · show

Large transformer models have shown extraordinary success in achieving state-of-the-art results in many natural language processing applications. However, training and deploying these models can be prohibitively costly for long sequences, as the standard self-attention mechanism of the Transformer uses $O(n^2)$ time and space with respect to sequence length. In this paper, we demonstrate that the self-attention mechanism can be approximated by a low-rank matrix. We further exploit this finding to propose a new self-attention mechanism, which reduces the overall self-attention complexity from $O(n^2)$ to $O(n)$ in both time and space. The resulting linear transformer, the \textit{Linformer}, performs on par with standard Transformer models, while being much more memory- and time-efficient.

1 Introduction

Transformers deliver strong NLP performance but their self-attention creates a quadratic sequence-length bottleneck. Linformer exploits the low-rank structure of self-attention to achieve linear complexity while maintaining comparable task performance.

  • Motivation: O(n^2) self-attention complexity makes Transformer training and deployment slow for long sequences.The bottleneck arises because each token attends to all other tokens.
  • Related approaches: Existing sparse attention reduces complexity to O(n√n), but can incur a 2% performance drop for only 20% speedup.This limitation motivates alternatives to sparsifying attention.
  • Approach: Linformer approximates the self-attention matrix with a low-rank matrix and decomposes attention through linear projections.The projected attentions combine into a low-rank factorization of the original attention.
  • Approach: The proposed mechanism reduces self-attention complexity to O(n) in both time and space.The paper reports this reduction through theoretical and empirical analysis.
  • Results: Linformer shows similar pretraining performance to the standard Transformer and comparable or slightly better performance on downstream tasks.The evaluated tasks include three GLUE tasks and IMDB sentiment analysis.

2 Backgrounds and Related works

Transformer self-attention forms a context-mapping matrix whose computation scales quadratically with sequence length. Prior efficiency methods trade accuracy, speed, memory, or applicability, motivating direct optimization of the attention layer.

  • Transformer self-attention: Multi-head self-attention jointly attends to information at different positions from multiple representation subspaces.It uses input embedding matrices Q, K, and V and combines multiple attention heads.
  • Transformer self-attention: The context-mapping matrix P captures sequence-wide context but requires O(n^2) time and space to compute.The quadratic cost comes from multiplying two n × d matrices.
  • Efficiency methods: Mixed-precision training improves numerical efficiency but is orthogonal to Linformer’s attention optimization.The paper uses mixed-precision training by default.
  • Efficiency methods: Knowledge distillation speeds inference through a smaller student model but does not accelerate teacher training and can reduce performance.Distilling 12-layer BERT to 6-layer BERT produced an average 2.5% performance drop on several benchmark tasks.
  • Efficiency methods: Sparse attention and LSH attention reduce attention complexity but have limited practical gains or performance degradation.Sparse attention reported a 2% drop with 20% speedup, while LSH attention is most effective only for extremely long sequences.
  • Efficiency methods: Microbatching and gradient checkpointing trade time for memory and do not speed up inference.These techniques optimize training memory rather than the attention operation itself.
  • Motivation: These limitations motivate optimizing self-attention layers directly to reduce both training and inference time and memory consumption.The paper introduces its approach after surveying these efficiency techniques.

3 Self-Attention is Low Rank

The paper shows that Transformer self-attention’s context mapping matrix P is low-rank, theoretically and empirically, enabling approximation with substantially lower complexity.

  • 3 Self-Attention is Low Rank: Self-attention’s context mapping matrix P is demonstrated to be low-rank.The section introduces this as the central structural property of self-attention.
  • 3 Self-Attention is Low Rank: Spectrum analysis across RoBERTa models, Wiki103, and IMDB shows a clear long-tail distribution of singular values.The analysis averages normalized cumulative singular values across layers and heads over 10k sentences.
  • 3 Self-Attention is Low Rank: Theorem 1 states that an approximate context mapping matrix can achieve ϵ-relative error with rank Θ(log(n)).The theorem gives the probability guarantee Pr(∥P̃w^T−Pw^T∥<ϵ∥Pw^T∥)>1−o(1).
  • 3 Self-Attention is Low Rank: The theoretical proof constructs an approximate low-rank matrix using the distributional Johnson–Lindenstrauss lemma and random projections.The construction uses a diagonal normalization matrix and a projection matrix with i.i.d. Gaussian entries.
  • 3 Self-Attention is Low Rank: SVD can approximate P with O(nk) time and space complexity, but computing an SVD for every self-attention matrix adds complexity.This motivates a different low-rank approximation approach that avoids per-matrix SVD decomposition.

4 Model

The proposed linear self-attention mechanism projects keys and values into lower-dimensional spaces, reducing sequence-length dependence while preserving an efficient multihead design.

  • 4 Model: The proposed mechanism computes contextual mapping through linear projections, achieving linear time and memory complexity in sequence length.Projection matrices E_i and F_i have shape n×k and are applied when computing keys and values.
  • 4 Model: Linformer’s architecture uses two linear projection matrices for each attention head’s key and value computations.The model decomposes the original attention computation by projecting the original n×d representations.
  • 4 Model: Figure 2 compares the proposed architecture with inference time versus sequence length for Linformer and standard Transformer models.The figure holds the total number of tokens fixed in its speed comparison.
  • 4 Model: Standard Transformer inference becomes slower at longer sequence lengths, whereas Linformer speed remains relatively flat and is significantly faster there.This comparison is reported from the top-right plot of Figure 2.
  • 4 Model: Projection parameters can be shared across heads, key-value pairs, or layers to further optimize efficiency.These sharing schemes reduce the number of distinct projection matrices used by the model.
  • 4 Model: Projected dimension k can vary across heads and layers, with higher layers potentially using smaller k because their spectra are more skewed.The paper also discusses pooling and convolution as alternative projection methods.

5 Experiments

The experiments evaluate Linformer’s pretraining, downstream-task performance, and inference efficiency against Transformer-based models. Across these settings, Linformer retains comparable performance while improving efficiency and supporting longer sequences.

  • Pretraining Perplexities: Linformer’s validation perplexity is nearly on par with the original Transformer at k = 128 for n = 512 and k = 256 for n = 1024.Increasing the projected dimension improves performance.
  • Pretraining Perplexities: With k fixed at 256, final Linformer perplexities remain about the same as sequence length increases from 512 to 4096.This empirically supports the paper’s assertion that Linformer is linear-time.
  • Downstream Results: Linformer has comparable downstream performance to RoBERTa at n = 512 and k = 128, and slightly outperforms it at k = 256.Layerwise sharing achieves the best accuracy among the three parameter-sharing strategies.
  • Downstream Results: The n = 1024, k = 256 Linformer has similar downstream results to the n = 512, k = 256 model.These results support performance being determined mainly by projected dimension k rather than the ratio n/k.
  • Inference-time Efficiency Results: At n = 512 and k = 128, Linformer provides 1.5× faster inference and a 1.7× larger maximum batch size than the Transformer.Inference speed-ups and memory savings become more dramatic as sequence length increases.

6 Conclusion

The paper identifies quadratic self-attention as a source of Transformer inefficiency and proposes a theoretically and empirically supported linear-time alternative.

  • Conclusion: Self-attention has O(n^2) time and space complexity, motivating a mechanism with O(n) complexity in sequence length.The paper demonstrates the linear complexity through theoretical and empirical analysis.
  • Conclusion: The paper demonstrates that the stochastic matrix formed by self-attention is low-rank and leverages this observation to design a more efficient mechanism.

Broader Impact

The paper frames efficient Transformers as potentially broadening model accessibility and enabling applications involving very long sequences.

  • Broader Impact: Efficient Transformers could increase model accessibility for deployment on devices and for training in research.
  • Broader Impact: Linear-time attention could support training Transformers on images by handling very long sequences.
  • Broader Impact: The paper identifies potential positive environmental benefits associated with more efficient Transformer computation.

A Proof of Theorem 1

The proof constructs a low-rank approximation using a random projection and invokes the distributional Johnson–Lindenstrauss lemma to control approximation error.

  • Proof of Theorem 1: The proof uses the distributional Johnson–Lindenstrauss lemma as its main proof idea.
  • Proof of Theorem 1: A random k × n matrix R with i.i.d. N(0, 1/k) entries is used in the constructive approximation.
  • Proof of Theorem 1: The constructed approximation has rank bounded by k, while the proof sets k = 5 log(n)/(ϵ^2 − ϵ^3) to establish the theorem.

B Proof of Theorem 2

The proof uses Johnson–Lindenstrauss-based bounds to show that a low-dimensional projection can approximate the relevant self-attention products, with projection dimension independent of sequence length when exploiting rank.

  • Proof of Theorem 2: The proof bounds the approximation error for products involving rows of QK^T and columns of V using norm inequalities and the Johnson–Lindenstrauss lemma.The argument also uses Lipschitz continuity of the exponential function after a small scaling δ = θ(1/n).
  • Proof of Theorem 2: The proof applies the approximation bound to every row of A and every column of V to establish the desired matrix-level approximation.The cited steps combine the row-wise result with the construction based on the rank-d submatrix.
  • Proof of Theorem 2: The row-wise approximation result initially requires k = 5 log(nd)/(ϵ2 − ϵ3), which depends on sequence length n.This bound does not use the low-rank property rank(A) = d.
  • Proof of Theorem 2: Using rank(A) = d, the proof selects a 2d × d row submatrix As with rank d and applies the approximation result to obtain k = 9 log(d)/(ϵ2 − ϵ3).This choice removes the direct dependence of k on sequence length n.
Loading 2006.04768v3…