When Gradients Misbehave

In Module 2, Lesson 2.3 §9 you computed one gradient norm after loss.backward(). That number is useful as a quick check that some gradient exists, but it combines information from all parameters.

That combination can hide the problem we care about. Imagine the first layer has gradient norm 1e-10 while the last layer has gradient norm 1.0. The total can still look perfectly ordinary because the large gradient from the last layer dominates the summary.

So we need a more informative measurement: the gradient norm for each layer separately.

The important idea is simple:

• loss.backward() computes how much the loss would change if each parameter changed.
• p.grad stores that gradient for parameter p.
• The norm, such as p.grad.norm(), gives one number describing the overall size of that parameter's gradient.
• Looking at those numbers from the first layer to the last lets us see whether the learning signal is reaching every part of the network.

We are not measuring whether a layer is "good" in isolation. We are measuring whether the backward signal reaching different depths is on a comparable scale.

On a shallow, healthy network, the gradient norms of the layers are usually in roughly the same order of magnitude. That does not mean they must be identical. It means there is no enormous systematic gap where one part of the network receives a useful learning signal and another receives almost none.

Think of this function as an instrument rather than as part of the model. We will use it throughout the notebook to answer a specific diagnostic question:

Where in the network is the learning signal becoming too small or too large?

Lesson 3.4 will eventually connect this probe to the config-driven runner from Module 2 so that the same measurement can be recorded automatically during experiments.

Before training begins, every trainable weight needs an initial numerical value. Initialisation is the rule used to choose those starting values.

It is tempting to regard this as a minor implementation detail because the optimizer will change the weights anyway. But the starting values determine the scale of the first activations and gradients, and they can also determine whether different hidden units begin with different behaviour.

We will compare three schemes while keeping the architecture and training procedure fixed:

1. sensible framework-style initialisation;
2. all-zero weights;
3. deliberately enormous weights.

The point is not to memorise one magic initialisation formula. The point is to understand what a broken initialisation looks like and why it breaks.

Two different failures appear here, and they have different mechanisms.

If every weight in a layer starts at exactly the same value, the hidden units start in exactly the same state. They therefore produce the same activations, receive the same gradients, and make the same updates. Nothing distinguishes one unit from another, so they remain identical.

This is called symmetry. A 64-unit layer may contain 64 parameter sets, but if all those units behave identically, the network is not getting 64 independent pieces of representation.

Random initialisation is therefore not mainly about adding randomness for its own sake. Its important job is to break symmetry, giving different units slightly different starting points so they can specialise.

With a standard deviation of 3.0, the initial weights are deliberately far too large. That produces very large pre-activations and a huge starting loss. For ReLU networks, many units can also be pushed into the negative region, where ReLU outputs zero and has zero derivative.

The network is no longer failing because its units are identical. It is failing because the numerical scale of the computation is badly behaved.

Kaiming-style initialisation chooses the scale of weights using the number of inputs to a layer. The aim is to keep the variance of activations from changing wildly as the signal moves forward through the network.

PyTorch's standard Linear initialisation is already designed to give sensible starting scales, so you usually do not need to hand-design initialisation for ordinary layers.

What you do need to recognise are the signatures of a broken start:

• a loss that begins at an absurdly large value;
• a metric that remains at chance;
• gradients that are essentially zero;
• or activations that immediately become numerically extreme.

Now we reach a failure that becomes much more important as a network gets deeper.

Backpropagation uses the chain rule. To find how the loss changes with respect to an early layer, the gradient is passed backward through every operation between that layer and the loss. In simplified form, those local derivatives are multiplied together.

That creates a simple numerical effect:

• if many factors are smaller than 1, their product becomes very small;
• if many factors are larger than 1, their product can become very large.

For example, multiplying by 0.3 twelve times gives roughly 5e-7. The exact network is more complicated than this toy calculation, but the principle is the same.

The sigmoid activation makes the first case especially visible. Its derivative is never larger than 0.25, and it becomes much smaller when the sigmoid is saturated near 0 or 1.

Let's measure what happens as the gradient travels through a deep network.

The numbers are stark. In the 12-layer sigmoid network, the gradient reaching the first layer is roughly nine orders of magnitude smaller than the gradient reaching the last layer.

What does that mean in practice?

The layers near the output are receiving a usable learning signal. The layers near the input are receiving such a tiny signal that their weights may barely change at all. The network can therefore reduce the loss while a large portion of its depth contributes very little learning.

This is an important diagnostic lesson: a falling loss does not prove that every layer is training properly.

ReLU improves the situation because its derivative is 1 whenever its input is positive. In that active region, ReLU itself does not shrink the gradient. That is one major reason ReLU became a standard hidden activation instead of sigmoid.

But ReLU alone does not guarantee healthy gradient flow. In a deep network, the product of the linear-layer derivatives can still create a large difference between early and late layers. Sections 6 and 7 introduce two structural tools that address this more directly: normalisation and residual connections.

The same chain-rule issue also helps explain why very deep recurrent networks historically had difficulty learning long-range dependencies. An RNN unrolled over many time steps creates a long sequence of transformations through which the gradient must travel. LSTMs introduced gating mechanisms that provided additional ways for information and gradients to persist. Attention takes a different route by allowing positions to interact without requiring every dependency to pass through one long sequential chain.

