Module 4 — Vanishing and exploding gradients

Module 3 wrote the BPTT chain rule as a product of Jacobians. Everything unpleasant about training simple RNNs comes from that product. This module makes the failure concrete on a 50-step example, then covers the two mitigations that keep raw RNNs alive — gradient clipping and orthogonal initialisation — before setting up the case for LSTM and GRU in the next two modules.

Why a product of matrices is the problem

Inside BPTT, the gradient of the loss with respect to a hidden state $k$ steps in the past is:

$$\frac{\partial \mathcal{L}}{\partial h_{t-k}} = \frac{\partial \mathcal{L}}{\partial h_t} \cdot \prod_{i=t-k+1}^{t} \frac{\partial h_i}{\partial h_{i-1}}$$

Each Jacobian $\partial h_i / \partial h_{i-1}$ is roughly $W_h^\top \operatorname{diag}(\tanh'(\cdot))$. Multiply $k$ of them together and the norm of the product behaves like the $k$-th power of the largest singular value $\sigma_{\max}$ of $W_h$:

$$\left\| \prod \frac{\partial h_i}{\partial h_{i-1}} \right\| \approx \sigma_{\max}^{\,k}$$

Three regimes appear immediately:

• $\sigma_{\max} < 1$ → the gradient shrinks exponentially with $k$: vanishing gradient. The network cannot learn dependencies more than a few dozen steps back.
• $\sigma_{\max} > 1$ → the gradient grows exponentially: exploding gradient. A single batch can produce NaN losses and destroy weeks of training.
• $\sigma_{\max} \approx 1$ → the gradient stays roughly stable, and the RNN can learn moderately long dependencies. This is the razor's edge that orthogonal initialisation tries to sit on.

Add the $\tanh'$ terms, which are always in $(0, 1]$ and often much smaller when the state saturates, and even a well-initialised network drifts toward vanishing rather than explosion.

A 50-step numerical demonstration

The theory is easier to trust after seeing the numbers. The following snippet builds a fresh recurrent weight matrix with a chosen spectral radius, then applies it 50 times to a random vector and prints the norm at every step.

import numpy as np

def simulate(spectral_radius, steps=50, dim=32, seed=0):
    rng = np.random.default_rng(seed)
    W = rng.standard_normal((dim, dim))
    # Rescale to a chosen spectral radius (largest absolute eigenvalue)
    W = W / max(abs(np.linalg.eigvals(W))) * spectral_radius
    h = rng.standard_normal(dim)
    norms = []
    for _ in range(steps):
        h = np.tanh(W @ h)
        norms.append(np.linalg.norm(h))
    return norms

for r in (0.5, 0.9, 1.0, 1.1, 1.5):
    n = simulate(r)
    print(f"radius={r}  step 1={n[0]:.3f}  step 25={n[24]:.3e}  step 50={n[-1]:.3e}")

Typical output on a laptop:

radius=0.5  step 1=1.980   step 25=2.510e-08   step 50=6.301e-16
radius=0.9  step 1=3.271   step 25=2.114e-01   step 50=1.033e-02
radius=1.0  step 1=3.487   step 25=3.108e+00   step 50=3.081e+00
radius=1.1  step 1=3.605   step 25=5.128e+00   step 50=5.132e+00
radius=1.5  step 1=4.078   step 25=5.632e+00   step 50=5.630e+00

The saturation at high radii comes from $\tanh$; the exponential decay at low radii comes straight from the eigenvalues. The band around 1 is narrow — miss it by 10 %, and 50 steps later the signal is nine orders of magnitude too small.

Gradient clipping: the cheapest fix for explosion

Exploding gradients are easier to defend against than vanishing ones because they announce themselves as NaN losses or huge parameter jumps. Gradient clipping rescales the gradient whenever its global norm exceeds a threshold:

from tensorflow.keras.optimizers import Adam
opt = Adam(learning_rate=1e-3, clipnorm=1.0)

Or, more explicitly with a hand-written loop:

import torch
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

A threshold around 1.0 or 5.0 is standard. Clipping never helps with vanishing gradients — the norm is already small — but it prevents a rare bad batch from destroying the run. Every serious recurrent training pipeline has it.

A loss that oscillates and occasionally becomes NaN is almost always an exploding gradient

Reaching for a smaller learning rate is not the right first move: it slows learning without curing the cause. Add clipnorm=1.0 first, watch the training resume, then tune the learning rate normally.

Orthogonal initialisation: sit on the razor's edge

An orthogonal matrix has all its singular values equal to 1 by construction. Initialising $W_h$ from an orthogonal distribution therefore puts the network exactly at $\sigma_{\max} = 1$ at step zero, which gives training a fighting chance to keep it there.

from tensorflow.keras import layers

layer = layers.SimpleRNN(
    64,
    kernel_initializer="glorot_uniform",      # input-to-hidden
    recurrent_initializer="orthogonal",       # hidden-to-hidden
    activation="tanh",
)

Keras uses orthogonal as the default recurrent initialiser for SimpleRNN, LSTM and GRU. PyTorch does not: nn.RNN initialises with a uniform distribution scaled by the hidden size, which is a decent choice but not orthogonal. On PyTorch this is worth setting by hand:

import torch.nn as nn
rnn = nn.RNN(input_size=1, hidden_size=64, nonlinearity="tanh")
for name, p in rnn.named_parameters():
    if "weight_hh" in name:
        nn.init.orthogonal_(p)

Why LSTM and GRU exist

Neither trick removes the fundamental problem. Clipping bounds the norm from above; orthogonal initialisation bounds it near 1 at step zero, but training pushes the weights away from orthogonality. On sequences longer than roughly 100 to 200 steps, a SimpleRNN reliably fails to learn dependencies that span the whole window.

The solution, introduced by Hochreiter and Schmidhuber in 1997, is architectural: replace the multiplicative product $\prod \partial h_i / \partial h_{i-1}$ with a mostly additive path along a separate cell state. That path has a Jacobian close to the identity, so the product stays close to one over hundreds of steps. This is the LSTM, and its simpler cousin the GRU, and they are the subject of modules 5 and 6.

A quick diagnostic: are my gradients dead?

Before adopting a more complex architecture, always check that the simple RNN really is the bottleneck.

import tensorflow as tf

@tf.function
def grad_norms(model, x, y):
    with tf.GradientTape() as tape:
        loss = tf.reduce_mean((model(x) - y) ** 2)
    grads = tape.gradient(loss, model.trainable_variables)
    return [tf.norm(g).numpy() for g in grads]

print(grad_norms(model, x_train[:32], y_train[:32]))

A recurrent kernel whose gradient norm is $10^{-8}$ while the dense output kernel is $10^{-1}$ is a textbook vanishing gradient: only the last dense layer is learning. The fix is not more units or more epochs — it is an LSTM or a GRU.

Sanity-check the spectral radius during training

np.max(np.abs(np.linalg.eigvals(model.layers[0].get_weights()[1]))) returns the spectral radius of the recurrent kernel. If it drifts below 0.5 or above 1.5 during training, the model is either about to fall silent or about to explode, and the next problem is architectural, not numerical.

In summary

• The BPTT gradient is a product of Jacobians whose norm grows or decays like $\sigma_{\max}^{\,k}$; simple RNNs live on a razor's edge.
• Gradient clipping prevents explosion by rescaling any gradient whose global norm exceeds a threshold, typically 1.0; it does nothing for vanishing.
• Orthogonal initialisation starts the recurrent kernel with singular values equal to 1 and buys the network the best possible chance to preserve the gradient early in training.
• Neither trick fixes the underlying product-of-matrices problem beyond about 100 to 200 steps, which is exactly what LSTM and GRU solve architecturally.

Next module: the LSTM cell, its three gates, and why the cell state gives the gradient a mostly additive highway through time.