def wall_clock(f, *args, reps=10):
f(*args) # Warm up
torch.cuda.synchronize()
start = time.time()
for _ in range(reps):
f(*args)
torch.cuda.synchronize()
return (time.time() - start) / reps
for n in [2048, 4096, 8192, 16384]:
Q = torch.randn(1, n, d_h, device=d2l.try_gpu())
with torch.no_grad():
t = wall_clock(attention, Q, Q, Q)
print(f'n = {n:5d}: {t*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:
torch.manual_seed(0)
n, d_h = 2048, 64
Q, K, V = (torch.randn(n, d_h, device=d2l.try_gpu()) for _ in range(3))
err = (chunked_attention(Q, K, V) - causal_attention(Q, K, V)).abs().max()
print(f'maximum deviation: {float(err):.2e}')maximum deviation: 4.77e-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.