One Gradient That Must Be Exactly Zero: Eight Steps of softmax Backprop, by Hand

The previous part ended on this: next comes the twenty-odd steps, by hand. The tools are in order.

This one uses them eight times, working from loss all the way back to logits — exactly the stretch that lives inside F.cross_entropy.

But first, why derive any of it by hand.

backward() never gets it wrong, and never tells you anything

What loss.backward() promises is “write the forward pass and leave the gradients to me”. That promise holds numerically, always — the gradients PyTorch computes are correct. What leaks is the other end:

backward() computes your gradients; it does not judge whether they are reasonable. Judging them means knowing where they came from.

This is what Karpathy calls a leaky abstraction — a layer that spares you the details right up until something breaks, at which point not knowing the details means not being able to fix it. Parts 14 to 17 taught reading the instrument panel; this one opens the engine.

Part 4 put the other half of it in one line: automatic differentiation guarantees mathematical correctness, not numerical stability.

The forward pass, flattened into 18 lines

Deriving by hand needs somewhere to point. What used to be three or four lines becomes one atomic operation per line, each with its shape written beside it — and as part 18 said, the shape is the main clue for the gradient.

Two things about this code run against instinct, and both exist for the sake of marking:

① Back down to 2 layers, not part 17’s 6. With n_embd=10, n_hidden=64 and batch=32, the seven parameter tensors come to 4,137 numbers. Deriving twenty-odd intermediates by hand means keeping the depth as small as possible — one more BN layer is ten more backward steps.

② b1 is deliberately initialised non-zero.

b1 = torch.randn(n_hidden, generator=g) * 0.1     # not zeros

In ordinary training a zero bias is standard practice (with a BN after it, the bias can go entirely). But this section is running automated marking, and the rules are different:

If some quantity happens to be zero, a formula you derived wrongly may come out zero too, and the check hands you a false pass. With non-zero random numbers throughout, any mistake shows up immediately. Perturbing bngain and bnbias is the same idea.

Why split it this finely

Part 4 covered this trade-off once already — back then it was about splitting tanh apart inside micrograd. The bigger the operator you package, the less there is to write and the faster it runs, but a failure cannot be located to a line. There it was optional; here it is mandatory:

Write the forward as one F.cross_entropy(logits, Yb) and the only thing verifiable is the final dlogits; get something wrong and there is no telling which step. Flattened, every intermediate has its own .grad:

logprobs        | exact: True
probs           | exact: True
counts_sum_inv  | exact: False    ← located, right here

Flattening is what lets the marking point at a line.

Why there is not one parameter in these eight steps

Parts 14 to 17 kept coming back to which layers have parameters, but never cut across the question. Sorted by what each kind of layer does, there are only three kinds:

Nearly all parameters live on the mixing layers. Of this network’s 4,137 parameters, mixing takes 4,009 (96.9%) and normalisation the 128 γβ values (3.1%); the non-linearity has none at all.

And the eight steps derived here belong to none of the three — from logits onward it is all plain arithmetic: subtract the max, exp, sum, reciprocal, multiply, log, pick one column. Past that point there are no neurons and no parameters left, which is why no dW appears anywhere in the eight. Parameter gradients start with dW2 and db2, in the next section.

cmp: how to mark your own work

The answer key comes from one loss.backward() run before any of this. The intermediate tensors have to be given retain_grad() first — the same mechanism part 17’s plot ② ran into: by default PyTorch discards intermediate gradients as soon as they are used.

def cmp(s, dt, t):
    ex = torch.all(dt == t.grad).item()        # equal bit for bit
    app = torch.allclose(dt, t.grad)           # equal within floating-point tolerance
    maxdiff = (dt - t.grad).abs().max().item()
    print(f'{s:15s} | exact: {ex} | approximate: {app} | maxdiff: {maxdiff}')

❗ The two columns disagreeing does not mean you are wrong

The reason is that floating-point addition is not associative: (a+b)+c ≠ a+(b+c), differing in the last few bits. Your summation order and PyTorch’s need not match, and the results drift slightly apart.

💡 Run cmp after every single step; do not save up five and check them together. Gradients travel backward, so an error in step 3 makes steps 4 and 5 wrong too — saving them up only hides where it started.

The eight steps

The four shapes were covered in the previous part, so here they are named rather than explained. What is worth looking at in each step is the thing it brings that is new.

① dlogprobs: what the forward picked out, the backward puts back

The starting point, needing no upstream gradient at all.

The forward line is loss = -logprobs[range(n), Yb].mean() — out of a 32×27 table it picks one number per row, the correct answer’s column. Change any cell it did not pick and the loss does not move, so those gradients are zero.

