Module 6 — Weight initialization and vanishing gradients

Module 4 hinted at a problem: if a deep weight's gradient is a product of local derivatives, that product can collapse toward zero or diverge. This module addresses that problem and the set of answers that, between 2010 and 2015, made genuinely deep networks trainable.

Initialization is not a detail

Two naive choices fail, for instructive reasons.

Initializing all weights to zero makes every neuron in a layer identical: they receive the same gradient, update the same way, and remain indistinguishable forever. The whole layer behaves like a single neuron. Symmetry must be broken, and that is what randomness is for.

Initializing with values that are too large or too small produces a drift in the scale of activations from layer to layer. If the variance of the outputs is multiplied by a factor slightly above 1 at each layer, it explodes after twenty layers; slightly below, it vanishes. In both cases training is lost before it begins.

The objective is therefore precise: keep the variance of activations roughly constant across layers. Two formulas address this, depending on the activation used.

Xavier initialization, also called Glorot, suits symmetric activations such as tanh:

$$\text{Var}(W) = \frac{2}{n_{\text{in}} + n_{\text{out}}}$$

He initialization suits ReLU:

$$\text{Var}(W) = \frac{2}{n_{\text{in}}}$$

The difference comes from a simple argument: ReLU zeroes half its inputs, so it roughly halves the output variance. The factor 2 exactly compensates that loss. This is why using Xavier with ReLUs in a deep network leads to a progressive extinction of the signal.

In practice, libraries already apply the right default — kaiming_uniform_ in PyTorch for linear layers. The topic becomes visible when writing custom initialization, or when a deep network refuses to start learning.

The vanishing gradient

Recall the formula from module 4. Moving up the layers, the error is multiplied at each step by $f'(z)$ and by the weights. With a sigmoid, the derivative never exceeds 0.25. Over ten layers, the first layer's gradient is therefore attenuated by at most $0.25^{10} \approx 10^{-6}$.

Early layers no longer receive any usable signal. They stay at their initialization values while only the last layers learn. This is why, despite the universality theorem, deep networks were reputed untrainable until the 2000s: the problem was not theoretical, it was numerical.

Four answers, all either already met or forthcoming:

• ReLU, whose derivative is exactly 1 on the positive side, so attenuates nothing;
• suitable initialization, which prevents scale drift;
• normalization by batch or by layer, the subject of the next module;
• residual connections, which are the most radical answer.

Residual connections

Introduced by ResNet in 2015, they consist of adding a block's input to its output:

$$y = f(x) + x$$

The effect on the gradient is direct and worth seeing. The derivative of this sum with respect to $x$ is $f'(x) + 1$: the $+1$ term creates a direct path along which the gradient flows back without attenuation, whatever the depth. The network no longer learns the full transformation but only the deviation from identity, which is both easier and more stable.

It is this idea that allowed the move from a few dozen layers to more than a hundred, and it is present in almost every modern architecture — transformers included.

The exploding gradient

The symmetric problem exists: when the factors exceed 1, the gradient grows exponentially. The symptoms are unambiguous — a loss going NaN, or weights diverging within a few iterations. The recurrent networks of course 11 are particularly exposed.

The remedy is simple and effective: gradient clipping, which bounds its norm before the update.

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

The gradient's direction is preserved, only its magnitude is capped. It is an inexpensive precaution that is reasonable to enable by default on recurrent architectures or unstable training runs.

Diagnose before changing architecture

Facing a network that will not learn, inspect the per-layer gradient norms before anything else. Norms that decay strongly toward the input signal vanishing: check the activation, the initialization, and consider residual connections. Norms that explode call for clipping and a lower learning rate. That measurement takes a few lines of code and replaces hours of guesswork.

Summary

• Initializing to zero prevents breaking symmetry; a badly chosen scale makes activation variance drift layer after layer.
• Xavier for symmetric activations, He for ReLU, whose factor 2 compensates for the half of inputs zeroed out.
• The vanishing gradient comes from a product of derivatives below 1 — at most 0.25 for the sigmoid; ReLU, initialization, normalization and residual connections answer it.
• The $+1$ term of a residual connection creates a direct path for the gradient; conversely, the exploding gradient is treated by norm clipping.

Next module: regularization, to stop a now-trainable network from simply memorizing its data.