Xavier (Glorot) Initialization

As introduced, initializing the weights of a deep network requires careful consideration. Simply drawing weights from a standard normal or uniform distribution often leads to unstable training dynamics. Gradients can shrink exponentially as they propagate backward (vanishing gradients), or grow exponentially (exploding gradients), preventing the model from learning effectively.

The Xavier initialization method, proposed by Xavier Glorot and Yoshua Bengio in 2010, directly addresses this by aiming to keep the variance of activations and gradients roughly constant across layers. Maintaining consistent variance helps ensure that the signal (activations during the forward pass and gradients during the backward pass) doesn't vanish or explode as it traverses the network's depth.

The Core Idea: Variance Preservation

Consider a single linear layer in a neural network performing the operation $y = Wx + b$, where $W$ is the weight matrix, $x$ is the input vector, $b$ is the bias vector, and $y$ is the output before the activation function. Xavier initialization analyzes how the variance propagates through this layer.

The core assumptions are:

1. The inputs $x$ are independent and identically distributed (i.i.d.) with zero mean and variance $Var(x)$.
2. The weights $W$ are independent and identically distributed (i.i.d.) with zero mean and variance $Var(W)$.
3. The weights and inputs are independent of each other.
4. The activation function $f$ is symmetric around zero (like $tanh$) and its derivative is approximately 1 near zero ($f'(0) \approx 1$).

Under these assumptions, the variance of the layer's output $y_i$ (for a single neuron $i$) can be related to the input variance and weight variance:

$Var(y_i) = n_{in} Var(W_{ij}) Var(x_j)$

Here, $n_{in}$ represents the number of input connections to the neuron (the "fan-in"). To keep the output variance $Var(y_i)$ equal to the input variance $Var(x_j)$, we need:

$n_{in} Var(W_{ij}) = 1$ $Var(W_{ij}) = \frac{1}{n_{in}}$

This condition ensures that the variance of the activations doesn't significantly change during the forward pass.

Similarly, considering the backward pass, we analyze the variance of the gradients. The variance of the gradient with respect to the activations of the previous layer depends on the fan-out $n_{out}$ (the number of neurons the current layer's output connects to). To maintain gradient variance, the condition becomes:

$n_{out} Var(W_{ij}) = 1$ $Var(W_{ij}) = \frac{1}{n_{out}}$

Xavier initialization seeks a compromise between these two conditions (maintaining activation variance forward and gradient variance backward) by averaging the constraints:

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

This variance forms the basis for the initialization strategy.

Xavier Initialization Formulas

Weights are typically drawn from either a uniform or a normal distribution scaled according to this derived variance.

1. Xavier Uniform Initialization: Weights are sampled from a uniform distribution $U[-a, a]$, where $a$ is chosen such that the variance is $\frac{2}{n_{in} + n_{out}}$. The variance of $U[-a, a]$ is $\frac{(2a)^2}{12} = \frac{a^2}{3}$. Setting this equal to the target variance: $\frac{a^2}{3} = \frac{2}{n_{in} + n_{out}}$ $a^2 = \frac{6}{n_{in} + n_{out}}$ $a = \sqrt{\frac{6}{n_{in} + n_{out}}}$

   So, the weights are initialized using: $W \sim U\left[-\sqrt{\frac{6}{n_{in} + n_{out}}}, \sqrt{\frac{6}{n_{in} + n_{out}}}\right]$

2. Xavier Normal Initialization: Weights are sampled from a normal distribution $N(0, \sigma^2)$ with mean 0 and the target variance $\sigma^2$: $\sigma^2 = \frac{2}{n_{in} + n_{out}}$

   So, the weights are initialized using: $W \sim N\left(0, \sqrt{\frac{2}{n_{in} + n_{out}}}\right)$

In both cases, $n_{in}$ is the number of inputs to the layer (fan-in), and $n_{out}$ is the number of outputs from the layer (fan-out). Biases are typically initialized to zero.

Suitability and Implementation

Xavier initialization works well for layers followed by activation functions that are roughly linear around zero and symmetric, such as $tanh$ and the logistic sigmoid. The assumption $f'(0) \approx 1$ holds reasonably well for these functions. However, it is less suited for Rectified Linear Units (ReLU) and its variants, which are asymmetric and have a derivative of 0 for negative inputs. This limitation motivated the development of Kaiming initialization, which we will discuss next.

In PyTorch, you can easily apply Xavier initialization using the torch.nn.init module.

import torch
import torch.nn as nn

# Define example layer dimensions
fan_in, fan_out = 512, 256 # Example dimensions for a linear layer

# Create a linear layer
linear_layer = nn.Linear(fan_in, fan_out, bias=True)

# --- Xavier Uniform Initialization ---
# Apply Xavier uniform initialization to the weights
nn.init.xavier_uniform_(linear_layer.weight)

# Typically initialize biases to zero
if linear_layer.bias is not None:
    nn.init.constant_(linear_layer.bias, 0)

print(f"Layer: Linear({fan_in}, {fan_out})")
print("Initialization: Xavier Uniform")
print(f"Weight Mean: {linear_layer.weight.mean():.4f}, Std: {linear_layer.weight.std():.4f}")
# Check theoretical standard deviation for uniform
# Var = (sqrt(6/(fan_in+fan_out)))^2 / 3 = 2 / (fan_in+fan_out)
# Std = sqrt(2 / (fan_in+fan_out))
theoretical_std_uniform = (2.0 / (fan_in + fan_out))**0.5
print(f"Theoretical Std Dev (Uniform based): {theoretical_std_uniform:.4f}\n")


# --- Xavier Normal Initialization ---
# Re-create the layer or re-initialize
linear_layer_norm = nn.Linear(fan_in, fan_out, bias=True)

# Apply Xavier normal initialization to the weights
nn.init.xavier_normal_(linear_layer_norm.weight)

# Initialize biases to zero
if linear_layer_norm.bias is not None:
    nn.init.constant_(linear_layer_norm.bias, 0)

print("Initialization: Xavier Normal")
print(
    f"Weight Mean: {linear_layer_norm.weight.mean():.4f}, "
    f"Std: {linear_layer_norm.weight.std():.4f}"
)
# Theoretical standard deviation for normal is sqrt(2 / (fan_in+fan_out))
theoretical_std_normal = (2.0 / (fan_in + fan_out))**0.5
print(f"Theoretical Std Dev (Normal): {theoretical_std_normal:.4f}")

Example output showing weight statistics after Xavier initialization. The actual standard deviation of the initialized weights should be close to the theoretical value derived from the fan-in and fan-out.

Xavier initialization provides a principled way to set initial weights, promoting more stable signal propagation in deep networks, particularly those using symmetric activation functions. While it represented a significant improvement over naive random initialization, its assumptions don't perfectly match all modern network architectures, especially those heavily reliant on ReLU activations. This leads us to the next technique, Kaiming initialization, designed specifically for such cases.