0Pricing
Machine Learning Academy · Lesson

Weight Initialisation: Xavier and He Initialisation

Learners will apply Xavier uniform and He normal initialisation, observe how they prevent vanishing/exploding gradients in deep networks compared to default random init.

Weight Initialisation: Xavier and He Initialisation is a free Machine Learning Academy lesson on CoddyKit — lesson 4 of 4. You can read the complete lesson below for free — then practise it hands-on in the browser with a built-in code editor and a 24/7 AI tutor. It is part of the Machine Learning Academy learning path, one of 4 lessons in the course, and your progress syncs across the web and the CoddyKit app.

Why Initialisation Matters

The weights of a neural network must be initialised to non-zero values before training — but the choice of how to initialise them profoundly affects training dynamics. Poor initialisation causes vanishing gradients (weights shrink to near zero, gradients become negligible) or exploding gradients (weights grow unboundedly, gradients become NaN). Good initialisation keeps activations and gradients in a healthy range from the very first batch, enabling stable and fast training.

import torch
import torch.nn as nn

# All-zeros init: disaster! All neurons compute the same
# gradient (symmetry breaking fails)
model_bad = nn.Linear(4, 4)
nn.init.zeros_(model_bad.weight)
print('All-zero gradients:', model_bad.weight.grad)

# Constant init: same problem
# Random init from N(0,1): works for shallow, fails deep
# Xavier / He: designed for deep networks

The Symmetry Breaking Problem

If all weights are initialised to the same value (including zero), every neuron in a layer computes exactly the same output and receives exactly the same gradient. All neurons learn the same feature — the hidden layer collapses to a single neuron for all practical purposes. This symmetry problem is why random initialisation is necessary: each neuron must start with a different random weight to break symmetry and learn different representations.

import torch
import torch.nn as nn

# Demonstrate symmetry breaking failure
model = nn.Linear(3, 4, bias=False)
nn.init.constant_(model.weight, 0.1)  # all same

x = torch.randn(5, 3)
y = model(x)

# All 4 neurons produce identical outputs!
print('All neurons identical:', torch.allclose(y[:, 0], y[:, 1]))
# True -- the 4 output neurons are indistinguishable

Naive Normal Init and Its Failure

Initialising weights from a standard normal distribution N(0, 1) seems reasonable but causes problems in deep networks. For a layer with fan-in (number of input connections) of 1000, the weighted sum of 1000 normally distributed values has variance 1000 — causing exploding activations in deep networks. Conversely, very small random values (e.g., N(0, 0.001)) cause vanishing activations. Neither extreme allows gradients to flow through many layers.

import torch
import torch.nn as nn

# Track activation variance through 10 deep layers
def test_deep_init(std):
    x = torch.randn(1, 256)
    for i in range(10):
        W = torch.randn(256, 256) * std
        x = torch.tanh(x @ W)
    return x.std().item()

print(f'std=1.0:   activation_std={test_deep_init(1.0):.6f}')
# Huge -> saturation
print(f'std=0.01:  activation_std={test_deep_init(0.01):.6f}')
# Near zero -> vanishing

Xavier Glorot Initialisation

Xavier initialisation (Glorot and Bengio, 2010) was designed for networks using tanh or sigmoid activations. The insight is to choose weights so that the variance of activations and gradients remains roughly constant across layers. Weights are drawn from a uniform or normal distribution with variance 2 / (fan_in + fan_out). This is the default initialisation for nn.Linear (uniform variant) in PyTorch.

import torch
import torch.nn as nn

layer = nn.Linear(256, 128)

# Xavier uniform: default for nn.Linear
nn.init.xavier_uniform_(layer.weight)
print('Xavier uniform std:', layer.weight.std().item())
# Approximately sqrt(2 / (256 + 128)) = 0.081

# Xavier normal: Gaussian version
nn.init.xavier_normal_(layer.weight)
print('Xavier normal std:', layer.weight.std().item())

He (Kaiming) Initialisation for ReLU

He initialisation (He et al., 2015) was designed specifically for networks using ReLU activations. Because ReLU zeros out half of its inputs (the negatives), the effective variance after activation is halved. He init compensates by using variance 2 / fan_in — twice what Xavier uses. Using Xavier with ReLU causes gradients to vanish in deep networks; using He initialisation enables training networks with 100+ layers.