dlogprobs = torch.zeros_like(logprobs)   # lay down a table of zeros with the same shape
dlogprobs[range(n), Yb] = -1.0/n         # write only into the positions that were picked

🔑 The thing worth remembering is the indexing expression on the second line — it is word for word the one on the forward line.

Forward “picks a subset out of A”, so backward is “put the gradients back where they came from and pad the rest with zeros”. The index expression is reused verbatim; nothing else needs inventing.

Zeros first and values second, rather than the other way round, because the overwhelming majority of positions are zero anyway. And zeros_like rather than torch.zeros(32, 27): no shape written out by hand, so a different batch size cannot break it.

Three things that are easy to get wrong:

One useful self-check: dlogprobs.sum() should come to exactly −1.0 — .mean() makes each one −1/32, and 32 of those add to −1.

② dprobs: through the log

Shape A, local derivative 1/x, one line: dprobs = (1.0 / probs) * dlogprobs.

What is worth a second look is that this step pulls apart gradients that were all identical. Step ① wrote −1/32 everywhere, and 1/p blows up in the small-value region:

A factor of 99. The worse the guess, the harder the shove.

Nobody designed that rule; it falls out of 1/x exploding near zero. Cross-entropy naturally concentrates on the samples that are badly wrong and barely pushes the ones already right. Part 7 derived cross-entropy from maximum likelihood without this being visible — it only shows up going backward.

③ dcounts_sum_inv: the first broadcast

The forward is probs = counts * counts_sum_inv, shaped (32,27) = (32,27) * (32,1) — so counts_sum_inv was broadcast 27 times. By the previous part’s test: the variable is one index c short of the output, which is shape B, so sum c away.

dcounts_sum_inv = (counts * dprobs).sum(1, keepdim=True)     # (32,27) → (32,1)

Why the axis is 1 and not 0 took the previous part a whole section; here is a faster version of the same test: whichever of the variable’s axes has length 1 is the axis to sum over. counts_sum_inv is (32, 1), so axis 1, so sum(1). Get it backwards and cmp reports a shape mismatch at once — it fails loudly rather than quietly.

④⑤⑥ dcounts: the fork

counts appears on two right-hand sides in the forward pass (the counts_sum line and the probs line), so it has two paths. The previous part worked that through with this very set of numbers — path one being zero while path two is not, which is why dropping the second is so hard to catch — and it is not repeated here.

What is worth adding are the three traps in the ordering:

It also explains something in passing: why dcounts stops being sparse. The second path adds a non-zero value to every cell — up to here each row had exactly one non-zero, and now it is full.

Step ⑤ crosses counts_sum**-1, whose local derivative is −S⁻², and the sign flips positive. The previous part has a whole subsection on it (including “do not write S⁻¹ = S⁻¹, that says nothing”), so it is not re-derived here.

Also: the forward could simply have said counts / counts_sum. Splitting it into reciprocal and multiply is an arrangement made for deriving by hand — division needs the quotient rule and both branches are awkward, while split apart each step has one simple derivative. The bnvar_inv = (bnvar + 1e-5)**-0.5 in part 17’s network is the same design.

⑦ dnorm_logits: through the exp

Shape A, local derivative exp(x) — which is counts, already computed on the forward pass, so there is nothing to recompute:

dnorm_logits = counts * dcounts

⚠️ On condition that dcounts is complete. Run this after step ④ but before step ⑥ and it uses half a gradient, and everything after it is wrong.

❗ And every row of the result sums to exactly zero

print(dnorm_logits.sum(1).abs().max())    # around 1e-9

The notes reason backward from a property of softmax: its output rows always sum to 1, and since that sum is a constant, adding the same amount to a whole row cannot change any probability — so the row’s gradients must sum to zero.

That can also be computed directly, without relying on the observation. Chaining the previous steps together:

Σ_c e·(T·dp + dS)                    ← one row of dnorm_logits, summed
  = T·Σ_c e·dp  +  dS·Σ_c e          ← split apart
  = T·dT        +  (−S⁻²·dT)·S       ← the first term is step ③, the second is step ⑤
  = dT·(1/S − 1/S)
  = 0

Identically zero, independent of the numbers. Two angles confirming each other: one a structural property of softmax, the other three steps’ formulas substituted in and cancelled. That 1e-9 is not “approximately zero” — it is floating point computing an exact zero and landing there.

⑧ dlogits and dlogit_maxes: the section’s finale

This step spans two forward operations:

logit_maxes = logits.max(1, keepdim=True).values   # operation 1
norm_logits = logits - logit_maxes                 # operation 2

Drawn out, the dependencies look like this:

logits ──①──→ logit_maxes ──┐
   │                        ③
   └────────②───────────────┴──→ norm_logits

