Efficient Transformers

Chapter Overview

Standard transformers have $O(n^2)$ complexity in sequence length, limiting their application to long sequences. This chapter covers efficient attention mechanisms that reduce complexity: sparse attention, linear attention, low-rank methods, and kernel-based approaches.

Learning Objectives

1. Understand the quadratic bottleneck in standard attention
2. Implement sparse attention patterns (sliding window, strided, global)
3. Apply Linformer and Performer for linear complexity
4. Use Flash Attention for memory-efficient computation
5. Compare trade-offs: accuracy vs efficiency vs memory
6. Deploy long-context models (Longformer, BigBird)

The Quadratic Bottleneck

Complexity Analysis

The standard self-attention mechanism computes attention scores between all pairs of tokens in a sequence, leading to computational and memory requirements that scale quadratically with sequence length. The attention operation is defined as:

The computational bottleneck arises from computing the attention matrix $\mQ \mK\transpose \in \R^{n \times n}$, which requires $O(n^2 d)$ floating-point operations. For each of the $n$ queries, we compute dot products with all $n$ keys, where each dot product involves $d$ multiplications and additions. The subsequent softmax normalization adds $O(n^2)$ operations, and the final multiplication with values $\mV$ requires another $O(n^2 d)$ operations. The memory bottleneck is equally severe: storing the attention matrix requires $O(n^2)$ memory, which must be materialized before the softmax operation and retained for the backward pass during training.

This quadratic scaling becomes prohibitive for long sequences. Consider a BERT-base model with $d = 768$ and 12 attention heads processing a sequence of length $n = 4096$. Each attention head must store an attention matrix of size $4096 \times 4096$, requiring $4096^2 \times 4 = 67$ MB in FP32 format. Across all 12 heads, this amounts to 804 MB just for attention weights in a single layer. With 12 layers, the total memory for attention matrices alone reaches 9.6 GB, nearly filling an NVIDIA V100 GPU with 16 GB memory before accounting for activations, gradients, or model parameters.

Sparse Attention Patterns

Efficiency Taxonomy

Efficient attention mechanisms can be categorized into five main approaches, each targeting different aspects of the quadratic bottleneck. Sparse attention methods reduce the number of attention connections by restricting each query to attend to only a subset of keys, achieving $O(n \times k)$ complexity where $k \ll n$. Linear attention methods use mathematical approximations to avoid computing the full attention matrix, achieving $O(n)$ complexity in sequence length. Low-rank methods project keys and values to lower-dimensional spaces, reducing the effective size of the attention computation. Kernel-based methods reformulate attention using kernel functions and random features to enable linear-time computation. Finally, recurrent methods process sequences in chunks with recurrent connections, trading parallelism for reduced memory.

Each approach involves different trade-offs between computational efficiency, memory usage, approximation quality, and implementation complexity. Sparse methods maintain exact attention within their connectivity pattern but may miss important long-range dependencies. Linear methods achieve impressive speedups but introduce approximation errors that can degrade model quality. Low-rank methods work well when attention patterns have inherent low-rank structure but may fail for complex attention distributions. Understanding these trade-offs is essential for selecting the appropriate efficient attention mechanism for a given application.

Fixed Sparse Patterns

Sparse attention restricts each query to attend to only a subset of keys, dramatically reducing both computation and memory requirements. The fundamental idea is to identify which attention connections are most important and compute only those, setting all other attention weights to zero or negative infinity before the softmax operation.

The choice of sparsity pattern $\mathcal{S}$ determines which information can flow through the network. Three fundamental patterns have emerged as particularly effective building blocks for sparse attention.

The sliding window or local attention pattern restricts each token to attend only to nearby tokens within a fixed window. Formally, $\mathcal{S}_{\text{local}}(i) = \{j : |i-j| \leq w\}$ where $w$ is the window size. Each token attends to $2w+1$ tokens: itself, $w$ tokens before, and $w$ tokens after. This pattern is motivated by the observation that in many domains, particularly natural language, nearby tokens are more relevant than distant ones. For a window size $w = 256$ and sequence length $n = 4096$, each query attends to only 513 keys instead of 4096, reducing computation by 8× and memory by the same factor. The limitation is that information can only propagate $w$ positions per layer, requiring $\lceil n/w \rceil$ layers for full sequence communication.

The strided or dilated attention pattern samples tokens at regular intervals: $\mathcal{S}_{\text{strided}}(i) = \{j : (i-j) \bmod s = 0\}$ where $s$ is the stride. This pattern allows each token to attend to distant tokens, enabling faster information propagation across the sequence. With stride $s = 64$, a token at position 1024 can attend to positions 0, 64, 128, ..., 1024, ..., 4032, providing long-range connectivity with only $n/s$ connections per query. However, strided attention alone misses local context, so it is typically combined with local attention in alternating layers.

