The Math We Waved Our Hands At

Four experiments in, we can describe a transformer fluently — attention reads the room, heads specialise, the mask blindfolds the future, positions are sinusoids. So here’s a comfortable hypothesis worth puncturing: we understand it now; the equations are just bookkeeping for the intuition.

They are not. The math is where the constraints live — the ones you can’t feel from a diagram but that decide every real design choice: why √d_k is mandatory, why context length is so expensive, why you can fine-tune a 70B model on a laptop, why the thing trains at all. Four rigorous passes — backprop, complexity, linear algebra, stability — and the verdict on “just bookkeeping”: FAR FROM IT.

In EXP-013 we said a saturated softmax “kills gradients.” Let’s earn that claim. Attention learns by pushing a loss gradient back into W_Q and W_K, and the chain runs through the softmax:

$\frac{\partial \mathcal{L}}{\partial W_Q} = \frac{\partial \mathcal{L}}{\partial p}\cdot\frac{\partial p}{\partial s}\cdot\frac{\partial s}{\partial q}\cdot\frac{\partial q}{\partial W_Q}$

Every gradient to the query weights must pass through ∂p/∂s — the softmax Jacobian. And that Jacobian has a specific, revealing shape:

$\frac{\partial p_i}{\partial s_j} = p_i\,(\delta_{ij} - p_j)$

The diagonal is p_i(1 − p_i), the off-diagonal −p_i p_j. When softmax is confident — one p near 1, the rest near 0 — every one of those terms collapses to ≈ 0. The Jacobian is the only door the gradient can walk through, and saturation slams it shut.

↑ Drag the logit scale up (what happens WITHOUT ÷√d_k). The whole Jacobian heatmap fades to zero and ‖J‖ collapses — no gradient reaches W_Q or W_K. Scaling by √d_k is what keeps this door open.

Self-attention compares every token with every token, so it builds an n × n score matrix:

$S = \frac{QK^\top}{\sqrt{d_k}} \in \mathbb{R}^{n\times n}, \qquad \text{time, memory} = O(n^2 d)$

That n² is the single most consequential number in the architecture. Double the context and you quadruple the attention cost. The KV cache — the keys and values you store to avoid recomputing them each step — grows only linearly, O(n). So as context stretches, the quadratic term is what runs away from you.

↑ Drag the context length. The coral (attention, n²) curve pulls away from the blue (KV cache, n) one. This gap is the whole reason FlashAttention (never materialise the matrix) and linear-attention exist.

Strip the story away and attention is linear algebra. W_Q, W_K, W_V are learned projections — changes of basis that re-view each embedding as a query, a key, a value. And the score matrix hides a structural fact:

$S = QK^\top, \qquad Q \in \mathbb{R}^{n\times d_k},\ K \in \mathbb{R}^{n\times d_k} \;\Rightarrow\; \operatorname{rank}(S) \le d_k$

When d_k < n — almost always — the attention matrix is low-rank. It literally cannot express an arbitrary n × n pattern; it’s a rank-d_k object wearing an n × n costume. That’s a real bottleneck (and the premise every linear-attention method leans on). Multi-head is a further restriction — a block-diagonal carve-up of the projections — and W_O is the matrix that mixes the concatenated heads back into one space.

Hold onto “weight matrices are low-rank objects in disguise” — the whole of next quest’s LoRA is built on it.

Backprop multiplies one Jacobian per layer. If each has norm below 1, the gradient shrinks geometrically — by 32 layers, 0.7³² ≈ 10⁻⁵, and the bottom of the network gets no signal at all. Deep stacks should be untrainable. Two tricks save them.

Residual connections. Wrap each sub-layer as x + F(x) instead of F(x), and its Jacobian becomes:

$\frac{\partial}{\partial x}\big(x + F(x)\big) = I + F'(x), \qquad \lVert I + F'(x)\rVert \approx 1$

The I is an identity “highway” the gradient can ride straight down, so its norm stays near 1 no matter the depth. LayerNorm is the second trick — it renormalises each token so the F′ term stays well-scaled and the residual sum never blows up:

$\text{LN}(x) = \gamma \odot \frac{x - \mu}{\sigma} + \beta$

↑ Residual OFF and drag the depth up — layer 0 goes dark (vanished gradient). Flip it ON and the whole stack stays lit at any depth. Putting the norm inside the residual branch (pre-norm) is what makes very deep training stable.

That closes the from-scratch build: attention, multi-head, masking, cross-attention, position — and now the math beneath all of it. But one result here is a door into a whole new quest. We just saw that the attention matrix is low-rank; it turns out the weight updates a model needs during fine-tuning are low-rank too. So you can adapt a frozen 70-billion-parameter model by training two tiny matrices instead of all of them. That’s LoRA — and the opening of Quest-05: Fine-Tuning (EXP-017).