Loss & Stability Primer

A classifier's final linear layer hands back one real number per class — a logit. Negative, positive, huge, tiny: nothing constrains them. They are the output of a matrix multiply, and a matrix multiply has never heard of probability.

Here are five of them. Drag the slider to move the third score and watch what the row does — and what it never does:

Notice the total at the bottom wandering as you drag the third score. At rest it reads 1.8; at a third score of 6 it reads 5.8. Nothing is pulling it anywhere, and nothing could be — the layer that produced these numbers has no normalising step in it.

To make a distribution we need every value positive and the total pinned at one. Exponentiating buys the first half, at a price — raise the third score and watch the second row:

Watch how fast that row loses its shape. At a score of 2 the largest exponential is 7.39 against a total of 12.38; at 6 it is 403.43 out of 408.42 and the other four bars have collapsed onto the baseline. Exponentiating does not merely make numbers positive: it turns a difference of scores into a ratio, and that ratio multiplies by e for every unit of score.

One division buys the second half. Divide each exponential by their total and the row becomes a probability distribution, whatever the scores were:

That is softmax, complete: exponentiate, then divide by the sum. Σp reads 1.0000 at every slider position, because the denominator is the sum of exactly the numbers in the numerator. The invariant is not approximate and it is not learned — it is arithmetic, and it holds on an untrained network too.

Divide the scores by a constant before exponentiating and you get a dial for how sharp the answer is. Drag the temperature across its range while the scores hold still:

At T = 0.1 the winner reads 1.0000 to four decimals and the figure is doing argmax; at T = 4 the largest probability is 0.2905 and the five are nearly level. The scores never move. Temperature is not a property of the model — it is a reading of one, chosen at sampling time.

So a logit is a score, and a probability is a reading of a whole set of scores. Everything below is about the gap between those two sentences, because how you cross it decides whether a training run produces a number or a NaN.

Squared error is the loss everyone meets first, and for regression it is the right one: it is the maximum-likelihood loss under Gaussian noise, its gradient is 2(ŷ − y), and it is convex in the prediction.

Drag the marker along the bowl and read the loss off it — the target is the dashed line at 1:

Notice that the shape does all the work. The penalty grows with the square of the miss, so a prediction 2 away costs four times one that is 1 away, and the gradient grows linearly with the error: the further out you are, the harder you are pushed back.

Classification asks a different question. There the model reports a probability for the true class, and what we want to punish is confident wrongness. Drag the marker toward zero:

Watch the wall on the left. −log p is 0 when the model is certain and right, 0.69 at a coin flip, 4.61 at p = 0.01, and it diverges as p → 0. That is cross-entropy against a one-hot target: the surprise of the answer the model actually gave, unbounded on purpose.

Point squared error at the same question and the difference stops being subtle. Both curves read the same probability; drag the marker left, to where the model is confident and wrong:

At p = 0.05 cross-entropy charges 3.00 and squared error charges 0.90. Worse, squared error is bounded by 1 however wrong the model gets, so the two disagree most sharply exactly where it matters most.

Training does not use the loss value, though. It uses the gradient, and that is where squared error stops working at all. Drag the marker to a badly wrong score:

Because squared error is composed with a sigmoid, its gradient carries a σ(1 − σ) factor that dies exactly where the model is most wrong. At a score of −6 with the target 1, cross-entropy pushes with 0.9975 and squared error with 0.00246 — 405 times weaker. The saturated unit stops learning, and the run looks like it has converged.

That cancellation is the whole reason classification uses cross-entropy: the σ′ term the chain rule contributes is exactly cancelled by the 1/p term in the derivative of −log p. What is left is p − y and nothing else, and we will watch it fall out of the algebra in §06.

Every number in a training step lives in an IEEE-754 float. binary32 holds magnitudes up to 3.4028235e38 and, below 1.4012985e-45, nothing at all. Those two constants are not trivia — they are the walls exp runs into.

The ladder below is one decade a rung, from the smallest subnormal to past the largest finite float32. Drag the slider and watch e to the score climb it:

Each rung up is a factor of ten, so the column's height is the number's exponent rather than its size. exp(2) sits just above 1; exp(89) is off the top, and past the ceiling every value is the same value: inf. There is no gradual degradation — exp(88.72) is a finite 3.39e38 and exp(88.73) is infinity.

Mixed-precision training moves most tensors to sixteen bits, and float16 spends five bits on the exponent against float32's eight. Push the score up and watch which ceiling arrives first:

float16 tops out at 65,504, which is exp(11.09) — a score of eleven, not eighty-nine. That is why every autocast policy keeps softmax, log_softmax and cross_entropy in float32 while the matmuls run in half precision: the multiply is safe at sixteen bits and the exponential is not.

Now drive the pipeline from §01 into that wall. Raise the largest score past 88.72 and read all three rows at once:

Because the numerator and the denominator overflow together, the winning class computes inf ÷ inf and gets NaN, while every other class computes a finite number ÷ inf and gets a clean, plausible 0. The loss is NaN. One backward pass later so is every parameter that touched it — and nothing threw.

This is the failure mode that matters, because it is the quiet one. No exception, no warning, no partial result: the run keeps going, the loss prints nan for the rest of training, and the checkpoint is worthless. The two lines of the next section are what stand between you and it.

softmax(z) = softmax(z − c) for every scalar c. The proof is a line: multiplying numerator and denominator by e^−c changes nothing, and e^z−c = e^ze^−c. It is worth seeing rather than believing.

Slide every one of the five scores by the same amount and watch the probability row underneath it:

