🎓 A1Learner
🔍

PyTorch Weight Initialization

Weight initialization is the choice of starting values for a model's weights. The starting scale decides whether signals survive a deep stack or fade out inside it.

The wrong scale, shown

The drift table in the closing example is the pain here. A poor start shrinks the signal before training even begins.

h = torch.relu(layer(h))  # with the default start, std fell under 0.05 by layer 5

Xavier weights

Xavier draws weights from a uniform range set by the layer's input and output counts. The drawn weights have this spread.

import torch

torch.manual_seed(0)
layer = torch.nn.Linear(16, 16)
torch.nn.init.xavier_uniform_(layer.weight)
print(f"Xavier weight std: {layer.weight.std():.4f}")

Kaiming weights

Kaiming draws from a wider range for the same layer. ReLU clips half the signal, so the start compensates with bigger weights.

import torch

torch.manual_seed(0)
layer = torch.nn.Linear(16, 16)
torch.nn.init.kaiming_uniform_(layer.weight, nonlinearity="relu")
print(f"Kaiming weight std: {layer.weight.std():.4f}")

Example

import torch


def build(use_kaiming):
    torch.manual_seed(0)
    layers = [torch.nn.Linear(64, 64) for _ in range(20)]
    if use_kaiming:
        for layer in layers:
            torch.nn.init.kaiming_uniform_(
                layer.weight, nonlinearity="relu")
            torch.nn.init.zeros_(layer.bias)
    return layers


def spreads(layers, x):
    h = x
    result = {}
    for i, layer in enumerate(layers):
        h = torch.relu(layer(h))
        if i + 1 in (1, 5, 10, 15, 20):
            result[i + 1] = h.std().item()
    return result


torch.manual_seed(2)
x = torch.randn(128, 64)
default = spreads(build(False), x)
kaiming = spreads(build(True), x)

for depth in (1, 5, 10, 15, 20):
    print(f"Layer {depth}: default std {default[depth]:.4f}, "
          f"Kaiming std {kaiming[depth]:.4f}")

Key points

  • Initialization: the starting values given to weights before training begins
  • Use torch.nn.init to set starting weights with a chosen rule
  • Xavier initialization: a rule that scales weights by the number of inputs and outputs, built for tanh-style activations
  • Kaiming initialization: a rule that scales weights for ReLU, allowing for the half of values ReLU clips to zero
  • Default init: the starting rule a layer applies on its own when none is given
  • A starting scale matched to the activation keeps activation spread steady with depth

In this example, two 20-layer stacks differ only in their starting weights. The input batch and the seed are shared, and each stack is built from seed 0. The default stack starts at std 0.3343 and falls to 0.0458 by layer 5. It stays near 0.045 for the rest of the stack. The Kaiming stack starts at 0.8376 and holds between 0.6527 and 0.8376 at every measured depth. Layer 20 still reads 0.7135. The matched start keeps the signal alive through all 20 layers.

Note

Initialization matters most in deep plain stacks, and normalization layers can hide a poor start.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Why does Kaiming initialization use a wider spread than Xavier for ReLU layers?