Module 5: The other fixes and one classic bug

Module 2 reduced the disease to one line: each layer multiplies the backward signal by (slope × weight scale) ≈ 0.1, and sub-unity factors compound geometrically. Every fix below is an attack on one ingredient of that product. That framing turns a zoo of tricks into a checklist: for each one, ask "which factor does it push toward 1?"

Fix 1: make the slope 1 (ReLU)

The sigmoid's slope tops out at 0.25 (Module 0). The modern default activation, ReLU (max(0, z)), simply refuses to have a small slope where it is active:

The load-bearing property is the orange line: for any z > 0 the slope is exactly 1 — a gradient passing backward through an active ReLU is not scaled down at all. (For z < 0 the slope is exactly 0: that neuron passes nothing, for that sample.) Swapping torch.sigmoid for torch.relu in the hidden layers — same weights, same everything else — moves the at-init gradient ratio from 3.32e-04 to 2.26e-01 and the final accuracy from 51.5% to 93.0% (results.log line 92; scoreboard below).

Why not 100%? An honest wrinkle worth knowing: a ReLU neuron whose input is negative for every sample has slope 0 for every sample — it is dead, and receives no gradient to revive itself. With only 4 neurons per layer, losing even one hurts. Wider layers make death survivable, which is one quiet reason real ReLU networks are wide.

Fix 2: pick the weight scale on purpose (Xavier init)

