1.8 — Multi-class cross-entropy and the ŷ − y gradient

1. Why this session exists

Softmax + cross-entropy is the output head of almost every modern classifier and the entire next-token loss of every LLM. Underneath, it admits a single backward identity:

This is the same identity we saw in 1.2 for sigmoid + BCE (the $K = 2$ case). It is the algebraic miracle that makes softmax + cross-entropy the gold-standard pairing — no $\sigma'$ tail, no $K \times K$ matrix to multiply during backward, just a single subtraction. This session is the careful derivation.

2. Setup

A single training example. Logits $z \in \mathbb{R}^K$, one-hot label $y \in \{0, 1\}^K$ with $\sum_k y_k = 1$. Predicted distribution:

Cross-entropy loss:

Because $y$ is one-hot, $L = -\log \hat y_c$ where $c$ is the index of the true class. Equivalently, $L = -z_c + \log \sum_j e^{z_j}$ (the log-sum-exp form from 0.10 / 1.7).

3. The Jacobian of softmax

We need $\partial \hat y_k / \partial z_j$ for all $k, j$. Two cases.

Case $j = k$:

Case $j \ne k$:

Combining,

where $\delta_{kj}$ is the Kronecker delta. The Jacobian matrix is

It is symmetric, PSD, and singular — its null space is $\mathbf{1}$, reflecting the simplex constraint $\sum_k \hat y_k = 1$.

4. The gradient of cross-entropy: derivation 1 (via the Jacobian)

By the chain rule,

Each factor:

Multiply and sum:

The $\hat y_k$ in numerator and denominator cancel. Distribute the sum:

Writing this as a vector identity,

5. Derivation 2 (via the log-sum-exp form)

Same result, no Jacobian. Start from

For the first piece, $\partial(-z_c) / \partial z_j = -\delta_{jc} = -y_j$ (since $y$ is one-hot with $y_c = 1$).

For the second piece,

(A side identity worth committing to memory: the gradient of log-sum-exp is the softmax.)

Adding,

The second derivation is shorter because the log-sum-exp form bypasses the softmax Jacobian entirely. Internally, PyTorch's F.cross_entropy uses something equivalent.

6. Gradient w.r.t. the parameters

With logits $z = W x + b$ in $\mathbb{R}^K$, the chain rule extends through the linear layer.

$\nabla_W L$. $z_k = w_k^\top x + b_k$, so $\partial z_k / \partial w_k = x^\top$ (and $\partial z_k / \partial w_l = 0$ for $l \ne k$). Therefore

which stacks into the matrix outer-product

(Residual column times input row.)

$\nabla_b L$. $\partial z_k / \partial b_k = 1$, so

Averaged over a mini-batch of $m$ examples,

Same "residual × input" pattern as 1.3, just with vector-valued residuals now.

7. Why this identity is the most-used in deep learning

Three reasons.

1. It is the backward signal at the loss layer of every classifier. Every CNN classifier, every Transformer LM head, every reward model in RLHF (Chapter 9) — they all backpropagate $\hat y - y$ out of the softmax + CE head.
2. It composes cleanly with everything upstream. Whatever module produced the logits $z$ — a linear layer, a deep stack of attention blocks — receives a "tell me where you went wrong" signal of the right shape (a $K$-dim residual) at the right scale (no $\sigma'$ damping).
3. It is interpretable. $\hat y - y$ is the error in the predicted distribution. Moving $z$ a step in the direction $-(\hat y - y)$ shifts probability mass from the wrong classes onto the right class. The size of the move is how off the model currently is.

8. Takeaways

• $\partial L_{\text{CE}} / \partial z = \hat y - y$. Memorise this.
• Two derivations: through the softmax Jacobian, or by differentiating the log-sum-exp form.
• Gradient w.r.t. the weight matrix is the outer product $(\hat y - y)\, x^\top$, averaged over the batch.
• Sigmoid + BCE is the $K = 2$ special case of this identity (1.2).

9. Where we go next

Chapter 0 gave us mathematical fluency. Chapter 1 used it to build the simplest end-to-end trainable models — linear regression, binary logistic regression, the perceptron, multi-class softmax regression — and articulated the central obstacle (XOR) that motivates depth.

Chapter 2 introduces the multilayer perceptron and backpropagation: the same chain rule from 0.3 and the same $\hat y - y$ gradient sink from this session, run through more than one matrix multiplication. The result will be every neural network in this course, all the way to the Transformer.