I used to think cross-entropy was just “the loss PyTorch picks when you’re doing classification.” Fair enough — nn.CrossEntropyLoss() is the default for a reason. But when I wired up pretraining for AL-1, I had to implement the backward pass myself. That’s when I realized cross-entropy isn’t an arbitrary choice. It’s maximum likelihood in disguise, and for next-token prediction it’s the only loss that penalizes wrong probability mass in the right way.
If you’ve got working autograd and you’ve run numerical gradient checking, this is the next piece: turn model outputs into a scalar you can minimize, derive the gradient, and see the famous cancellation that makes softmax + cross-entropy so clean to implement.
Why doesn’t MSE work for classification?
Mean squared error on one-hot targets sounds reasonable until you try it.
Say three classes with logits $z = [2.0,\, 1.0,\, 0.1]$ and true label $y = 0$. One-hot target is $t = [1, 0, 0]$. Softmax probabilities:
MSE loss:
The model already assigns 66% to the correct class and MSE is tiny. MSE treats probability $0.66$ and $0.99$ as nearly equivalent errors from $1.0$ — both are “close” in L2 sense. For classification I care about confidence on the right class and suppressing wrong classes. Cross-entropy punishes low probability on the true label much harder:
Still not catastrophic, but the gradient structure is what matters. MSE + softmax gives messy, vanishing gradients when the model is confidently wrong. Cross-entropy + softmax gives a clean $p - t$ residual. That’s why every language model uses it.
flowchart LR
LOG["Logits z"] --> SM["softmax"]
SM --> P["probs p"]
P --> CE["−log p_y"]
TGT["target y"] --> CE
CE --> L["scalar loss L"]
What is cross-entropy, intuitively?
Cross-entropy measures how surprised I am by reality, given my model’s distribution.
If my model says $p(\text{cat}) = 0.9$ and the true label is cat, surprise is low: $-\log 0.9 \approx 0.105$. If my model says $p(\text{cat}) = 0.01$ but the label is cat, surprise is huge: $-\log 0.01 \approx 4.605$.
For a single example with true class $y$ and predicted probabilities $p$:
Log does two jobs: turns multiplication of probabilities into addition (handy for long sequences), and blows up as $p_y \to 0$, which is exactly the penalty I want.
How do you derive softmax?
The model outputs raw scores $z \in \mathbb{R}^C$ — one logit per class. I need a probability vector $p$ that sums to 1.
Exponentiate (makes everything positive):
Normalize:
Numerical trick I always use: subtract the max before exponentiating so nothing overflows.
For $C = 3$ logits $z = [2.0,\, 1.0,\, 0.1]$:
How do you combine softmax and cross-entropy?
For one example with true index $y$:
Expand:
The second term is the log-sum-exp (LSE). This form is what I implement in code — no need to materialize $p$ in the forward pass unless I want probabilities for inspection.
For a batch of $N$ examples, mean loss:
Language models are the same thing with more dimensions. Logits shape $(B, T, V)$, targets shape $(B, T)$ — each position is an independent classification over vocabulary size $V$. GPT-2 small: $V = 50257$. Average the loss over all non-masked token positions.
What is the gradient? (The cancellation)
This is the part that made implementation click. Define $p = \mathrm{softmax}(z)$. For one example:
Read that: gradient is predicted probability minus one-hot target. Correct class gets $p_y - 1$ (negative when $p_y < 1$, pushing logits up). Wrong classes get $p_i$ (positive, pushing logits down).
Quick sanity check with $p = [0.659, 0.242, 0.099]$ and $y = 0$:
The true class logit should increase; the others should decrease. The gradient says exactly that.
Why does this fall out? Chain rule through $L = -z_y + \mathrm{LSE}(z)$:
No separate softmax Jacobian multiplied by cross-entropy gradient — they collapse. In autograd I can fuse the ops or implement one backward kernel. I fused them in AL-1’s cross_entropy and gradcheck passed first try.
flowchart TB
subgraph "Backward at one logit z_i"
PY["p_y"] --> GY["∂L/∂z_y = p_y − 1"]
PI["p_i (i≠y)"] --> GI["∂L/∂z_i = p_i"]
end
A full numeric walk-through (3 classes)
Logits $z = [1.0,\, 3.0,\, 0.5]$, true label $y = 1$ (middle class).
Step 1 — stable softmax
Step 2 — loss
Step 3 — gradient
Model was already pretty confident on class 1 (82%). Gradient magnitude on the true class is modest. If instead $p_1 = 0.05$:
Much steeper pull. Cross-entropy automatically scales the correction by how wrong I am.
What about perplexity?
Perplexity is just cross-entropy dressed up for language modeling:
If average CE is $2.3$ nats per token, perplexity is $e^{2.3} \approx 9.97$ — “as confused as if uniformly guessing among ~10 tokens.” GPT-2 at convergence on web text lands around 20–30 depending on dataset; random over 50k vocab would be ~50,000. Perplexity is monotonic with CE, so minimizing one minimizes the other.
From-scratch implementation sketch
Forward (mean over $N$ tokens):
import numpy as np
def cross_entropy_forward(logits, targets):
N, V = logits.shape
m = logits.max(axis=1, keepdims=True)
shifted = logits - m
log_sum_exp = m[:, 0] + np.log(np.exp(shifted).sum(axis=1))
rows = np.arange(N)
nll = log_sum_exp - logits[rows, targets]
loss = nll.mean()
return loss, shifted
Backward (given stored shifted logits from forward):
def cross_entropy_backward(shifted, targets, grad_out=1.0):
N, V = shifted.shape
exp = np.exp(shifted)
probs = exp / exp.sum(axis=1, keepdims=True)
grad = probs.copy()
grad[np.arange(N), targets] -= 1.0
grad *= grad_out / N
return grad
PyTorch’s F.cross_entropy expects raw logits and integer targets — it applies log-softmax + NLL internally. Don’t softmax first; you’ll double-softmax and get garbage.
ignore_index and masked positions
Instruction fine-tuning masks loss on padding or prompt tokens. Same formula, but zero out positions where target == ignore_index and divide by the count of kept positions, not total length.
For AL-1 pretraining I use every token in the packed window. For instruction tuning I’ll mask user prompt positions and only backprop on assistant completions — same cross_entropy, different mask.
Where does this sit in the training loop?
flowchart LR
BATCH["batch x, y"] --> GPT["GPT forward"]
GPT --> LOG["logits (B,T,V)"]
LOG --> CE["cross_entropy"]
CE --> L["loss"]
L --> BW["backward"]
BW --> OPT["AdamW step"]
Cross-entropy is the objective. Autograd computes $\partial L / \partial \theta$. AdamW uses those gradients to update weights. Embeddings, attention, LayerNorm — none of it learns without a scalar loss telling the network what “better” means.
For language models, “better” means assign high probability to the true next token at every position. Cross-entropy is that requirement written as calculus.
Common mistakes I hit
| Mistake | Symptom |
|---|---|
Softmax before cross_entropy |
Loss stuck, gradients near zero |
log(softmax) manually without max trick |
inf / nan on large logits |
| Mean vs sum mismatch in backward | Gradients off by factor of $N$ |
| Wrong target dtype (float one-hot vs int indices) | Silent wrong loss |
| Forgetting to divide masked CE by kept count | Loss scale drifts with sequence length |
Run gradcheck on the fused op with random logits shape $(2, 3, 7)$ before trusting the training loop. Takes five seconds, saves five hours.
I used to treat cross-entropy as a black-box API call. Implementing it once — forward, backward, numeric check — fixed that. The loss isn’t magic; it’s “make $p_y$ big” expressed as $-\log p_y$, with a gradient so clean ($p - t$) that the whole training stack downstream gets simpler. Next up in the series: actually moving the weights with AdamW.
Support the writing
If this post helped, a coffee keeps the deep dives, paper picks, and news digests coming.
Buy Me a Coffee