Weight Initialization for Neural Networks

Hello! Let's dive into the fourth lesson of our module on Deep Neural Network Fundamentals.

Introduction

In our last lesson, we explored the critical role of activation functions and saw how choices like Sigmoid and Tanh can lead to vanishing gradients, while ReLU and its variants can suffer from "dying neurons." These problems are not caused by the activation functions alone; they arise from the interaction between the activations and the magnitude of the weights in the network. A poor choice of initial weights can cause the signal to either explode to infinity or vanish to zero within just a few layers, halting the learning process before it even begins.

This lesson directly addresses that challenge. Our goal is to implement Xavier and He weight initialization strategies to improve convergence. These are not just arbitrary tricks; they are principled methods derived from a careful analysis of how signals and gradients propagate through a network. By setting the initial weights to an appropriate scale, we can create the conditions for stable and efficient training.

We'll start by understanding the core principle of variance preservation, then derive the formulas for both Xavier and He initialization, and finally, implement them in code. This will connect the theory from our previous lessons on backpropagation and activations to the practical reality of building a network that can actually learn.

1. The Problem: Why Naive Initialization Fails

Before we dive into the solutions, let's solidify our understanding of the problem. We know from the previous lesson that we must initialize weights randomly to break symmetry—if all weights start at zero, all neurons in a layer will learn the same thing. But how should we scale this randomness?

To see why this is a critical question, let's watch the beginning of this video from DeepLearningAI.

As the video explains, the pre-activation $z$ in a neuron is a sum of $n$ terms: $z={\sum }_{i=1}^n{w}_i{x}_i$. If we're summing up many terms, the variance of the sum can grow very quickly. If the weights $w_i$ are too large, $z$ will have a large variance, pushing the activations into the saturated regions of Sigmoid/Tanh or leading to huge values with ReLU. Conversely, if the weights are too small, $z$ will have a tiny variance, and the activations will shrink towards zero layer after layer.

This leads to the core idea behind modern initialization: we must preserve the variance of the activations as they flow forward through the network, and the variance of the gradients as they flow backward.

2. The Principle of Variance Preservation

To find the right scale for our weights, we can analyze how variance propagates through a layer. This is where your mathematical background will be very useful.

Let's consider the pre-activation for a single neuron:

$$z=\sum _{i=1}^{n_{in}}{w}_i{x}_i$$

Here, $n_{in}$ is the number of input neurons (the "fan-in"). We want to find the relationship between the variance of the output, $Var(z)$, and the variances of the inputs, $Var(x_i)$, and weights, $Var(w_i)$.

The following article provides an exceptionally clear, step-by-step derivation.

The key result from that reading is:

$$Var(z)=n_{in}\cdot Var(w)\cdot Var(x)$$

This equation assumes that the inputs and weights are independent and have zero mean. It tells us that the output variance is scaled by the fan-in ($n_{in}$) and the weight variance ($Var(w)$).

For the signal to propagate without exploding or vanishing, we ideally want the output variance to equal the input variance, i.e., $Var(z)\approx Var(x)$. If we set $Var(z)=Var(x)$ in our equation and solve for $Var(w)$, we get:

$$1=n_{in}\cdot Var(w)⟹Var(w)=\frac{1}{n_{in}}$$

This simple but powerful result is the foundation of principled weight initialization. It tells us that the variance of our weights should be inversely proportional to the number of inputs to the layer.

3. Xavier (Glorot) Initialization

The first popular method to use this principle was proposed by Xavier Glorot and Yoshua Bengio. Their key insight was that we need to preserve variance not only in the forward pass (for activations) but also in the backward pass (for gradients).

The derivation for the backward pass is analogous to the forward pass, but now the gradients flow from $n_{out}$ neurons (the "fan-out"). This leads to a requirement of $Var(w)=\frac{1}{n_{out}}$.

We now have two conflicting goals:

• Forward pass: $Var(w)=\frac{1}{n_{in}}$
• Backward pass: $Var(w)=\frac{1}{n_{out}}$

Glorot and Bengio's solution was a simple compromise: average the two!

As the article explains, the final formula for Xavier initialization is:

$$Var(w)=\frac{2}{n_{in}+n_{out}}$$

To draw weights from a distribution with this variance, we use one of the following:

• Normal Distribution: Sample from $\mathcal{N}(0,{\sigma }^2)$ with $\sigma =\sqrt{\frac{2}{n_{in}+n_{out}}}$.
• Uniform Distribution: Sample from $U[-a,a]$ with $a=\sqrt{\frac{6}{n_{in}+n_{out}}}$. (The $\sqrt{6}$ comes from the fact that $Var(U[-a,a])=a^2/3$).

This initialization was derived assuming a linear activation function. It works well in practice for Tanh and Sigmoid because they are roughly linear around zero.

Test your understanding!

You are initializing a linear layer in a neural network that connects a layer with 512 neurons to a layer with 128 neurons. You are using a Tanh activation function.

1. What are the fan_in and fan_out for this layer?
2. If you use Xavier normal initialization, what should the standard deviation of the weights be?

Show answer

