The Same Deviation Is Both Numerator and Denominator: Thirteen Steps Back Through BatchNorm

The previous part’s eight steps stopped at logits. That is the edge of the neural network — those eight were all softmax and logarithms, and not one of them touched a parameter.

This one goes thirteen steps further in: four derivations in section 4, nine in section 5. Parameters appear for the first time.

From here on, the gradients are worth keeping

The previous eight steps all computed gradients of intermediates — dprobs, dcounts, dnorm_logits — and threw them away. Their only reason for existing was to pass the gradient further downstream.

dW2 is different:

W2.data -= lr * dW2       # section 9 will use it like this

This is the first thing that can go straight into updating the model.

The four thousand parameters do not arrive all at once; three sections take a share each:

W2 is 64×27 = 1728 plus b2’s 27, so 1755, which is 42% of every parameter in the network — one output layer eats four tenths of it, because the 27-class dimension is only half the size of the hidden layer’s 64 and sits directly behind it.

Which weight produces h, and which one uses it

The forward lines are these two:

h = torch.tanh(hpreact)      # (32,64)
logits = h @ W2 + b2         # (32,27)

❗ The weight that produces h is W1, not W2

“h is a weight times the previous layer’s output” is half right — that weight is W1:

🔑 h has two identities: it is the first layer’s output (produced by W1) and at the same time the second layer’s input (used by W2). Section 4 deals with the second identity.

The h family: three names, three positions

All three shapes are (32,64). Neither BatchNorm nor tanh changes the dimension — which is exactly what the previous part’s table of the three kinds of layer concluded: only mixing layers change the dimension (30→64, 64→27); non-linearities and normalisation never do.

Matmul backward: shape-matching is fast, the indices tell you why

Part 18 finished deriving the test for shape D, so it is not repeated. What is worth recording is how to be fast in practice.

Written out with indices, the forward is:

logits[b,j] = Σ_i h[b,i] · W2[i,j] + b2[j]        i from 0 to 63

Matching shapes: three seconds to an answer

given: dlogits (32,27), h (32,64), W2 (64,27)

for dh, target (32,64):
  (32,64) = (32,27) @ (27,64)
     dh      dlogits    W2.T        ✅ the only way

  W2 @ dlogits    → (64,27) @ (32,27)   inner 27≠32   ✗
  dlogits.T @ W2  → (27,32) @ (64,27)   inner 32≠64   ✗

for dW2, target (64,27):
  (64,27) = (64,32) @ (32,27)
    dW2       h.T      dlogits        ✅

There are very few candidates, and matrix multiplication is fussy about shapes — a mismatched inner dimension errors out rather than quietly computing something wrong. So matching shapes is reliable here.

But shapes do not tell you which index the sum runs over

dh  = dlogits @ W2.T      # (32,27) @ (27,64)
dW2 = h.T @ dlogits       # (64,32) @ (32,27)

The two expressions look alike and mean opposite things:

This is exactly the previous part’s hook: part 13 said “no need to memorise it, just match the shapes”, part 18 answered “it can actually be derived”, and only after deriving it is it clear why the thing transposed in dW2 is h and not dlogits — because the dimension being contracted is the sample dimension. Matching shapes cannot give you that.

db2: no sample index, so sum along samples

b2[j] is missing the index b, so shape B:

db2 = dlogits.sum(0)

⚠️ This is sum(0), not sum(1). Every sum in the previous part’s eight steps was sum(1) (along classes), because those variables were shaped (32,1). b2 is shaped (27,), and what it is missing is the sample dimension.

⚠️ And no keepdim — b2 is already one-dimensional (27,), not (1,27).

Through tanh

The local derivative is 1 - h**2, using tanh’s output:

dhpreact = (1.0 - h**2) * dh

This line is now its third appearance in the series — part 14 used it to explain why saturation kills the gradient and part 17 used it to compute a quantity, and this time it is derived by hand (the quotient rule, with f = eˣ − e⁻ˣ and g = eˣ + e⁻ˣ, noting that f' = g and g' = f). The table of 1−h² values is in part 14 and is not repeated.

🔑 This is the second time the forward output serves as the local derivative: the previous part’s step ⑦ used counts for exp, and tanh uses h here. Both save a recomputation.

BatchNorm’s nine steps

The seven forward lines:

bnmeani   = 1/n * hprebn.sum(0, keepdim=True)        # (1,64)   column mean μ
bndiff    = hprebn - bnmeani                         # (32,64)  centred x−μ    ⚠️fork
bndiff2   = bndiff**2                                # (32,64)  squared
bnvar     = 1/(n-1) * bndiff2.sum(0, keepdim=True)   # (1,64)   variance σ², Bessel n−1
bnvar_inv = (bnvar + 1e-5)**-0.5                     # (1,64)   1/σ, ε stops div-by-zero
bnraw     = bndiff * bnvar_inv                       # (32,64)  standardisation done
hpreact   = bngain * bnraw + bnbias                  # (32,64)  γβ puts it back   ★the only parameter line

What BatchNorm does, why it is done and what γβ buys are the whole of part 16; only its code is needed here. The two rules E[ax+b] and Var(ax+b) are in part 15.

Two details of notation:

(bnvar + 1e-5)**-0.5 does two things with one exponent — x^0.5 is a square root and x^-1 is a reciprocal, so x^-0.5 is 1/√x. Check: 4.0 ** -0.5 = 0.5 ✓.

The + 1e-5 prevents division by zero: if some neuron outputs the same value on all 32 samples its variance is 0, and 1/√0 blows up into nan. With ε the result is 1/√1e−5 ≈ 316 — very large, but at least finite. PyTorch’s nn.BatchNorm1d also defaults to 1e-5.

