attention
**Attention is a softmax-weighted average, and the softmax — not the query-key-value metaphor — determines everything that matters about training stability, context scaling, and inference cost.** The standard explanation describes attention as "queries looking up relevant keys to retrieve values," which is a mnemonic, not a mechanism. The mechanism is: project each token into three vectors, compute pairwise dot products between one set and another, push those products through a softmax to get a probability distribution, then take the weighted average of the third set. Every property of attention — its capacity, its failure modes, its cost — follows from the interaction between the dot products and the softmax.
**The $\sqrt{d_k}$ divisor is not a cosmetic normalisation; it is the difference between a functioning network and a collapsed one.** If queries and keys are drawn independently with unit variance, each dot product $q \cdot k$ has mean zero and variance $d_k$. At $d_k = 512$, the unscaled logits have standard deviation $\sqrt{512} \approx 22.6$, which pushes softmax outputs toward one-hot: the winning token gets weight 0.932 and entropy drops to 0.173 nats, far below the 4.159-nat uniform ceiling. Dividing by $\sqrt{d_k}$ restores variance to 1.0 regardless of dimension, keeping entropy at 3.684 nats and the maximum weight at 0.107. This is not a training trick to be applied and forgotten — it is a structural requirement. Any mechanism that lets the magnitude of logits grow (deeper networks, longer training, larger weight norms) recreates the same pathology: at weight scale 16$\times$, even with $\sqrt{d_k}$ scaling, logit variance reaches 63,556 and the maximum attention weight climbs to 0.994. This is why QK-norm and logit capping exist.
| Head dimension $d_k$ | Unscaled entropy (nats) | Unscaled max weight | Scaled entropy (nats) | Scaled max weight |
|---|---|---|---|---|
| 8 | 3.268 | 0.204 | 3.699 | 0.114 |
| 32 | 1.723 | 0.534 | 3.692 | 0.107 |
| 64 | 0.987 | 0.713 | 3.685 | 0.109 |
| 128 | 0.533 | 0.844 | 3.685 | 0.107 |
| 256 | 0.311 | 0.903 | 3.677 | 0.109 |
| 512 | 0.173 | 0.932 | 3.684 | 0.107 |
**Multi-head attention does not split the model's capacity; it raises the rank of the output.** A single attention head produces a weighted average of value vectors, which is a convex combination — rank at most $\min(n, d_h)$. With 32 heads each operating in a 16-dimensional subspace of a 512-dimensional model, the concatenated output reaches effective rank 64.0 (equal to the sequence length), while a single 512-dimensional head reaches only 59.5 under the same conditions. The gain is not from "attending to different things" in a vague sense; it is from the fact that $h$ independent rank-$d_h$ matrices, concatenated and projected, span a subspace of dimension up to $h \cdot d_h = d$, which is strictly larger than any single rank-$d$ attention matrix can produce after passing through softmax's convex-combination constraint. This is measurable: at $d = 512$ and $n = 64$, going from 1 head to 8 lifts effective rank from 59.5 to 64.0, saturating the sequence-length ceiling. The diminishing returns beyond 8 heads are not a sign that more heads are useless — they are a sign that the rank ceiling at $n = 64$ has already been hit.
**The "quadratic cost" of attention is a statement about sequence length, not about the operation as a whole, and it is misleading below the crossover.** A single attention layer at model dimension $d = 4{,}096$ with 32 heads performs two kinds of work: four weight-matrix projections ($W_Q$, $W_K$, $W_V$, $W_O$) costing $8nd^2$ FLOPs total, and the score computation ($QK^\top$ plus attention-weighted value summation) costing $4n^2 d$ FLOPs. Setting $8nd^2 = 4n^2 d$ gives the crossover at $n = 2d = 8{,}192$: below that length, projections dominate and doubling the sequence merely doubles total cost (linear); above it, score computation dominates and doubling the sequence quadruples it. At $n = 256$, scores account for 3.0% of FLOPs. At $n = 4{,}096$, they account for 33.3%. The score fraction does not reach 50% until $n = 2d$, and 94.1% at $n = 131{,}072$. Calling attention "quadratic" without stating the crossover leaves the impression that halving the sequence halves the cost, when in practice it saves 1.5% at typical fine-tuning lengths.
```svg
```
**Without a residual connection, repeated self-attention collapses token representations to rank one.** This is a structural property, not an empirical observation about specific trained models. Simulating $L$ layers of self-attention (no learned weights, just the softmax-weighted averaging) on 64 random token vectors shows entropy falling from 3.698 nats after one layer to 0.032 nats after 64 layers — near-perfect convergence to a single representation. The mechanism is straightforward: softmax produces a convex combination, and iterating convex combinations is a contraction. Adding a residual connection ($X \leftarrow X + \mathrm{Attn}(X)$) with layer normalisation preserves entropy at 4.159 nats (the uniform maximum) even after 64 layers. The residual stream is not a training convenience; it is the structural element that prevents attention from destroying the information it is supposed to route. This is also why the "attention sink" phenomenon — early tokens accumulating disproportionate weight in autoregressive models — is a consequence of the residual stream: the model needs a no-op attention pattern, and concentrating weight on a token whose value vector is already in the residual stream accomplishes exactly that.
**The KV cache is an inference artefact, not an architectural feature, and its memory cost is determined entirely by the number of KV heads.** During autoregressive generation, each new token must attend to all previous tokens, requiring their key and value vectors. Recomputing them would cost $O(n)$ per step and $O(n^2)$ total; caching them costs $O(1)$ per step and $O(n)$ total in compute, but $2 \cdot L \cdot n_{\mathrm{kv}} \cdot d_h \cdot 2$ bytes per token in memory (the factor 2 covers keys and values; the final 2 is fp16). For a 32-layer, 4096-dimension model at 32K context length, multi-head attention (MHA) stores 32 KV heads and consumes 16.00 GB. Grouped-query attention with groups of 4 (GQA, 8 KV heads) consumes 4.00 GB. Groups of 8 (4 KV heads) consume 2.00 GB. Multi-query attention (MQA, 1 KV head) consumes 0.50 GB — a 32$\times$ reduction from MHA, linear in the head ratio.
| Variant | KV heads | Bytes per token | Cache at 32K context | Ratio vs MHA |
|---|---|---|---|---|
| Multi-Head (MHA) | 32 | 524,288 | 16.00 GB | 1.00 |
| Grouped-Query (G=4) | 8 | 131,072 | 4.00 GB | 0.25 |
| Grouped-Query (G=8) | 4 | 65,536 | 2.00 GB | 0.125 |
| Multi-Query (MQA) | 1 | 16,384 | 0.50 GB | 0.03 |
**FlashAttention does not change the FLOPs of attention; it changes the memory hierarchy level at which the work happens.** Standard attention materialises the $n \times n$ score matrix in HBM (GPU global memory), requiring $O(n^2)$ memory and $O(n^2 + nd)$ HBM read/write operations. FlashAttention tiles the computation into SRAM (on-chip memory, roughly 20 MB on an A100), computing softmax in an online fashion without ever materialising the full matrix. The memory footprint drops from $O(n^2)$ to $O(n)$: at $n = 4{,}096$ and $d = 128$, the attention matrix would consume 32 MB in fp16, while FlashAttention uses 1.0 MB — a 32$\times$ reduction. At $n = 32{,}768$, the reduction is 256$\times$ (2,048 MB to 8.0 MB). The HBM access reduction follows from Dao et al.'s Theorem 2: total accesses scale as $\Theta(n^2 d^2 / M)$ where $M$ is the SRAM capacity, giving a ratio of $M/d^2 \approx 640$ at $M = 10\text{M}$ elements and $d = 128$. The wall-clock improvement is smaller (2–4$\times$) because FlashAttention is compute-bound rather than IO-bound, but the memory savings are exact and are what enable long-context training without gradient checkpointing.
**Attention is one-third of a Transformer layer's parameters and carries none of its nonlinearity.** At $d = 4{,}096$ with SwiGLU MLP (the standard choice in modern LLMs), attention contributes 67.1 million parameters per layer — the four projection matrices $W_Q$, $W_K$, $W_V$, $W_O$, each $d \times d$. The MLP contributes 134.2 million — the gate, up, and down projections at $d \times \frac{8d}{3}$. Attention is 33.3% of layer parameters and 100% of the token-mixing computation; the MLP is 66.7% of parameters and 100% of the channel-mixing computation. The softmax is the only nonlinearity in the attention sublayer, but it acts on the scores, not on the representations — the output is a linear function of the value vectors. All feature transformation happens in the MLP. This division matters for efficiency: linear attention variants that replace softmax with a kernel approximation ($\phi(Q)\phi(K)^\top V$) can reduce score computation to $O(nd^2)$ but sacrifice the adaptive sparsity that softmax provides, which is why no linear-attention model has matched softmax attention at scale despite eliminating the quadratic term.
**The causal mask halves the average context and exactly halves the score FLOPs, but its deeper effect is forcing every token to make predictions from a different-sized context.** In a causal (autoregressive) model, token $i$ can attend to positions $0$ through $i$ only. The average number of visible tokens is $(n+1)/2$: at $n = 4{,}096$, each token sees on average 2,048 predecessors, and the first token sees only itself. The triangular mask zeros out half the $n \times n$ score matrix, reducing score FLOPs to $n(n+1)/2 \approx n^2/2$. But the asymmetry is the important part: the model must produce a useful representation at position 1 (1 token of context) and at position 4,096 (4,096 tokens of context) using the same weight matrices. Positional encoding, whether learned, sinusoidal, or rotary (RoPE), exists to let the model distinguish these situations — without it, self-attention is permutation-equivariant and cannot tell position 1 from position 4,096. Cross-attention, used in encoder-decoder models, drops the causal constraint entirely: queries from the decoder attend to all encoder positions, with no mask.
**Through the lens of hardware, attention is a memory-bandwidth problem dressed as an arithmetic one, and every successful optimisation since 2020 has targeted the memory side.** The attention score matrix at $n = 32{,}768$ is 2 GB in fp16 — larger than the entire KV cache of a GQA-8 model at the same context length. FlashAttention eliminates it. GQA and MQA shrink the KV cache by 4–32$\times$ with minimal quality loss. Sliding-window attention (Mistral-style) bounds the matrix to $n \times w$ with window size $w$, trading global context for linear memory and enabling million-token contexts. Ring attention distributes the sequence across devices, keeping each device's memory proportional to $n/P$ rather than $n$. Paged attention (vLLM) eliminates memory fragmentation in the KV cache. In every case, the FLOPs are the same or nearly so — what changes is where and how the bytes move. The arithmetic intensity of attention (FLOPs per byte of memory traffic) is approximately $d/4$ for the score computation, which at $d = 128$ is 32 — well below the ~300 ratio needed to saturate an A100's compute. Attention is, and has always been, an IO-bound operation, and the $\sqrt{d_k}$ divisor, the causal mask, the KV cache, and FlashAttention are all responses to the same underlying constraint: the softmax requires materialising or simulating a structure that grows as $n^2$, and the only question is which level of the memory hierarchy pays for it.