Global attention designates certain tokens as global tokens that attend to all positions and are attended to by all positions. These tokens act as information hubs, aggregating information from the entire sequence and broadcasting it back. In practice, special tokens like [CLS] in BERT or separator tokens are often designated as global. For $g$ global tokens in a sequence of length $n$, each global token requires $O(n)$ computation, and each regular token requires $O(g)$ additional computation to attend to globals, adding $O(ng)$ total cost.

BigBird: Random + Window + Global

BigBird extends sparse attention by combining three complementary patterns: local windows for nearby context, random connections for long-range dependencies, and global tokens for information aggregation. This combination provides both theoretical guarantees and practical efficiency for processing sequences up to 4096 tokens or longer.

The random attention component is crucial for BigBird's theoretical properties. While local windows provide nearby context and global tokens enable information aggregation, random connections create shortcuts across the sequence that allow information to propagate efficiently. The random graph formed by these connections has high probability of being well-connected, ensuring that any two positions are connected by a short path through the attention graph. This property enables BigBird to approximate full attention's expressiveness while maintaining linear complexity.

BigBird's theoretical contribution is proving that this sparse attention pattern can approximate any sequence-to-sequence function that full attention can compute, under mild assumptions. Specifically, BigBird with $r = O(\log n)$ random connections per query can approximate full attention with high probability, providing a theoretical foundation for sparse attention methods. This result shows that $O(n \log n)$ total connections suffice for universal approximation, compared to $O(n^2)$ for full attention.

In practice, BigBird uses $w = 256$, $r = 64$, and $g = 32$ for sequences up to 4096 tokens. Each regular token attends to $2 \times 256 + 64 + 32 = 608$ keys instead of 4096, reducing computation by 6.7×. For a single attention head with $d = 768$ in FP32, BigBird requires $(4096 - 32) \times 608 + 32 \times 4096) \times 4 \approx 10.4$ MB compared to 67 MB for full attention, a 6.4× memory reduction. With 12 heads and 12 layers, total attention memory decreases from 9.6 GB to 1.5 GB.

The performance benefits are substantial on modern hardware. On an NVIDIA A100 GPU, BigBird processes 4096-token sequences in approximately 15 milliseconds per layer compared to 98 milliseconds for full attention, a 6.5× speedup. The speedup is slightly less than the theoretical 6.7× due to overhead from irregular memory access patterns in the random attention component. For sequences of 8192 tokens, BigBird takes 30 milliseconds per layer while full attention would require approximately 390 milliseconds, a 13× speedup that makes previously impractical sequence lengths feasible.

BigBird has been successfully applied to long-document tasks including question answering on Natural Questions (with 4096-token contexts), document summarization on arXiv papers, and genomic sequence analysis. On the Natural Questions benchmark, BigBird achieves 79.2

Linear Attention Methods

Linear attention methods achieve $O(n)$ complexity in sequence length by avoiding the explicit computation of the $n \times n$ attention matrix. These methods use mathematical reformulations or approximations that allow attention to be computed through matrix operations with different associativity, reducing the dominant term from $O(n^2 d)$ to $O(nd^2)$ or even $O(nd)$ in some cases.

Linformer

Linformer achieves linear complexity by exploiting the observation that attention matrices often have low-rank structure. Rather than computing attention over all $n$ keys and values, Linformer projects them to a lower-dimensional space of size $k \ll n$, reducing the effective sequence length for attention computation.

The key insight is that the attention matrix $\mA = \text{softmax}(\mQ \mK\transpose / \sqrt{d})$ often has low-rank structure, meaning it can be well-approximated by a rank-$k$ matrix with $k \ll n$. Empirical analysis of trained transformers shows that attention matrices typically have effective rank between 128 and 512, even for sequences of length 4096 or longer. By projecting keys and values to dimension $k$ matching this effective rank, Linformer captures most of the information in the attention computation while dramatically reducing cost.

The projection matrices $\mE$ and $\mF$ can be implemented in several ways. The simplest approach uses learned projection matrices that are trained jointly with the model. Alternatively, fixed projections such as max pooling or average pooling can be used, where $\mE$ and $\mF$ partition the sequence into $k$ segments and pool within each segment. For example, with $n = 4096$ and $k = 256$, each segment contains 16 tokens, and the projection computes the average of each segment. Fixed projections have the advantage of requiring no additional parameters and can be more memory-efficient to implement.

