Introduction

Weight initialization has a direct effect on how easily a neural network trains. In PyTorch, you can rely on sensible defaults for many layers, but understanding when to apply Xavier, Kaiming, or custom initialization will help you avoid unstable gradients and slow convergence.

Why Initialization Matters

When weights start too small, signals shrink as they pass through layers and gradients can vanish. When weights start too large, activations and gradients can explode. Good initialization keeps the scale of activations in a useful range so optimization begins from a stable point.

PyTorch modules already initialize parameters when you create them, but those defaults are generic. If your model architecture or activation functions have specific needs, it is common to override them after constructing the model.

1import torch
2import torch.nn as nn
3
4model = nn.Sequential(
5    nn.Linear(32, 64),
6    nn.ReLU(),
7    nn.Linear(64, 10),
8)
9
10for name, param in model.named_parameters():
11    print(name, param.shape)

The weights exist as soon as the layers are created. Initialization means replacing those values with a strategy suited to the network.

Common Initialization Strategies

PyTorch exposes initialization helpers in torch.nn.init. Two of the most common choices are Xavier initialization and Kaiming initialization.

Xavier, also called Glorot initialization, works well for activations that keep values roughly centered, such as tanh. Kaiming, also called He initialization, is usually a better fit for ReLU-style networks because it preserves variance more effectively when half the activations may become zero.

1import torch
2import torch.nn as nn
3
4
5class Classifier(nn.Module):
6    def __init__(self):
7        super().__init__()
8        self.fc1 = nn.Linear(100, 128)
9        self.fc2 = nn.Linear(128, 64)
10        self.fc3 = nn.Linear(64, 3)
11        self.relu = nn.ReLU()
12
13        self.apply(self._init_weights)
14
15    def _init_weights(self, module):
16        if isinstance(module, nn.Linear):
17            nn.init.kaiming_uniform_(module.weight, nonlinearity="relu")
18            nn.init.zeros_(module.bias)
19
20    def forward(self, x):
21        x = self.relu(self.fc1(x))
22        x = self.relu(self.fc2(x))
23        return self.fc3(x)
24
25
26model = Classifier()

This example initializes every nn.Linear layer with Kaiming uniform weights and zero bias terms. Zero bias is a common default because it does not create symmetry problems the way zero weights would.

If you are using tanh or sigmoid, Xavier is often a better choice:

1import torch.nn as nn
2
3
4def init_for_tanh(module):
5    if isinstance(module, nn.Linear):
6        nn.init.xavier_uniform_(module.weight)
7        nn.init.zeros_(module.bias)

Initializing Different Layer Types

Convolutional and recurrent layers also benefit from deliberate initialization. For example, convolutional layers in ReLU networks commonly use Kaiming initialization, while embeddings may use a small normal distribution.

1import torch.nn as nn
2
3
4def init_model(module):
5    if isinstance(module, nn.Conv2d):
6        nn.init.kaiming_normal_(module.weight, nonlinearity="relu")
7        if module.bias is not None:
8            nn.init.zeros_(module.bias)
9    elif isinstance(module, nn.Linear):
10        nn.init.xavier_uniform_(module.weight)
11        nn.init.zeros_(module.bias)
12    elif isinstance(module, nn.Embedding):
13        nn.init.normal_(module.weight, mean=0.0, std=0.02)

The important part is matching the initialization to the layer behavior and activation pattern rather than applying one rule blindly everywhere.

Verifying That Initialization Happened

It is worth checking statistics after initialization, especially in experimental models:

1import torch
2
3layer = nn.Linear(16, 8)
4nn.init.xavier_normal_(layer.weight)
5
6print(layer.weight.mean().item())
7print(layer.weight.std().item())

You do not need exact target values, but wildly unexpected means or standard deviations usually signal a bug such as forgetting to call apply or accidentally reinitializing the wrong module.

Common Pitfalls

• Initializing all weights to zero prevents neurons in the same layer from learning different features.
• Using Xavier for a deep ReLU stack can work, but Kaiming is usually the better default.
• Forgetting biases when writing a custom initializer leaves part of the layer at its original default values.
• Reinitializing a pretrained model destroys learned weights, so custom initialization should happen before training, not after loading a checkpoint.
• Applying initialization rules by layer name rather than module type can break when the architecture changes.

Summary

• PyTorch provides initialization helpers in torch.nn.init for common strategies.
• Kaiming initialization is a strong default for ReLU networks.
• Xavier initialization often fits tanh and similar activations better.
• Use model.apply to walk the module tree and initialize each layer consistently.
• Verify parameter statistics when debugging training instability or suspicious convergence.