Sparse & sub-quadratic attention

FlashAttention made full attention cheaper to run without changing the math — but it’s still O(N²) in compute: every token attends to every other. At 100k+ tokens that quadratic wall returns. This lesson is about genuinely sub-quadratic attention — changing what gets computed, not just how memory is touched.

Escape 1 — sparsify the attention matrix

The N×N score matrix is mostly wasted: a token usually cares about its recent neighbours and a few special tokens, not all 100,000 others. So don’t compute the cells you don’t need. Each pattern below keeps only a fraction of the causal matrix lit — the lit cells (query rows down, key columns across) are the ones you actually pay for:

The workhorse patterns:

• Sliding window (Mistral): each token attends to the last w tokens — O(N·w). One layer sees w back; stack L layers and the effective receptive field grows to L·w, exactly like depth in a CNN.
• Global + local (Longformer, BigBird): most tokens stay local, but a handful of global tokens attend to (and are attended by) everyone — cheap long-range shortcuts. O(N).
• Dilated: attend to every k-th token for wide reach at coarse resolution.
• Block-sparse: attend within your block and the previous one — the GPU-friendly version, since blocks map cleanly to tiles.

The catch: a fixed sparsity pattern can miss a genuinely long-range dependency it didn’t budget for. That’s why the next escape is more radical.

Escape 2 — change the operator

Instead of pruning attention, replace it with a sequence mixer that’s linear by construction.

Linear attention rewrites softmax attention with a kernel feature map so you can use associativity: compute KᵀV first (a small d×d matrix), then Q(KᵀV). That reorder turns O(N²·d) into O(N·d²) — linear in sequence length. See the associativity that makes it possible:

import numpy as np
rng = np.random.default_rng(0)
N, d = 2000, 64
Q = rng.standard_normal((N, d)); K = rng.standard_normal((N, d)); V = rng.standard_normal((N, d))

# Quadratic order: (Q Kᵀ) V  ->  builds an N×N matrix
A = (Q @ K.T) @ V                 # O(N²·d) work, O(N²) memory

# Linear order: Q (Kᵀ V)  ->  builds only a d×d matrix
B = Q @ (K.T @ V)                 # O(N·d²) work, O(d²) memory

print("same result? ", np.allclose(A, B))
print("N×N matrix floats :", N*N, " vs  d×d floats:", d*d)
# Identical answer, but the linear order never forms the N×N matrix.
# (Real linear attention applies a feature map φ so softmax ≈ φ(Q)φ(K)ᵀ — same trick.)

same result?  True
N×N matrix floats : 4000000  vs  d×d floats: 4096

The two orders give the identical result — but the linear one never materialises the N×N matrix (here 4,000,000 floats versus just 4,096). That reordering is the whole game: O(N²·d) work becomes O(N·d²), linear in sequence length.

State-space models (SSMs) — Mamba. SSMs treat the sequence like a linear recurrence (equivalently a long convolution): a hidden state is carried forward and updated at each step. Mamba’s contribution is making that update selective — input-dependent gates decide what to keep or forget, recovering much of attention’s content-awareness. The payoffs are large:

• O(N) time, and at inference a fixed-size state per step — there is no KV cache that grows with context. Constant memory, high throughput.
• The cost: history is compressed into that fixed state. Attention can look up any past token exactly; an SSM must have chosen to remember it. So SSMs can lag on tasks needing precise long-range recall (e.g. “copy this exact ID from 40k tokens ago”).

How to choose

In one breath

• Full attention is O(N²) in compute; FlashAttention only makes that cheaper to run, so past ~100k tokens the quadratic wall is still there.
• Escape 1 — sparsify: compute only a structured slice of the matrix. Sliding-window (O(N·w), reach grows with depth like a CNN), global+local (O(N)), dilated, and GPU-friendly block-sparse — trading guaranteed reach for cost.
• Escape 2 — change the operator: linear attention reorders Q(KᵀV) to skip the N×N matrix (O(N·d²)); SSMs (Mamba) carry a fixed selective state forward — O(N) time, constant per-step memory, no growing KV cache.
• The trade is recall: a fixed state or pattern must choose what to remember, so exact long-range lookup can suffer — which is why hybrids (Jamba) mix cheap linear/SSM layers with a few full-attention layers, the pragmatic long-context default.

Check yourself

What to remember

• Full attention is O(N²); FlashAttention makes it cheaper but not sub-quadratic. Two real escapes change the computation.
• Sparse patterns (sliding-window, global+local, block) skip most of the matrix — O(N·w) to O(N) — trading guaranteed reach for cost.
• Linear attention reorders the product to avoid the N×N matrix; SSMs (Mamba) carry a fixed selective state — O(N) time, no growing KV cache, at the cost of exact recall.
• Hybrids (Jamba) mix cheap linear/SSM layers with a few attention layers — the pragmatic recipe behind much of today’s long context.