Normalization & Initialization

Start here — the thing nobody warns you about

You've learned what a neural network is: a stack of layers, each one multiplying by a matrix and bending the result with a nonlinearity. You've learned that gradient descent turns the knobs and backpropagation works out which way to turn them.

So here's a reasonable thing to try: stack thirty of those layers instead of three, and let it rip.

It won't work. The loss will sit there, flat, forever — or it will explode to NaN in the first fifty steps. Not because your code is wrong. Not because thirty layers is too many for the problem. It fails for a reason that is baked into what "stacking layers" means, and for about twenty-five years that reason was the single biggest obstacle in the whole field.

This lesson is about that obstacle and the three fixes that removed it. Every deep network you will ever use — every CNN, every Transformer, every LLM — is built out of those three fixes:

1. Initialization — start the weights at exactly the right scale.
2. Normalization — re-center and re-scale the signal at every layer, continuously, during training.
3. Residual connections — give the signal a shortcut that skips layers entirely.

The core problem — every layer multiplies

Here is the whole disease in one sentence: a deep network is a long chain of multiplications, and long chains of multiplications go to zero or to infinity.

Two names for the two failure modes. When the signal (or the gradient flowing back) shrinks toward zero layer after layer, that's the vanishing gradient problem. When it grows without bound, that's the exploding gradient problem.D2L

See the collapse happen

This is not a hand-wave — you can compute the exact scale at every depth. Below, a 20-layer network with 256 units per layer and ReLU activations, under three different weight scales. The vertical axis is a log scale, so a straight line means a fixed multiplication per layer.

Fix 1 — Initialization: start at the right scale

Initialization is the choice of what the weights are before step one. It sounds like a footnote. It is not: it decides whether the network trains at all.

Why not just set everything to zero?

This is the symmetry problem, and randomness is what breaks it. The weights must be random. The only question is how big.

Why not just "small random numbers"?

The obvious guess — draw from a Gaussian with a small standard deviation like 0.01 — is what everyone tried first, and it fails for a reason you can now compute yourself.

Now do the same arithmetic with the right scale. He initialization says $\sigma = \sqrt{2/n_{\text{in}}}$, so for $n_{\text{in}} = 512$: $\sigma = \sqrt{2/512} = \sqrt{0.00390625} = 0.0625$. Only six times larger than the naive 0.01 — and that six-times difference is the difference between a network that trains and one that doesn't.

The two formulas you'll actually use

There are two standard recipes, and which you pick depends on your activation function.

Run the numbers yourself. This cell propagates a signal through 20 layers under all three scales and prints what survives:

The other way signals die — saturation

Bad scaling is one killer. There's a second, and it comes from the activation function itself.

That flatness is called saturation, and it interacts horribly with depth. Play with the curves below: drag along the sigmoid and watch its derivative curve collapse to nothing on either side.

Here is the killer arithmetic. The sigmoid's derivative never exceeds 0.25, even at its very best point. So in the most favourable case imaginable, a 10-layer sigmoid network multiplies the gradient by at most $0.25^{10} \approx 10^{-6}$ on the way back. In the realistic case, where activations have wandered into the flat regions, the factor is far smaller still.

Fix 2 — Normalization: reset the scale, continuously

Good initialization gets you a network whose signal is well-scaled on step zero. But training moves the weights. By step 1,000 the careful scaling has drifted, and by step 100,000 it's gone.

Here is what that does to a layer's output distribution. Before normalization, a layer deep in a drifting network might be producing values centered around 3 with a spread of 2.5; after, they're centered on 0 with a spread of 1 — every single time, regardless of what the previous layers did.

BatchNorm — normalize each feature across the batch

The first version to work, and the one that changed everything, is BatchNorm.Ioffe & Szegedy 2015 Its choice is: for each feature (each neuron), compute the mean and spread across the examples in the current mini-batch, and standardize using those.

Step 5 deserves its own explanation, because it looks like it undoes the whole thing.

Compute it yourself and confirm the hand-worked numbers:

LayerNorm — normalize each example across its own features

BatchNorm has a structural weakness: it needs a batch. Its statistics are computed across examples, which makes one example's output depend on its batch-mates, breaks with tiny batches, and gets awkward when examples have different lengths — as sentences do.

LayerNorm makes the opposite choice.Ba et al. 2016 Same arithmetic — subtract a mean, divide by a spread — but computed in the perpendicular direction.

Fix 3 — Residual connections: give the gradient a shortcut

Initialization sets the scale right at the start. Normalization keeps it right throughout. Yet even with both, networks past about 20 layers were still getting worse as you added depth — and not from overfitting: their training error went up too. Adding layers was making the model worse at the data it could see.

The change is one addition. Instead of a block computing $\mathbf{y} = F(\mathbf{x})$, it computes:

• $\mathbf{x}$ — what came into the block, passed through completely untouched. This is the escalator.
• $F(\mathbf{x})$ — whatever the block's layers compute: the residual, the change being proposed.
• $+$ — plain addition. That is the entire architectural innovation, and it made 100-plus-layer networks routine.He et al. 2015 (ResNet)
• Why it matters for training: if the block should do nothing, it only has to drive $F(\mathbf{x})$ toward zero — which is trivially easy, since shrinking weights toward zero is exactly what gradient descent does naturally. Compare that with making a stack of matrices reproduce the identity exactly. The residual form makes "do nothing" the default and improvement the deviation.

Putting the three together — the modern block

Every serious architecture since about 2016 combines all three fixes, and they compose in a specific order. In a Transformer the pattern is:

import torch
import torch.nn as nn

class TransformerBlock(nn.Module):
    """The pre-LN block used by essentially every modern LLM."""

    def __init__(self, d_model, n_heads):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)          # fix 2: normalize BEFORE the sublayer
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.norm2 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, 4 * d_model),
            nn.GELU(),
            nn.Linear(4 * d_model, d_model),
        )
        # fix 1: PyTorch already applies a sensible init; be explicit for the MLP
        for layer in self.mlp:
            if isinstance(layer, nn.Linear):
                nn.init.normal_(layer.weight, std=0.02)   # the GPT-2 convention
                nn.init.zeros_(layer.bias)

    def forward(self, x):
        h = self.norm1(x)
        x = x + self.attn(h, h, h, need_weights=False)[0]   # fix 3: residual
        x = x + self.mlp(self.norm2(x))                     # fix 3: residual
        return x

Read the two lines in forward. Each is x = x + something(norm(x)): normalize, transform, add back to the untouched input. That one line, repeated 32 or 80 or 120 times, is the entire skeleton of a large language model.

Practice — and how to make it stick

1. By hand: a layer has 2048 inputs and 512 outputs with ReLU. Compute the He standard deviation, then the Xavier one. By what factor do they differ, and what would that do to the signal over 20 layers?
2. From scratch: implement LayerNorm's backward pass in NumPy and check it against torch.nn.LayerNorm with gradcheck. The mean and variance both depend on every input, which makes this a genuinely instructive derivative.
3. Experiment: train a 20-layer MLP on MNIST three times — no normalization, BatchNorm, LayerNorm — and plot the loss curves together. Then delete the residual connections from a small ResNet and watch the training error get worse with depth.
4. Read the source: skim §3 of the ResNet paper and find Figure 1 — the plot of a 56-layer plain network training worse than a 20-layer one. That single figure is the motivation for the entire architecture.

Try it right here — edit and run the code, and if you get stuck or hit an error, ask Ada on the right: she can see your code and terminal output.

Next: with deep networks now trainable, learn how to stop them memorizing the training set in Regularization — then see all three fixes assembled at scale in the Transformer Block.