Linear Attention › Summary

Linear attention transforms the attention mechanism from quadratic to linear complexity by exploiting the associative property of matrix multiplication. Instead of computing the full $n \times n$ attention matrix, it precomputes a $d \times d$ key-value aggregate that can be queried by each position independently.

The key concepts covered in this chapter:

• Kernel decomposition: Replacing softmax with a kernel function $\kappa(q, k) = \phi(q)^T\phi(k)$ enables reordering $(QK^T)V$ to $Q(K^TV)$, eliminating the quadratic dependency on sequence length.

• Feature maps: Different choices of $\phi$ trade off between computational efficiency and how well they approximate softmax attention. Simple maps like ELU are fast but imprecise; random feature methods like FAVOR+ provide unbiased approximation with controllable variance.

• Causal formulation: Linear attention naturally supports autoregressive modeling through a running sum formulation that incrementally accumulates key-value aggregates, enabling $O(1)$ per-token generation.

• Performer and variants: The FAVOR+ approach uses positive random features to provide an unbiased estimator of softmax attention, with approximation quality improving as the number of random features increases.

• Trade-offs: Linear attention sacrifices the ability to express arbitrarily sharp attention patterns, introduces approximation variance with random features, and may exhibit position biases in causal settings. These limitations make it best suited for long-document tasks where the quadratic bottleneck is the primary constraint.

Linear attention changes how we think about the attention mechanism: from an $n \times n$ interaction matrix to a fixed-size state that summarizes the key-value relationships. This perspective has influenced subsequent architectures like state space models and linear recurrent units, which further develop the idea of maintaining a compact state that captures sequence information.

For practitioners, linear attention is most valuable when sequence length is the binding constraint. At shorter lengths where $n < d$, the overhead of feature transformation may not be worthwhile. But for tasks involving documents, books, or streaming data where sequences stretch to tens of thousands of tokens, linear attention's $O(nd^2)$ complexity makes previously intractable problems accessible.