Transformers & Attention: The Full Derivation

The √d_k, derived

Assume the components of $q$ and $k$ are independent with mean 0 and variance 1. Then for $q\cdot k={\sum }_{i=1}^{d_k}{q}_i{k}_i$:

So the logits have standard deviation $\sqrt{d_k}$. At $d_k=128$ that is ±11 — and softmax over logits spread that wide is essentially a hard argmax. Its gradient is $p_i({\delta }_{ij}-p_j)$, which is ~0 when one $p$ is ~1 and the rest ~0. Attention stops learning. Dividing by $\sqrt{d_k}$ restores unit variance regardless of head dimension.

Causal masking

Complexity: $Q{K}^{\top }$ is $O(n^{2}d)$ in time and, naively, $O(n^2)$ in memory for the score matrix — per head, per layer. That quadratic memory term is what FlashAttention eliminates, and what makes long context expensive.

Crucially, $d_k=d_{model}/h$, so total FLOPs are unchanged versus a single head of full width. Multi-head attention is free expressiveness — you get several independent attention patterns for the same compute.

The KV cache is why MQA and GQA exist

During autoregressive decoding, the K and V for every previous token are recomputed at every step unless you cache them. Caching turns each step from $O(n^2)$ to $O(n)$ — mandatory. But the cache itself grows:

Worth naming if the conversation goes deeper: Multi-head Latent Attention (MLA) compresses K and V into a shared low-rank latent that is cached instead of the full K/V, achieving roughly a 10× cache reduction while matching MHA quality — the approach used in DeepSeek's models and a good signal that you follow current architecture work.

RoPE, the current default

RoPE treats each pair of dimensions as a 2D plane and rotates it by $m{\theta }_i$, where $m$ is the position. The essential property is that the dot product of two rotated vectors depends only on their relative offset:

Extending context beyond the trained length

1. 1

   Position interpolation (PI)

   Divide positions by a factor s so that position 8192 maps to 2048 in the model's trained range. Cheap, needs a short fine-tune, and works — but compresses high-frequency dimensions and degrades local resolution.
2. 2

   NTK-aware scaling

   Instead of scaling positions uniformly, increase the RoPE base so low-frequency (long-range) dimensions stretch while high-frequency (local) ones stay intact. Often works with no fine-tuning at all.
3. 3

   YaRN

   Combines NTK-aware interpolation with a temperature correction on the attention logits. The current best-in-class method; extends context 8–32× with a small fine-tune.
4. 4

   Long-context continued pre-training

   What frontier labs actually do: extend RoPE base and continue pre-training on genuinely long documents. Expensive but the only approach that produces reliable long-range retrieval rather than merely non-broken perplexity.

ALiBi deserves a mention as the elegant alternative: no embedding at all, just subtract $m\cdot |i-j|$ from the pre-softmax score with a per-head slope $m$. Because it is a monotone recency bias rather than a learned encoding, it extrapolates to sequences far longer than training with no modification — the trade-off being that it hard-codes 'closer is more relevant', which is not always true for retrieval-style tasks.

SRAM is roughly 10× faster than HBM. Standard attention does: write $S=Q{K}^{\top }$ to HBM (n² elements), read it back for softmax, write $P$ back, read it again for $PV$. For n = 4096 with 32 heads in bf16, that score matrix is 4096² × 32 × 2 bytes ≈ 1 GB per layer moved several times. The matmuls themselves are trivial by comparison — the kernel spends its life waiting on memory.

The online softmax trick

Softmax appears to need the whole row before it can normalise. The insight is that you can keep a running maximum and a running sum, and rescale previous partial results when a larger maximum appears:

The backward pass uses the same trick with one extra idea: instead of storing the $n\times n$ attention matrix for the backward pass, it recomputes each tile from the saved $m$ and $\ell$ statistics. Recomputation is cheap because it is compute-bound; storing would be expensive because it is memory-bound. That trade is the essence of the whole method.

• FlashAttention-2 improves work partitioning across thread blocks and warps, roughly doubling throughput again.
• FlashAttention-3 targets Hopper, overlapping matmul with softmax and exploiting FP8.
• PagedAttention (vLLM) applies an analogous idea to the KV cache — virtual-memory-style paging that removes fragmentation and enables far higher batch concurrency.
• In PyTorch you get it for free: F.scaled_dot_product_attention dispatches to a Flash kernel when shapes and dtypes allow. Hand-writing attention forfeits this.

Parameter accounting

Design decisions to be able to justify

• Pre-norm over post-norm — keeps the residual stream unnormalised so gradients reach layer 0 directly; removes warmup fragility.
• RMSNorm over LayerNorm — drops mean-centring and bias for ~10–15% faster normalisation at identical quality.
• SwiGLU over ReLU MLP — a learned gate; FFN hidden is set to 8d/3 to keep the parameter count matched.
• GQA over MHA — 8× smaller KV cache, the dominant serving cost at long context.
• RoPE over learned positions — relative by construction and extendable beyond the trained length.
• Weight tying — sharing embedding and output head saves V·d parameters (0.5 B at V = 128k, d = 4096) and usually improves quality slightly.
• No biases in linear layers — they contribute almost nothing at scale and removing them simplifies sharding and normalisation interactions.