Off by 0.0001: Twenty-Six Steps Compressed into Six, and Hand-Derived Gradients Training the Same Model

The previous three parts walked the whole path from loss to every parameter — twenty-six intermediates and parameters, each one checked once by cmp.

But that was opening the black box and walking through it. This part does the opposite: it puts the box back together.

Cross-entropy’s eight steps compress into one line and BatchNorm’s ten into one, by the same method. Then loss.backward() gets commented out and the model trains for 200,000 steps on the hand-derived gradients — to see whether it trains at all.

One thing to put on the table first

The one thing worth taking from this figure: parameters have no batch axis.

W2 is (64,27), a single copy; h is (32,64), 32 rows. Thirty-two samples share one set of parameters, so the backward pass has to pool 32 opinions into one — which is why sum(0) is everywhere in the backward pass, all through part 20’s nine steps.

💡 On any network diagram from now on, write the real shape beside each node and then ask one question: does this layer’s parameter have a batch axis? The answer is always no ⟹ the backward pass must collapse that axis.

The task changed: simplify on paper, then translate to code

The previous three parts flattened the forward into atomic operations and walked back one at a time. Ex2 and Ex3 go the other way:

The method is one sentence:

🔑 Count how many times a variable appears in the forward pass; the backward has that many terms. Ex2 is twice → two terms, Ex3 is three times → three terms. Not guessed at — counted.

This is really part 19’s third operating rule (“to get a variable’s gradient, search for the right-hand sides it appears on”) applied after simplification — there the search was over lines of code, here it is over the places it occurs within a single expression.

Eight steps into one line

Why there is no starting from logit_maxes

Ex1’s first step was logit_maxes. In Ex2 it does not exist — F.cross_entropy still subtracts the maximum internally to stop exp overflowing, but that intermediate is never exposed.

🔑 Backpropagation walks the forward dependency graph in reverse, so however many intermediates the forward has is how many gradients the backward computes. Merge the forward and the intermediates are gone, leaving only dlogits.

And part 19’s step ⑧ already proved that dlogit_maxes is identically zero (softmax’s shift-invariance). Since that path contributes zero, skipping it cannot change the answer — which is exactly why the two algorithms agree: Ex1’s route through logit_maxes was adding a zero.

Simplifying: log turns the division into a subtraction

For one sample (dropping the sample index b, with y the correct answer):

L = −log p_y = −log( e^(z_y) / Σ_k e^(z_k) )

Expand with log(a/b) = log a − log b:

L = −z_y + log Σ_k e^(z_k)
    └─┬─┘   └──────┬──────┘
   first term   second term

Check it numerically (cut down to 4 classes, correct answer y=2, z = [1.2, 1.9, 3.0, 1.2]): −3.0 + log(33.42) = −3.0 + 3.51 = 0.51 ✓

🔑 After simplification the forward goes from one fraction to two terms added together. The division is gone, so no quotient rule; the intermediates are gone, so nothing to step through. Ex1’s eight steps were all spent computing that fraction the hard way.

Differentiating term by term

First term −z_y: when j = y the derivative is −1; when j ≠ y the expression contains no z_j at all, so it is 0. Together that is −1[j=y].

Second term log Σ_k e^(z_k): it holds all 27, so it is non-zero whatever you differentiate by. Using d(log f)/dx = f'/f, and noting that only the k=j term of the sum contains z_j (the rest are constants), the numerator is e^(z_j):

∂/∂z_j · log Σ_k e^(z_k) = e^(z_j) / Σ_k e^(z_k) = p_j

💡 This step is neat: after differentiating, that pile of exp turns itself back into p_j — softmax reappears. Which is why the final formula needs nothing but one F.softmax.

Adding the two terms:

∂L/∂z_j = p_j − 1[j=y]

Two terms forward, two terms back, one to one: forward “pick exactly one” → backward “act at exactly that one”; forward “pool everything” → backward “every position contributes”.