🔑 Three edges, so three derivations — but only two gradients come out, because two of the edges both feed dlogits.

The only possible order is ② → ③ → ①: ① needs dlogit_maxes as its upstream, and ③ is what produces it.

② Subtraction with respect to logits: the .clone() is not optional

norm_logits = logits - logit_maxes differentiated by logits gives a coefficient of +1, shape A, the gradient passing straight through:

dlogits = dnorm_logits.clone()

⚠️ Writing dlogits = dnorm_logits is wrong — that just gives the same block of memory a second name. Adding the second path to dlogits below would modify dnorm_logits along with it.

The symptom is a misleading one: cmp('norm_logits', ...) is green before this line and red if you come back to it afterwards — without having touched it. This is Python’s reference semantics rather than anything to do with gradients, but in a long string of d variables it is unusually hard to find.

Part 10 mentioned it once while covering how view shares storage: “if debugging leads you to change a value by hand, know that it may have changed something else with it.” Same thing, this time in the backward pass.

③ Subtraction with respect to maxes: shape B

Same line, differentiated by logit_maxes. It is missing the index c, so c gets summed away, and the coefficient is −1:

dlogit_maxes = (-dnorm_logits).sum(1, keepdim=True)    # (32,27) → (32,1)

① max’s gradient is a routing

logit_maxes = logits.max(1).values. Nudge the largest cell up a little and m follows by exactly as much; nudge any other cell (without overtaking the maximum) and m does not move at all. So:

∂m[b]/∂logits[b,c] = 1 where c holds that row’s maximum, and 0 everywhere else.

Only the winner receives gradient; the rest get zero. That is what routing means — the gradient is not divided up, it is forwarded to one recipient.

dlogits += F.one_hot(logits.max(1).indices, num_classes=logits.shape[1]) * dlogit_maxes

What one_hot does here is turn “which one is the maximum” into a table of 0s and 1s. In part 7, one-hot was a way of encoding the input; this time it is a switch.

❗ And dlogit_maxes must be exactly zero

logit_maxes | exact: False | approximate: True | maxdiff ~1e-9

Part 10 already derived the reason: softmax is shift-invariant, softmax(z) = softmax(z − c) holds for any constant, because the exp(−c) picked up by numerator and denominator cancels. Since logit_maxes has no effect on the output at all, its gradient is of course zero.

It also follows directly from step ⑦ — that step just verified “each row sums to zero”, and this step is exactly summing each row and negating it. Two routes, one conclusion.

That it is not exactly zero is once again floating-point addition not being associative: 27 numbers cancel and the last few bits remain. This is precisely why cmp has two columns.

⚠️ cmp('logits', ...) still has to wait — logits has an upstream of its own, logits = h @ W2 + b2, not yet connected. Numerically dlogits is already complete, so checking early does work.

Three operating rules worth taking away

The eight expressions will fade. These three will not.

① To get a variable’s gradient, search for the right-hand sides it appears on

counts_sum_inv = counts_sum**-1      # where it is produced — does not count
probs = counts * counts_sum_inv      # where it is used — this is the battlefield

Because gradients flow back from downstream, only the places it gets used matter.

That one question answers two: however many right-hand sides it appears on is how many paths it has (and therefore whether += is needed), while the index relationship on that line decides which shape it is. The first step of the previous part’s flow chart is asking exactly this.

It also clears up a naming confusion: d + the variable’s name is that variable’s gradient. Step ③ is called dcounts_sum_inv because it is counts_sum_inv’s gradient — even though that one line mentions counts, dprobs and counts_sum_inv all at once, in three completely different roles: one is what is being differentiated, one is the upstream gradient, one is just the value of the local derivative.

② “How many derivations” is how many edges the dependency graph has, not how many variables

Step ⑧ is the worked example: three variables, three edges, three derivations — but only two gradients out. Counting variables gives two; counting edges gives the right answer.

③ “This gradient should be zero” is itself a checkable property

This is the most valuable line in the part, and it transfers directly.

Subtracting the maximum is a purely numerical device — its whole purpose is to stop exp from overflowing, and mathematically it changes nothing. And “changes nothing” has an executable test: its gradient must be zero.

Turned around:

You add something to your own code that “should not affect the result” (a clamp, a rescaling, a stabilisation, a reparameterisation). Run the backward pass and look at its gradient — zero → it really is a pure device, and the model is unchanged; not zero → you quietly changed the model, and that is not a device, it is a bug.

backward() will never tell you what a gradient ought to be. But when you know, it becomes a ruler.

Closing out

All five backward patterns turned up across these eight steps:

There will not be a sixth. Everything left is these five recombined — except that in the next section’s BatchNorm the same patterns stack ten deep, and that is the one place where samples affect one another at all (part 16 did that accounting).