def wall_clock(f, *args, reps=10):
f(*args).block_until_ready() # Warm up (and compile)
start = time.time()
for _ in range(reps):
f(*args).block_until_ready()
return (time.time() - start) / reps
layer = jax.jit(attention_layer)
for n in [2048, 4096, 8192, 16384]:
Q = jax.random.normal(jax.random.key(0), (n, d_h))
print(f'n = {n:5d}: {wall_clock(layer, Q, Q, Q)*1e3:6.2f} ms')The bargain and the bill
Map n tokens of dimension d to n tokens of dimension d:
| layer | complexity | sequential ops | max path length |
|---|---|---|---|
| convolution (k) | \mathcal{O}(knd^2) | \mathcal{O}(1) | \mathcal{O}(n/k) |
| recurrence | \mathcal{O}(nd^2) | \mathcal{O}(n) | \mathcal{O}(n) |
| self-attention | \mathcal{O}(n^2d) | \mathcal{O}(1) | \mathcal{O}(1) |
Full parallelism and constant path length — paid for by the one quadratic entry.
Online softmax: the matrix never needs to exist
Per query, carry a running max, normalizer, and output; process keys in blocks:
m' = \max(m, \max_j a_j), \quad s' = s\,e^{m-m'} + \sum_j e^{a_j - m'}, \quad \mathbf{o}' = \mathbf{o}\,e^{m-m'} + \sum_j e^{a_j-m'}\mathbf{v}_j
\mathbf{o}/s is exactly \mathrm{softmax}(\mathbf{a})\mathbf{V} — raising the max retroactively rescales the past.
Chunked attention: exact, in linear memory
Twenty lines carry the update over all queries, one key block at a time:
n, d_h = 2048, 64
Q, K, V = (jax.random.normal(k, (n, d_h))
for k in jax.random.split(jax.random.key(0), 3))
with jax.default_matmul_precision('highest'):
err = jnp.abs(chunked_attention(Q, K, V)
- causal_attention(Q, K, V)).max()
print(f'maximum deviation: {float(err):.2e}')maximum deviation: 4.62e-07
- Same answer, same FLOPs; the n^2 footprint is gone.
Attention through a window
Let each query attend to the w most recent positions — the causal mask becomes a band:
- \mathcal{O}(nw) work instead of \mathcal{O}(n^2).
- Deployed: Mistral 7B (w = 4096, 32 layers); Longformer adds global tokens; NSA learns the sparsity pattern itself.
Depth restores the reach
Information hops w-1 positions per layer: receptive field 1 + L(w-1) after L layers — computed, not asserted:
depth 1: last query reaches 8 positions (formula: 8)
depth 2: last query reaches 15 positions (formula: 15)
depth 4: last query reaches 29 positions (formula: 29)- The CNN bet (ch. 7), re-made: locality compounds through depth.
The price of attention, measured
- Dense bends to slope 2; windowed hugs the launch floor; linear pays nd_h^2 memory and bandwidth-bound cumsums — its decisive win is the constant-memory recurrent mode.
- The approximation zoo (Performer, Linformer, Reformer) is history: what survived is exact-but-clever, windowed, and linear-as-recurrence.