Notice that the bottom row is completely still. Every bar above it moves by forty and the third probability holds at 0.5967. Softmax reads only the differences between scores; the absolute level is information it does not have.

So c is ours to choose, and one choice is special. Drag the shift up to the largest score and watch the exponentials:

At c = max the largest shifted score is exactly 0, so its exponential is exactly 1and every other one lands in (0, 1]. Nothing can overflow, because the biggest thing exp is ever handed is zero — and the denominator is at least 1, so nothing can divide by zero either.

Put the two routes on one ladder. Raise the raw score as far as the slider goes and watch which column stays on the rungs:

The raw column leaves the top of the ladder at 89 and never comes back. The shifted one reads 1 at every score, because after the shift the largest exponent is always e⁰. The guarantee costs one extra pass to find the maximum: three passes over the row instead of two.

This is why you never write softmax as its own definition. exp(z) / exp(z).sum() is a correct formula and a broken program; exp(z − z.max()) / exp(z − z.max()).sum() is the same formula and a program that cannot fail. Every framework's softmax is the second one.

Cross-entropy is −log p, so what training actually wants is a log probability. Computing softmax and then taking its log is one line and one bug: the probability has to survive the round trip through the float, and for a confident model it does not.

First the object that replaces it. log Σ exp is a soft maximum — the largest score plus a correction that depends only on the gap. Drag the marker and watch the correction:

At a gap of 0 the correction is ln 2 = 0.6931: two equal scores, so the sum is twice the max. By a gap of 8 it is 0.0003. log-sum-exp is max with the corner rounded off, and the rounding is confined to the region where the top two scores are close.

Now the failure the log inherits. Widen the gap between two scores and watch the runner-up's probability walk down the ladder:

Below 1.1754944e-38 float32 has left its normal range and is spending mantissa bits to keep going; below 1.4012985e-45 it has nothing left and returns exactly 0. That happens at a gap of 104, which is nothing — a trained language model routinely puts a hundred nats between its top and bottom logits.

A probability of exactly zero has a log of −inf, and that is what the naive routereports. Both rows below compute the same quantity; drag the gap past 104:

Because z − log Σ exp never forms the probability, it never has to represent it. At a gap of 140 it returns −140.0 where log(softmax) returns −inf, and −140 is a perfectly ordinary float32. The subtraction is exact; the exponential was the lossy step, and we skipped it.

So log_softmax is not a convenience wrapper around log(softmax(x)). It is a different computation with a different numeric range, and that difference is the entire reason it ships as its own kernel.

Cross-entropy from logits is logΣexp(z) − z[y]: one reduction and one subtraction, no probability materialised, no division. Its gradient is where the design pays for itself.

The derivative with respect to every logit is p − y —the probability row minus the one-hot target row. Drag the true class's score and read the third row off the first two:

Notice that the bottom row really is the difference of the two above it, component by component: at rest, 0.5967 − 1 = −0.4033 under the target and the raw probabilities everywhere else. No σ′ survived, because the 1/p from the log and the p from the exponential cancelled exactly. Every component sits in [−1, 1] whatever the scores are, which is why cross-entropy gradients do not explode on their own.

Hand that loss a probability row instead of a score row and it will softmax it again — silently, because a probability is a perfectly valid float. Switch what the loss is fed:

Watch the ceiling appear. The second softmax sees inputs already squeezed into [0, 1], so the widest gap it can ever see is 1 and the winner is capped at e / (e + n − 1) = 0.4046 for five classes. The model still trains and the loss still falls; it simply cannot go below 0.9048 here, or below 9.82 for a 50,257-token vocabulary — against 10.82 for guessing uniformly.

The other half of the bargain is that the fused form has no overflow to have. The loss cannot depend on a shift, so it must be flat; drag the largest score past 88.72 and see which route agrees:

Because logΣexp shifts internally, the fused route still reads 0.5163 at a top score of 120, while −log(softmax) has been NaN since 88.73. Below the ceiling the two agree to four decimals; above it, one of them is a number and the other is not.

There is a memory argument on top. A materialised probability tensor for 8 × 1,024 positions over a 50,257-token vocabulary is 1.53 GiB in float32, and its gradient is another 1.53 GiB — allocated, written and freed on every single step. The fused kernel forms neither. That is why the API takes the logits.

All of this page is one exponential meeting one finite format. Drag the score and watch which formats can still hold the result:

float16 gives out at a score of 12, float32 at 88.72, and float64 at 709.78 — off the top of this ladder entirely. bfloat16 shares float32's eight exponent bits, so it reaches 88.72 too, with about three decimal digits of precision instead of seven.

Five lines carry everything above. The first is the only one that is not the textbook definition, and it is the one that makes the other four safe:

m    = z.max()                  # the shift
lse  = m + log(exp(z - m).sum())
logp = z - lse                  # log_softmax
loss = lse - z[y]               # cross-entropy
grad = softmax(z) - onehot(y)   # p - y

Fusing them is a memory decision as much as a numeric one, and the cost scales with the vocabulary. Drag it from a small classifier to a language model and read the two totals:

At a 50,257-token vocabulary the fused step holds 1.53 GiB of logits; the unfused one holds 4.60 GiB, because the probabilities and their gradient are each another copy of the same shape. Both axes are log, so the two lines are parallel: the ratio is 3 at every vocabulary, and only the absolute cost moves.

Four ways this fails quietly. Exponentiating raw scores returns NaN rather than raising. Taking log of a softmax output returns −inf once a probability underflows. Feeding probabilities to a loss that expects logits trains a systematically under-confident model that never errors. And moving any of it to float16 to save memory drops the ceiling from 88.72 to 11.09 without changing a line of your code.