Computational Complexity of Attention

Standard Self-Attention has $O(n^2)$ time and memory in sequence length $n$ . This is the fundamental bottleneck of transformers.

Where the $n^2$ comes from:

Each of $n$ tokens attends to all $n$ tokens → the attention matrix $\text{softmax}(QK^\top / \sqrt{d_k})$ is $n \times n$
Full cost per layer: $O(n^2 d)$ FLOPs, $O(n^2)$ memory for the score matrix

The scaling problem:

Approaches to break the quadratic wall:

IO-aware exact attention:

FlashAttention — computes exact attention in $O(n^2 d)$ FLOPs but $O(n)$ memory by tiling and avoiding materializing the full $n \times n$ matrix. Doesn't change asymptotics but 2–4x faster in practice

Sparse attention:

Local / sliding window — each token attends to the last $w$ tokens → $O(nw)$ , and the KV cache becomes independent of sequence length (tokens outside the window are discarded). Used in Mistral, Longformer
Dilated / strided patterns — attend to every $k$ -th token for long-range
Block-sparse — hand-crafted or learned sparsity patterns
Common design: interleave local attention with periodic global attention layers

Linear attention:

Replace $\text{softmax}(QK^\top)V$ with $\phi(Q)(\phi(K)^\top V)$ using the associativity trick → $O(nd^2)$ instead of $O(n^2 d)$
Trade-off: weaker in practice, though active research (RetNet, RWKV, Mamba)

State-space models (SSMs):

Mamba, S4 — $O(n)$ sequence processing via recurrence, competitive with transformers on long sequences

KV-cache for inference:

At generation time, cache $K, V$ from previous tokens → each new token costs $O(nd)$ instead of recomputing $O(n^2 d)$
Memory grows as $O(n \cdot d \cdot L)$ where $L$ = layers. This is why long-context inference is memory-bound

Linked from