---
layout: post
title: "Cross-Entropy Loss Explained: From Logits to Language Model Training"
date: 2026-06-14
tags: ["Deep Learning", "Loss Functions", "Softmax", "GPT", "From Scratch", "Python"]
category: "ai"
description: "What is cross-entropy loss? Derive softmax + negative log-likelihood step by step, work a 3-class numeric example, see why MSE fails for classification, and build the loss GPT training actually uses."
keywords: "cross entropy loss explained, softmax cross entropy, negative log likelihood, classification loss, language model loss, cross entropy from scratch, why not MSE for classification"
image: /assets/images/og-cross-entropy-loss.png
order: 7
---

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](/blog/build-autograd-from-scratch/), 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](/blog/numerical-gradient-checking-explained/), 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:

$$
p = \mathrm{softmax}(z) \approx [0.659,\, 0.242,\, 0.099]
$$

MSE loss:

$$
L_{\mathrm{MSE}} = \frac{1}{3}\sum_i (p_i - t_i)^2
$$

$$
L_{\mathrm{MSE}} \approx \frac{(0.659-1)^2 + (0.242)^2 + (0.099)^2}{3} \approx 0.058
$$

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:

$$
L_{\mathrm{CE}} = -\log p_y = -\log(0.659) \approx 0.417
$$

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.

```mermaid
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$:

$$
L = -\log p_y
$$

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):

$$
\tilde{p}_i = e^{z_i}
$$

Normalize:

$$
p_i = \frac{e^{z_i}}{\sum_j e^{z_j}}
$$

Numerical trick I always use: subtract the max before exponentiating so nothing overflows.

$$
m = \max_j z_j
$$

$$
p_i = \frac{e^{z_i - m}}{\sum_j e^{z_j - m}}
$$

For $C = 3$ logits $z = [2.0,\, 1.0,\, 0.1]$:

$$
m = 2.0
$$

$$
e^{z - m} = [e^0,\, e^{-1},\, e^{-1.9}] = [1.0,\, 0.368,\, 0.150]
$$

$$
\sum = 1.518
$$

$$
p = [0.659,\, 0.242,\, 0.099]
$$

---

## How do you combine softmax and cross-entropy?

For one example with true index $y$:

$$
L = -\log p_y = -\log \frac{e^{z_y}}{\sum_j e^{z_j}}
$$

Expand:

$$
L = -z_y + \log\sum_j e^{z_j}
$$

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:

$$
L = \frac{1}{N}\sum_{n=1}^{N} \left( \log\sum_j e^{z_{nj}} - z_{n,y_n} \right)
$$

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:

$$
\frac{\partial L}{\partial z_i} = p_i - \mathbb{1}[i = y]
$$

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$:

$$
\frac{\partial L}{\partial z} = [0.659 - 1,\, 0.242,\, 0.099] = [-0.341,\, 0.242,\, 0.099]
$$

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)$:

$$
\frac{\partial L}{\partial z_i} = -\mathbb{1}[i=y] + \frac{e^{z_i}}{\sum_j e^{z_j}} = p_i - \mathbb{1}[i=y]
$$

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](/blog/numerical-gradient-checking-explained/) passed first try.

```mermaid
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**

$$
m = 3.0
$$

$$
e^{z-m} = [e^{-2},\, e^{0},\, e^{-2.5}] = [0.135,\, 1.0,\, 0.082]
$$

$$
Z = 1.217
$$

$$
p = [0.111,\, 0.822,\, 0.067]
$$

**Step 2 — loss**

$$
L = -\log p_1 = -\log(0.822) \approx 0.196
$$

**Step 3 — gradient**

$$
\frac{\partial L}{\partial z} = [0.111,\, 0.822 - 1,\, 0.067] = [0.111,\, -0.178,\, 0.067]
$$

Model was already pretty confident on class 1 (82%). Gradient magnitude on the true class is modest. If instead $p_1 = 0.05$:

$$
L = -\log(0.05) \approx 2.996
$$

$$
\frac{\partial L}{\partial z_1} = 0.05 - 1 = -0.95
$$

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:

$$
\mathrm{PPL} = e^{L}
$$

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):

```python
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):

```python
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?

```mermaid
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](/blog/build-autograd-from-scratch/) computes $\partial L / \partial \theta$. [AdamW](/blog/adam-optimizer-explained/) 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](/blog/numerical-gradient-checking-explained/) 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](/blog/adam-optimizer-explained/).