Unit 05.02: Initialization and vanishing/exploding gradients

In a deep network, the gradient reaching the first layer has passed through every layer above it. Small errors in scale compound multiplicatively.

Watching the signal die

If each layer shrinks the gradient slightly, twelve layers shrink it by that factor twelve times over — the early layers receive almost nothing and stop learning. If each layer amplifies it, the gradient explodes into NaN.

The fix is to initialise weights so that variance is preserved through depth. Xavier does this for symmetric activations like Tanh:

import torch
from torch import nn

torch.manual_seed(0)


def gradient_scale(init, depth=12, width=64):
    layers = []
    for _ in range(depth):
        layer = nn.Linear(width, width)
        init(layer.weight)
        nn.init.zeros_(layer.bias)
        layers += [layer, nn.Tanh()]
    model = nn.Sequential(*layers)
    out = model(torch.randn(16, width))
    out.pow(2).mean().backward()
    first = model[0].weight.grad.abs().mean().item()
    last = model[-2].weight.grad.abs().mean().item()
    return first, last


for name, init in [
    ("too small (x0.01)", lambda w: nn.init.normal_(w, std=0.01)),
    ("too large (x1.0)", lambda w: nn.init.normal_(w, std=1.0)),
    ("Xavier", nn.init.xavier_uniform_),
]:
    first, last = gradient_scale(init)
    print(f"{name:18} grad at layer 1: {first:.3e}   at layer 12: {last:.3e}")

# With a bad initialisation the gradient reaching the first layer is orders of
# magnitude smaller or larger than at the last. Those early layers either stop
# learning or blow up. Xavier keeps the scale roughly stable through depth.

Compare the gradient at layer 1 against layer 12 in each row. With weights too small, the first layer receives orders of magnitude less signal. With Xavier, the two stay within a reasonable factor of each other.

The mistake this prevents

Assuming default initialisation is always right. PyTorch's defaults are sensible for common layers, but a custom initialisation copied from a blog post can silently disable the first half of your network.

Takeaway

Initialisation controls whether gradients survive depth. Check the scale at both ends before blaming the architecture.

Mark complete