The axis changed: every sum here is sum(0)

bnmeani.shape     # (1, 64)   ← dim=0 is 1, so sum(0)
counts_sum.shape  # (32, 1)   ← previous part, dim=1 is 1, so sum(1)

🔑 The test has not changed by a word: sum over whichever of the variable’s axes has length 1. It is just that the 1 is now at the front. Put mathematically, μ_j is missing the sample index b, so the sum runs over b.

The previous part’s SumAxis figure drew precisely this, and is not redrawn.

🔑 Two forks, and what all normalisation shares

Why bndiff must fork takes only the mathematical definition to see:

σ²_j = (1/(n−1)) · Σ_b (x[b,j] − μ_j)²        ← the bracket is bndiff
x̂[b,j] = (x[b,j] − μ_j) / σ_j                  ← the numerator is bndiff too

The same deviation is the numerator, and is also squared and sent off to build the denominator. This is not a question of how the code is written — BatchNorm’s mathematical definition itself requires it to appear twice.

And the previous part has already met this once:

🔑 This is what all normalisation shares, not a defect of either operator. Whenever a batch of numbers is divided by a quantity computed from that same batch, the numerator has to go build the denominator first, and the backward pass must fork. softmax does it, BatchNorm does it, and so do LayerNorm and RMSNorm.

The previous part’s third operating rule was “to get a variable’s gradient, search for the right-hand sides it appears on”. This section is that rule used in reverse: see a normalisation and you know without searching that its numerator has at least two paths.

★ The first sum-backward whose coefficient is not 1

Step ⑥: the forward is a sum, so the backward is a broadcast (shape C). But the local derivative is not 1:

dbndiff2 = (1.0/(n-1)) * torch.ones_like(bndiff2) * dbnvar

The previous part’s step ⑥ looked like this:

dcounts += torch.ones_like(counts) * dcounts_sum        # coefficient 1

The difference is in the forward: there it was a plain counts.sum(1), coefficient 1; here the forward reads 1/(n-1) * bndiff2.sum(0), and that 1/(n−1) has to come back through with it.

This is where Bessel’s correction shows up — using n−1 rather than n for a sample variance does not affect the method of derivation at all; it changes this one coefficient.

⚠️ The mistake that is only 3% wide

BatchNorm’s backward has two different coefficients:

The mean divides by n, the variance by n−1. What happens if they get swapped:

1/32 = 0.031250
1/31 = 0.032258
relative difference = 3.23%

A 3% relative error on gradients whose magnitudes are already tiny means maxdiff comes out as a very small non-zero number — looking much like the 1e-9 the previous part attributed to floating point addition not being associative.

It is not. The previous part’s cmp table gets its first real test here:

🔑 A small maxdiff does not mean correct. The arbiter is always the approximate column. torch.allclose defaults to rtol=1e-5 — a 3% relative error does not get past that, and approximate reports False honestly. Watch only the absolute size of maxdiff and this mistake walks straight through.

It is also why the previous part insisted on running cmp after every single step: save them up and this 3% propagates downstream, so what you end up looking at is a whole run of wrong answers rather than a wrong step ⑥.

.clone() appears a second time

Step ⑨:

dhprebn = dbndiff.clone()
dhprebn += (1.0/n) * torch.ones_like(hprebn) * dbnmeani

The reasoning behind .clone() was covered in the previous part’s step ⑧ (Python’s reference semantics, nothing mathematical), and the conditions here are identical: the right-hand side is a single already-existing variable name, and this gradient still has a += coming.

The test: = dbndiff → an alias, so .clone(); = a * b → multiplication makes a new tensor, so no. Which means only a fork’s first path ever needs it.

The symptom is identical too: cmp('bndiff', ...) is green before this line and red after step ⑨ — without bndiff having been touched.

Where that 1/n comes from

Expand the sum and fix attention on one term; the other 31 are constants as far as it is concerned:

∂μ_j / ∂x[5,j] = 1/n = 1/32 = 0.03125

What it means: make any single number in hprebn larger and it lifts that column’s mean by 1/32, which then moves the bndiff of all 32 samples in that column.

Why nine steps, when Linear takes three lines

Linear’s backward is three lines with nothing crossing between samples. BatchNorm is nine steps. There is exactly one reason:

BatchNorm is the only operation in the network that lets samples affect one another.

hprebn[5][j] ──→ bnmeani[j] ──→ moves the bndiff of all 32 samples in that column
     └──────────────────────→ affects only its own cell

Part 16 did that accounting (its SampleCoupling figure is about the same thing: put one row among two different sets of neighbours and the output changes), and part 11, while defining “evaluation”, added a clause for it: BatchNorm is the first thing in the series that makes the forward pass itself differ between training and inference. This part is the same fact seen from the backward side: hprebn’s gradient has the contributions of all 32 samples mixed into it.

❗ Three paths in the formula, only two in the code

∂L/∂hprebn[5][j] = direct path + via μ_j (from 32 samples) + via σ_j (from 32 samples)

Yet step ⑨ writes only two. Because the one via σ is hidden inside dbndiff — it is step ⑦’s path B (through bndiff2 → bnvar → bnvar_inv). Once merged, only two remain from hprebn’s point of view.

💡 This is exactly where backpropagation saves effort: there is no need to compute a whole long path in one go. Each step handles only its own hop, and the number handed down from upstream already contains everything behind it.

Closing out

Thirteen steps done, and 1883 of the four thousand parameters are in hand:

The remaining 2254 — W1, b1 and C — are section 6’s, and that section’s first three lines are a copy of this part’s matmul-backward section with h swapped for embcat and W2 for W1. Only the last two steps are genuinely new.