Training Stability Techniques

By the end of this lesson you'll be able to spot an unstable training run, clip gradients to a safe norm, warm up and schedule the learning rate, and make every run reproducible — so deep learning actually converges instead of blowing up.

Learn Training Stability Techniques in our free AI & Machine Learning course — a beginner-friendly interactive lesson with worked examples, a practice…

Part of the free AI & Machine Learning course at LearnCodingFast — hands-on lessons with examples you run in your browser, plus practice exercises and a quick quiz.

Think of a deep network as a very tall tower of building blocks, one block per layer. Training nudges the blocks to make the tower taller and straighter. The taller it gets, the easier it is to topple — and the same techniques a builder uses map exactly onto stable training:

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:

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:

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:

Fill in the two blanks so the clipper scales [6, 8] (norm 10) down to norm 5.

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:

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).

Detecting trouble early is half the battle. Fill in the blanks so the loop reports the first NaN or Inf in the loss list.

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):

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

Support is faded now — only a comment outline is given. Write the whole function yourself.

❌ No gradient clipping → loss explodes to NaN

An unclipped gradient spikes, the weights jump, and the loss goes inf then NaN .

✅ Fix: add nn.utils.clip_grad_norm_(model.parameters(), 1.0) after backward() and before optimizer.step() .

❌ Learning rate too high → loss spikes or diverges

The loss oscillates or climbs instead of falling — the steps are too big.

✅ Fix: lower the LR (try 10x smaller) and add warmup so early steps stay small.

❌ Bad initialisation → dead neurons or instant explosion

randn * 1.0 explodes through deep layers; randn * 0.01 makes activations vanish and ReLUs go dead.

✅ Fix: use He init for ReLU ( kaiming_normal_ ), Xavier for tanh/sigmoid.

You can't reproduce a result or tell whether a tweak actually helped.

✅ Fix: call seed_everything(42) once at the very top, before building the model.

Once a NaN enters the weights it poisons all subsequent updates, but training keeps "running".

✅ Fix: check each step with loss != loss or math.isnan(loss) and stop or skip immediately.

You can now keep a tall network from toppling: you read a loss curve for spikes and NaN, clip gradients to a safe norm (by hand and with clip_grad_norm_ ), warm up and schedule the learning rate, pick init/normalisation that hold signals steady, and seed every RNG for reproducible runs.

🚀 Up next: Generative Models — now that training is stable, build models that create entirely new data.

Practice quiz

• Browse all free courses
• Cheat sheets & quizzes
• Coding dictionary