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')Sequence-layer complexity
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}(nd^2+n^2d) | \mathcal{O}(1) | \mathcal{O}(1) |
Sequence-layer connectivity for a convolution, a recurrence, and self-attention.
Self-attention combines full parallelism and constant path length with quadratic complexity in sequence length.
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 equals \mathrm{softmax}(\mathbf{a})\mathbf{V}. When the running maximum increases, the factors e^{m-m'} rescale earlier partial sums.
Online softmax carries a running maximum, normalizer, and weighted-value sum across key blocks.
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
If each query attends 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, as the following check confirms:
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.
Measured time and memory
- Dense approaches slope 2; windowed remains near the launch-overhead floor; linear attention uses nd_h^2 memory and bandwidth-bound cumulative sums. Its principal advantage is the constant-memory recurrent mode.
- Current approaches emphasize memory-efficient exact attention, windowed attention, and linear attention in recurrent form.