Vanishing and exploding gradients

Measuring it in a real network

Now watch it happen in PyTorch. The code below builds four networks, each 20 layers deep. They differ only in the activation and in how the starting weights are chosen:

• sigmoid
• tanh
• ReLU with PyTorch's default starting weights
• ReLU with He initialization, which sizes the starting weights to suit ReLU (the weight initialization lesson shows where that size comes from)

Each network looks at the same 256 small images of handwritten digits. The code runs the backward pass once. Then it measures the size of each layer's weight gradient.

"Size" here means the gradient's norm. It's one number for how big a whole list of numbers is. Square every number, add them up, and take the square root. Think of it as the length of an arrow.

The chart uses a log scale. Each step up the side means 10 times bigger. A flat line means every layer gets a gradient of about the same size.

import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
from sklearn.datasets import load_digits

digits = load_digits()
X = torch.tensor(digits.data[:256] / 16.0, dtype=torch.float32)
y = torch.tensor(digits.target[:256])

def deep_net(act, he_init=False, depth=20):
    layers = []
    for _ in range(depth):
        lin = nn.Linear(64, 64)
        if he_init:
            nn.init.kaiming_normal_(lin.weight, nonlinearity="relu")
            nn.init.zeros_(lin.bias)
        layers += [lin, act()]
    return nn.Sequential(*layers, nn.Linear(64, 10))

for name, act, he in [("sigmoid", nn.Sigmoid, False), ("tanh", nn.Tanh, False),
                      ("ReLU, default init", nn.ReLU, False), ("ReLU, He init", nn.ReLU, True)]:
    torch.manual_seed(0)
    model = deep_net(act, he)
    nn.CrossEntropyLoss()(model(X), y).backward()
    norms = [m.weight.grad.norm().item() for m in model if isinstance(m, nn.Linear)]
    print(f"{name:<19} layer 1: {norms[0]:.1e}   layer 10: {norms[9]:.1e}   last: {norms[-1]:.1e}")
    plt.semilogy(range(1, len(norms) + 1), norms, marker="o", markersize=3, label=name)
plt.xlabel("layer")
plt.ylabel("weight gradient size (log scale)")
plt.legend()

Read the printed numbers the same way as before. For example, 1.0e-05 means 0.00001.

With sigmoid, the first layer's gradient is about $10^{17}$ times smaller than the last layer's. That's a 1 followed by 17 zeros. The early layers barely learn at all. Tanh, and ReLU with PyTorch's default starting weights, do better. They still lose several factors of 10 on the way down. ReLU with He initialization keeps every layer's gradient on about the same scale. (The torch here is the course's in-browser PyTorch subset.)

Five fixes, one goal

Every fix in this table works on the same product of factors. Each one keeps the factors close to 1.

A few of these need a word of explanation:

• An active unit is a ReLU neuron whose input is above zero. Its output equals its input, so its slope is 1. It passes the gradient back unchanged.
• A normalization layer works like a volume knob that resets the level at every layer. It has its own lesson.
• A residual connection adds a layer's input straight onto its output. The gradient then has two routes back: through the layer, and around it. The route around it has a factor of exactly 1. So even if the layer's own part is small, say 0.1, the total factor is 1.1, close to 1. Without the shortcut it would be 0.1. (That's a simplified picture. The next lesson covers it properly.)

Modern large language models use all five at once. They have GELU or SwiGLU activations (close relatives of ReLU), scaled initialization, and a normalization layer in every block. They also have residual connections around every block, and they clip gradients during training.

Exploding gradients and clipping

Exploding gradients show up as sudden spikes in the loss. They can also show up as nan, short for "not a number". That's what you get when a calculation breaks down, often after numbers grow too large to store.

Fixing the network's design helps. The standard safety net on top is gradient clipping. Think of a speed limiter on a car. You can still steer anywhere, but you can't go faster than the limit.

Here's how it works. Measure the combined size (the norm) of all the gradients together. If it's bigger than a threshold you choose, called max_norm, scale every gradient down by the same amount. Afterward, the combined size equals the threshold. The direction stays the same. Only the length of the step is capped.

The combined size uses the norm from earlier. Write the gradient numbers as $g_1$, $g_2$, and so on. The last one is $g_k$, where $k$ is how many numbers there are.

In plain words: square every gradient number, add them all up, and take the square root. If that total is above max_norm, multiply every gradient by max_norm ÷ total.

It's the same rule as measuring a walk on a map. Go 6 blocks east and 8 blocks north, and you end up 10 blocks away in a straight line. The worked example below clips exactly that gradient.

In PyTorch, clipping is one line, between the backward pass and the optimizer step:

loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)   # typical LLM setting
optimizer.step()

import math

grads = [6.0, 8.0]
max_norm = 2.0
total = math.sqrt(sum(g * g for g in grads))
print("total:", total)
if total > max_norm:
    grads = [g * max_norm / total for g in grads]
print("clipped:", grads)
print("new total:", math.sqrt(sum(g * g for g in grads)))

grads = [9.0, 12.0]
max_norm = 5.0
total = (9.0**2 + 12.0**2) ** 0.5
print(total, [g * max_norm / total for g in grads])

Recurrent networks: the same problem over time

A recurrent network reads a sequence one step at a time, such as one word at a time. It applies the same layer at every step. So a sequence of 100 words acts like a 100-layer network, where every layer shares the same weights.

The gradient from the end of a long sequence has to pass back through all those steps. It vanishes before it reaches the beginning. That's why plain recurrent networks (RNNs) struggled to remember anything far back. It's also why gated versions (LSTMs) were invented, and later attention. The recurrent networks lesson covers them.