The second factor in the product was the weight scale (W shrink ≈ 0.5 — from PyTorch's default init, std 0.31). Xavier initialization chooses the scale deliberately: set Var(W) = 2/(fan_in + fan_out), which for our 4-into-4 layers gives std = √(2/8) = 0.5, so that each layer approximately preserves the variance of the signal passing through it instead of shrinking it.

for lin in list(xavier.hidden) + [xavier.out]:
    nn.init.xavier_uniform_(lin.weight)
    nn.init.zeros_(lin.bias)

xavier layer-2 weight std = 0.5155 (default was 0.3114)

(Captured in captures/results.log line 87 — the measured 0.5155 is the theory's 0.5 plus sampling noise.) That ~1.7x scale-up per layer, compounded over the depth, moves the at-init gradient ratio from 3.32e-04 to 1.59e-02 — 48x healthier — and the network trains to 100%. Note what this fix is: a one-time calibration of the same product BN re-tunes every step. It works beautifully at this depth; Module 4's answer key explains why the one-time version is the more fragile of the two.

Fix 3: give the gradient a bypass (residual connections)

Instead of shrinking the per-layer factor less, residual connections change the topology: each layer computes h + f(h) — its input, plus a learned correction — so where dimensions match (our layers 2–4):

a = self.act(z)
if self.residual and a.shape == h.shape:
    a = a + h

The layer-crossing rate is then ∂(h + f(h))/∂h = I + f′(h): the identity term is a lane the backward signal crosses with factor exactly 1, no matter how small f′ is. The shrinkage is not fixed — it is made additive on top of 1 instead of multiplicative below 1. Measured: at-init ratio 3.24e-02 (nearly 100x better than plain), final accuracy 93.0%. This is the trick that lets ResNets and transformers stack tens to hundreds of layers.

Fix 4: normalize along a different axis (LayerNorm)

LayerNorm does the same standardize-then-γ/β dance as batch norm, but computes μ and σ across the features of each individual sample, not across the batch. Consequences: no batch statistics, so no train/eval distinction and no batch-size sensitivity (the bug demonstrated below cannot happen), which is why it took over in transformers, where batch composition varies wildly.

Measured on our net: final accuracy 100% — and an instructive at-init ratio of just 2.76e-03, only ~8x better than plain. Yet it trains perfectly. This is a useful honest calibration for the whole workshop: the at-init gradient ratio is a strong diagnostic, not a sufficient statistic — normalizers keep re-tuning the gain throughout training, and that ongoing correction can rescue a mediocre starting ratio. (BN's advantage at init comes from dividing by the small per-neuron batch spread ~0.35; LN divides by a per-sample spread across only 4 features, a noisier, generally larger quantity here.)

The scoreboard

All six variants, identical starting weights (Module 1's checksum), identical recipe (SGD, lr = 0.5, 4,000 steps):

variant           init grad ratio L1/out  final test acc
plain sigmoid                   3.32e-04           51.5%
batch norm                      2.29e-01          100.0%
ReLU                            2.26e-01           93.0%
Xavier init                     1.59e-02          100.0%
residual                        3.24e-02           93.0%
layer norm                      2.76e-03          100.0%

(Captured in captures/results.log lines 89–95.)

Every attack on the compounding factor clears the bar; the untreated network alone sits at the coin-flip line. In practice these compose: a modern network typically uses several at once (ReLU-family activation + principled init + residual connections + a normalizer), because each covers the others' weaknesses.

The classic bug: forgetting model.eval()

From Module 4: BN normalizes with batch statistics μ_B, σ_B during training, and keeps running averages to use at inference, when batches may not exist. PyTorch switches between the two behaviors with model.train() / model.eval(). The bug is deploying a BN model without calling model.eval() — silently leaving batch-statistics mode on. The script measures the damage on the trained BN network, predicting the same 200 test points three ways:

bn.eval()                      # correct: normalize with running averages
... bn(X_test) ...
bn.train()                     # bug: normalize each batch by its own stats
... bn(X_test) ...             #   (a) all 200 test points as one batch
... bn(X_test[i:i+2]) ...      #   (b) batches of 2

test accuracy, model.eval() (correct)          = 100.0%
test accuracy, train mode, full test batch     = 100.0%
test accuracy, train mode, batches of 2        = 71.5%

(Captured in captures/results.log lines 100–102.) Read it with me: the correct path scores 100%. The buggy path happens to also score 100% when all 200 points arrive as one batch — with 200 samples, batch statistics approximate the running averages well, which is exactly what makes this bug so treacherous: it can pass your offline evaluation. Then in production, where requests arrive in dribs and drabs, the same model in batches of 2 scores 71.5% — each pair of samples is normalized by its own mean and spread, so each prediction depends on which other sample shared its batch. Two rules fall out: always model.eval() before inference, and be suspicious of any accuracy that changes with batch size — that is BN statistics talking.

You made it

The whole workshop in four sentences. Backprop multiplies one factor per layer, and for plain sigmoid stacks that factor is (slope ≤ 0.25) × (weight scale ≈ 0.5) ≈ 0.1, so the learning signal decays geometrically with depth — measured here as a 3,015x gradient gap at just 4 layers, and a flatlined, coin-flip network. The failure is about scale, not saturation — the dying network's sigmoids were measurably at their steepest. Batch norm fixes the mechanism, not the symptom: by re-standardizing each layer's pre-activations every step it pins the per-layer factor near 1 (measured 0.7–0.97), in both directions, at any depth — same weights, 100% accuracy, while brute-force learning rates die at depth 8. The other standard fixes — ReLU, Xavier init, residual connections, LayerNorm — each attack the same product a different way, and all of them work where the untreated network fails.

Reproducing this workshop

Every number and every plot on these pages comes from one deterministic script:

cd workshops/batchnorm-vanishing-gradients/experiments
python3 experiment.py    # ~1 minute on a plain CPU

Environment used for the captured run: Python 3.10.12, torch 2.12.1+cpu, numpy 2.2.6, matplotlib 3.10.6, on x86-64 Linux; seed 0 (set for torch and numpy at the top of the script, and re-set inside every TinyNet constructor so all variants share weights). The run writes captures/results.log and all measurement figures in diagrams/; line numbers cited throughout the workshop refer to that log. The three structural diagrams were rendered from the .d2 sources in diagrams/ with D2 v0.7.1 (ELK layout) and rsvg-convert 2.52.5. Exact figures may drift with major PyTorch version changes (different default-init RNG streams); the qualitative results — the ~10x-per-layer cliff, the flatline, the rescue — are robust to that.