Vanishing and exploding gradients, measured

Deep networks were hard to train for years for one reason: the signal telling the early layers how to change either shrank to nothing or blew up on the way back. This section measures that directly, layer by layer, instead of describing it. Then it shows what each of the standard fixes does to those numbers — the activation you choose, how the weights are initialised, normalisation, and residual connections — which is why modern architectures look the way they do.

Measure it#

import torch, torch.nn as nn

def deep_mlp(act, depth=20, width=256):
    layers = []
    for _ in range(depth):
        layers += [nn.Linear(width, width), act()]
    return nn.Sequential(*layers, nn.Linear(width, 1))

for act in (nn.Sigmoid, nn.Tanh, nn.ReLU):
    torch.manual_seed(0)
    model = deep_mlp(act)
    model(torch.randn(64, 256)).pow(2).mean().backward()
    norms = [m.weight.grad.norm().item() for m in model if isinstance(m, nn.Linear)]
    print(act.__name__, f"first layer {norms[0]:.3e}", f"last layer {norms[-1]:.3e}")

# Sigmoid  first layer 4.237e-18  last layer 1.324e+00
# Tanh     first layer 3.593e-06  last layer 7.948e-02
# ReLU     first layer 1.330e-09  last layer 1.901e-02

With sigmoid, gradients in the first layers are typically many orders of magnitude smaller than in the last, because each layer multiplies the gradient by a derivative of at most 0.25. The early layers barely learn. ReLU passes gradients through unchanged where active, and the gap shrinks dramatically.

Notice that ReLU shrinks the gap by nine orders of magnitude and still leaves seven. The activation is not the whole story, which is the honest version of an argument usually stated as though it were. What is left is the size of the weights, and that is set by the initialisation.

Initialisation#

If weights are too large, activations and gradients grow at every layer; too small, they shrink. He initialisation, with standard deviation equal to the square root of 2 divided by the number of inputs, keeps signal size roughly steady through ReLU layers. PyTorch's defaults are chosen with the same goal in mind.

Set the initialisation deliberately and run the same measurement. With He initialisation the first layer's gradient is 4.84 and the last is 24.9 — a factor of five apart rather than ten million, on the identical architecture. Then set every weight from a normal distribution with standard deviation 0.01, which looks harmlessly small, and every gradient in the network reads exactly zero: the signal has underflowed to nothing on the way forward, and there is nothing to send back.

def measure(init_fn):
    torch.manual_seed(0)
    model = deep_mlp(nn.ReLU)
    for m in model:
        if isinstance(m, nn.Linear):
            init_fn(m.weight)
            nn.init.zeros_(m.bias)
    model(torch.randn(64, 256)).pow(2).mean().backward()
    norms = [m.weight.grad.norm().item() for m in model if isinstance(m, nn.Linear)]
    return norms[0], norms[-1]

he = lambda w: nn.init.kaiming_normal_(w, nonlinearity="relu")
tiny = lambda w: nn.init.normal_(w, 0.0, 0.01)

print("%.3e %.3e" % measure(he))      # 4.843e+00 2.486e+01
print("%.3e %.3e" % measure(tiny))    # 0.000e+00 0.000e+00

Residuals and normalisation#

Adding each block's input to its output gives gradients a direct path back through the network, and normalisation keeps each layer's inputs in a steady range. Add both to the experiment above and the first-layer gradient becomes comparable to the last even at depth 100. This is why every modern deep architecture uses them.

class Block(nn.Module):
    def __init__(self, width):
        super().__init__()
        self.norm = nn.LayerNorm(width)
        self.lin = nn.Linear(width, width)
        self.act = nn.ReLU()

    def forward(self, x):
        return x + self.act(self.lin(self.norm(x)))     # the addition is the whole trick

for depth in (20, 100):
    torch.manual_seed(0)
    model = nn.Sequential(*[Block(256) for _ in range(depth)], nn.Linear(256, 1))
    model(torch.randn(64, 256)).pow(2).mean().backward()
    norms = [m.lin.weight.grad.norm().item() for m in model if isinstance(m, Block)]
    print(depth, "%.3e %.3e" % (norms[0], norms[-1]))

# 20   3.456e+00 1.754e+00
# 100  1.388e+02 2.635e+01

The reason is visible in the forward expression. Because each block computes x plus something, the derivative of its output with respect to its input is 1 plus something. The chain rule multiplies those together, and a product of terms near 1 does not vanish the way a product of terms near 0.25 does. The gradient has a road straight back to the first layer that does not pass through any weight matrix.

How to use this on your own model#

• Log the gradient norm per layer, not just the total. A single total hides the shape entirely; the per-layer profile is what tells you where learning has stopped.
• A healthy profile is flat to within an order of magnitude or two across layers. A ratio of 1e6 between first and last means the early layers are effectively frozen.
• Exactly zero anywhere means underflow or a dead path, not a small gradient. Look at initialisation, at a ReLU whose inputs are all negative, and at anything you detached.
• Do this once on a few batches at the start of a run, not continuously. It costs a backward pass you were doing anyway plus one norm per layer.