Machine Learning Foundations

Loss functions in depth: MSE, cross-entropy, softmax

The objective the whole model chases

~12 minLesson 10 of 57Written by Claude Opus 4.8

The loss is the single number that defines what "good" means - and its shape, not just its value, decides how fast and where the network learns.

A neural network never sees the word "wrong." It only sees a number - the loss - and the slope of that number, which it slides down with gradient descent. So the loss is the definition of good you hand the optimizer, and choosing it is a real decision. Use the mean squared error when the target is a quantity you want close; use cross-entropy when the target is a probability you want the model to believe. They are not interchangeable: swap them and the training either crawls or converges to the wrong thing. The reason lives in the shape of each curve - how steeply it punishes a confident mistake - which is exactly what you will drag around below.

The one-line idea

A loss is a scorecard for one prediction: hand it the truth and the model's guess, get back a penalty. Averaged over the data it becomes the cost the optimizer minimises. Squared error scores numbers by their distance; cross-entropy scores probabilities by their surprise. The optimizer then flows downhill on whichever scorecard you picked - so the scorecard's slope is the lesson the model actually learns.

MSE: penalise the gap, squared

For a regression task the model outputs a raw number $\overset{y}{^}$ and the truth is another number $y$ . The natural score is the gap, squared, averaged over the $n$ examples:

MSE = \frac{1}{n} i = 1 \sum n (y_{i} - \overset{y}{^}_{i})^{2} .

Squaring does two jobs. It makes every error positive, so overshoots and undershoots cannot cancel; and it punishes a big miss far more than a small one - double the error, quadruple the penalty. As a function of a single prediction it traces a gentle parabola, a bowl, with its lowest point sitting exactly on the true value. This is the cost function of Lesson 17, drawn for numbers. Its crucial property for later: the penalty is bounded and shallow - even a wildly wrong prediction on a $[0, 1]$ probability can cost at most $1$ .

Cross-entropy: penalise surprise

Classification is a different animal. The model outputs a probability $p$ (or a whole vector of them), and "how far off" is not a distance but a surprise. Cross-entropy measures how surprised the true label is by the model's stated probabilities. For a binary label $y \in {0, 1}$ with predicted $p = Pr (y = 1)$ :

L_{BCE} = - [y ln p + (1 - y) ln (1 - p)] .

Only the term for the true class survives: if the truth is $y = 1$ the loss is simply $- ln p$ . For more than two classes, the categorical version sums over classes $c$ against a one-hot truth $y_{c}$ , and again collapses to a single log:

L_{CE} = - c \sum y_{c} ln p_{c} = - ln p_{true} .

That logarithm is the entire story. Right and confident ( $p_{true} \to 1$ ) costs almost $0$ . But wrong and confident ( $p_{true} \to 0$ ) sends the loss racing to $+ \infty$ - an unbounded penalty for asserting a falsehood with certainty. Lesson 63 built cross-entropy out of entropy; here that same quantity becomes a training objective.

Softmax: logits become a distribution

One catch: a network's final layer emits unbounded real numbers called logits. They can be negative; they sum to nothing in particular. Cross-entropy needs genuine probabilities, so softmax squeezes a logit vector $z$ into a distribution:

p_{i} = softmax (z)_{i} = \frac{e ^{z_{i}}}{\sum _{j} e ^{z_{j}}} .

Exponentiating forces every entry positive; dividing by the total forces them to sum to exactly $1$ . The largest logit wins the largest share, yet every class keeps a nonzero slice. Feed those $p_{i}$ into categorical cross-entropy and you get the workhorse classification loss - "softmax cross-entropy" - sitting under every image classifier and every language model's next-token head (Lesson 37). Drag both pieces below.

Two losses, one confident wrong answer drag the sliders

The true label is 1 (say, "this image is a cat"). Drag the model's predicted probability $p$ of that true class and watch both penalties. squared error $(1 - p)^{2}$ cross-entropy $- ln p$

Predicted p0.900

true label y = 1 | at p = 0.900 → squared error = 0.010 (never exceeds 1), cross-entropy = 0.105.

As you drag $p$ toward $0$ - the model insisting the cat is not a cat - the teal squared-error curve flattens against its ceiling of $1$ , while the red cross-entropy curve climbs off the top of the chart toward $+ \infty$ . Same mistake, wildly different punishment.

Softmax: raw scores → a probability distribution

Drag any logit. Softmax re-normalises all three so the bars still add to exactly 100% - the biggest score takes the biggest share, but nobody hits zero.

z_cat2.0z_dog1.0z_bird0.1

Push one logit far above the rest and softmax sharpens toward a near-one-hot pick; nudge them together and it spreads toward a uniform $\frac{1}{3}$ each. These $p_{i}$ are precisely what feed cross-entropy above.

Why the shape is the point

