Weight Initialization

Learning Objectives

By the end of this section, you will be able to:

1. Understand why initialization matters: Grasp how poor initialization causes vanishing/exploding gradients that prevent deep networks from learning
2. Derive Xavier initialization: Mathematically derive the variance-preserving initialization for linear and tanh activations
3. Derive He initialization: Understand why ReLU networks need different initialization and derive the formula
4. Implement initialization in PyTorch: Apply proper initialization to your networks using PyTorch's built-in functions
5. Choose the right strategy: Select appropriate initialization based on your network architecture and activation functions

Why This Matters: Before Xavier and He initialization were discovered, training deep networks was extremely difficult. These insights about variance preservation enabled the training of networks with hundreds of layers. Understanding initialization is essential for training modern deep learning models.

The Big Picture

The Training Disaster

Imagine spending weeks designing a sophisticated 50-layer neural network, only to find that training fails completely—the loss doesn't decrease, gradients are either zero or infinity, and no learning occurs. This was a common experience in the early days of deep learning, and the culprit was often something as simple as how the weights were initialized.

Weight initialization might seem like a minor implementation detail, but it's actually one of the most critical factors determining whether a deep network can be trained at all. The values we assign to weights before training begins set the stage for everything that follows.

The Core Insight

The fundamental insight is this: as signals (activations during forward pass, gradients during backward pass) flow through a deep network, their magnitudes must remain approximately constant. If signals shrink at each layer, they vanish by the time they reach deep layers. If they grow, they explode into numerical overflow.

The key breakthrough was realizing that we can control this variance flow by carefully choosing the variance of the initial weights.

Historical Context

Two seminal papers solved this problem:

These discoveries, along with batch normalization and residual connections, enabled the training of very deep networks that power modern AI.

The Symmetry Breaking Problem

Why Not Initialize All Weights to the Same Value?

Consider initializing all weights to zero (or any constant). What happens?

In the forward pass, every neuron in a layer receives the same weighted sum of inputs (since all weights are identical). They all produce the same output. During backpropagation, they all receive the same gradient. And after the weight update, they all have the same new weight value.

This symmetry persists throughout training. The network effectively collapses to having just one neuron per layer, wasting all the extra capacity. No matter how long you train, the neurons never learn different features.

Random Initialization: The First Attempt

The obvious solution is random initialization. But what distribution should we sample from? A naive approach is to use standard normal: $W_{ij} \sim \mathcal{N}(0, 1)$.

This breaks symmetry, but creates a new problem. Consider a fully-connected layer with $n$ inputs:

If $\text{Var}(w_i) = 1$ and $\text{Var}(x_i) = 1$ (with everything independent and zero-mean):

The variance grows by a factor of $n$ at each layer! In a network with 256 neurons per layer, the variance explodes to astronomical values within just a few layers.

Variance Flow Through Networks

The Mathematical Framework

Let's carefully analyze how variance propagates through a single layer. Consider a fully-connected layer with:

• $n_{\text{in}}$ input neurons (fan-in)
• $n_{\text{out}}$ output neurons (fan-out)
• Weights $W_{ij}$ drawn i.i.d. with mean 0 and variance $\sigma^2$
• Inputs $x_i$ with mean 0 and variance $v_{\text{in}}$

For a single output neuron (before activation):

Using the property that the variance of a sum of independent terms is the sum of variances:

The Variance Preservation Condition

For stable forward propagation, we want the output variance to equal the input variance:

This gives us the condition:

This is the key insight: the variance of the weights should be inversely proportional to the number of input connections.

Backward Pass: Gradient Variance

The same analysis applies to gradients flowing backward. During backpropagation, the gradient with respect to an input is:

For gradients to maintain variance:

We have a conflict! Forward pass wants $\sigma^2 = 1/n_{\text{in}}$, but backward pass wants $\sigma^2 = 1/n_{\text{out}}$.

Interactive: Variance Propagation

Experiment with different initialization strategies and see how variance changes as signals propagate through a deep network. Observe how naive initialization causes variance to explode or vanish, while proper initialization maintains stability:

Xavier (Glorot) Initialization

Resolving the Conflict

