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, so 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 93; 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 gain ≈ 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.

# Xavier: choose the weight scale so each layer preserves signal variance,
# std = sqrt(2 / (fan_in + fan_out)) = 0.5 for our 4-into-4 layers.
xavier = TinyNet(norm="none")
for lin in list(xavier.hidden) + [xavier.out]:
    nn.init.xavier_uniform_(lin.weight)
    nn.init.zeros_(lin.bias)

Section 8's pass of the experiment script (the part that builds and races every alternative fix) logs the resulting scale next to the default (line 88):

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

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):

            # Skip connection -- only Module 5's residual variant takes this.
            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). The script measures each one's at-init gradient ratio, then trains it:

for name, net in candidates.items():
    ratio = init_grad_ratio(net)
    hist = train(net)
    variants[name] = (ratio, hist["test_acc"][-1])
    log(f"  {name:<16} {ratio:>23.2e} {hist['test_acc'][-1]*100:>14.1f}%")

When Section 8's pass runs this loop over all six contenders, it logs the scoreboard (captures/results.log lines 90–96):

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%

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. Section 9's pass of the experiment script measures the damage on the trained batch-norm network, predicting the same 200 test points three ways:

bn.eval()    # correct: inference normalizes with the running averages
with torch.no_grad():
    acc_eval = ((bn(X_test) > 0).float() == y_test).float().mean().item()

bn.train()   # the bug: still normalizing every batch by its own statistics
with torch.no_grad():
    acc_train_full = ((bn(X_test) > 0).float() == y_test).float().mean().item()
    # batch-of-2 predictions in train mode (BN needs >1 sample to compute a std)
    preds = []
    for i in range(0, len(X_test) - 1, 2):
        preds.append((bn(X_test[i:i+2]) > 0).float())
    preds = torch.cat(preds)
    acc_train_b2 = (preds == y_test[:len(preds)]).float().mean().item()

and logs the three verdicts (lines 101–103):

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%

Let's go deeper: 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 gain ≈ 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, giving the 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, rendered in full, ready to copy, on the Source page:

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.