Efficient Transformers Survey

Why This Matters

Standard self-attention computes all $n^2$ pairwise interactions for a sequence of length $n$, costing $O(n^2 d)$ time and $O(n^2)$ memory. For long sequences (documents, genomes, high-resolution images), this quadratic cost is prohibitive.

Between 2019 and 2022, dozens of papers proposed sub-quadratic attention variants. Most are now rarely used. FlashAttention (2022) showed that the bottleneck was memory bandwidth, not FLOPs, and that exact attention with IO-aware tiling is faster in practice than most approximate methods up to sequences of length 8K-16K. Understanding why this happened teaches more about systems-level thinking than about any single algorithm.

The Quadratic Bottleneck

For queries $Q$, keys $K$, values $V \in \mathbb{R}^{n \times d}$:

$\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d}}\right) V$

The matrix $QK^\top \in \mathbb{R}^{n \times n}$ has $n^2$ entries. Storing it costs $O(n^2)$ memory. For $n = 100{,}000$ and $d = 128$, the attention matrix alone requires 40 GB in float32. The question: can we avoid materializing this matrix?

Linear Attention

Proposition

Linear Attention via Kernel Decomposition

Statement

Replace the softmax kernel $\kappa(q, k) = \exp(q^\top k / \sqrt{d})$ with a feature map $\phi: \mathbb{R}^d \to \mathbb{R}^m$ such that $\kappa(q, k) \approx \phi(q)^\top \phi(k)$. Then:

$\text{Attention}(Q, K, V)_i = \frac{\phi(Q_i)^\top \sum_j \phi(K_j) V_j^\top}{\phi(Q_i)^\top \sum_j \phi(K_j)}$

The sums $S = \sum_j \phi(K_j) V_j^\top \in \mathbb{R}^{m \times d}$ and $z = \sum_j \phi(K_j) \in \mathbb{R}^m$ can be precomputed in $O(nmd)$ time. Each attention output costs $O(md)$, giving total cost $O(nmd)$, which is linear in $n$ if $m$ is fixed.

Intuition

Instead of computing the full $n \times n$ attention matrix and then multiplying by $V$, change the order of operations. First aggregate key-value pairs into a compact summary (the matrix $S$), then query this summary. This is the associativity trick: $(QK^\top)V = Q(K^\top V)$ up to normalization.

Proof Sketch

Write $\text{Attn}_i = \sum_j \kappa(Q_i, K_j) V_j / \sum_j \kappa(Q_i, K_j)$. Substitute $\kappa(Q_i, K_j) = \phi(Q_i)^\top \phi(K_j)$ and factor out $\phi(Q_i)^\top$ from both sums.

Why It Matters

Linear attention reduces the complexity from $O(n^2 d)$ to $O(n m d)$. For $m \ll n$, this is a large speedup for very long sequences. The idea also enables recurrent computation: process tokens one by one, updating $S$ and $z$ incrementally, giving constant memory per step.

Failure Mode

The approximation quality depends on $\phi$. No known feature map perfectly approximates softmax attention. The approximation error degrades attention's ability to do sharp, sparse lookups (e.g., copying a specific token). In practice, linear attention models underperform standard attention on tasks requiring precise retrieval from context.

Sparse Attention Methods

Rather than approximating the softmax, sparse methods compute exact attention on a subset of the $n^2$ pairs.

Longformer (Beltagy et al., 2020): Each token attends to a fixed-size local window (e.g., 512 tokens around it) plus a small set of global tokens. Cost is $O(nw)$ where $w$ is the window size. Good for document-level tasks where local context dominates.

BigBird (Zaheer et al., 2020): Combines local window attention, global tokens, and random attention edges. Proved that this sparse pattern is a universal approximator of sequence-to-sequence functions and is Turing complete, matching the theoretical expressiveness of full attention.

Block-sparse attention: Partition tokens into blocks and attend within blocks. Efficient on GPUs because block operations map well to hardware. This is the foundation of FlashAttention's tiling strategy.

Low-Rank Approximations

Linformer (Wang et al., 2020): Project keys and values to a lower dimension $k \ll n$ using learned projection matrices $E, F \in \mathbb{R}^{k \times n}$:

$\text{Attention}(Q, EK, FV) = \text{softmax}\!\left(\frac{Q(EK)^\top}{\sqrt{d}}\right) FV$

Cost is $O(nkd)$ instead of $O(n^2d)$. The claim: the attention matrix is approximately low-rank for most practical inputs, so projecting to rank $k$ loses little information.

Performer (Choromanski et al., 2021): Uses random orthogonal features to approximate the softmax kernel:

$\exp(q^\top k / \sqrt{d}) \approx \phi(q)^\top \phi(k)$

where $\phi(x) = \frac{1}{\sqrt{m}} \exp(-\|x\|^2/2) [\exp(\omega_1^\top x), \ldots, \exp(\omega_m^\top x)]$ with random $\omega_i \sim \mathcal{N}(0, I_d / \sqrt{d})$.

