10 · Weight initialisation: Xavier & He
How you initialise weights determines whether gradients vanish, explode, or train cleanly — Xavier for sigmoid/tanh, He for ReLU.
How you initialise weights determines whether gradients vanish, explode, or train cleanly — Xavier for sigmoid/tanh, He for ReLU.
Without this:
Random init from N(0, 1) explodes a 10-layer net within one batch.
Weight initialisation is one of those topics that looks like a footnote until you watch your 10-layer net produce NaN on the very first forward pass. The problem is deceptively simple: if initial weights are too large, activations blow up exponentially layer by layer. If they're too small, signals shrink to zero and gradients vanish before they reach the early layers.
The key insight is variance preservation: we want the variance of the activation signal to remain roughly constant as it flows through each layer. If layer l has n_in incoming connections and we draw weights from a distribution with standard deviation σ, the output variance scales as n_in × σ². To keep output variance ≈ input variance, we need σ = sqrt(1/n_in).
That reasoning leads directly to the two schemes every practitioner uses:
- Xavier / Glorot (2010) — designed for sigmoid and tanh:
W ~ N(0, sqrt(1/n_in)) - He / Kaiming (2015) — designed for ReLU:
W ~ N(0, sqrt(2/n_in))
The factor of 2 in He init exists because ReLU zeroes out roughly half of all activations (the negative half), which halves the effective variance. Doubling the initialisation variance compensates exactly.
Python (in browser)
Python runs entirely in your browser via Pyodide (~6 MB on first Run, cached after).
Python (in browser)
Python runs entirely in your browser via Pyodide (~6 MB on first Run, cached after).
Python (in browser)
Python runs entirely in your browser via Pyodide (~6 MB on first Run, cached after).
Deriving Xavier vs He: why the factor of 2?
PyTorch read-along: kaiming_normal_ (He) and xavier_uniform_ (Xavier)
Why use He init for ReLU layers instead of Xavier?
- Bad init (N(0,1)) explodes activations within a few layers; the fix is to scale weights by sqrt(1/n_in) or sqrt(2/n_in).
- Xavier (Glorot): sigma = sqrt(1/n_in) — designed for sigmoid and tanh, preserves variance through symmetric activations.
- He (Kaiming): sigma = sqrt(2/n_in) — designed for ReLU; the factor of 2 compensates for ReLU zeroing ~50% of activations.
- Never initialise all weights to zero — symmetry never breaks, all neurons in a layer learn identically.
Every deep net you train uses Kaiming/He init for hidden layers and Xavier for output layers — `nn.Linear` and `nn.Conv2d` default to these in PyTorch.
If you remove it: Use N(0,1) random init and your 10-layer net either explodes to NaN on step 1 (large weights) or flatlines with zero gradient (small weights) — both untrainable.