The same chain-rule mechanism can fail in the opposite direction.

If the factors being multiplied through the backward pass are repeatedly greater than 1, the product grows rapidly with depth. The result is an exploding gradient: an unusually large gradient reaches some parameters, producing an unusually large parameter update.

A common symptom is a training run that appears normal and then suddenly:

• the loss jumps to an enormous value;
• the loss becomes NaN or Inf;
• activations or weights become numerically extreme;
• or the next few updates are completely unusable.

The important point is that one oversized update can be enough to destabilise the rest of training.

Gradient clipping is a direct way to limit that update. The next experiment shows what it does and, just as importantly, what it does not do.

Gradient clipping rescales the gradient when its norm is larger than a chosen threshold.

Suppose the gradient vector has norm 10 and the threshold is 1. Clipping scales the whole vector down so that its norm becomes 1. The direction stays the same; only the magnitude changes.

That gives clipping a very specific role: it prevents a single update from being catastrophically large. It does not discover why the gradient became large in the first place.

The placement is important enough to memorise:

backward() → clip gradients → step()

Clipping before backward() would not clip the gradients produced by the current loss. Clipping after step() would be too late because the oversized parameter update has already happened.

A threshold such as 1.0 is a common starting point, but there is no universal threshold that is correct for every model. The right lesson is not "always use 1.0." It is:

Use clipping as protection against unusually large updates, then investigate the cause if clipping is happening constantly.

If gradients repeatedly hit the clipping threshold, possible causes include poor initialisation, an excessive learning rate, or badly scaled inputs. Continually lowering the threshold can hide the symptom without fixing the underlying problem.

ReLU is defined as max(0, x).

That simple definition gives it a useful property: for positive inputs, its derivative is 1, so active units do not shrink the gradient. But it also creates a hard failure mode.

For negative inputs, ReLU outputs exactly 0, and its derivative is 0. If a hidden unit ends up producing negative pre-activations for every training example, that unit receives no useful gradient from those examples. It can become permanently inactive — a dead ReLU.

This is different from ordinary overfitting or poor performance. The problem is local: a particular unit has effectively stopped participating in learning.

Lesson 3.1 gave us a suspected cause: the very large learning rate that caused the model to collapse to a constant prediction. Here we will measure whether dead units are the mechanism behind that collapse.

There is the causal chain that Lesson 3.1 could only suggest.

A learning rate that is too large can move the first-layer weights so far in one update that many ReLU units are pushed permanently into their zero-output region. Once a unit is dead, its gradient is zero, so later updates cannot bring it back. If enough units die, the layer loses most of its useful representation capacity and the whole model can collapse toward a constant prediction — hence AUC near 0.5.

But do not turn this into the rule "any dead ReLU is bad." A modest fraction of inactive units can be completely harmless. Neural networks are often over-parameterised, so the surviving units can carry the useful computation.

What matters is the fraction of a layer that has become inactive and whether that loss of capacity is associated with degraded model behaviour.

This failure is permanent for the current training run. Lowering the learning rate after the units have died does not give them a gradient to recover with. In practice, you normally restart training after fixing the cause.

Useful responses include:

• Lower the learning rate when large updates are the cause.
• Use LeakyReLU or GELU when you want a non-zero gradient on the negative side.
• Check input scaling. Features with very different or very large scales can create extreme pre-activations and gradients. The StandardScaler in this pipeline is therefore part of the numerical-stability story, not merely cosmetic preprocessing.

The measurement is simple enough to keep in your diagnostic toolkit: ask what fraction of units in a layer never activate on the data being inspected.

The previous sections showed two opposite gradient problems: gradients can become extremely small or extremely large as they pass through depth.

A useful way to understand both is to focus on scale. If the activations entering each layer become progressively too large or too small, the local derivatives in the chain rule can also become badly scaled. That makes the product through many layers unstable.

Normalisation layers try to control that scale by rescaling activations during the forward pass.

The two normalisation schemes we need to distinguish are based on different sets of values:

• BatchNorm normalises each feature using statistics computed across the examples in the current batch. For a particular feature, it asks: "How are the values for this feature distributed across these rows?"
• LayerNorm normalises each individual example across its features. For one row, it asks: "How are this example's feature activations distributed?"

So the difference is not simply two names for the same operation. They use different axes and therefore have different behaviour with respect to batch size and evaluation mode.

Let's measure how these choices affect gradient flow.

This experiment shows the structural effect clearly. In the plain 12-layer ReLU network, the gradient reaching the last layer can be thousands of times larger than the gradient reaching the first. Adding LayerNorm reduces that gap to single digits, while BatchNorm brings it close to one in this particular experiment.

The important lesson is not that BatchNorm will always beat LayerNorm numerically. The lesson is that normalisation can keep the scale of intermediate activations and gradients under control, making deep stacks easier to train.

BatchNorm computes statistics across examples in a batch. That creates a dependency between rows: during training, the representation of one row is partly determined by the other rows that happen to be in the same batch.