import torch
import torch.nn as nn

layer = nn.Linear(512, 256)

# He (Kaiming) uniform: designed for ReLU
nn.init.kaiming_uniform_(layer.weight,
                          nonlinearity='relu')
print('Kaiming uniform std:', layer.weight.std().item())
# Approximately sqrt(2/512) * sqrt(3) = 0.108

# He (Kaiming) normal: Gaussian variant
nn.init.kaiming_normal_(layer.weight,
                         nonlinearity='relu')
print('Kaiming normal std:', layer.weight.std().item())

Comparing Init Methods on a Deep Network

The effect of initialisation becomes visible when you track activation statistics across layers of a deep network. With Xavier init and tanh, activation variance stays near 1 through all layers. With He init and ReLU, the same holds for ReLU networks. Using the wrong combination (Xavier + ReLU, or He + sigmoid) leads to systematic activation collapse or explosion, confirming that initialisation choice must match the activation function.

import torch
import torch.nn as nn

def track_activation_std(init_fn, activation, n_layers=10):
    x = torch.randn(64, 256)
    stds = []
    for _ in range(n_layers):
        W = torch.empty(256, 256)
        init_fn(W)
        x = activation(x @ W.T)
        stds.append(x.std().item())
    return stds

xavier = lambda W: nn.init.xavier_uniform_(W)
he     = lambda W: nn.init.kaiming_uniform_(W, nonlinearity='relu')

xavier_stds = track_activation_std(xavier, torch.tanh)
he_stds     = track_activation_std(he, torch.relu)

print('Xavier+tanh layer stds:', [f'{s:.2f}' for s in xavier_stds])
print('He+ReLU layer stds:', [f'{s:.2f}' for s in he_stds])

Applying Custom Init to a Full Model

You can apply a custom initialisation to an entire model using model.apply(init_fn), which recursively visits every module. The function receives each module and can apply different initialisations based on the layer type. A common pattern is to apply He init to Linear and Conv2d layers, Xavier init to embedding layers, and set biases to zero. This single call replaces PyTorch's defaults across the entire network.

import torch.nn as nn

def init_weights(module):
    if isinstance(module, nn.Linear):
        nn.init.kaiming_normal_(module.weight,
                                nonlinearity='relu')
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Conv2d):
        nn.init.kaiming_normal_(module.weight,
                                nonlinearity='relu')

model = nn.Sequential(
    nn.Linear(64, 128), nn.ReLU(),
    nn.Linear(128, 64), nn.ReLU(),
    nn.Linear(64, 10)
)
model.apply(init_weights)
print('Custom He init applied to all layers')

Orthogonal Initialisation for RNNs

Orthogonal initialisation sets weight matrices to be orthogonal (Q from QR decomposition), which preserves gradient norms during backpropagation through time. This is particularly useful for recurrent networks, where the same weight matrix is multiplied repeatedly (once per timestep). Orthogonal weights prevent gradients from exploding or vanishing as they are propagated back through many timesteps in long sequences.

import torch
import torch.nn as nn

# Orthogonal init: columns are orthonormal
layer = nn.Linear(64, 64)
nn.init.orthogonal_(layer.weight)

# Verify: W @ W.T should be identity (approx)
I_approx = layer.weight @ layer.weight.T
print('Close to identity:', torch.allclose(
    I_approx,
    torch.eye(64),
    atol=1e-5
))
# True -- orthogonal matrices preserve vector norms

Default PyTorch Initialisations

PyTorch applies sensible defaults automatically: nn.Linear uses Kaiming uniform for weights and uniform distribution for biases. nn.Conv2d also uses Kaiming uniform. nn.Embedding uses standard normal N(0, 1). nn.LSTM uses uniform in [-1/sqrt(hidden), 1/sqrt(hidden)]. In many cases the defaults work well, but for very deep networks or non-standard activations, explicit initialisation with the formulas above gives better results.

import torch.nn as nn

# Check PyTorch defaults
linear = nn.Linear(256, 128)
print('Linear weight std:', linear.weight.std().item())
# ~0.088 = Kaiming uniform for fan_in=256

conv = nn.Conv2d(3, 64, kernel_size=3)
print('Conv2d weight std:', conv.weight.std().item())
# Kaiming uniform based on receptive field size