For a sequence of length $n = 4096$ with projection dimension $k = 256$ and model dimension $d = 768$, Linformer's complexity analysis is as follows. Computing $\bar{\mK} = \mE \mK$ requires $O(nkd) = 4096 \times 256 \times 768 \approx 805$ million FLOPs. Computing $\mQ \bar{\mK}\transpose$ requires $O(nkd) = 805$ million FLOPs. The softmax over the $n \times k$ matrix requires $O(nk) = 1$ million operations, and the final multiplication with $\bar{\mV}$ requires another $O(nkd) = 805$ million FLOPs. The total is approximately 2.4 billion FLOPs compared to $O(n^2d) = 4096^2 \times 768 \approx 12.9$ billion FLOPs for full attention, a 5.4× reduction.

Memory requirements are similarly reduced. The attention matrix $\mQ \bar{\mK}\transpose \in \R^{n \times k}$ requires $4096 \times 256 \times 4 = 4.2$ MB in FP32 compared to 67 MB for the full $n \times n$ matrix, a 16× reduction. With 12 heads and 12 layers, total attention memory decreases from 9.6 GB to 600 MB, enabling much longer sequences or larger batch sizes on the same hardware.

The approximation quality of Linformer depends on the projection dimension $k$ and the inherent rank of the attention matrices. Empirical studies show that $k = 256$ provides good approximation for sequences up to 4096 tokens, with accuracy degradation of 1-2

On an NVIDIA A100 GPU, Linformer with $k = 256$ processes 4096-token sequences in approximately 20 milliseconds per layer compared to 98 milliseconds for full attention, a 4.9× speedup. The speedup is less than the theoretical 5.4× due to the overhead of the projection operations and less efficient memory access patterns. For sequences of 8192 tokens, Linformer takes 40 milliseconds per layer while full attention would require 390 milliseconds, a 9.8× speedup that enables processing of very long documents.

Performer (Kernel-based)

Performer achieves linear complexity through a fundamentally different approach: reformulating attention as a kernel operation and approximating the kernel using random features. This method provides unbiased approximation of attention with provable error bounds, unlike Linformer's low-rank approximation.

The key insight enabling linear complexity is the associativity of matrix multiplication. In standard attention, we compute $(\mQ \mK\transpose) \mV$, which requires first computing the $n \times n$ matrix $\mQ \mK\transpose$ at cost $O(n^2 d)$. By approximating the attention kernel with feature maps $\phi$, we can instead compute $\phi(\mQ) (\phi(\mK)\transpose \mV)$, where the parentheses indicate we first compute $\phi(\mK)\transpose \mV \in \R^{m \times d}$ at cost $O(nmd)$, then multiply by $\phi(\mQ) \in \R^{n \times m}$ at cost $O(nmd)$. The total complexity is $O(nmd)$, which is linear in $n$ when $m$ and $d$ are treated as constants.

Performer uses the FAVOR+ (Fast Attention Via Orthogonal Random features) algorithm, which constructs the feature map $\phi$ using random projections. For a query or key vector $\vx \in \R^d$, the feature map is defined as:

FAVOR+ improves upon basic random features by using orthogonal random features, where the random vectors $\vw_1, \ldots, \vw_m$ are orthogonalized using Gram-Schmidt or similar procedures. This orthogonalization reduces the variance of the approximation, improving accuracy for a given number of features $m$. Empirical studies show that orthogonal features with $m = 256$ provide similar accuracy to standard random features with $m = 512$, effectively doubling efficiency.

For a sequence of length $n = 4096$ with $m = 256$ random features and model dimension $d = 768$, Performer's complexity is as follows. Computing $\phi(\mQ)$ and $\phi(\mK)$ requires $O(nmd) = 4096 \times 256 \times 768 \approx 805$ million FLOPs each. Computing $\phi(\mK)\transpose \mV$ requires $O(nmd) = 805$ million FLOPs, and multiplying by $\phi(\mQ)$ requires another $O(nmd) = 805$ million FLOPs. The total is approximately 3.2 billion FLOPs compared to 12.9 billion for full attention, a 4× reduction. The memory requirement is $O(nm + md) = 4096 \times 256 + 256 \times 768 \approx 1.2$ million elements or 4.8 MB in FP32, compared to 67 MB for full attention.

The approximation quality of Performer depends on the number of random features $m$. With $m = 256$, Performer typically achieves accuracy within 2-3

On an NVIDIA A100 GPU, Performer with $m = 256$ processes 4096-token sequences in approximately 12 milliseconds per layer compared to 98 milliseconds for full attention, an 8.2× speedup. This speedup exceeds the theoretical 4× reduction in FLOPs because Performer's computation is more memory-bandwidth efficient—it never materializes the large $n \times n$ attention matrix, reducing memory traffic. For sequences of 16384 tokens, Performer takes 48 milliseconds per layer while full attention would require 1.5 seconds, a 31× speedup that enables processing of extremely long sequences.

