🎓 A1Learner
🔍

PyTorch Vanishing and Exploding Gradients

A gradient norm is the length of a gradient treated as one arrow, showing how strongly a layer learns. In a deep plain network, that norm can shrink to zero or blow up as it travels back through layers.

Gradients fade in a small-weight stack

Five tanh layers start with weight std 0.10. Read the norms from the last layer back to the first, the way gradients travel.

import torch

torch.manual_seed(0)
layers = [torch.nn.Linear(8, 8, bias=False) for _ in range(5)]
with torch.no_grad():
    for layer in layers:
        torch.nn.init.normal_(layer.weight, std=0.10)
x = torch.randn(16, 8)
target = torch.randn(16, 8)
h = x
acts = []
for layer in layers:
    h = torch.tanh(layer(h))
    h.retain_grad()
    acts.append(h)
loss = torch.nn.MSELoss()(h, target)
loss.backward()
print([f"{a.grad.norm():.1e}" for a in acts])

The same stack, oversized weights

Bigger starting weights push the other way. Each layer multiplies the travelling gradient instead of shrinking it.

torch.nn.init.normal_(layer.weight, std=0.50)  # one change: the starting scale

Gradients hold with a matched start

The same five layers with Kaiming init and ReLU tell a calmer story. The norms stay in one small band from last layer to first.

import torch

torch.manual_seed(0)
layers = [torch.nn.Linear(8, 8, bias=False) for _ in range(5)]
with torch.no_grad():
    for layer in layers:
        torch.nn.init.kaiming_uniform_(layer.weight, nonlinearity="relu")
x = torch.randn(16, 8)
target = torch.randn(16, 8)
h = x
acts = []
for layer in layers:
    h = torch.relu(layer(h))
    h.retain_grad()
    acts.append(h)
loss = torch.nn.MSELoss()(h, target)
loss.backward()
print([f"{a.grad.norm():.1e}" for a in acts])

Example

import torch


def grad_norms(mode):
    torch.manual_seed(0)
    layers = [torch.nn.Linear(16, 16, bias=False) for _ in range(20)]
    with torch.no_grad():
        for layer in layers:
            if mode == "small":
                torch.nn.init.normal_(layer.weight, std=0.10)
            elif mode == "large":
                torch.nn.init.normal_(layer.weight, std=0.50)
            else:
                torch.nn.init.kaiming_uniform_(
                    layer.weight, nonlinearity="relu")
    torch.manual_seed(1)
    x = torch.randn(64, 16)
    target = torch.randn(64, 16)
    h = x
    activations = []
    for layer in layers:
        if mode == "small":
            h = torch.tanh(layer(h))
        else:
            h = torch.relu(layer(h))
        h.retain_grad()
        activations.append(h)
    loss = torch.nn.MSELoss()(h, target)
    loss.backward()
    return [a.grad.norm().item() for a in activations]


for name, mode in [("Small weights", "small"),
                   ("Large weights", "large"),
                   ("Kaiming init", "kaiming")]:
    norms = grad_norms(mode)
    picked = [norms[i] for i in (0, 4, 9, 19)]
    text = ", ".join(f"{v:.2e}" for v in picked)
    print(f"{name}: layers 1, 5, 10, 20: {text}")

Key points

  • Gradient norm: the square root of the sum of squared gradient values, read with norm()
  • Vanishing: gradients shrinking toward zero as they travel back, so early layers barely learn
  • Exploding: gradients growing huge as they travel back, so one step can wreck the weights
  • Use retain_grad() to keep a gradient on an activation, so its norm can be read after backward()
  • Gradients travel from the last layer back to the first, so read the tables in that direction
  • A matched initialization holds gradient norms in one band across all layers, as the Kaiming table shows

In this example, one function builds a 20-layer plain network and returns the gradient norm at each layer. Small weights with tanh give a norm of 6.37e-02 at layer 20 and only 7.31e-10 at layer 1. The gradient all but vanishes before it reaches the first layers. Large weights with ReLU give 7.65e+01 at layer 20 and 3.17e+05 at layer 1. The gradient explodes on the same trip back. Kaiming init holds every measured layer between 3.89e-02 and 6.39e-02. The starting scale alone decides which story the stack tells.

Note

Gradient clipping caps an exploding norm during training, while good initialization stops it from forming.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

A first layer's gradient norm is 7e-10 while the last layer's is 6e-02. What is happening?