1. fan_in = 512 (number of inputs) and fan_out = 128 (number of outputs).
2. The variance is $Var(w)=\frac{2}{512+128}=\frac{2}{640}=\frac{1}{320}$. The standard deviation is the square root of the variance, so $\sigma =\sqrt{\frac{1}{320}}\approx 0.0559$.

4. He (Kaiming) Initialization

Xavier initialization was a huge step forward, but it has a problem when used with the ReLU activation function. In our last lesson, we saw that ReLU sets all negative inputs to zero. For a symmetric input distribution centered at zero, this means ReLU discards half of the information. This has a direct effect on the variance.

It turns out that passing a zero-centered signal through a ReLU function halves its variance.

This means our forward pass variance equation for ReLU networks becomes:

$$Var(a)=\frac{1}{2}\cdot n_{in}\cdot Var(w)\cdot Var(x)$$

where $a$ is the activation after ReLU.

To counteract this halving effect, we need to double our weight variance. This is the insight behind He initialization, proposed by Kaiming He et al.

To maintain variance ($Var(a)\approx Var(x)$), we set:

$$1=\frac{1}{2}\cdot n_{in}\cdot Var(w)⟹Var(w)=\frac{2}{n_{in}}$$

This is the formula for He initialization. Notice that it only depends on $n_{in}$. The authors found that for very deep ReLU networks, properly initializing the forward pass was more critical, and this simpler formula worked very well.

The corresponding distribution parameters are:

• Normal Distribution: Sample from $\mathcal{N}(0,{\sigma }^2)$ with $\sigma =\sqrt{\frac{2}{n_{in}}}$.
• Uniform Distribution: Sample from $U[-a,a]$ with $a=\sqrt{\frac{6}{n_{in}}}$.

For a visual demonstration of how effective these methods are, the following video provides excellent comparative plots.

5. Implementation in Python

Now let's put theory into practice. The learning outcome for this lesson is to implement these strategies. We'll first see how to code them from scratch using NumPy to solidify your understanding, and then see how to use the standard, built-in functions in PyTorch.

From-Scratch Implementation

The formulas we derived can be translated directly into code.

import numpy as np

def he_normal(shape):
    """He initialization with a normal distribution."""
    # shape is typically (fan_out, fan_in)
    fan_in = shape[1]
    std = np.sqrt(2 / fan_in)
    return np.random.randn(*shape) * std

def xavier_normal(shape):
    """Xavier initialization with a normal distribution."""
    fan_in, fan_out = shape[1], shape[0]
    std = np.sqrt(2 / (fan_in + fan_out))
    return np.random.randn(*shape) * std

# Example usage for a layer with 512 inputs and 256 outputs
shape = (256, 512)
weights_he = he_normal(shape)
weights_xavier = xavier_normal(shape)

print(f"He normal std: {weights_he.std():.4f}")
print(f"Xavier normal std: {weights_xavier.std():.4f}")

# Theoretical values:
# He std = sqrt(2/512) = 0.0625
# Xavier std = sqrt(2/(512+256)) = 0.0515

PyTorch Implementation

In practice, you will almost always use the built-in functions from your deep learning framework. PyTorch's torch.nn.init module makes this very easy.

As shown in the article, you can initialize a layer like this:

import torch
import torch.nn as nn

# Create a linear layer
layer = nn.Linear(in_features=512, out_features=256)

# Apply He normal initialization (for ReLU networks)
# The '_' at the end signifies an in-place operation
nn.init.kaiming_normal_(layer.weight, nonlinearity='relu')

# Apply Xavier uniform initialization (for Tanh/Sigmoid networks)
nn.init.xavier_uniform_(layer.weight)

# Biases are typically initialized to zero
nn.init.zeros_(layer.bias)

Note that PyTorch refers to He initialization as kaiming initialization, after the first author of the paper.

Conclusion

Congratulations on completing this lesson! You've moved from simply using random weights to understanding and implementing principled initialization strategies that are fundamental to training deep neural networks.

Key Takeaways:

• The Goal: Proper weight initialization aims to preserve the variance of signals and gradients to prevent them from vanishing or exploding.
• The Principle: The variance of the weights in a layer, $Var(w)$, should be scaled inversely with the number of connections.
• Xavier (Glorot) Initialization: Designed for symmetric, saturating activations like Tanh and Sigmoid. It uses $Var(w)=\frac{2}{n_{in}+n_{out}}$ to balance the forward and backward passes.
• He (Kaiming) Initialization: Designed for ReLU and its variants. It uses $Var(w)=\frac{2}{n_{in}}$ to compensate for the fact that ReLU halves the signal variance.
• Implementation: Modern frameworks like PyTorch provide convenient functions (nn.init.xavier_normal_, nn.init.kaiming_normal_) to apply these initializers.

Preview of the next lesson:
Weight initialization is our first and most important line of defense against unstable gradients. However, even with perfect initialization, the complex dynamics of training can sometimes still cause gradients to grow uncontrollably. In our next lesson, we will look at a complementary technique: diagnosing and mitigating vanishing and exploding gradients using techniques like gradient clipping. This is a reactive measure that monitors gradients during training and rescales them if they become too large, ensuring stability throughout the learning process.