Will Ahmad

Linear algebra for LLMsPart 8: Gradients through matrices

Softmax and cross-entropy

6 min

The loss on one position

p=softmax⁡(z),L=−log⁡py\displaystyle p = \operatorname{softmax}(z), \qquad \mathcal{L} = -\log p_y

zz are the logits for one position, yy is the index of the token that actually came next (Part 6).

Its gradient is p−onehot(y)p - \text{onehot}(y)

L=−zy+log⁡∑jezj\displaystyle \mathcal{L} = -z_y + \log \textstyle\sum_j e^{z_j}
∂L∂zk=−[k=y]+ezk∑jezj=pk−[k=y]\displaystyle \frac{\partial \mathcal{L}}{\partial z_k} = -[k = y] + \frac{e^{z_k}}{\sum_j e^{z_j}} = p_k - [k = y]

Write log⁡py\log p_y out, differentiate, and the softmax comes back. The gradient on the logits is the predicted distribution minus the one-hot target.

Write

∂L/∂z=p−onehot⁡(y)\partial \mathcal{L}/\partial z = p - \operatorname{onehot}(y).

On numbers

z=[2,1,0]  →  p≈z = [2, 1, 0] \;\to\; p \approxy=2nd token:p−e2≈y = 2\text{nd token}: \quad p - e_2 \approx

The correct logit gets −(1−py)=−0.755-(1 - p_y) = -0.755, so descent pushes it up. Every wrong logit gets +pk+p_k, pushed down in proportion to how much probability it stole. The entries sum to 0.

✏️ Your turn

Gradient by hand

p=[0.1,0.7,0.2],target=2nd token\displaystyle p = [0.1, 0.7, 0.2], \quad \text{target} = \text{2nd token}

Write ∂L/∂z\partial \mathcal{L} / \partial z.

Work it out on paper, then go to the next card to check.

Answer

[0.1, 0.7−1, 0.2]=[0.1,−0.3,0.2]\displaystyle [0.1,\ 0.7 - 1,\ 0.2] = [0.1, -0.3, 0.2]

A confident correct prediction (0.7) gets a small push; a perfect one (py=1p_y = 1) would get none.

The softmax Jacobian

∂p∂z=diag⁡(p)−p p⊤\displaystyle \frac{\partial p}{\partial z} = \operatorname{diag}(p) - p\,p^\top

On its own, softmax mixes every input into every output: ∂pi/∂zj=pi([i=j]−pj)\partial p_i/\partial z_j = p_i([i = j] - p_j). Combined with cross-entropy, most of it cancels into p−yp - y, which is why frameworks fuse the two (F.cross_entropy takes logits, not probabilities).

In a transformer

Inside attention there’s no cross-entropy to cancel against, so the backward pass multiplies by diag⁡(p)−pp⊤\operatorname{diag}(p) - pp^\top for every row of attention weights.

Write

Softmax Jacobian: diag⁡(p)−pp⊤\operatorname{diag}(p) - pp^\top.

In PyTorch

# mean over all B·T positions
loss = F.cross_entropy(logits.view(-1, V), targets.view(-1))
# never pass logits.softmax(-1): it expects raw logits

📌 Recap

Softmax and cross-entropy in one card

  • L=−log⁡py\mathcal{L} = -\log p_y, p=softmax⁡(z)p = \operatorname{softmax}(z).
  • ∂L/∂z=p−onehot⁡(y)\partial \mathcal{L}/\partial z = p - \operatorname{onehot}(y): sums to 0.
  • Softmax alone: diag⁡(p)−pp⊤\operatorname{diag}(p) - pp^\top, used inside attention.

Write

The three lines above.