Proposition

Performer Approximation Quality

Statement

The Performer approximation $\hat{\kappa}(q, k) = \phi(q)^\top \phi(k)$ satisfies:

$\mathbb{E}[\hat{\kappa}(q, k)] = \exp(q^\top k / \sqrt{d})$

The variance decreases as $O(1/m)$, so with $m = O(d \log d / \epsilon^2)$ random features, the relative approximation error is at most $\epsilon$ with high probability.

Intuition

Random Fourier features (Rahimi and Recht, 2007) applied to the softmax kernel. More random features give a better approximation, but the required $m$ can be large for high accuracy, reducing the practical speedup.

Proof Sketch

The unbiasedness follows from the identity $\exp(q^\top k) = \mathbb{E}_{\omega \sim \mathcal{N}}[\exp(\omega^\top q - \|q\|^2/2) \exp(\omega^\top k - \|k\|^2/2)]$. Concentration follows from standard random feature arguments.

Why It Matters

This gives a principled way to trade accuracy for speed. But the required number of features for high accuracy often makes the method no faster than exact attention for typical sequence lengths.

Failure Mode

For tasks requiring precise attention patterns (copying, exact retrieval), the approximation error causes significant quality degradation. The number of features needed for acceptable accuracy grows with the "peakiness" of the attention distribution.

Why FlashAttention Won

FlashAttention computes exact softmax attention but tiles the computation to minimize HBM (high bandwidth memory) reads and writes. The key insight: the attention bottleneck for sequences up to 16K is not FLOPs but memory IO. The $n \times n$ attention matrix does not need to be materialized; it can be computed and consumed block by block.

FlashAttention achieves 2-4x wall-clock speedup over standard attention for typical LLM training sequence lengths (2K-8K), while computing the exact same output. This made most approximate attention methods unnecessary for their target use case (moderate-length sequences on modern GPUs).

Sub-quadratic methods still matter when $n$ is truly enormous (100K+ tokens), where the $O(n^2)$ FLOP cost itself dominates. But for most current LLM training and inference, FlashAttention (and its successors) suffice.

Common Confusions

Watch Out

Sub-quadratic FLOPs does not mean faster wall-clock time

A method with $O(n \log n)$ FLOPs can be slower than $O(n^2)$ if it has poor memory access patterns, cannot use tensor cores, or requires custom kernels. On modern GPUs, memory bandwidth is often the bottleneck, not arithmetic. This is why FlashAttention (which uses more FLOPs due to recomputation) is faster than standard attention (which materializes the full attention matrix).

Watch Out

Linear attention is not a drop-in replacement

Linear attention changes the model's expressive power, not just its cost. The softmax attention matrix can be arbitrarily sparse and peaked. Linear attention with a fixed feature map cannot represent sharp attention patterns. Models trained with linear attention often require architectural changes (gating, decay factors) to match softmax attention quality.

Key Takeaways

• Standard attention costs $O(n^2 d)$ time and $O(n^2)$ memory
• Linear attention uses the kernel trick to achieve $O(nmd)$ but sacrifices sharp attention patterns
• Sparse methods (Longformer, BigBird) compute exact attention on a structured subset of pairs
• Low-rank methods (Linformer, Performer) approximate the attention matrix
• FlashAttention computes exact attention with IO-aware tiling, making approximate methods unnecessary for $n < 16K$
• For truly long sequences ($n > 100K$), sub-quadratic methods and state-space models remain relevant

Exercises

ExerciseCore

Problem

In linear attention with feature map $\phi: \mathbb{R}^d \to \mathbb{R}^m$, what are the shapes of the precomputed matrices $S = \sum_j \phi(K_j) V_j^\top$ and $z = \sum_j \phi(K_j)$? What is the total memory cost of storing $S$ and $z$ compared to storing the full $n \times n$ attention matrix?

Hint

Reveal Solution

ExerciseAdvanced

Problem

Explain why FlashAttention uses more FLOPs than standard attention (due to recomputation in the backward pass) yet achieves faster wall-clock time. Quantify the tradeoff: if HBM bandwidth is $B$ bytes/sec and compute throughput is $C$ FLOP/sec, under what condition on $n$ and $d$ is standard attention memory-bound?

Hint

Reveal Solution

References

Canonical:

• Vaswani et al., "Attention Is All You Need" (NeurIPS 2017)
• Katharopoulos et al., "Transformers are RNNs" (ICML 2020), linear attention

Current:

• Dao et al., "FlashAttention: Fast and Memory-Efficient Exact Attention" (NeurIPS 2022)
• Tay et al., "Efficient Transformers: A Survey" (ACM Computing Surveys 2022)
• Dao, "FlashAttention-2" (ICLR 2024)

Next Topics

• Mamba and state-space models: a recurrent architecture that avoids attention entirely for long sequences
• KV cache: memory optimization for autoregressive inference