Memory-Efficient Attention

Flash Attention

Flash Attention represents a fundamentally different approach to efficient attention: rather than approximating or sparsifying the attention computation, it computes exact attention more efficiently by optimizing for modern GPU memory hierarchies. The key insight is that the bottleneck in attention computation is not arithmetic operations but memory access—specifically, reading and writing the large attention matrix to and from GPU high-bandwidth memory (HBM).

Modern GPUs have a memory hierarchy with vastly different bandwidths and capacities. An NVIDIA A100 GPU has 40 GB of HBM with bandwidth 1.5 TB/s, and 20 MB of SRAM (shared memory) per streaming multiprocessor with bandwidth exceeding 19 TB/s—more than 12× faster. Standard attention implementations compute $\mS = \mQ \mK\transpose$, write it to HBM (consuming $n^2 \times 4$ bytes), read it back for softmax, write the result to HBM, read it back for multiplication with $\mV$, and finally write the output. For $n = 2048$, this involves reading and writing $2048^2 \times 4 = 16.8$ MB multiple times, totaling over 100 MB of memory traffic.

Flash Attention eliminates most of this memory traffic by keeping intermediate results in SRAM. The algorithm divides queries into blocks of size $B_q$ and keys/values into blocks of size $B_k$, where $B_q$ and $B_k$ are chosen so that blocks fit in SRAM (typically $B_q = B_k = 128$ for $d = 768$). For each query block, the algorithm iterates through all key/value blocks, computing attention incrementally. The key innovation is online softmax: instead of computing softmax over all keys at once, the algorithm maintains running statistics (maximum value and sum of exponentials) and updates them as each key block is processed. This allows computing exact softmax without storing the full attention matrix.

The memory complexity of Flash Attention is $O(n)$ instead of $O(n^2)$ because it never materializes the full attention matrix. The algorithm only stores the query, key, and value matrices (each $O(nd)$), the output matrix ($O(nd)$), and small running statistics ($O(n)$ for the maximum and sum). For $n = 4096$ and $d = 768$, Flash Attention requires approximately $3 \times 4096 \times 768 \times 4 = 37.7$ MB compared to $67 + 37.7 = 104.7$ MB for standard attention (attention matrix plus activations), a 2.8× memory reduction. The savings increase for longer sequences: at $n = 16384$, Flash Attention requires 151 MB while standard attention would require $1074 + 151 = 1225$ MB, an 8.1× reduction.

The computational complexity remains $O(n^2 d)$ since Flash Attention computes exact attention, but the wall-clock time is significantly reduced due to fewer memory accesses. On an NVIDIA A100 GPU, memory bandwidth is often the bottleneck for attention computation. Standard attention achieves only 30-40

Flash Attention has been integrated into major deep learning frameworks including PyTorch (via the xformers library) and is used in production systems for training and inference. The technique has been extended to Flash Attention 2, which provides additional optimizations including better parallelization across attention heads and improved handling of non-power-of-two sequence lengths, achieving up to 2× additional speedup over the original Flash Attention.

Memory-Efficient Transformers

Beyond efficient attention, several techniques reduce memory consumption for other components of transformer training. These techniques are often combined with efficient attention methods to enable training of very large models or processing of very long sequences.

Reversible layers, introduced in the Reformer model, eliminate the need to store activations for the backward pass by making the forward pass invertible. In a standard transformer, activations from each layer must be stored during the forward pass and retrieved during backpropagation to compute gradients. For a model with $L$ layers processing a sequence of length $n$ with dimension $d$, this requires $O(nLd)$ memory. Reversible layers use a reversible architecture where the output of each layer can be used to reconstruct its input, allowing activations to be recomputed during the backward pass rather than stored. This reduces activation memory from $O(nLd)$ to $O(nd)$, a factor of $L$ reduction. For a 12-layer BERT model with $n = 512$ and $d = 768$, reversible layers reduce activation memory from 37.7 MB to 3.1 MB per sequence.

Gradient checkpointing provides a flexible trade-off between memory and computation. Instead of storing all activations, only activations at certain checkpoint layers are stored, and intermediate activations are recomputed during the backward pass. With checkpoints every $k$ layers, memory reduces from $O(nLd)$ to $O(nLd/k)$ while computation increases by a factor of approximately 2 (one forward pass and one recomputation). For $k = 3$ in a 12-layer model, memory reduces by 3× while training time increases by only 20-30

Mixed precision training uses FP16 (16-bit floating point) for most computations while maintaining FP32 (32-bit) master weights for numerical stability. This reduces activation memory by 50

Comparison of Efficient Methods

Comprehensive Benchmarks

