Module 4: Batch norm brings it back

What batch normalization computes

Batch normalization (BN) is a small computation inserted between the linear layer and the sigmoid. For each neuron independently, it looks at the current batch of pre-activations z and re-standardizes them:

ẑ = (z − μ_B) / σ_B,  then  z_out = γ · ẑ + β

where μ_B and σ_B are the mean and standard deviation of that neuron's z across the batch. After the first step, the batch of ẑ values has mean 0 and spread 1 — always, no matter what the linear layer produced. The second step multiplies by γ and shifts by β, two learned parameters per neuron, initialized to γ = 1, β = 0 (so at the start, BN is pure standardization). Why hand back a scale and shift right after removing them? Hold that thought for the end of the page.

The diagram shows the one-block difference: in the plain layer, the Linear box feeds its output z straight into the sigmoid box; in the batch-norm layer, z first passes through the BatchNorm box and the sigmoid receives the standardized z-hat instead. In the code, the difference is a single line in the constructor you saw in Module 1:

self.norms = nn.ModuleList(nn.BatchNorm1d(WIDTH) for _ in range(depth))
...
z = nrm(lin(h))       # forward: Linear -> BatchNorm -> sigmoid
a = self.act(z)

Everything else — starting weights (checksum-verified identical, Module 1), data, loss, optimizer, learning rate, step count — is unchanged.

Re-measuring: the cliff is gone

Same one-pass measurement as Module 2, same untrained starting weights, batch norm inserted:

BATCHNORM (Linear -> BN -> sigmoid): loss at init = 0.7914
  layer    ||dL/dW||    ||dL/dh||  std(z_pre)   std(z)   std(h)  mean slope  W shrink
      1     5.39e-02     8.25e-03      0.3513   1.0001   0.2180      0.2025     0.767
      2     9.34e-02     1.01e-02      0.4858   0.9996   0.2192      0.2020     0.604
      3     1.15e-01     1.44e-02      0.5569   0.9999   0.2215      0.2009     0.623
      4     2.07e-01     1.48e-02      0.2728   0.9995   0.2195      0.2018     0.470
 output     2.36e-01
gradient ratio, layer 1 vs output: 2.29e-01  (4x smaller)
backward shrink factor per layer (||dL/dh_k|| / ||dL/dh_k+1||): 0.816, 0.703, 0.972

(Captured in captures/results.log lines 43–51.) Read it against Module 2's table, column by column. std(z_pre) — what the linear layers produce — is unchanged (0.23–0.56; same weights, so of course). But std(z) — what the sigmoid actually receives — is now pinned at 1.000 at every layer: that is BN doing its one job. Downstream of that, std(h) holds steady at ~0.22 per layer instead of sagging to 0.06, and the weight gradients ||dL/dW|| now span just 5.4e-2 to 2.4e-1 from bottom to top: a ratio of 4x, where the plain network had 3,015x. The bottom line shows why: the per-layer backward factors are 0.82, 0.70, 0.97 — pinned near 1. No sub-unity factor, no geometric decay, no cliff. That is the flat green line you saw in Module 2's gradient plot:

And the forward signal tells the same story from the other direction — the plain network's activation spread (red) decays toward "every input looks the same," while batch norm (green) delivers a constant-strength signal to every layer:

Training confirms what the measurements promise (results.log line 65): loss dives off the ln 2 plateau immediately, crosses 90% test accuracy at step 100, and finishes at 100% — the right-hand decision boundary of Module 3, the S-curve through the moons. And from Module 3's depth-8 experiment (line 79): at double the depth, BN's gradient ratio is still 0.34. The fix does not merely survive depth; it makes depth stop mattering.

Why it works — test the textbook story first

The most common explanation goes: sigmoids kill gradients when saturated; BN centers and scales the pre-activations, keeping the sigmoid in its steep, non-saturated regime; hence gradients flow. It sounds airtight. It is also checkable in one column of our tables — if BN worked by keeping sigmoids steep, the BN network's mean slope should be higher than the plain network's.

Measured, per layer:

(From results.log lines 35–38 and 45–48.) It is exactly backwards: batch norm made the sigmoids less steep, not more. The plain network was never saturated in the first place — Module 2 showed its z's huddle near zero, the steepest place — and BN, by widening the z spread to std 1, actually pushed more pre-activations toward the flatter shoulders. Yet its gradients are a thousand times healthier. For this network at initialization, the saturation-avoidance story is not just incomplete; it predicts the wrong sign.

The real mechanism: automatic gain control

What BN actually fixed is the thing Module 2 actually measured: scale. Recall the hop factor ≈ slope × weight scale ≈ 0.24 × 0.5 ≈ 0.1 — every ingredient is a size, not a saturation effect. BN attacks it in both directions:

• Forward. Whatever the small weights shrink the signal to, dividing by σ_B stretches it right back to spread 1. The sigmoid then compresses it (slope ≈ 0.2, std(h) ≈ 0.22) — but the next BN stretches again. The signal arrives at the output with the moons still distinguishable in it, instead of every input mapping to nearly the same activations.
• Backward. Differentiate ẑ = (z − μ_B)/σ_B and the same 1/σ_B appears in the gradient path: on the way down, the backward signal is multiplied by 1/σ_B. Since σ_B here is the small pre-activation spread (0.23–0.56 — the std(z_pre) column), that multiplication is a boost of roughly 2–4x per layer, canceling the shrinkage from the slope and the small weights. Net hop factor, measured: 0.82, 0.70, 0.97 — approximately 1.

The compact way to say it: BN is a per-layer automatic gain control. Each layer re-amplifies the signal to a standard strength — forward and backward — so no factor compounds, and the geometric decay that defines vanishing gradients never gets started. It is also self-tuning: σ_B is recomputed from the batch at every step, so if the weights grow during training (they do — Module 3's lr = 5 run exploited exactly that), the normalizer adapts instead of falling out of calibration. That is the property a fixed weight rescaling cannot give you, and it is why the fix held up unchanged at depth 8.

Why hand back gamma and beta?

If standardizing is the fix, why let the network immediately re-scale (γ) and re-shift (β) what BN standardized? Because the network — not the normalizer — should decide what distribution each sigmoid input ought to have. Forcing mean 0, spread 1 forever would forbid useful choices: a neuron that wants to act as a hard, nearly-binary switch needs its inputs spread wide into the saturated zones; one that wants to pass information gently needs them narrow. γ and β restore that freedom, while (unlike raw weights) starting from the perfectly conditioned 1-and-0 state. The trained network shows the freedom being exercised (results.log lines 103–104):

learned gamma (BN layer 1) = ['1.42', '4.70', '0.82', '0.87']
learned beta  (BN layer 1) = ['-0.10', '0.46', '0.16', '0.04']

Read it with me: after training, neuron 2 of layer 1 chose γ = 4.70 — spreading its inputs almost five times wider than standard, deliberately driving itself into saturation to act as a sharp threshold — while neurons 3 and 4 chose γ < 1, narrower than standard. Had we pinned γ = 1, all four would be stuck at one compromise. (And note the irony given the textbook story: given the choice, the trained network bought more saturation, not less.)

One flag before you use this in real code: μ_B and σ_B are batch statistics — they exist at training time, when batches exist. At inference you may see one sample at a time, so BN keeps running averages during training and uses those instead when the model is in eval mode. The switch is model.eval(), forgetting it is a classic bug, and Module 5 demonstrates the damage with a measurement.