That can work very well, particularly in convolutional vision models, but it becomes less convenient with very small or variable batches.

It also explains why model.train() and model.eval() matter. During training, BatchNorm uses batch statistics and updates running statistics. During evaluation, it uses the stored running statistics instead of recomputing them from the evaluation batch.

LayerNorm normalises each example independently across its features. There is therefore no dependence on the other rows in the batch.

That property is particularly useful for sequence models and Transformers, where batch composition should not determine how an individual sequence is normalised. LayerNorm is consequently a standard component of Transformer architectures.

You do not need to derive the backward pass here. What you should be able to explain is what is being normalised, over which dimensions, and why that choice matters.

Normalisation changes the scale of the signal. A residual connection changes the route the signal takes.

A plain block computes:

y = f(x)

A residual block computes:

y = x + f(x)

The original input x is therefore passed directly to the output as well as through the learned transformation f.

Why does this help gradients?

Differentiate the expression with respect to x:

dy/dx = 1 + f'(x)

The important part is the constant 1. Even if the derivative through f becomes very small, the gradient still has a direct path through the x term.

This is the core idea. You do not need a complicated derivation to remember it:

A residual connection gives the gradient a route that bypasses the transformation.

That route becomes extremely valuable when many blocks are stacked.

The experiment shows nearly three orders of magnitude of improvement — roughly a factor of 500 — simply by changing the block from f(x) to x + f(x).

The important mechanism is the direct path. In a plain deep network, the gradient must pass through every transformation. In a residual network, part of the gradient can travel along the identity path instead.

This is one reason very deep residual networks can be trained reliably. The problem with adding depth is not only that a larger model can overfit; before residual connections became standard, depth itself could make optimisation difficult because early layers stopped receiving a useful gradient.

There is another useful way to understand residual learning. A residual block does not have to learn the entire desired transformation. It only has to learn the difference from the identity. If leaving the representation unchanged is already close to the right thing, the block can make f(x) close to zero and let the skip path carry the representation through unchanged.

Real residual architectures sometimes change the width of the representation. In that case, x and f(x) cannot be added directly because their shapes differ. A common solution is a learned projection — often a linear layer — on the shortcut path so that the two tensors have matching shapes.

The experiments above are different demonstrations of one underlying idea: a training curve tells you what happened overall; internal diagnostics help explain why.

When an unfamiliar model is behaving strangely, you do not want to immediately change five hyperparameters and hope one helps. Start by measuring the internal state of the computation.

The following checklist packages the measurements from this notebook into a quick first pass.

1. Check the total gradient norm. The absolute value depends on the model and loss, so treat the ranges here as rough warning signs, not laws. A norm above about 1e3 can indicate explosion; a value below about 1e-6 can indicate that the learning signal is effectively disappearing.
2. Check the per-layer gradient profile. Look for a large systematic gap between early and late layers. A last/first ratio of more than roughly three orders of magnitude is a strong reason to investigate.
3. Check the fraction of dead ReLU units. A small inactive fraction can be normal. If roughly half or more of a layer is permanently inactive, most of that layer's capacity has been lost.
4. Check explicitly for NaN and Inf. These values can propagate through subsequent calculations and turn a useful numerical computation into nonsense.
5. Check activation scale. If pre-activation standard deviations are extremely large or extremely small, investigate initialisation, input scaling, learning rate, or normalisation.

These checks are diagnostic signals, not universal pass/fail thresholds. The exact healthy range depends on architecture, parameterisation, batch size, loss, and data. Their value is that they turn a vague statement such as "the model is not learning" into a more specific question such as "the early layers are receiving gradients 10,000 times smaller than the final layer."

A summed gradient norm hides structure. Per-layer measurements tell you where the learning signal is being lost or amplified.

Initialisation affects the starting numerical regime. All-zero weights preserve symmetry and prevent hidden units from specialising; excessively large weights can create extreme activations and gradients. Standard PyTorch layers already use sensible defaults, so the skill is recognising broken behaviour rather than memorising a custom initialisation recipe.

Vanishing and exploding gradients have the same mathematical family. Both arise from repeated multiplication through the chain rule. When the factors are effectively below one, gradients shrink; when they are effectively above one, they grow.

Dead ReLU units are a local, potentially permanent failure. A unit that receives zero gradient because it is always in ReLU's inactive region cannot repair itself during that run. Excessive learning rates are one important cause.

Normalisation and residual connections address gradient flow structurally. Normalisation controls activation scale; residual connections provide a shortcut path for information and gradients. These are not arbitrary tricks — each changes a specific part of the computation that was causing the problem.

You now have three kinds of diagnostic evidence:

• Lesson 3.1: learning-rate and optimisation signatures
• Lesson 3.2: generalisation and overfitting signatures
• Lesson 3.3: gradient-flow and internal-layer signatures

The remaining problem is procedural. When a new experiment fails, which measurement should you make first? And how do you avoid changing settings until something happens and then mistaking the change for a diagnosis?

Lesson 3.4 turns these measurements into a diagnostic protocol ordered by cost and evidence, then gives you six deliberately broken experiments to diagnose — including a case where the usual curves look healthy even though the model is still not useful.