Initialization: Xavier and He

Why initialization matters

If you initialize weights too small, signal vanishes through the network and gradients vanish back. Too large, signal explodes and gradients explode. The right scale keeps activation variance roughly constant from layer to layer, which keeps the gradient chain product well-conditioned at depth.

Xavier (a.k.a. Glorot) initialization: std = sqrt(2 / (fan_in + fan_out)). Designed for tanh/sigmoid activations. He initialization: std = sqrt(2 / fan_in). Designed for ReLU (which kills half the inputs, so you need slightly bigger weights to compensate). PyTorch's default for nn.Linear is a uniform variant of He.

Tip: PyTorch's defaults are sensible for ReLU/GELU networks. The first time you write a custom layer or use a non-standard activation, override the init explicitly — defaults can silently break.

The custom-init pattern

For non-default architectures, write a single init_weights function and apply it via model.apply(init_weights). This recursively walks the module tree and calls your function on every submodule. Standard practice in transformer code.

Special cases

Output layer — sometimes initialized at smaller scale (e.g., 0.02 std for transformer LM heads). Embeddings — Normal(0, 0.02) for transformers; uniform for word embeddings. BatchNorm — γ=1, β=0 by default; some recipes use γ=0 for the last BN in a residual block to make it start as identity.

Principle: Initialization is one of the things you tune until it works, then stop touching. Use He for ReLU/GELU/SiLU networks, Xavier for tanh/sigmoid, and PyTorch defaults for the rest until you have a reason.

Code

He init for a ReLU network·python

import torch.nn as nn
import math

def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.Conv2d):
        nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, (nn.LayerNorm, nn.BatchNorm2d)):
        nn.init.ones_(m.weight)
        nn.init.zeros_(m.bias)
    elif isinstance(m, nn.Embedding):
        nn.init.normal_(m.weight, mean=0.0, std=0.02)

model = MyModel()
model.apply(init_weights)

Why initialization matters

The custom-init pattern

Special cases

Code

External links

Exercise

Progress

Comments 0