Backward through attention

Now an entry off the diagonal. Its sign says which way $p_0$ goes when another score grows.

One more fact about softmax: adding the same number $c$ to every score changes nothing. $e^{z_i + c} = e^c \, e^{z_i}$, and the factor $e^c$ appears in the numerator and in every term of the denominator, so it cancels.

From dp to dz

Each score $z_j$ moves every $p_i$, so $dz_j$ collects one part from each of them: $dp_i \times J_{ij}$. This is the += of U3: a value used in several places gets the sum of all its parts.

When you put the formula for $J$ into the sum, the sum becomes short. $dp \; \text{diag}(p)$ is just $dp_j \, p_j$, and $dp \; p^\top p$ is $(dp \cdot p) \, p_j$, where $dp \cdot p$ is one number for the whole row. So

where $\odot$ means entry by entry (* in NumPy). Subtract the one number $dp \cdot p$ from every entry of $dp$, then multiply by $p$.

Check both ways on numbers, with $p = [0.5, 0.25, 0.25]$ and $dp = [1, 1, 0]$:

• the long way: column 0 of $J$ is $[0.25, -0.125, -0.125]$, so $dz_0 = 1 \times 0.25 + 1 \times (-0.125) + 0 \times (-0.125) = 0.125$;
• the short way: $dp \cdot p = 0.5 + 0.25 + 0 = 0.75$, so $dz_0 = p_0 (dp_0 - 0.75) = 0.5 \times 0.25 = 0.125$.

In NumPy, and checked

In attention, softmax runs along each row of a table, so dp·p is needed once per row: a sum along the last axis. A sum along an axis normally removes that axis: for a (2, 3) table, .sum(axis=-1) gives shape (2,). The argument keepdims=True keeps the axis with size 1, so the result is (2, 1), one number per row. Broadcasting can then subtract it from every entry of its own row.

How do you know a backward pass is right? Compare it with numeric slopes $(L(z + h) - L(z - h)) / 2h$, as in level 6 and U5. A loss that makes the check easy is $L = \sum dp \odot \text{softmax}(z)$: its gradient at $p$ is exactly $dp$. The template does this check for you.

2. One attention head, backward

The forward pass of one head (level 14) has three lines. Here are their shapes, with $L$ tokens:

Backward handles the lines in reverse order, starting with the last one. It starts with dO, which has the shape of O. Two rules do all the work:

• A matrix product $Y = X W$ (level 6): dX = dY @ W.T and dW = X.T @ dY.
• Softmax, each row: section 1’s formula, with $A$ in place of $p$.

The figure shows the order. Every gradient has the shape of its value, which makes the transposes easy to check:

An example by hand

Two tokens, $d_k = 4$ (so $\sqrt{d_k} = 2$), $d_v = 2$. The queries and keys are chosen so that every score is 0:

Every dot product of a row of $Q$ with a row of $K$ is 0, so $S$ is all zeros and every weight in $A$ is 0.5. Each output row is the average of the rows of $V$: $O = [[1.5, 0.5], [1.5, 0.5]]$. The gradient that arrives from the layers above is $dO = [[1, 0], [0, 2]]$.

The product $O = A V$ also gives $dA$: $dA = dO \; V^\top = [[2, 1], [0, 2]]$. Row 0 of $dA$ says how the loss changes when token 0 gives more weight to each token. Now the softmax step, one row at a time.

The full score gradient is $dS = [[0.25, -0.25], [-0.5, 0.5]]$. Each row adds up to 0, like the rows of the Jacobian. The last step is the backward pass of $S = Q K^\top / \sqrt{d_k}$. The forward pass divided the scores by $\sqrt{d_k}$, so both $dQ$ and $dK$ are divided by $\sqrt{d_k}$ too.

Now the keys. $S$ has one row per query and one column per key: $S[1][0]$ is query 1 against key 0. Look at where key $j$ appears in $S = Q K^\top / \sqrt{d_k}$.

The shapes give a quick check on any line you write. Try it with sizes that are all different.

The whole head

Now write the backward pass in NumPy: five lines, in the order of the second table, with the rules you just used by hand. The template runs it on the example above, then checks every gradient against numeric slopes, on random numbers with 3 tokens, $d_k = 2$ and $d_v = 4$.

3. What backward must keep in GPU memory

Look at the five backward lines again. Each one uses values from the forward pass: A, V, K, Q. Those values must still exist when the backward pass reaches them, so the forward pass keeps them. These kept values are called saved activations. In your code they were the tuple saved.

So one head keeps Q, K, V and A, and the other layers keep what their own backward rules need (the input of every matrix product, for example). Training uses GPU memory in two ways.

Per parameter: a fixed cost

Large models train in mixed precision (U2). The forward and backward passes use 16-bit numbers, bfloat16 (bf16 for short). The update is done on a 32-bit float32 copy (fp32 for short), so that small changes aren’t lost to rounding. With Adam (level 7), each parameter needs:

That is 16 bytes per parameter, whatever the batch or the length.

Per token: activations

Saved activations grow with the data in the batch. Most of them have one row per token, shape (B, L, d_model): with $B = 4$, $L = 1000$, d_model = 1000 and 2 bytes per number, one such tensor is 8 MB. A layer keeps several of them, for example the inputs of its matrix products. The attention weights A are different: they have one number for every pair of tokens, in every head. Their shape is (B, h, L, L).

How many (B, L, ·) tensors does one layer keep? Here is one pre-norm block (level 16), with an FFN of width 4 × d_model. The table lists what its backward rules need, step by step, in a plain implementation:

Now make the sequences longer and keep everything else.

The lab adds up everything training keeps, for one model. Move $L$ and watch which part grows fastest.