Attention Complexity › Efficient Attention Variants

The quadratic complexity of standard attention has motivated an entire subfield of research into more efficient alternatives. Understanding these options is important for two reasons: it shows you what trade-offs are available when quadratic attention becomes prohibitive, and it illustrates the creative approaches researchers have taken to break the $n^2$ barrier without abandoning the core insight that all-pairs interaction is useful.

The key challenge in designing efficient attention is that you cannot simply compute fewer interactions without risking the loss of important long-range dependencies. Different variants make different assumptions about which interactions are dispensable. Linear attention methods assume that the interaction pattern can be approximated using lower-rank representations. Sparse attention methods assume that each token only needs to attend to a small subset of positions. Local attention methods assume that most relevant context is nearby, with occasional global dependencies handled by special mechanisms. Each assumption is correct in some settings and wrong in others.

Think of the trade-off this way: standard attention is like an exhaustive library search that checks every single book for relevance. It never misses anything, but it is expensive. Efficient attention variants are like using an index, relying on subject headings, or trusting that nearby shelves are most relevant. Each strategy can dramatically reduce the search effort, at the risk of occasionally missing an important reference that the full search would have found.

While a detailed treatment of each variant deserves its own chapter, this comparison puts the complexity analysis in context and helps you evaluate which variant is appropriate for a given task.

Several methods approximate attention with linear complexity by avoiding the explicit $n \times n$ matrix. The core insight is that softmax attention can be written as a similarity function between queries and keys, and certain similarity functions can be decomposed into feature maps that allow the computation to be reordered. Instead of computing $(QK^T)V$ (which requires the $n \times n$ intermediate), these methods compute $Q(K^TV)$, where the parenthesized term has shape $d \times d$ rather than $n \times n$.

• Linformer: Projects keys and values to a fixed low-dimensional space before computing attention. If keys and values are projected to $k$ dimensions where $k \ll n$, the attention matrix becomes $n \times k$ rather than $n \times n$, achieving $O(nk)$ complexity with $k$ as a fixed hyperparameter.
• Performer: Uses random feature maps to approximate softmax attention in $O(n)$ time. The random features are chosen so that their inner product approximates the exponential kernel that softmax attention implicitly computes.
• Linear Transformers: Replace softmax with kernel functions (such as the ReLU feature map) that allow associative computation. Because the kernel decomposes as a product of query and key features, the computation can be reordered to avoid the $n \times n$ matrix.

These methods trade accuracy for speed, often working well in practice despite theoretical approximation. The main limitation is that removing the softmax normalization can destabilize training, and the approximation quality degrades for queries that are far from the key distribution. In practice, linear attention methods often underperform full attention on tasks requiring precise long-range matching, while performing comparably on tasks where approximate retrieval is sufficient.

Instead of attending to all positions, sparse methods attend to a carefully chosen subset. The key question is which subset to choose: a random subset, a structured subset based on proximity, or a learned subset based on content. Different choices lead to different architectural properties.

• Longformer: Combines local windowed attention (each token attends to nearby tokens within a window) with global attention on special tokens like [CLS]. The global tokens can attend to and be attended by all positions. This provides pathways for long-range information while keeping most attention local.
• BigBird: Uses a combination of random and window attention with global patterns. The theoretical motivation is that this combination is sufficient to preserve the expressive power of full attention for most practical tasks, at a fraction of the cost.
• Sparse Transformers: Learn which positions to attend to, adapting the attention pattern to the data rather than fixing it in advance. This is more powerful but introduces training complexity.

Sparse attention reduces complexity to $O(n \cdot k)$, where:

• $n$: sequence length
• $k$: number of positions each token attends to (typically $k \ll n$)

For example, with a local window of 256 tokens, $k = 256$ regardless of sequence length, making the complexity effectively linear in $n$. The key assumption behind sparse attention is that most relevant context for any given token is local, with occasional non-local dependencies that can be handled through global token mechanisms or random connections. This assumption holds well for natural language, where adjacent sentences are usually more relevant than distant ones, but may fail for tasks requiring systematic long-range reasoning such as mathematical proofs or algorithmic reasoning over large structures.

A third class of approaches does not approximate the attention function at all, but instead restructures the computation to be more hardware-efficient. These methods compute exactly the same attention as the standard algorithm but do so in a way that reduces the amount of data moved between different memory levels on the GPU.

• FlashAttention: Reorders operations to minimize memory transfers between GPU high-bandwidth memory (HBM) and the faster on-chip SRAM. Standard attention writes the $n \times n$ attention matrix to HBM and reads it back for the output computation. FlashAttention instead tiles the computation so that each block of the attention matrix is computed, used for the corresponding output block, and discarded without ever writing to HBM. The result is identical to full attention mathematically, but achieves 2 to 4x speedup due to reduced memory bandwidth, with memory usage $O(n)$ rather than $O(n^2)$.
• KV Caching: During autoregressive generation, caches key-value pairs from all previous tokens to avoid recomputing them at each step. The cache grows linearly with the number of generated tokens and must be stored in memory for the duration of generation. For long outputs, KV cache management (including eviction strategies when the cache exceeds memory) becomes a significant engineering challenge.
• Gradient Checkpointing: Trades computation for memory during training by discarding intermediate activations after the forward pass and recomputing them from saved checkpoints during backpropagation. This reduces activation memory from $O(n^2 L)$ (proportional to all layers and sequence length squared) to roughly $O(n^2 \sqrt{L})$ with some computational overhead.