Xavier Glorot and Yoshua Bengio proposed a compromise: use the average of the forward and backward requirements:

This ensures that variance is approximately preserved in both directions. The name "Xavier" comes from Xavier Glorot's first name.

Xavier Normal Initialization

Sample weights from a normal distribution with the computed variance:

Xavier Uniform Initialization

Alternatively, sample from a uniform distribution with the same variance:

The limit $a$ is derived from the uniform distribution's variance formula: $\text{Var}(\mathcal{U}(-a, a)) = a^2/3$. Setting this equal to $2/(n_{\text{in}} + n_{\text{out}})$ gives $a = \sqrt{6/(n_{\text{in}} + n_{\text{out}})}$.

When to Use Xavier

Xavier initialization was derived assuming linear activations or activations symmetric around zero like tanh. It works well for:

• Networks with tanh or sigmoid activations
• Linear layers without activation (like the final output layer)
• Attention mechanisms and transformers (which often use linear projections)

He (Kaiming) Initialization

The ReLU Problem

Consider ReLU activation: $\text{ReLU}(z) = \max(0, z)$. For a symmetric input distribution centered at zero, ReLU outputs zero for half the inputs. This halves the variance:

This derivation assumes $z$ is symmetric around zero. For the forward pass through a ReLU layer:

To preserve variance ($v_{\text{out}} = v_{\text{in}}$):

This is exactly twice the variance of Xavier initialization (considering only fan-in). Kaiming He and colleagues derived this in their 2015 paper on training very deep networks.

He Normal Initialization

He Uniform Initialization

Fan-in vs Fan-out Mode

PyTorch allows you to choose whether to use $n_{\text{in}}$ (fan-in) or $n_{\text{out}}$ (fan-out):

Interactive: Weight Distributions

Explore how different initialization strategies create different weight distributions. Adjust the fan-in and fan-out values to see how they affect the distribution's spread:

Interactive: Signal Propagation

Visualize how activations and gradients flow through a deep network with different initialization strategies. Watch how poor initialization causes signals to vanish or explode:

PyTorch Implementation

Built-in Initialization Functions

PyTorch provides initialization functions in torch.nn.init:

init.xavier_uniform_(layer.weight)

'relu', 'leaky_relu', 'tanh', 'sigmoid'

1import torch
2import torch.nn as nn
3import torch.nn.init as init
4
5# Create a linear layer
6layer = nn.Linear(512, 256)
7
8# Xavier initialization
9init.xavier_uniform_(layer.weight)  # Uniform version
10init.xavier_normal_(layer.weight)   # Normal version
11
12# He (Kaiming) initialization
13init.kaiming_uniform_(layer.weight, mode='fan_in', nonlinearity='relu')
14init.kaiming_normal_(layer.weight, mode='fan_in', nonlinearity='relu')
15
16# Zero-initialize biases (common practice)
17init.zeros_(layer.bias)
18
19# Constant initialization
20init.constant_(layer.weight, 0.01)  # Not recommended!
21
22# Check the resulting variance
23print(f"Weight variance: {layer.weight.var().item():.6f}")
24print(f"Expected (He): {2/512:.6f}")

Custom Initialization for a Model

1import torch.nn as nn
2import torch.nn.init as init
3
4class MLP(nn.Module):
5    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
6        super().__init__()
7
8        self.layers = nn.Sequential(
9            nn.Linear(input_dim, hidden_dim),
10            nn.ReLU(),
11            nn.Linear(hidden_dim, hidden_dim),
12            nn.ReLU(),
13            nn.Linear(hidden_dim, output_dim),
14        )
15
16        # Apply custom initialization
17        self._init_weights()
18
19    def _init_weights(self):
20        for module in self.modules():
21            if isinstance(module, nn.Linear):
22                # He initialization for ReLU layers
23                init.kaiming_normal_(
24                    module.weight,
25                    mode='fan_in',
26                    nonlinearity='relu'
27                )
28                if module.bias is not None:
29                    init.zeros_(module.bias)
30
31    def forward(self, x):
32        return self.layers(x)
33
34# Usage
35model = MLP(784, 256, 10)
36
37# Verify initialization
38for name, param in model.named_parameters():
39    if 'weight' in name:
40        print(f"{name}: mean={param.mean():.4f}, var={param.var():.4f}")

