Training Stability Techniques

During backpropagation the gradient is multiplied through every layer. Exploding gradients happen when those factors are larger than 1 on average — the product grows huge, the weights jump, and the loss blows up to inf then NaN. Vanishing gradients are the opposite: factors below 1 shrink the product toward zero, so the early layers barely update and the network stops learning.

You spot exploding gradients by watching the gradient norm climb and the loss spike; you spot vanishing gradients when the loss plateaus early and the first layers' weights hardly move. The cures in this lesson — init, normalisation, clipping, and warmup — all work by keeping that running product near 1.

The loss numbers tell you almost everything. A sudden jump (say 3x the previous step) is a loss spike — often a too-high learning rate or an unclipped gradient. A NaN ("not a number") means a value became undefined and the run is dead from that point on. The neat trick: NaN is the only value that is not equal to itself, so loss != loss is a one-line NaN test (no import needed).

Run this and watch it flag the spike at step 4 and the NaN at step 7:

# Worked example: read a loss curve and spot the trouble
# A training run that goes wrong usually leaves clues in the loss numbers.

losses = [2.40, 2.10, 1.85, 1.60, 9.80, 1.55, 1.50, float("nan"), 1.48]
#                                  ^^^^ a loss SPIKE         ^^^ then NaN

prev = losses[0]
for step, loss in enumerate(losses):
    note = ""
    if loss != loss:                 # NaN is the only value not equal to itself
        note = "<-- NaN! training is broken from here"
    elif loss > 
...

The cleanest fix for exploding gradients is to cap the norm (the length) of the gradient. If the gradient's L2 norm exceeds a threshold, you scale the whole vector down by one factor — max_norm / norm — so its direction stays the same and only its size is capped. That last point is why clipping by norm beats clipping each value independently: distorting direction sends the optimiser somewhere it never intended to go.

This worked example clips the vector [3, 4, 12] (norm 13) down to norm 5:

# Worked example: clip a gradient vector to a maximum norm
# Exploding gradients = the update vector gets huge and blows up the weights.
# The fix: if its length (L2 norm) is too big, scale the WHOLE vector down.

def l2_norm(vec):
    """Length of the vector = sqrt(sum of squares)."""
    return sum(v * v for v in vec) ** 0.5

def clip_grad_norm(grad, max_norm):
    """Scale grad so its norm is at most max_norm. Direction is preserved."""
    norm = l2_norm(grad)
    if norm > max_norm:
       
...

In real code you never hand-roll this — PyTorch's clip_grad_norm_ does it across every parameter at once and returns the pre-clip norm so you can log it:

import torch
import torch.nn as nn

# The real thing: PyTorch ships clip_grad_norm_ so you never hand-roll it.
# It clips across ALL the model's parameters at once, in place.

torch.manual_seed(0)                     # reproducibility: same run every time
model = nn.Linear(4, 1)
x = torch.randn(8, 4)
y = torch.randn(8, 1)

loss = ((model(x) - y) ** 2).mean()
loss.backward()                          # fills .grad on every parameter

before = torch.nn.utils.clip_grad_norm_(model.parameters(), max_
...

A fresh model has random weights, so the first gradients are loud. Warmup ramps the learning rate linearly from near zero up to your target over the first few steps, letting Adam's running statistics settle before you take big steps. After warmup you usually decay the rate — cosine or step schedules — for a cleaner final model.

Mixed precision (AMP) speeds training by computing in lower precision, but tiny gradients can underflow to zero. GradScaler multiplies the loss up before backward() and unscales afterwards; it also skips the step automatically when it detects inf/NaN gradients. Always unscale_ before clipping so the clip sees the true norms:

import torch
import torch.nn as nn
from torch.cuda.amp import autocast, GradScaler

# Production stability stack: warmup schedule + AMP loss scaling + clipping.
torch.manual_seed(0)

model = nn.Linear(4, 1)
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler()                    # AMP: scales loss so small grads don't underflow to 0
warmup_steps = 5
base_lr = 1e-3

def lr_at(step):
    # Linear warmup: ramp 0 -> base_lr over warmup_steps, then hold.
    if step < warmup_steps
...

Initialisation sets the starting scale of every weight so the signal neither shrinks nor blows up as it passes through layers. Use He (Kaiming) init for ReLU networks and Xavier (Glorot) for tanh/sigmoid. Never use plain randn * 1.0 or randn * 0.01 for a deep network — the first explodes, the second vanishes.

Normalisation re-centres each layer's activations during training so the next layer always sees a well-behaved distribution. Batch Norm normalises across the batch (great for CNNs with large batches); Layer Norm normalises across features (the default for Transformers and RNNs, and the safe choice for small batches).

nn.init.kaiming_normal_(w)   # He
nn.BatchNorm2d(channels)     # Batch Norm

nn.init.xavier_uniform_(w)   # Xavier
nn.LayerNorm(features)       # Layer Norm

If two runs give different results you can't tell whether a change helped or you just got lucky. Seed every source of randomness before you build the model — Python's random, NumPy, and PyTorch (CPU and GPU):

import random, numpy as np, torch

def seed_everything(seed=42):
    random.seed(seed)                 # Python's RNG
    np.random.seed(seed)              # NumPy
    torch.manual_seed(seed)           # PyTorch CPU
    torch.cuda.manual_seed_all(seed)  # PyTorch GPU
    # For bitwise-identical runs (slower):
    torch.use_deterministic_algorithms(True)

seed_everything(42)   # call ONCE, at the very top, before building the model

Call seed_everything once at the top of your script. Without it, weight init, shuffling, and dropout all differ between runs.