A loss is only useful through its gradient - the slope gradient descent rides downhill. Pair softmax with cross-entropy and the gradient with respect to each logit collapses to something astonishingly clean:

\frac{\partial L _{CE}}{\partial z _{i}} = p_{i} - y_{i} (predicted - true) .

A confident wrong answer drives $p_{i} - y_{i}$ toward $- 1$ right where the log wall is steepest, so the model gets shoved hard in the correct direction; a nearly-right answer barely moves. Try MSE on that same softmax output and its gradient vanishes exactly when the model is confidently wrong - the curve you saw flatten against its ceiling has almost no slope there, so learning stalls precisely where cross-entropy accelerates. That is why classifiers use cross-entropy, not MSE.

Match the loss to the task

The loss is a design choice, not a law of nature. Numbers you want close → squared error (or robust cousins like MAE). A single yes/no → binary cross-entropy. One-of-

K

classes → softmax + categorical cross-entropy. A probability you want calibrated → cross-entropy again. Pick the wrong one and the optimizer faithfully optimises the wrong thing: a classifier trained with MSE learns slowly and stays timid and under-confident. Choosing the loss is choosing what the model is even trying to be good at.

import numpy as np

# --- the two penalties ---
def squared_error(y, p):   return (y - p) ** 2      # regression penalty, bounded on [0,1]
def cross_entropy(p_true): return -np.log(p_true)     # classification penalty, unbounded

# demo case: true label y = 1, sweep the predicted probability p
for p in [0.9, 0.5, 0.1, 0.01, 0.001]:
    print(p, round(squared_error(1, p), 3), round(cross_entropy(p), 3))
# p=0.9    SE 0.010   CE 0.105
# p=0.5    SE 0.250   CE 0.693
# p=0.1    SE 0.810   CE 2.303
# p=0.01   SE 0.980   CE 4.605   <- SE nearly capped, CE still climbing
# p=0.001  SE 0.998   CE 6.908   <- confidently wrong: CE explodes

# --- softmax: raw logits -> a probability distribution ---
def softmax(z):
    e = np.exp(z - z.max())     # subtract max: numerically stable, same result
    return e / e.sum()          # every output > 0, and they all sum to 1

logits = np.array([2.0, 1.0, 0.1])
print(softmax(logits))          # [0.659 0.242 0.099]  -> sums to 1.000

The two losses, and softmax, in NumPy

Check yourself

MSE and cross-entropy differ most when a prediction is:

Softmax outputs are best described as:

Which loss for which task - and why does cross-entropy, not MSE, punish a confident wrong classification?

Try to state it, then check.

Regression → MSE

\frac{1}{n} \sum (y_{i} - \overset{y}{^}_{i})^{2}

: a gentle bowl, bounded (at most

1

on a probability). Classification → cross-entropy

- \sum_{c} y_{c} ln p_{c} = - ln p_{true}

: because of the logarithm, a confident wrong probability

p_{true} \to 0

sends the loss to

+ \infty

while MSE merely tops out. Softmax

e^{z_{i}} / \sum_{j} e^{z_{j}}

first turns raw logits into the

p_{c}

. That steep log wall gives gradient descent a strong push exactly where the model is badly wrong - the gradient is a clean

p_{i} - y_{i}

Lock it in

The loss is the definition of "good" you hand the optimizer; its shape, not just its value, decides how fast and where the network learns.
MSE penalises the squared gap - a gentle, bounded bowl (at most $1$ on a probability) whose minimum sits on the true value - the right choice for regression.
Cross-entropy penalises surprise through $- ln p_{true}$ : right-and-confident costs almost $0$ , wrong-and-confident races to $+ \infty$ - the right choice for probabilities.
Softmax turns raw logits into a distribution (every entry positive, summing to $1$ ); paired with cross-entropy the gradient collapses to a clean $p_{i} - y_{i}$ .
Match the loss to the task: MSE on a softmax output has almost no slope exactly where the model is confidently wrong, so learning stalls - which is why classifiers use cross-entropy, not MSE.

Primary source

For an unhurried, build-it-up derivation of cross-entropy - from entropy to the loss classifiers actually train on - watch StatQuest's Neural Networks / cross-entropy playlist. For the visual intuition of how a loss becomes the surface gradient descent walks down, see 3Blue1Brown's Neural networks series. Together they show the same softmax-then-cross-entropy pipeline drawn two complementary ways.

Ask your teacher

This lesson quietly unifies half the course. The cost function of Lesson 17, the LLM's next-token loss (Lesson 37), and maximum likelihood estimation (Lesson 62) are all the same move: pick a loss whose shape makes gradient descent (Lesson 11) learn the right thing - in fact minimising cross-entropy is maximising likelihood. Ask me why

- ln p

falls out of MLE, why we scale nothing but still avoid overflow with the max-subtraction trick, or when a task wants a bespoke loss (ranking, contrastive, focal) instead of these two staples.