emb = nn.Embedding(1000, 128)
print('Embedding weight std:', emb.weight.std().item())
# ~1.0 = N(0, 1) default

Practical Initialisation Guide

A practical guide for choosing initialisation: use He (Kaiming) normal or uniform for any network using ReLU or its variants (LeakyReLU, ELU, GELU). Use Xavier normal or uniform for tanh or sigmoid activations. Use orthogonal for recurrent weights. Set biases to zero in all cases. For Transformers with GELU, N(0, 0.02) is the empirical standard used in GPT-2 and subsequent models. Trust PyTorch's defaults for standard architectures; only override when training is unstable.

# Quick reference table
init_guide = {
    'ReLU (Linear, Conv)': 'kaiming_normal_ / kaiming_uniform_',
    'Tanh / Sigmoid':       'xavier_normal_ / xavier_uniform_',
    'RNN hidden matrix':    'orthogonal_',
    'Transformer (GELU)':   'normal_(mean=0, std=0.02)',
    'Embedding':            'normal_(mean=0, std=1)',
    'Biases':               'zeros_()'
}
for activation, method in init_guide.items():
    print(f'{activation}: {method}')

Verifying Initialisation Quality

After applying initialisation, verify it by checking activation statistics in the first forward pass. A healthy network should have activation standard deviations near 1.0 across all layers and gradient norms of similar magnitude across layers. Large discrepancies (e.g., std=10 in one layer, std=0.001 in another) indicate initialisation mismatch. This quick sanity check takes seconds and can save hours of debugging poor training dynamics.

import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(64, 256), nn.ReLU(),
    nn.Linear(256, 256), nn.ReLU(),
    nn.Linear(256, 64),  nn.ReLU(),
    nn.Linear(64, 10)
)
model.apply(lambda m: nn.init.kaiming_normal_(m.weight)
            if isinstance(m, nn.Linear) else None)

# Check activation std through the network
x = torch.randn(32, 64)
hooks = []
stds = []
for layer in model:
    x = layer(x)
    if hasattr(x, 'std'):
        stds.append(x.std().item())
print('Activation stds:', [f'{s:.2f}' for s in stds])

Quick Check

Test your understanding of Machine Learning with Python concepts from this lesson.

Lesson Recap

In this lesson you learned: Xavier initialisation is designed for tanh/sigmoid activations using variance 2/(fan_in + fan_out), He (Kaiming) initialisation is designed for ReLU using variance 2/fan_in to compensate for ReLU zeroing half its inputs, and model.apply(init_fn) applies a custom initialisation to every layer in the network. Next up we dive into convolutional neural networks starting with convolution and filter operations.

Frequently asked questions

Is the “Weight Initialisation: Xavier and He Initialisation” lesson free?

Yes — the full text of “Weight Initialisation: Xavier and He Initialisation” is free to read here on the web, and the Machine Learning Academy course includes 4 lessons in total. To practise it interactively (a built-in code editor and a 24/7 AI tutor) and unlock the rest of the Machine Learning Academy course, upgrade to CoddyKit PRO.

What will I learn in “Weight Initialisation: Xavier and He Initialisation”?

Learners will apply Xavier uniform and He normal initialisation, observe how they prevent vanishing/exploding gradients in deep networks compared to default random init. You practise Machine Learning Academy with hands-on code you run directly in the browser, and a 24/7 AI tutor answers your questions as you work through the lesson.

Do I need any experience to start Machine Learning Academy?

No prior experience is required. Machine Learning Academy on CoddyKit is structured for beginners through advanced learners; this is — lesson 4 of 4, so you can start here or from the beginning and move at your own pace.

How long does the “Weight Initialisation: Xavier and He Initialisation” lesson take?

Most CoddyKit lessons take about 5–10 minutes. Each one is bite-sized and interactive, so you make steady progress and pick up exactly where you left off across the web and the app.

Can I write and run code in this Machine Learning Academy lesson?

Yes. Every Machine Learning Academy lesson includes a built-in code editor, so you write and run real code right in your browser and get instant AI feedback — no local setup required.

All lessons in this course

  1. Learning Rate: The Most Important Hyperparameter
  2. Batch Normalisation: Stable and Faster Training
  3. Dropout Regularisation to Prevent Overfitting
  4. Weight Initialisation: Xavier and He Initialisation
← Back to Machine Learning Academy