Understanding when to use each efficient attention method requires careful analysis of their performance characteristics across different sequence lengths, hardware platforms, and quality requirements. This section provides detailed benchmarks on NVIDIA A100 GPUs with concrete memory and speed measurements.

Memory Scaling Analysis

Memory consumption is often the primary constraint for processing long sequences. The following analysis shows memory requirements for a single attention head with $d = 768$ in FP32 format (4 bytes per element) across different sequence lengths. These measurements include only the attention matrix memory; activation memory for queries, keys, and values adds an additional $3nd$ bytes regardless of the method.

For $n = 1024$ tokens, standard attention requires $1024^2 \times 4 = 4.2$ MB per head. Sparse methods with window $w = 256$ require $1024 \times 512 \times 4 = 2.1$ MB (50

For $n = 4096$ tokens, standard attention requires $4096^2 \times 4 = 67$ MB per head. Sparse methods with $w = 512$ require $4096 \times 1024 \times 4 = 16.8$ MB (75

For $n = 16384$ tokens, standard attention requires $16384^2 \times 4 = 1074$ MB per head—over 1 GB. Sparse methods with $w = 512$ require $16384 \times 1024 \times 4 = 67$ MB (94

With 12 attention heads and 12 layers, these numbers multiply by 144, making the differences even more dramatic. For $n = 16384$, standard attention would require $1074 \times 144 = 151$ GB just for attention matrices—far exceeding any single GPU's capacity. Sparse methods require 9.4 GB, linear methods require 2.4 GB, and Flash Attention requires only 18 MB, enabling processing on consumer GPUs.

Speed Benchmarks on A100 GPU

Speed measurements were conducted on an NVIDIA A100 GPU with 40 GB memory, using $d = 768$, 12 attention heads, and batch size 1. Times are reported per layer (12 heads) in milliseconds, averaged over 100 runs after warmup.

For $n = 1024$ tokens, standard attention takes 8 milliseconds per layer. Sparse attention with $w = 256$ (Longformer-style) takes 5 milliseconds (1.6× speedup). Linformer with $k = 256$ takes 4 milliseconds (2× speedup). Performer with $m = 256$ takes 3 milliseconds (2.7× speedup). Flash Attention takes 3 milliseconds (2.7× speedup). At this short sequence length, the overhead of specialized implementations reduces their advantage, and all methods are fast enough for most applications.

For $n = 4096$ tokens, standard attention takes 98 milliseconds per layer. Sparse attention with $w = 512$ (Longformer) takes 18 milliseconds (5.4× speedup). BigBird with $w = 256$, $r = 64$, $g = 32$ takes 15 milliseconds (6.5× speedup). Linformer with $k = 256$ takes 20 milliseconds (4.9× speedup). Performer with $m = 256$ takes 12 milliseconds (8.2× speedup). Flash Attention takes 25 milliseconds (3.9× speedup). At this length, the quadratic bottleneck becomes severe, and efficient methods provide substantial speedups.

For $n = 16384$ tokens, standard attention takes 1.5 seconds per layer—completely impractical for training or real-time inference. Sparse attention with $w = 512$ takes 72 milliseconds (21× speedup). BigBird takes 60 milliseconds (25× speedup). Linformer with $k = 256$ takes 80 milliseconds (19× speedup). Performer with $m = 256$ takes 48 milliseconds (31× speedup). Flash Attention takes 220 milliseconds (6.8× speedup). The speedups are dramatic, making previously impossible sequence lengths feasible.

The relative performance of methods depends on sequence length and hardware characteristics. Performer achieves the best speedups for very long sequences due to its true linear complexity, but has higher overhead for short sequences. Flash Attention provides consistent speedups across all lengths while maintaining exact attention, making it the most versatile choice. Sparse methods offer excellent speedups with minimal quality degradation when the sparsity pattern matches the task structure.

Quality Trade-offs

Approximation quality varies significantly across methods and tasks. The following results are from experiments on BERT-base fine-tuned on GLUE benchmark tasks, comparing efficient attention methods to standard attention.

Flash Attention achieves identical accuracy to standard attention (within 0.1

Sparse attention methods (Longformer, BigBird) typically show 0.5-1.5

Linformer shows 1-2

Performer shows 2-3

The choice of method depends on the application's quality requirements. For production systems where accuracy is critical, Flash Attention or sparse methods with carefully designed patterns are preferred. For research or applications where 2-3

When to Use Each Method

Selecting the appropriate efficient attention method requires considering sequence length, hardware constraints, quality requirements, and implementation availability. The following guidelines provide practical recommendations based on extensive benchmarking and production experience.

For sequences with $n < 512$ tokens, use standard attention. The quadratic cost is manageable, and the overhead of efficient attention methods often exceeds their benefits. Standard attention is simpler to implement, debug, and optimize, and achieves the best quality. Most BERT-style models and many GPT-style models fall in this regime.

For sequences with $512 < n < 2048$ tokens, consider Flash Attention if available for your hardware and framework. Flash Attention provides 2-4× speedups with no quality degradation, making it an ideal drop-in replacement for standard attention. If Flash Attention is not available, sparse attention with window size $w = 256$ provides good speedups (2-3×) with minimal quality loss (< 1

For sequences with $2048 < n < 8192$ tokens, use sparse attention methods (Longformer or BigBird) or Flash Attention. Sparse methods provide 5-10× speedups and are well-suited for tasks where local context is important. Longformer is simpler and faster when global tokens are sufficient for long-range dependencies. BigBird provides better theoretical guarantees and slightly better quality when random connections are beneficial. Flash Attention provides 3-5× speedups with exact attention, making it preferable when quality is critical and memory is the primary constraint.

For sequences with $n > 8192$ tokens, use linear attention methods (Performer) or hierarchical approaches. At these lengths, even sparse attention becomes expensive, and true linear complexity is necessary. Performer with $m = 256$ provides 20-30× speedups compared to full attention, making sequences of 16384 or 32768 tokens feasible. Accept 2-3

Hardware considerations also matter. Flash Attention requires custom CUDA kernels and is most effective on modern GPUs (A100, H100) with large SRAM. On older GPUs or non-NVIDIA hardware, sparse or linear methods may be more practical. For CPU inference, sparse methods are often fastest due to efficient sparse matrix libraries. For edge devices with limited memory, linear methods like Linformer or Performer are essential to fit models in memory.

Task structure should inform the choice of sparsity pattern. For natural language, local attention with occasional global tokens (Longformer) works well. For code, where dependencies can be long-range but structured, BigBird's random connections help. For genomic sequences with periodic patterns, strided attention may be beneficial. For images, local attention in spatial dimensions is natural. Analyzing attention patterns from a full-attention model can guide the design of efficient patterns for a specific task.

Long-Context Models

Longformer

Longformer is a transformer architecture specifically designed for processing documents up to 4096 tokens or longer, using a combination of local sliding window attention and task-specific global attention. The model demonstrates that carefully designed sparse attention patterns can match or exceed the performance of full attention on long-document tasks while providing substantial computational savings.

The Longformer attention pattern combines two components. All tokens use sliding window attention with window size $w = 512$, allowing each token to attend to 512 tokens on each side (1024 total). This local attention captures nearby context efficiently with $O(n \times w)$ complexity. Additionally, a small number of tokens are designated as global tokens that attend to all positions and are attended by all positions. For classification tasks, the [CLS] token is global. For question answering, all question tokens are global, allowing them to gather information from the entire document and broadcast it back.

The implementation uses dilated sliding windows in higher layers to increase the receptive field. In the first few layers, window size is $w = 512$ with no dilation. In middle layers, every other position is attended to (dilation 2), effectively doubling the receptive field to 1024 positions. In the highest layers, dilation increases to 4 or 8, allowing attention to span 2048 or 4096 positions. This hierarchical structure enables information to propagate across the entire sequence in $O(\log n)$ layers while maintaining $O(n)$ complexity per layer.

Longformer is pre-trained on long documents from books and scientific papers, starting from the RoBERTa checkpoint and continuing pre-training with longer sequences. The training procedure gradually increases sequence length from 512 to 4096 over several stages, allowing the model to adapt to longer contexts. Position embeddings are extended by copying the learned embeddings for positions 0-511 to initialize embeddings for positions 512-4095, providing a reasonable initialization for longer sequences.

On long-document tasks, Longformer achieves state-of-the-art results. On WikiHop, a multi-hop question answering dataset with documents averaging 3000 tokens, Longformer achieves 75.3

The computational efficiency enables practical deployment. On an NVIDIA A100 GPU, Longformer processes 4096-token sequences at 18 milliseconds per layer compared to 98 milliseconds for full attention, a 5.4× speedup. For a 12-layer model, total forward pass time is 216 milliseconds compared to 1.2 seconds, enabling real-time inference. Memory consumption is 2.4 GB for batch size 8 compared to 9.6 GB for full attention, allowing 4× larger batches or longer sequences on the same hardware.

Reformer

Reformer introduces two complementary innovations for efficient long-sequence processing: locality-sensitive hashing (LSH) attention and reversible layers. Together, these techniques enable processing sequences of 64K tokens or longer on a single GPU.

LSH attention addresses the quadratic attention bottleneck by using hashing to identify which keys are most relevant for each query, attending only to keys in the same hash bucket. The key insight is that attention weights are dominated by keys with high similarity to the query (large dot product $\vq\transpose \vk$). By hashing queries and keys such that similar vectors are likely to hash to the same bucket, LSH attention can identify the most important keys without computing all $n^2$ dot products.

The LSH attention algorithm works as follows. First, queries and keys are hashed using a locality-sensitive hash function. Reformer uses random projection LSH: $h(\vx) = \arg\max_i (\vr_i\transpose \vx)$ where $\vr_1, \ldots, \vr_b$ are random unit vectors defining $b$ hash buckets. Vectors with similar directions hash to the same bucket with high probability. Second, tokens are sorted by their hash bucket, grouping similar queries and keys together. Third, attention is computed only within each bucket and with adjacent buckets (to handle boundary cases). Fourth, the output is reordered to the original sequence order.

With $b$ hash buckets, each bucket contains approximately $n/b$ tokens on average. Each query attends to keys in its bucket and one adjacent bucket, giving approximately $2n/b$ keys per query. The complexity is $O(n^2/b \times d)$, providing a factor of $b$ speedup. With $b = 8$ buckets, LSH attention is 8× faster than full attention. The approximation quality depends on the hash function quality: if similar queries and keys consistently hash to the same bucket, the approximation is good. Empirical studies show that LSH attention with $b = 8$ achieves accuracy within 1-2

Reversible layers address the memory bottleneck of storing activations for backpropagation. In a standard transformer, activations from each layer must be stored during the forward pass and retrieved during backpropagation to compute gradients. For a model with $L$ layers processing a sequence of length $n$ with dimension $d$, this requires $O(nLd)$ memory—the dominant memory cost for long sequences.

Reversible layers use a reversible architecture inspired by RevNets. Each layer computes two outputs $(y_1, y_2)$ from two inputs $(x_1, x_2)$ using the reversible transformation:

This transformation is invertible: given $(y_1, y_2)$, we can recover $(x_1, x_2)$ by:

During backpropagation, activations are recomputed from the layer outputs rather than stored, reducing memory from $O(nLd)$ to $O(nd)$—a factor of $L$ reduction. For a 12-layer model, this reduces activation memory by 12×. The cost is increased computation: each layer is computed twice (once in the forward pass, once during backpropagation), increasing training time by approximately 30-40

Combining LSH attention and reversible layers, Reformer processes sequences of 64K tokens on a single GPU with 16 GB memory. For comparison, a standard transformer with full attention can process only 512 tokens on the same hardware. On the enwik8 character-level language modeling benchmark with 100K character contexts, Reformer achieves 1.05 bits per character, matching transformer-XL while using 16× less memory. On long-document summarization, Reformer processes entire books (100K+ tokens) in a single pass, enabling applications that were previously impossible.

Exercises

Solutions

Full solutions for all exercises are available at \url{https://deeplearning.hofkensvermeulen.be}.

import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import seaborn as sns

def create_sliding_window_mask(n, window_size):
    """Create attention mask for sliding window"""
    mask = torch.zeros(n, n, dtype=torch.bool)
    
    for i in range(n):
        # Each position attends to window_size tokens on each side
        start = max(0, i - window_size // 2)
        end = min(n, i + window_size // 2 + 1)
        mask[i, start:end] = True
    
    return mask

def sliding_window_attention(Q, K, V, window_size):
    """Compute sliding window attention"""
    n, d = Q.shape
    
    # Create mask
    mask = create_sliding_window_mask(n, window_size)
    
    # Compute attention scores
    scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d, dtype=torch.float32))
    
    # Apply mask (set masked positions to -inf)
    scores = scores.masked_fill(~mask, float('-inf'))
    
    # Softmax
    attn_weights = torch.softmax(scores, dim=-1)
    
    # Apply attention to values
    output = torch.matmul(attn_weights, V)
    
    return output, attn_weights, mask

# Example with n=1024, w=256
n = 1024
d = 64
window_size = 256

# Random Q, K, V
torch.manual_seed(42)
Q = torch.randn(n, d)
K = torch.randn(n, d)
V = torch.randn(n, d)

# Compute sliding window attention
output_sw, attn_sw, mask = sliding_window_attention(Q, K, V, window_size)

# Compute full attention for comparison
scores_full = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d, dtype=torch.float32))
attn_full = torch.softmax(scores_full, dim=-1)
output_full = torch.matmul(attn_full, V)

print(f"Sequence length: {n}")
print(f"Window size: {window_size}")
print(f"Output shape: {output_sw.shape}")
print(f"Attention weights shape: {attn_sw.shape}")

# Visualize attention patterns
fig, axes = plt.subplots(1, 2, figsize=(15, 6))

# Full attention (sample 256x256 for visibility)
sample_size = 256
axes[0].imshow(attn_full[:sample_size, :sample_size].numpy(), cmap='viridis', aspect='auto')
axes[0].set_title('Full Attention Pattern')
axes[0].set_xlabel('Key Position')
axes[0].set_ylabel('Query Position')

# Sliding window attention
axes[1].imshow(attn_sw[:sample_size, :sample_size].numpy(), cmap='viridis', aspect='auto')
axes[1].set_title(f'Sliding Window Attention (w={window_size})')
axes[1].set_xlabel('Key Position')
axes[1].set_ylabel('Query Position')

plt.tight_layout()
plt.savefig('sliding_window_attention.png', dpi=150)

# Visualize mask pattern
plt.figure(figsize=(10, 10))
plt.imshow(mask[:sample_size, :sample_size].numpy(), cmap='binary', aspect='auto')
plt.title(f'Sliding Window Mask (w={window_size})')
plt.xlabel('Key Position')
plt.ylabel('Query Position')
plt.colorbar(label='Attention Allowed')
plt.savefig('sliding_window_mask.png', dpi=150)

import torch
import numpy as np

def generate_random_features(d, m, seed=42):
    """Generate random projection matrix for Performer"""
    torch.manual_seed(seed)
    # Gaussian random features
    omega = torch.randn(d, m) / np.sqrt(d)
    return omega

def phi_features(x, omega):
    """Compute random feature map phi(x)"""
    # x: (n, d), omega: (d, m)
    # phi(x) = exp(x @ omega) / sqrt(m)
    projection = torch.matmul(x, omega)  # (n, m)
    features = torch.exp(projection) / np.sqrt(omega.shape[1])
    return features

def performer_attention(Q, K, V, m=256):
    """Compute Performer attention using random features"""
    n, d = Q.shape
    
    # Generate random projection
    omega = generate_random_features(d, m)
    
    # Compute feature maps
    phi_Q = phi_features(Q, omega)  # (n, m)
    phi_K = phi_features(K, omega)  # (n, m)
    
    # Compute attention: phi(Q) @ (phi(K)^T @ V)
    # This is O(nmd) instead of O(n^2d)
    KV = torch.matmul(phi_K.T, V)  # (m, d)
    output = torch.matmul(phi_Q, KV)  # (n, d)
    
    # Normalize
    normalizer = torch.matmul(phi_Q, phi_K.sum(dim=0, keepdim=True).T)  # (n, 1)
    output = output / (normalizer + 1e-6)
    
    return output, phi_Q, phi_K

def standard_attention(Q, K, V):
    """Standard softmax attention"""
    d = Q.shape[1]
    scores = torch.matmul(Q, K.T) / np.sqrt(d)
    attn_weights = torch.softmax(scores, dim=-1)
    output = torch.matmul(attn_weights, V)
    return output, attn_weights

# Example with d=64, m=256
n = 512
d = 64
m = 256

torch.manual_seed(42)
Q = torch.randn(n, d)
K = torch.randn(n, d)
V = torch.randn(n, d)

# Compute both attentions
output_performer, phi_Q, phi_K = performer_attention(Q, K, V, m)
output_standard, attn_weights = standard_attention(Q, K, V)

print(f"Sequence length: {n}")
print(f"Hidden dimension: {d}")
print(f"Random features: {m}")
print(f"\nPerformer output shape: {output_performer.shape}")
print(f"Standard output shape: {output_standard.shape}")

# Measure approximation error
mse_error = torch.mean((output_performer - output_standard) ** 2)
relative_error = mse_error / torch.mean(output_standard ** 2)
cosine_sim = torch.nn.functional.cosine_similarity(
    output_performer.flatten(), 
    output_standard.flatten(), 
    dim=0
)

print(f"\nApproximation Quality:")
print(f"MSE: {mse_error.item():.6f}")
print(f"Relative Error: {relative_error.item():.4f}")
print(f"Cosine Similarity: {cosine_sim.item():.4f}")

# Test with different numbers of random features
m_values = [32, 64, 128, 256, 512, 1024]
errors = []

for m in m_values:
    output_perf, _, _ = performer_attention(Q, K, V, m)
    error = torch.mean((output_perf - output_standard) ** 2).item()
    errors.append(error)
    print(f"m={m:4d}: MSE={error:.6f}")

# Plot error vs number of features
import matplotlib.pyplot as plt

plt.figure(figsize=(10, 6))
plt.plot(m_values, errors, 'o-', linewidth=2, markersize=8)
plt.xlabel('Number of Random Features (m)')
plt.ylabel('Mean Squared Error')
plt.title('Performer Approximation Error vs Random Features')
plt.xscale('log')
plt.yscale('log')
plt.grid(True, alpha=0.3)
plt.savefig('performer_error.png', dpi=150)