Why does weight initialization matter in neural networks, and what is the difference between Xavier and He initialization?

If weights are initialized too small, activations and gradients shrink layer by layer — a form of vanishing gradient from the start. If too large, they explode. The goal of principled initialisation is to keep the variance of activations and gradients roughly constant across all layers at the start of training.

Xavier (Glorot) initialisation draws weights from a distribution with variance 2/(fan_in + fan_out). It was derived assuming linear activations (or tanh in the original paper) by requiring that the variance of the layer's output equals the variance of its input. He (Kaiming) initialisation uses variance 2/fan_in, derived for ReLU activations specifically — since ReLU zeroes out half the input on average, the variance of the output is halved, so doubling the initial weight variance compensates for this. Using Xavier with ReLU causes variance to shrink by roughly half per layer, eventually vanishing.

import torch
import torch.nn as nn

# Default PyTorch Linear layers use Kaiming Uniform initialisation
layer = nn.Linear(256, 128)
print(layer.weight.std())  # approximately sqrt(2/256) ≈ 0.088

# Explicit initialisation
def init_weights(m):
    if isinstance(m, nn.Linear):
        # Xavier: good for sigmoid/tanh activations
        nn.init.xavier_uniform_(m.weight)
        # He/Kaiming: good for ReLU activations (default in PyTorch)
        # nn.init.kaiming_uniform_(m.weight, nonlinearity='relu')
        nn.init.zeros_(m.bias)

model = nn.Sequential(
    nn.Linear(784, 256), nn.ReLU(),
    nn.Linear(256, 128), nn.ReLU(),
    nn.Linear(128, 10)
)
model.apply(init_weights)  # apply init_weights to every sub-module

# Verifying activation variance stays stable across layers:
x = torch.randn(100, 784)
for layer in model:
    x = layer(x)
    print(f'{layer.__class__.__name__}: std={x.std():.3f}')
# With He init + ReLU: std should remain near 1.0 throughout