Source: Difficulty of training deep neural networks by Glorot et. al.
Introduction
Xavier Initialization, also known as Glorot Initialization, is a weight initialization technique for deep learning models that was introduced by Xavier Glorot and Yoshua Bengio in their 2010 paper “Understanding the difficulty of training deep feedforward neural networks”.
It aims to maintain the variance of activations and gradients across layers in deep neural networks and prevents the vanishing or exploding gradient problem during training.
Why Do We Need Xavier Initialization?
When weights are initialized poorly (e.g., too large or too small), it can lead to:
Vanishing Gradients (approaching zero)：
- In deep neural networks, gradients can become extremely small as they are backpropagated through many layers. This causes the early layers to learn very slowly or stop learning altogether, since their weight updates become negligible. The weight update rule during training is typically:
- where $w$ is the weight, $\eta$ is the learning rate, and $\frac{\partial L}{\partial w}$ is the gradient of the loss with respect to the weight. When the gradient approaches zero, the update term becomes negligible, and the weight essentially remains unchanged.
Exploding Gradients (becoming too large)
- When gradients grow exponentially during backpropagation, leading to extremely large weight updates. This can make the model unstable, causing the loss to oscillate or even diverge.
Activation Saturation:
- Saturation occurs when activation functions like sigmoid or tanh receive inputs that are too large (positive or negative), causing their outputs to approach their maximum or minimum values (e.g., sigmoid approaches 0 or 1, tanh approaches -1 or 1).
- In saturated regions, the gradients become very small (close to zero) because the derivative of the activation function flattens out. For example:
- Sigmoid: when the input is very large/small, the derivative approaches 0
Source: Sigmoid function
- Tanh: similar behavior, derivatives vanish at extreme values
Source: Tanh function
- This leads to the vanishing gradient problem, where backpropagated gradients become negligible, effectively stopping learning in deeper layers.
These issues hinder effective training, especially in deep neural networks. Xavier Initialization addresses this by keeping the variance of the weights balanced across layers, helping gradients stay within a reasonable range and preventing activations from entering saturation regions.
Xavier Initialization Formula
Xavier Initialization sets weights using a uniform or normal distribution such that:
where:
- $n_{\text{in}}$: number of input units to the layer
- $n_{\text{out}}$: number of output units from the layer
Motivation behind Xavier Initialization
In deep networks, each neuron’s output becomes the input to the next layer. If the output values grow too large or shrink too small, the network becomes unstable.
To ensure stability, Xavier Initialization aims to keep the variance of activations roughly constant across layers. That means:
where $z = \sum w_i x_i$ is the output of a neuron, and $x$ is the input. Assuming $x_i$ and $w_i$ are independent and have a zero-mean:
Here n is the number of input connections to the neuron (fan-in) — or the, i.e., how many w_i x_i terms are being summed. If we want Var(z) to be same as Var(x), Var(w) needs to be:
This works for forward propagation. Similarly, for backward propagation, where gradients flow in reverse, the variance condition implies:
To balance both, Xavier Initialization uses the average:
This variance ensures that the signal magnitude remains stable from layer to layer during both forward and backward passes. If the variance is too high, the network might explode; if too low, it might vanish. Xavier keeps the flow of information smooth.
Importantly, the initialization is computed per layer, using the layer’s fan-in and/or fan-out — i.e., the number of input and output connections to determine the appropriate scale for the weights.
Code Example:
The following example demonstrates how to apply Xavier Initialization in a simple multilayer perceptron (MLP) using PyTorch. This model consists of three fully connected (Linear) layers. The weights of each layer are initialized using nn.init.xavier_uniform_, which ensures that the variance of activations is preserved across layers.
import torch
import torch.nn as nn
import torch.nn.functional as F
class SimpleMLP(nn.Module):
def __init__(self):
super(SimpleMLP, self).__init__()
self.fc1 = nn.Linear(784, 256)
self.fc2 = nn.Linear(256, 128)
self.fc3 = nn.Linear(128, 10)
# Apply Xavier Initialization
nn.init.xavier_uniform_(self.fc1.weight)
nn.init.xavier_uniform_(self.fc2.weight)
nn.init.xavier_uniform_(self.fc3.weight)
def forward(self, x):
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
return self.fc3(x)
# Example usage
model = SimpleMLP()Note:
- nn.init.xavier_uniform_() uses a uniform distribution
- You can also use nn.init.xavier_normal_() for a normal distribution. While both work well in practice, uniform is sometimes preferred when avoiding outliers is important, whereas normal may introduce more variance due to its heavier tails.
When to Use Xavier Initialization?
Xavier Initialization is most appropriate when:
- You are using activation functions like tanh or sigmoid: These functions are symmetric around zero and are prone to saturation (i.e., gradients becoming very small), especially when inputs are too large or too small. Xavier Initialization helps maintain a balanced variance of activations across layers, reducing the chance of vanishing or exploding gradients.
- The number of input and output units are roughly equal: Xavier assumes a balance between the fan-in (number of input neurons) and fan-out (number of output neurons). When this holds, it helps stabilize training by keeping signal variance consistent across layers.
Avoid Xavier Initialization when:
- You’re using ReLU or variants (Leaky ReLU, GELU, etc.):
ReLU activation zeroes out negative values, breaking the zero-mean assumption made by Xavier. In such cases, He (Kaiming) Initialization is more suitable, as it accounts for the asymmetric behavior of ReLU.
- Your architecture involves layers with very different fan-in and fan-out:
Since Xavier Initialization averages fan-in and fan-out in its variance calculation, layers with high mismatch might benefit from more tailored schemes (like using only fan-in in He initialization).
Video Explanation
- Andrew Ng explains Xavier Initialization in this video using a single-neuron example with the tanh activation function, providing clear intuition behind its design and effectiveness (Runtime: 6 mins).
- In the video, Sebastian Raschka discusses Xavier Initialization and the rationale behind it, providing insights into its mathematical foundation and practical benefits (Runtime: 12 mins)
Related Articles: