07. Weight Initialization Strategies (Xavier & He/Kaiming)
Preventing vanishing and exploding gradients through principled weight initialization: Xavier/Glorot, He/Kaiming, and variance preservation analysis.
Weight Initialization: Complete Notes (Beginner to Advanced)
Introduction#
Weight initialization is the process of assigning starting values to a network's weights before training begins. This might seem like a minor detail, but the choice of initial weight values has a major effect on whether a network trains successfully at all, especially as networks get deeper.
Why this matters so much: during the forward pass, the output of each layer becomes the input to the next layer. If the initial weights cause activations to shrink toward zero or grow toward infinity as they pass through each layer, the same problem compounds during the backward pass with gradients (since gradients depend on activations through the chain rule). A poor initialization strategy can cause vanishing or exploding gradients (covered in the Backpropagation notes) right from the very first training step, before the network has had any chance to learn anything.
Without a well-chosen weight initialization strategy: deep networks can fail to train entirely, since activations and gradients can vanish to near-zero or explode to very large values purely as a mathematical consequence of how the initial weights interact with matrix multiplication across many layers, regardless of how good the architecture or optimizer is.
With a well-chosen weight initialization strategy: activations and gradients maintain a stable, reasonable scale as they pass through the network's layers, giving the optimizer a good starting point from which it can actually make effective progress.
1. Random Initialization#
Random initialization means assigning weights small random values, typically drawn from a uniform or normal (Gaussian) distribution, rather than a fixed or zero value.
W = random values from a distribution (e.g., Normal(mean=0, std=some_small_value))
A very basic version:
🐍 PythonInteractive WebAssemblyimport numpy as np
# Random weights drawn from a standard normal distribution, scaled down
W = np.random.randn(input_size, output_size) * 0.01
Why the scaling factor (e.g., 0.01) matters: if weights are drawn from a distribution with too large a standard deviation, the weighted sums computed in the forward pass can become very large in magnitude. Combined with activation functions like sigmoid or tanh, this pushes activations into the "saturated" regions of those functions (where the function's output flattens out and its derivative approaches zero), causing vanishing gradients right from the start of training. If the standard deviation is too small, activations can shrink toward zero as they pass through each layer, again losing useful signal.
Without any randomness (all weights initialized to the exact same value): every neuron in a layer would compute the exact same output for the exact same input, and consequently receive the exact same gradient during backpropagation. This causes what's called the symmetry problem: all neurons in a layer update identically forever, meaning the layer effectively behaves as if it only had a single neuron, no matter how many neurons it actually contains. See Zero Initialization below for a concrete look at this exact failure.
With proper random initialization: each neuron starts with different weight values, so each neuron in a layer computes different outputs and receives different gradients, allowing the layer to actually learn a diverse and useful set of features rather than many redundant, identical copies of the same feature.
The remaining challenge with naive random initialization: simply picking small random numbers without accounting for the size of the layer (number of inputs/outputs) can still lead to activations and gradients that shrink or grow as the network gets deeper, since the variance of the weighted sum naturally scales with the number of inputs being summed. This is exactly the problem that Xavier/Glorot and He initialization (covered below) are specifically designed to solve.
2. Zero Initialization#
Zero initialization means setting every weight in the network to exactly 0 at the start of training.
🐍 PythonInteractive WebAssemblyimport numpy as np
W = np.zeros((input_size, output_size))
Why this fails completely, demonstrated concretely:
Consider a hidden layer with multiple neurons, all initialized with weight 0. During the forward pass, since every weight is 0, every neuron in that layer computes the exact same weighted sum (0, before adding bias), and therefore produces the exact same activation output, regardless of how many neurons the layer has.
During the backward pass, since every neuron produced the same output, every neuron also receives the exact same gradient with respect to its weights. This means every weight in that layer gets updated by the exact same amount in the exact same direction.
The result: after the update, all the weights in that layer are still identical to each other (just no longer zero). This pattern repeats at every single training step, forever. No matter how many neurons a layer has, they will always behave identically, as if the layer had collapsed down to a single effective neuron repeated multiple times.
Without breaking this symmetry (using zero, or any other identical-value initialization for every weight): a network with hidden layers becomes no more powerful than a network with just one neuron per layer, completely defeating the purpose of having multiple neurons, since none of them can specialize into detecting different features.
With random initialization instead of zero initialization: this symmetry is broken from the very first step, since different neurons start with different weights, compute different outputs, and receive different gradients, allowing them to specialize and learn distinct, useful features.
Important nuance: zero initialization for biases (as opposed to weights) is generally fine and is in fact the common default in most frameworks, since the weights are already randomly initialized and break the symmetry; the bias terms don't need randomness for this specific reason.
Code example (demonstrating the symmetry problem concretely):
🐍 PythonInteractive WebAssemblyimport numpy as np
def relu(z):
return np.maximum(0, z)
np.random.seed(0)
X = np.array([[1.0, 2.0, 3.0]]) # 1 sample, 3 features
# Zero-initialized weights for a hidden layer with 4 neurons
W_zero = np.zeros((3, 4))
b_zero = np.zeros((1, 4))
Z = np.dot(X, W_zero) + b_zero
A = relu(Z)
print("All neuron outputs identical (zero init):", A)
# Every value in A will be identical, since every neuron computed the same weighted sum
3. Xavier/Glorot Initialization#
Xavier initialization (also called Glorot initialization, named after its author Xavier Glorot) is a strategy specifically designed to keep the variance of activations roughly consistent as data passes forward through the network's layers, and correspondingly keep the variance of gradients roughly consistent as they pass backward. It was developed with sigmoid and tanh activation functions in mind.
The underlying reasoning: when you compute a weighted sum z = w1*x1 + w2*x2 + ... + wn*xn, the variance of z grows in proportion to the number of inputs n being summed (assuming the inputs and weights are independent random variables). If the weights aren't scaled down to compensate for this, activations tend to grow (or shrink) layer after layer as the network gets deeper, since each layer's output variance compounds with the next layer's.
Xavier initialization formulas:
For a normal distribution:
W ~ Normal(mean=0, variance = 2 / (n_in + n_out))
For a uniform distribution (the original formulation from Glorot's paper):
W ~ Uniform(-limit, +limit), where limit = sqrt(6 / (n_in + n_out))
Where:
n_in= number of input units to the layer (fan-in)n_out= number of output units from the layer (fan-out)
Why both n_in and n_out are used: Xavier initialization balances two goals simultaneously: keeping the variance of activations stable during the forward pass (which depends on n_in) and keeping the variance of gradients stable during the backward pass (which depends on n_out). Averaging both n_in and n_out gives a single scaling factor that reasonably balances both goals at once.
Without Xavier initialization (using a fixed, layer-size-independent scale like a plain Normal(0, 0.01) regardless of layer size): in a deep network with sigmoid/tanh activations, activations can progressively saturate (get pushed toward the flat regions of the sigmoid/tanh curve where the derivative is near zero) as they pass through many layers, since the naive fixed-scale initialization doesn't account for how variance compounds with layer size and depth. This directly leads to vanishing gradients starting from the very first training step.
With Xavier initialization: the scale of the initial weights is automatically adjusted based on each specific layer's number of inputs and outputs, keeping activation variance roughly stable across layers and helping sigmoid/tanh-based networks avoid immediate saturation and vanishing gradients from the start of training.
🐍 PythonInteractive WebAssemblyimport numpy as np
def xavier_init(n_in, n_out):
limit = np.sqrt(6 / (n_in + n_out))
return np.random.uniform(-limit, limit, size=(n_in, n_out))
W = xavier_init(n_in=784, n_out=128)
print("Weight shape:", W.shape)
print("Weight std:", W.std())
🐍 PythonInteractive WebAssembly# Using PyTorch's built-in Xavier/Glorot initialization
import torch
import torch.nn as nn
layer = nn.Linear(784, 128)
nn.init.xavier_uniform_(layer.weight) # uniform version
# or: nn.init.xavier_normal_(layer.weight) # normal distribution version
🐍 PythonInteractive WebAssembly# Using TensorFlow/Keras (Glorot is actually the DEFAULT initializer for Dense layers)
import tensorflow as tf
layer = tf.keras.layers.Dense(128, activation='tanh', kernel_initializer='glorot_uniform')
4. He Initialization#
He initialization (named after its author Kaiming He) is a variant designed specifically for ReLU and its variants (Leaky ReLU, ELU, GELU), addressing a specific mismatch between Xavier initialization's assumptions and how ReLU actually behaves.
Why Xavier initialization isn't ideal for ReLU: Xavier initialization's derivation assumes an activation function that is roughly linear and symmetric around zero (like tanh near the origin). ReLU, however, zeroes out all negative inputs entirely (ReLU(z) = max(0, z)), meaning roughly half of the neurons' pre-activations get set to exactly zero at any given layer, on average. This effectively halves the variance of the activations compared to what Xavier's formula assumes, so using Xavier's scaling with ReLU tends to make activations shrink layer after layer, gradually leading toward vanishing gradients as depth increases, though less severely than with a naive fixed-scale initialization.
He initialization formula:
For a normal distribution:
W ~ Normal(mean=0, variance = 2 / n_in)
For a uniform distribution:
W ~ Uniform(-limit, +limit), where limit = sqrt(6 / n_in)
Where n_in is the number of input units to the layer (fan-in only, unlike Xavier which uses both fan-in and fan-out).
Why the factor of 2 (compared to Xavier's more balanced formula) and why only n_in: He initialization compensates directly for the fact that ReLU zeroes out roughly half of its inputs, doubling the variance scale to counteract that reduction and keep the overall activation variance stable across layers. It uses only n_in (fan-in) because the primary goal is to preserve the variance of the forward pass activations specifically, which is what matters most for ReLU-based networks in practice.
Without He initialization (using Xavier initialization in a deep ReLU-based network instead): activation variance can still shrink progressively through the layers because Xavier's scale doesn't account for ReLU zeroing out negative values, contributing to a milder but still meaningful vanishing gradient effect in very deep ReLU networks.
With He initialization: the scale is specifically calibrated to compensate for ReLU's zeroing-out behavior, keeping activation variance much more stable across many layers, which is a major reason very deep ReLU-based networks (including architectures like ResNet, which can have over 100 layers) are able to train successfully at all.
🐍 PythonInteractive WebAssemblyimport numpy as np
def he_init(n_in, n_out):
std = np.sqrt(2 / n_in)
return np.random.randn(n_in, n_out) * std
W = he_init(n_in=784, n_out=128)
print("Weight shape:", W.shape)
print("Weight std:", W.std())
🐍 PythonInteractive WebAssembly# Using PyTorch's built-in He (Kaiming) initialization
import torch
import torch.nn as nn
layer = nn.Linear(784, 128)
nn.init.kaiming_normal_(layer.weight, nonlinearity='relu')
# or: nn.init.kaiming_uniform_(layer.weight, nonlinearity='relu')
🐍 PythonInteractive WebAssembly# Using TensorFlow/Keras
import tensorflow as tf
layer = tf.keras.layers.Dense(128, activation='relu', kernel_initializer='he_normal')
Xavier vs He: Quick Comparison#
| Aspect | Xavier/Glorot | He |
|---|---|---|
| Designed for | Sigmoid, Tanh | ReLU, Leaky ReLU, ELU, GELU |
| Variance formula (normal) | 2 / (n_in + n_out) | 2 / n_in |
| Accounts for | Symmetric, roughly-linear activations | ReLU zeroing out ~half of pre-activations |
| Uses fan-out? | Yes (balances forward and backward pass) | No (fan-in only) |
5. Weight Initialization and Gradient Flow#
This section ties weight initialization directly back to the gradient flow concepts (vanishing and exploding gradients) covered in the Backpropagation notes, making the connection explicit.
The core relationship: during the forward pass, each layer's output variance depends on the incoming activations' variance combined with the layer's weight variance and the number of inputs being summed. During the backward pass, by the chain rule, each layer's gradient is the product of the local derivatives (which depend on the activation function and its input's scale) and the weights, propagated backward layer by layer. If the initial weight scale is wrong for a given depth and activation function, both the forward-pass activations and the backward-pass gradients can compound multiplicatively across layers, either shrinking toward zero (vanishing) or growing without bound (exploding), purely due to the initialization choice, independent of anything the optimizer does later.
Concretely, how poor initialization causes each problem:
-
Weights initialized too small (variance too low) relative to the network's depth and activation function: activations shrink layer after layer during the forward pass. During backpropagation, the local derivatives being multiplied together (via the chain rule) are correspondingly small, so gradients shrink exponentially with depth, contributing to vanishing gradients, particularly severe when combined with sigmoid/tanh activations that already have small maximum derivatives.
-
Weights initialized too large (variance too high) relative to the network's depth: activations grow layer after layer during the forward pass (and can also push sigmoid/tanh into their saturated, near-zero-derivative regions). During backpropagation, when weights are large, the products in the chain rule compound to become very large as well, contributing to exploding gradients, where the loss can become unstable or turn into
NaN.
Why He and Xavier initialization directly address this: by deliberately scaling the initial weight variance based on the number of inputs and/or outputs of each specific layer (rather than using a single fixed scale everywhere), these methods are mathematically derived to keep the variance of activations, and consequently the variance of gradients, approximately constant as data and gradients pass through each layer, regardless of how many layers the network has. This doesn't eliminate vanishing/exploding gradients entirely on its own for extremely deep networks (which is why techniques like Batch Normalization and residual/skip connections were later developed as additional, complementary solutions), but proper initialization is the essential first line of defense and makes a substantial practical difference, especially for moderately deep networks.
Without matching your initialization strategy to your activation function (e.g., using Xavier initialization with ReLU activations, or He initialization with sigmoid/tanh activations): you lose the specific variance-preserving guarantee each method was mathematically derived for, since Xavier assumes activations that don't zero out half their inputs like ReLU does, while He's stronger scaling factor is calibrated specifically for that zeroing-out behavior and would be an unnecessarily strong, mismatched scale for sigmoid/tanh.
With matching your initialization strategy to your activation function (Xavier/Glorot for sigmoid/tanh, He for ReLU-family activations): you give your network the best possible starting point for stable gradient flow across its layers from the very first training step, directly reducing the risk of vanishing or exploding gradients purely due to a poor initial weight scale.
Code example (demonstrating the practical effect of matching initialization to activation function):
🐍 PythonInteractive WebAssemblyimport torch
import torch.nn as nn
torch.manual_seed(0)
def build_and_check_gradient_flow(activation_fn, init_fn, num_layers=15):
layers = []
for _ in range(num_layers):
linear = nn.Linear(64, 64)
init_fn(linear.weight) # apply the chosen initialization strategy
layers.append(linear)
layers.append(activation_fn())
model = nn.Sequential(*layers)
X = torch.randn(1, 64)
y_true = torch.randn(1, 64)
output = model(X)
loss = nn.MSELoss()(output, y_true)
loss.backward()
first_layer_grad = model[0].weight.grad.abs().mean().item()
return first_layer_grad
# ReLU network with MATCHING He initialization
grad_relu_he = build_and_check_gradient_flow(
activation_fn=nn.ReLU,
init_fn=lambda w: nn.init.kaiming_normal_(w, nonlinearity='relu')
)
# ReLU network with MISMATCHED Xavier initialization (for comparison)
grad_relu_xavier = build_and_check_gradient_flow(
activation_fn=nn.ReLU,
init_fn=nn.init.xavier_normal_
)
print("Gradient magnitude, first layer, ReLU + He init (matched):", grad_relu_he)
print("Gradient magnitude, first layer, ReLU + Xavier init (mismatched):", grad_relu_xavier)
# The mismatched combination typically shows a smaller first-layer gradient,
# reflecting the initial stages of a vanishing gradient trend as depth increases
Summary Table: Initialization Strategies#
| Strategy | Formula (Normal Dist.) | Best Suited For | Key Risk If Misused |
|---|---|---|---|
| Zero Initialization | All weights = 0 | Never suitable for weights (biases only) | Symmetry problem; network can't learn diverse features |
| Naive Random Initialization | Small fixed-scale random values (e.g., N(0, 0.01)) | Breaking symmetry only; not depth-aware | Activations/gradients shrink or grow with depth |
| Xavier/Glorot Initialization | variance = 2 / (n_in + n_out) | Sigmoid, Tanh | Suboptimal (activations may shrink) if used with ReLU |
| He Initialization | variance = 2 / n_in | ReLU, Leaky ReLU, ELU, GELU | Unnecessarily large scale if used with Sigmoid/Tanh |
Quick Recap (Beginner to Advanced Flow)#
- Random initialization (as opposed to using identical values for every weight) is necessary to break symmetry between neurons so they can each learn distinct features.
- Zero initialization is a specific failure case of non-random initialization: every neuron in a layer computes identically and updates identically forever, effectively collapsing the layer to one useful neuron regardless of its actual size.
- Xavier/Glorot initialization scales the initial weight variance based on both the number of inputs and outputs of a layer, specifically designed to keep activations and gradients stable for sigmoid and tanh activations.
- He initialization is a variant that scales weight variance based only on the number of inputs, using a larger scaling factor to compensate for ReLU-family activations zeroing out roughly half their inputs.
- Weight initialization directly affects gradient flow: a poorly-scaled initialization can cause vanishing or exploding gradients purely from how the initial weight variance compounds across layers during both the forward and backward pass, making the choice of initialization strategy (and matching it to your activation function) an essential first step in successfully training deep networks.
07. Weight Initialization Checkpoint
Finished studying this notebook?
Mark this guide as completed to update your course progress roadmap.