1[·] is the indicator function, 1 when the condition holds and 0 otherwise. Laid out over all j, it is the one-hot vector — its job is to squeeze a piecewise definition into one expression.

Part 10 already gave this result (∂L/∂logits = (prob − onehot(y)) / N), along with the intuition that the correct answer’s gradient is negative so its logit is pushed up while the rest are pushed down, and it closed with “V5 is the episode that has you derive this simplification yourself”. This section is that promise delivered — what is new is the derivation, not the result.

Where that 1/n comes from

What was derived above is one sample’s loss, while the loss in the code is the mean over 32. z[b,c] occurs only inside L_b (sample b’s logits do not affect any other sample’s loss), and the other 31 terms are constants:

∂loss/∂z[b,c] = (1/n) · (p[b,c] − 1[c = y_b])

That is shape C, and this 1/n local derivative has now appeared twice before:

It has the same origin as the −1/n in part 19’s step ①; only the timing differs. Ex1 multiplies at the first step and carries it through the remaining eight; Ex2 skips the middle and multiplies once at the end. Multiplication commutes.

★ What if the loss used .sum() instead of .mean()

Without that 1/n the gradient is just p − onehot, 32 times larger.

The consequence is not only the magnitude: change the batch size and the gradient magnitude changes with it, so the learning rate has to be retuned. With .mean() there is no such problem — the gradient’s size is independent of batch size. Which is why nearly every framework’s loss defaults to mean rather than sum.

Formula → code

dlogits = F.softmax(logits, 1)      # ① compute p
dlogits[range(n), Yb] -= 1          # ② subtract 1 at the correct answer
dlogits /= n                        # ③ divide by n

Three symbols, three lines. The index [range(n), Yb] on line ② is word for word part 19’s step ① — both mean “pick the correct answer’s column out of each row”. The code never builds a one-hot matrix to subtract; it just subtracts where the subtraction belongs.

⚠️ Here exact: False is right, and True would be suspicious

logits | exact: False | approximate: True | maxdiff ~8e-9

The two algorithms are mathematically equivalent but take different floating-point routes: Ex1 does dozens of operations, Ex2 does three. The rounding accumulates differently, so the last few bits must differ.

🔑 This is a different False from part 19’s step ⑧: there the gradient was supposed to be exactly zero and two numbers cancelled down to a remainder; here both are genuinely non-zero, both routes are correct, and neither is exact.

⚠️ Put the other way round: exact: True here would be suspicious — it would suggest Ex1’s result had been reused by accident rather than computed independently.

That makes three uses of the cmp table, no two alike:

Each row still sums to zero

Σ_c ( p[b,c] − 1[c=y_b] ) = 1 − 1 = 0

Because probabilities sum to 1 and so does one-hot. This is the same fact as part 19’s step ⑦ (“every row of dnorm_logits sums to zero”) — there it was observed and then proved, and here it is visible directly in the formula.

Ten steps into one line

Same method, different target. BatchNorm merged into one expression:

y[b,j] = γ_j · ( x[b,j] − μ_j ) / σ_j + β_j

(with √(σ² + ε) abbreviated to σ to keep it short.)

Count how many times x appears

Three appearances ⟹ three paths ⟹ three terms in the backward formula.

This is the previous part’s “two forks” said differently: there the forks were counted in the flattened code, here the appearances are counted in the merged expression. Same thing counted.

❗ First, a proof that three is all of them

“Count the appearances” needs one hole plugged first: the expression for σ² also contains μ, so is x → μ → σ² a fourth path?

∂σ²/∂μ = (1/(n−1)) Σ_k 2(x_k − μ)·(−1) = −2/(n−1) · Σ_k (x_k − μ)

And Σ_k (x_k − μ) ≡ 0 — that is the definition of μ (subtract the mean and the total is zero). So this path is identically zero.

🔑 Three really is all of them — not an approximation, not something small being dropped. Without this step, “count the appearances” is a slogan; with it, it is a test.

