Hugging Face Reads: Long-range Transformers

TL;DR

Hugging Face has analyzed four primary architectural approaches—Longformer, Compressive Transformer, Linformer, and Performer—to solve the quadratic memory and time complexity bottleneck of standard Transformers. These methods enable the processing of sequences far beyond the traditional 512 or 1024 token limits, which is critical for document-level NLP, speech, and protein modeling.

Overcoming the Quadratic Bottleneck

Standard Transformer self-attention scales quadratically ($O(n^2)$) with sequence length $n$, creating a significant memory and compute bottleneck for long documents. To address this, researchers have developed "Efficient Transformers" that aim to reduce this complexity to linear ($O(n)$). These approaches generally fall into four categories: custom attention patterns, recurrence, low-rank approximations, and kernel approximations.

Architectural Approaches to Long-Range Modeling

Longformer: Custom Attention Patterns

Longformer replaces standard self-attention with a combination of windowed (local) and global attention, allowing it to scale linearly with sequence length.

  • Mechanism: It uses dilated windowed self-attention for autoregressive language modeling and a mix of local windowed and global bi-directional attention for encoder pre-training. Global attention is applied to task-specific tokens (e.g., the [CLS] token or question tokens in QA) to allow information to flow across the entire sequence.
  • Key Advantage: The self-attention layer is a drop-in replacement, meaning pre-trained checkpoints can be adapted to long-range inputs without requiring costly pre-training from scratch.
  • Trade-off: The sliding window attention relies on indexing operations that can be slow on certain hardware, such as TPUs.

Compressive Transformer: Recurrence

Building on Transformer-XL, the Compressive Transformer introduces a compressed memory to store past activations that would otherwise be discarded.

  • Mechanism: It uses a compression function (such as max/mean pooling or 1D convolution) to compress past activations by a factor $c$. This allows the model to attend to both a regular memory of recent tokens and a compressed memory of much older tokens.
  • Key Advantage: It significantly improves perplexity on long-range language modeling benchmarks like enwik8 and WikiText-103, particularly for rare words occurring over long distances.
  • Trade-off: Training is brittle and requires a specialized optimization schedule where the effective batch size is progressively increased.

Linformer: Low-Rank Approximations

Linformer reduces complexity by projecting the sequence length into a smaller dimension, based on the observation that the self-attention matrix is low-rank.

  • Mechanism: Utilizing the Johnson-Lindenstrauss lemma, Linformer learns a low-rank decomposition of the attention context matrix. This ensures that no $n \times n$ matrix is ever computed or stored.
  • Key Advantage: Inference speed (time-clock) is not affected by increasing sequence length, and convergence speed remains stable compared to standard Transformers.
  • Trade-off: The decomposition is designed for a fixed context length determined at training and does not generalize to longer sequences without adaptation.

Performer: Kernel Approximations

Performer uses the FAVOR+ (Fast Attention Via Orthogonal Random positive features) algorithm to approximate the softmax attention kernel without relying on sparsity or low-rank priors.

  • Mechanism: It uses random feature maps to approximate the softmax function, allowing the matrix multiplication $K \times V$ to be performed before the query multiplication. This effectively bypasses the computation of the $n \times n$ attention matrix.
  • Key Advantage: Because it makes no assumptions about the attention matrix's structure, it is highly applicable across different modalities, including speech and protein sequences.
  • Trade-off: Small approximation errors can propagate through multiple Transformer layers, potentially impacting the stability of fine-tuning pre-trained networks.

Comparative Analysis and Discussion

Inductive Biases and Trade-offs

The choice of architecture depends on the specific requirements of the task and the available data:

  • Longformer vs. Linformer: Longformer uses fixed sparse patterns, while Linformer learns a low-rank factorization. Longformer is generally less efficient than Linformer but more flexible regarding sequence length (though Linformer is limited to its training context length).
  • Performer: Differs from both by approximating the kernel itself, making it a versatile drop-in replacement that does not assume the attention matrix is sparse or low-rank.

Positional Embeddings

Positional encoding is a critical factor in long-range efficiency:

  • Relative Positional Embeddings: Used in Transformer-XL and Compressive Transformers; they extend easily to unseen sequence lengths but are computationally expensive.
  • Absolute Positional Embeddings: Used in Longformer and Linformer; they are more computationally efficient but less flexible for sequences longer than those seen during training.
  • Position-Infused Attention: An alternative introduced by Shortformer that adds positional information to queries and keys rather than token embeddings.

Training Strategies

Evidence from models like Shortformer and Longformer suggests that training on short sequences and gradually increasing the length leads to faster training and stronger downstream performance, preventing the model from relying on spurious correlations in the data.

Sources