Initialization with apply()

PyTorch also provides apply() for recursively applying a function to all modules:

1def init_weights(m):
2    """Initialize weights for different layer types."""
3    if isinstance(m, nn.Linear):
4        init.kaiming_normal_(m.weight, nonlinearity='relu')
5        if m.bias is not None:
6            init.zeros_(m.bias)
7    elif isinstance(m, nn.Conv2d):
8        init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
9        if m.bias is not None:
10            init.zeros_(m.bias)
11    elif isinstance(m, nn.BatchNorm2d):
12        init.ones_(m.weight)
13        init.zeros_(m.bias)
14
15# Apply to model
16model = MyNetwork()
17model.apply(init_weights)

Practical Guidelines

Choosing the Right Initialization

Special Cases

Residual Networks

For residual networks with skip connections, it's common to initialize the last layer of each residual block to zero. This makes the network initially behave like a shallower network:

1class ResidualBlock(nn.Module):
2    def __init__(self, dim):
3        super().__init__()
4        self.layers = nn.Sequential(
5            nn.Linear(dim, dim),
6            nn.ReLU(),
7            nn.Linear(dim, dim),  # Initialize this to zero
8        )
9        # Zero-initialize the final layer
10        init.zeros_(self.layers[-1].weight)
11        init.zeros_(self.layers[-1].bias)
12
13    def forward(self, x):
14        return x + self.layers(x)  # Skip connection

Transformers

Modern transformers often scale initialization by layer depth:

1# GPT-style initialization
2# Scale residual layers by 1/sqrt(num_layers)
3scale = 1 / math.sqrt(num_layers)
4init.normal_(layer.weight, mean=0.0, std=0.02 * scale)

Output Layers

For classification tasks, initializing the output layer with smaller weights often helps:

1# Smaller initialization for output layer
2init.normal_(output_layer.weight, mean=0.0, std=0.01)
3init.zeros_(output_layer.bias)

Summary

Weight initialization is crucial for training deep neural networks. Here are the key takeaways:

Quick Reference: Initialization Formulas

1Xavier Normal:   W ~ N(0, sqrt(2 / (n_in + n_out)))
2Xavier Uniform:  W ~ U(-sqrt(6/(n_in+n_out)), sqrt(6/(n_in+n_out)))
3
4He Normal:       W ~ N(0, sqrt(2 / n_in))
5He Uniform:      W ~ U(-sqrt(6/n_in), sqrt(6/n_in))

Knowledge Check

Test your understanding of weight initialization concepts:

Exercises

Conceptual Questions

1. Explain why initializing all weights to small constants (e.g., 0.01) is better than zeros but still problematic for deep networks.
2. A layer has 1024 input neurons and 512 output neurons. Calculate the variance and uniform initialization bounds for both Xavier and He initialization.
3. Why might you choose fan_out mode for He initialization in a generative model?
4. How does batch normalization reduce the dependence on initialization?

Coding Exercises

1. Implement from scratch: Write your own xavier_normal and kaiming_normal functions without using PyTorch's built-in versions. Verify they produce the same variance.
2. Variance tracking: Create a neural network that records the variance of activations at each layer during a forward pass. Compare different initializations.
3. Training comparison: Train the same model architecture on MNIST with zeros, random N(0,1), Xavier, and He initialization. Plot loss curves and final accuracy for each.
4. Gradient analysis: Extend exercise 2 to also track gradient variance during backward pass. Verify that He initialization preserves gradient magnitudes better than Xavier for ReLU networks.

1class VarianceTracker(nn.Module):
2    def __init__(self, layers):
3        super().__init__()
4        self.layers = nn.ModuleList(layers)
5        self.activation_vars = []
6
7    def forward(self, x):
8        self.activation_vars = [x.var().item()]
9        for layer in self.layers:
10            x = layer(x)
11            self.activation_vars.append(x.var().item())
12        return x

In the next section, we'll explore regularization techniques—methods like dropout, weight decay, and data augmentation that prevent overfitting and improve generalization.