Deriving term by term

Fixing the notation first (column j, n samples, the index j dropped throughout):

μ   = (1/n) Σ_k x_k
σ²  = (1/(n−1)) Σ_k (x_k − μ)²
v   = (σ² + ε)^(−1/2)          ← bnvar_inv in the code
x̂_b = (x_b − μ) · v            ← bnraw in the code
y_b = γ · x̂_b + β

What comes down from upstream is dy_b = ∂L/∂y_b.

A　through γ — shape A, elementwise:

∂L/∂x̂_b = γ · dy_b

B　∂L/∂v — v is (1,64), shared by the whole column of n x̂ values, missing the sample index ⟹ shape B, so sum over k:

∂L/∂v = Σ_k (∂L/∂x̂_k)(∂x̂_k/∂v) = γ Σ_k dy_k (x_k − μ)

Here the key substitution x_k − μ = x̂_k / v gets used once:

      = (γ/v) Σ_k dy_k · x̂_k

bndiff has become bnraw — the swap matters because bnraw is on hand in the code, while bndiff is not computed at all once things are compressed.

C　v to σ² — the power rule with n = −1/2:

∂v/∂σ² = −½ (σ² + ε)^(−3/2) = −½ v³

∂L/∂σ² = (γ/v)(Σ_k dy_k x̂_k) · (−½ v³) = −½ γ v² Σ_k dy_k x̂_k

D　σ² to x_b — only the k=b term of the sum contains x_b:

∂σ²/∂x_b = 2(x_b − μ)/(n−1)

⚠️ This is where that n−1 first appears. It rides all the way to the end as the n/(n−1) in the final formula.

E　∂L/∂μ — μ is shared by the whole column too, so shape B again:

∂L/∂μ = Σ_k (∂L/∂x̂_k)(∂x̂_k/∂μ) = γ Σ_k dy_k · (−v) = −γ v Σ_k dy_k

F　Adding the three paths

∂L/∂x_b =   γ·dy_b · v                                      ← ① direct
          + (−γ v Σ_k dy_k) · (1/n)                         ← ② via μ
          + (−½ γ v² Σ_k dy_k x̂_k) · 2(x_b − μ)/(n−1)       ← ③ via σ²

The third term takes x_b − μ = x̂_b / v once more, cancelling one of the v²:

③ = −γ v · x̂_b · Σ_k dy_k x̂_k / (n−1)

All three now carry γv, so pull it out:

∂L/∂x_b = γv · [ dy_b − (1/n) Σ_k dy_k − x̂_b Σ_k dy_k x̂_k / (n−1) ]

Pull 1/n outside the bracket as well, and term ③’s 1/(n−1) becomes n/(n−1):

∂L/∂x_b = γv/n · [ n·dy_b − Σ_k dy_k − n/(n−1) · x̂_b · Σ_k dy_k x̂_k ]

Which is those three lines of code, with v as bnvar_inv and x̂ as bnraw:

dhprebn = bngain * bnvar_inv / n * (              # common factor γv/n
    n * dhpreact                                  # ① itself
    - dhpreact.sum(0)                             # ② the mean's correction
    - n/(n-1) * bnraw * (dhpreact * bnraw).sum(0) # ③ the variance's correction
)

Three appearances forward, three terms back, three lines of code. The two .sum(0) calls are exactly the cross-sample sums from steps B and E.

✅ This chain was checked on its own: computing ∂L/∂x from the formula above in pure Python and comparing against central-difference numerical differentiation on the same random inputs gives a maximum relative error of 9.4e−09 at n=32; the intermediate ∂L/∂v, ∂L/∂μ, ∂L/∂σ² and ∂σ²/∂x_b each agree with numerical differentiation to six decimal places.

🔑 That n/(n−1) was derived, not fitted — it originates in step D’s ∂σ²/∂x_b denominator and becomes n/(n−1) when 1/n is pulled out. This is Bessel’s correction showing through: with the textbook 1/n variance the coefficient would be 1 and the expression would be otherwise identical. The previous part’s step ⑥, “the first sum-backward whose coefficient is not 1”, has become a factor inside this expression.

Why the last two terms are subtracted

Look at the forward, x̂ = (x − μ)/σ: when x[b,j] grows, μ_j grows with it, so part of the numerator’s increase is cancelled; and σ_j grows too, so the denominator grows and compresses it further.

🔑 BN resists changes to its own input, so the gradient has to subtract that resistance. The direct effect is weakened by the knock-on changes in the mean and the variance, and the backward pass has to reflect that faithfully.

★ Two quantities that must be zero: BN’s backward is a projection

Reorganising the three terms (ignoring coefficients):

dx ∝ dy − (dy's batch mean) − x̂ · (correlation of dy with x̂)
        └──── term ② ────┘   └──────── term ③ ────────┘

Each of those guarantees one thing:

Why they must be zero: BN’s output is always mean 0 and variance 1 — so adding a constant to the whole column of hprebn, or multiplying it by one, is cancelled out and the output does not change at all. Since neither direction affects the loss, the gradient’s component along each must be zero.

🔑 This is the same phenomenon as part 19’s “dlogit_maxes must be exactly zero”, one level up. There a single scalar was identically zero (subtracting the maximum is a pure numerical device with no effect on the output); here two directional components are.

Part 19’s operating rule said that “this gradient should be zero” is itself a checkable property. This is its strongest form: however many degrees of freedom the forward has that do not affect the output, that is how many directions the backward gradient must vanish along.

Measured:

print(dhprebn.sum(0).abs().max())              # 2.6e-09
print((dhprebn * bnraw).sum(0).abs().max())    # 2.9e-07

⚠️ The second being two orders of magnitude larger than the first is normal — it does one extra elementwise multiply before summing, so the rounding error is amplified once by bnraw’s magnitude (around ±2). The test is not which is smaller but how each compares with the values themselves: the numbers in dhprebn sit around 1e−3, so 1e−7 is already full cancellation.

⚠️ unbiased=True cannot be dropped

When checking the merged forward:

hpreact_fast = bngain * (hprebn - hprebn.mean(0, keepdim=True)) / \
               torch.sqrt(hprebn.var(0, keepdim=True, unbiased=True) + 1e-5) + bnbias

unbiased=True is what corresponds to 1/(n−1). Write False and the forward no longer matches Ex1’s — it is off by several percent, not by 1e−7. This is the other face of the previous part’s “1/n or 1/(n−1)” trap.

maxdiff depends on the magnitude of the intermediates, not the number of steps

Ex2’s maxdiff was 8.1e−09; here Ex1 takes ten steps and yet comes out at only 9.3e−10.

Because the size of the error depends on the magnitude of the intermediates — BN’s (bnvar, bnvar_inv) all sit near 1 and the rounding is barely amplified, whereas softmax’s counts_sum can be in the tens.

Removing backward()

Every check so far has been on one fixed batch, for one step. Ex4 asks a different question: do the hand-derived gradients actually train the model?

From twenty-six steps down to six groups

None of bndiff, counts_sum and the rest needs computing any more:

💡 Group ④ cannot be dropped. Ex3 only merged the path through BN down to hprebn; BN’s own γβ still have to be derived separately — they are parameters, not part of the chain that got compressed.

Three things to watch

① The whole thing is wrapped in with torch.no_grad():, whose job is not to build a computation graph inside that block:

y1 = x * 3
print(y1.grad_fn)      # <MulBackward0>   ← remembers “I came from a multiply”

with torch.no_grad():
    y2 = x * 3
    print(y2.grad_fn)  # None             ← nothing remembered

🔑 This confirms the observation from part 19’s step ⑦: the backward pass needs the forward’s intermediates, so they have to stay in memory. Turn the backward off and they do not. Forgetting it still produces correct results — this is not a correctness problem but a memory and speed one.

② p.data += -lr * grad — the gradient points along the steepest increase in the loss, hence the minus. The learning rate is 0.1 for the first 100,000 steps and 0.01 after: walk in large steps, then refine in small ones.

③ The order of grads must line up with parameters.

parameters = [C, W1, b1, W2, b2, bngain, bnbias]
grads      = [dC, dW1, db1, dW2, db2, dbngain, dbnbias]

⚠️ A misalignment does not raise an error (only a shape mismatch would), but training goes completely off the rails.

Why running statistics are necessary

bnmean_running = 0.999 * bnmean_running + 0.001 * bnmean

An exponential moving average, updating with 0.1% of the new value each step, which amounts to remembering roughly the last thousand steps.

🔑 At inference the batch is a single row, x.std(0) is nan, and μσ simply cannot be computed. This is the only reason model.train() / model.eval() exist: what they switch between is “compute μσ on the spot” and “use what was accumulated during training”. Part 16 implemented it once; this is its second outing.

The evaluation forward differs from the training forward by exactly one line:

# training: compute this batch's μσ on the spot
hpreact = bngain * (hprebn - bnmean) * (bnvar + 1e-5)**-0.5 + bnbias
# evaluation: use what training accumulated
hpreact = bngain * (hprebn - bnmean_running) * (bnvar_running + 1e-5)**-0.5 + bnbias

The structure is identical; only the source of μσ changes. If evaluation used the current batch’s μσ, the same model would predict differently on 22,655 rows than on one — the model would stop being well defined.

★ The verdict

train 2.0713
val   2.1105

Off by 0.0001 — already down at the level of floating-point noise.

❗ The parameter counts differ by 200, but the two models are the same one. Those 200 are b1: part 16 did that accounting and a bias on the layer in front of BN has no effect, so it was dropped there. V5’s starting code keeps it — and as part 20 explained, it is even deliberately initialised non-zero, to stop a wrong derivation passing by accident.

Keeping it changes nothing: hprebn = embcat @ W1 + b1 adds the same constant to a whole column, and on the next line bnmeani in bndiff = hprebn - bnmeani grows by the same amount, so they cancel exactly. That is a direct consequence of this part’s previous section — b1 lies along one of those “did nothing” directions. So the two rows above compare the same model.

🔑 Why this is more convincing than cmp

cmp verifies one batch, one step; it is static. If some gradient carries a systematic bias, a single-step cmp may well pass, and yet after 200,000 steps the loss curve diverges and the val loss comes out clearly higher.

The previous part’s “mistake only 3% wide” is the ready-made example: swap 1/n and 1/(n−1) and maxdiff stays small enough to look like floating-point error — easy to wave through when looking at that one number. Two hundred thousand steps of accumulation would not wave it through.

🔑 cmp checks that this step’s expression is right; Ex4 checks that twenty-six steps together stay on course for 200,000 steps. They are different questions, and neither replaces the other.

train and val differ by only 0.039

No noticeable overfitting. For 180,000 training samples, 12,297 parameters is still on the small side, and the model sits on the underfitting side — wider and deeper would keep driving the loss down. (If train fell to 1.5 while val stayed at 2.1, that would be overfitting, the model memorising the training set.)

The ruler part 12 set up — train/val/test before any tuning — gives its first reading comparable across parts here.

The end of V5

From part 18’s tools to this part’s 0.0001, the episode is done. What is worth keeping is not those twenty-six expressions:

backward() is no longer a black box. Next time a gradient looks wrong, you know which step to open.

The next episode (V6, a WaveNet-style hierarchy) is being skipped. It makes this MLP deeper and wider and the loss falls further, but it does so with the same parts this article just took apart — so, on value for time, straight on to V7. The next part starts from Shakespeare and a tokenizer and builds up to a Transformer; what is inside backward(), you already know.