Weight Initialization: Xavier, He & Variance Preservation for Deep Networks

Training a neural network means finding good values for its weights, but where do we start? The answer matters more than you might expect. If we initialize weights poorly, gradients can explode to infinity or vanish to zero, making learning impossible before it even begins. The right initialization sets the stage for stable, efficient training.

This chapter explores why weight initialization matters and how techniques like Xavier and He initialization solve the problems that plagued early deep networks. We'll derive these methods from first principles, implement them from scratch, and see how they enable training networks that would otherwise fail to learn.

Consider what happens when we initialize all weights to zero. Each neuron in a layer receives the same input and produces the same output. During backpropagation, every neuron receives identical gradients. They all update identically, remaining the same as each other forever. The network has many neurons but effectively behaves as if it has one. This is the symmetry problem: if neurons start identical, they stay identical.

The following visualization contrasts zero initialization (where all weights look identical) with random initialization (where each weight has a distinct value):

We need randomness to break symmetry. But random initialization introduces a new challenge: how do we choose the scale of our random values? Too large, and activations explode. Too small, and gradients vanish.

Let's see this problem in action:

With scale 0.01, activations shrink to nearly zero. With scale 2.0, they saturate at the extremes of tanh ($\pm 1$). Neither is good for learning. Scale 0.1 and 1.0 preserve more reasonable activation magnitudes, but which is optimal?

The visualizations reveal two failure modes. With small initialization, activations progressively shrink layer by layer, eventually collapsing to near-zero values. The network loses the ability to distinguish different inputs. With large initialization, activations immediately saturate at the extremes of the activation function. Since the derivative of tanh near $\pm 1$ is close to zero, gradients vanish here too.

We need an initialization that preserves the variance of activations across layers, keeping them in the useful range where gradients flow effectively.

The experiments above reveal that initialization scale matters, but they don't tell us what the right scale is. To find it, we need to shift from empirical observation to mathematical analysis. The key question is: how does the variance of activations change as we pass through each layer?

If we can answer this question, we can work backward to determine what weight variance will preserve activation variance across layers. This is the central insight behind principled initialization: treat variance preservation as a design constraint, then solve for the weights that satisfy it.

Consider a single layer with $n_{\text{in}}$ input neurons and $n_{\text{out}}$ output neurons. Each output neuron computes a weighted sum of its inputs before applying the activation function. For a single output neuron, this pre-activation is:

where:

• $z$: the pre-activation value (before applying the activation function)
• $w_i$: the weight connecting input neuron $i$ to this output neuron
• $x_i$: the activation from input neuron $i$
• $n_{\text{in}}$: the number of input neurons (fan-in)

This equation is the starting point for our analysis. The pre-activation $z$ is a sum of many random terms, and we want to understand how its variance relates to the variances of the weights and inputs.

Assuming $w_i$ and $x_i$ are independent random variables with zero mean (a reasonable assumption when weights are initialized symmetrically around zero and inputs are centered), we can derive the variance of the pre-activation step by step.

Step 1: Variance of a sum. Since variance is additive for independent random variables:

Step 2: Variance of a product. For two independent, zero-mean random variables $A$ and $B$, we have $\text{Var}(AB) = \text{Var}(A) \cdot \text{Var}(B)$. This follows because $\mathbb{E}[AB] = \mathbb{E}[A]\mathbb{E}[B] = 0$ and $\mathbb{E}[(AB)^2] = \mathbb{E}[A^2]\mathbb{E}[B^2] = \text{Var}(A) \cdot \text{Var}(B)$.

Step 3: Combine. Assuming all weights have the same variance $\text{Var}(w)$ and all inputs have variance $\text{Var}(x)$:

where:

• $\text{Var}(z)$: variance of the pre-activation output
• $n_{\text{in}}$: number of input neurons (fan-in)
• $\text{Var}(w)$: variance of each weight (assumed identical)
• $\text{Var}(x)$: variance of each input activation (assumed identical)

This result tells us something useful: the output variance equals the input variance multiplied by both the number of inputs and the weight variance. It immediately suggests why naive initialization fails. With 256 input neurons and unit-variance weights, the output variance would be 256 times larger than the input. Each layer amplifies the signal by a factor of $n_{\text{in}}$, leading to exponential growth.

Now we can solve for the weight variance that preserves signal magnitude. To maintain stable variance across layers, we want $\text{Var}(z) = \text{Var}(x)$. Starting from our variance equation:

Setting $\text{Var}(z) = \text{Var}(x)$ and dividing both sides by $\text{Var}(x)$:

Solving for the weight variance:

This is the key insight: the weight variance should scale inversely with the number of input connections. Intuitively, when we sum $n_{\text{in}}$ random terms, the total variance grows proportionally to $n_{\text{in}}$. To counteract this growth and keep the output variance equal to the input variance, we must shrink each weight's variance by the factor $1/n_{\text{in}}$.

Theory is only useful if it matches reality. Let's verify our analysis by comparing different weight variances and observing how activation variance propagates through a deep network:

With weight variance of $1/n$, the output variance remains close to the input variance, confirming our analysis. The other choices cause variance to explode or collapse exponentially with depth.

Our variance analysis derived the optimal weight variance for the forward pass. But training a neural network involves both forward and backward propagation. Xavier Glorot and Yoshua Bengio, in their influential 2010 paper, recognized that we must consider both directions of signal flow.

Their insight was that gradients during backpropagation face the same variance amplification problem as activations during forward propagation, but in reverse. If we only optimize for forward variance, gradients might explode or vanish. The challenge is finding a weight variance that works reasonably well for both directions.

As we derived above, to maintain variance during forward propagation:

where $n_{\text{in}}$ is the fan-in (number of input connections to each neuron).

During backpropagation, gradients flow in the opposite direction. Each input neuron receives gradient contributions from all $n_{\text{out}}$ output neurons. By similar reasoning to the forward pass, to maintain gradient variance:

where $n_{\text{out}}$ is the fan-out (number of output neurons that receive input from each input neuron).

We now face a dilemma. The forward pass demands $\text{Var}(w) = 1/n_{\text{in}}$. The backward pass demands $\text{Var}(w) = 1/n_{\text{out}}$. These two requirements conflict unless $n_{\text{in}} = n_{\text{out}}$, which is rare in practice.

Since we cannot satisfy both simultaneously, Glorot and Bengio proposed a compromise that balances both requirements. Rather than favoring one direction over the other, they take the average of the two fan values:

where:

• $\text{Var}(w)$: the variance of each weight in the layer
• $n_{\text{in}}$: fan-in (input dimension)
• $n_{\text{out}}$: fan-out (output dimension)

This is Xavier initialization (also called Glorot initialization). To sample weights from this distribution, we need to convert the variance to distribution parameters.

For a uniform distribution $U[-a, a]$, the variance is $a^2/3$. Setting this equal to our target variance:

Solving for $a$:

So we sample:

For a normal distribution $\mathcal{N}(0, \sigma)$, the variance is $\sigma^2$. Setting this equal to our target variance:

Taking the square root:

So we sample:

With the formulas in hand, implementing Xavier initialization is straightforward. We need to sample weights from a distribution with the correct variance, choosing either uniform or normal based on preference.

The following visualization shows the resulting weight distributions for both uniform and normal variants:

Xavier initialization maintains activation variance much better than naive approaches. The standard deviation stays relatively stable across layers, preventing both vanishing and exploding activations.

Xavier initialization was a breakthrough for networks using tanh or sigmoid activations. But by the mid-2010s, ReLU (Rectified Linear Unit) had become the dominant activation function due to its simplicity and effectiveness in avoiding vanishing gradients. Unfortunately, Xavier initialization doesn't work well for ReLU networks.

The problem is that Xavier's derivation assumes the activation function is approximately linear around zero. For tanh and sigmoid, this is reasonable: near zero, they behave almost like the identity function. But ReLU is decidedly non-linear: it sets all negative values to zero, keeping only the positive half of the distribution.

Consider a pre-activation $z$ drawn from a symmetric distribution centered at zero with variance $\sigma^2$. ReLU sets all negative values to zero while keeping positive values unchanged. Since the distribution is symmetric, approximately half the values are negative, so half become zero.

The variance of the output can be computed as:

For a symmetric zero-mean distribution, the expected value of the positive half is small, and the key term is $\mathbb{E}[\text{ReLU}(z)^2]$. Since only positive values contribute, this equals approximately half of $\mathbb{E}[z^2] = \sigma^2$:

where:

• $\text{ReLU}(z) = \max(0, z)$: the rectified linear unit function
• $\sigma^2$: variance of the input pre-activation $z$

This variance halving has severe consequences for deep networks. After $L$ layers, the variance becomes approximately $\sigma^2 / 2^L$, causing activations to shrink exponentially. A 10-layer network would see activations reduced by a factor of $2^{10} = 1024$. A 20-layer network would see a reduction of over a million.

The following visualization makes this concrete by showing how ReLU transforms a symmetric input distribution:

Let's see this variance collapse in action by comparing Xavier initialization on tanh versus ReLU networks:

With tanh, Xavier initialization maintains activations reasonably well across 10 layers. With ReLU, however, activations decay dramatically. The decay ratio shows that ReLU activations shrink to a tiny fraction of their initial magnitude, confirming that Xavier is unsuitable for ReLU networks.

Kaiming He and colleagues addressed this problem in their 2015 paper by modifying the variance analysis to account for ReLU's behavior. The derivation follows the same logic as Xavier, but incorporates the factor of 1/2 from ReLU's variance reduction.

Starting from our original variance equation and incorporating the ReLU effect:

where:

• $\text{Var}(a)$: variance of the activation output (after ReLU)
• $\frac{1}{2}$: the variance reduction factor from ReLU
• $n_{\text{in}}$: fan-in (number of input connections)
• $\text{Var}(w)$: variance of the weights
• $\text{Var}(x)$: variance of the input activations

To maintain $\text{Var}(a) = \text{Var}(x)$, we set:

Dividing both sides by $\text{Var}(x)$ and solving for $\text{Var}(w)$:

This is He initialization (also called Kaiming initialization). The factor of 2 in the numerator compensates exactly for ReLU's variance-halving effect.

Just as with Xavier, we need to convert our target variance into the parameters of either a normal or uniform distribution.

For a normal distribution, we set the standard deviation to the square root of the target variance:

where $\mathcal{N}(\mu, \sigma)$ denotes a normal distribution with mean $\mu$ and standard deviation $\sigma$.

For a uniform distribution $U[-a, a]$, using the same derivation as Xavier (variance of uniform is $a^2/3$), we solve $a^2/3 = 2/n_{\text{in}}$ to get:

Notice that He initialization uses only $n_{\text{in}}$ (fan-in) rather than averaging $n_{\text{in}}$ and $n_{\text{out}}$ like Xavier. This is because He et al. found that matching forward pass variance was more important for training deep ReLU networks. The backward pass still works reasonably well because ReLU's simple gradient (either 0 or 1) doesn't introduce the same variance distortion as its forward computation.

Let's implement He initialization and compare it directly to Xavier on a ReLU network:

He initialization maintains much more stable activation magnitudes in ReLU networks compared to Xavier.

We've now seen how Xavier initialization works for symmetric activations and He initialization works for ReLU. But the deep learning toolkit includes many more activation functions, each with its own variance characteristics. The general principle remains the same: analyze how the activation affects variance and adjust the initialization accordingly.

The following table summarizes the recommended initialization for common activations:

For Leaky ReLU with negative slope $\alpha$, the situation is more nuanced than standard ReLU. Instead of zeroing negative inputs, Leaky ReLU scales them by a small factor $\alpha$ (typically 0.01 or 0.2). This preserves some information from the negative half of the distribution, meaning the variance reduction is less severe than standard ReLU.

The adjusted variance formula is:

where:

• $\alpha$: the negative slope of Leaky ReLU (typically 0.01 or 0.2)
• $n_{\text{in}}$: fan-in (number of input connections)
• $(1 + \alpha^2)$: a correction factor that accounts for variance from both positive inputs (coefficient 1) and negative inputs (coefficient $\alpha^2$)

The denominator $(1 + \alpha^2)$ arises because variance scales with the square of the coefficient. When $\alpha = 0$ (standard ReLU), the factor becomes $(1 + 0) = 1$, giving us $\text{Var}(w) = 2/n_{\text{in}}$, which is He initialization. When $\alpha = 1$ (linear activation), the factor becomes $(1 + 1) = 2$, giving $\text{Var}(w) = 1/n_{\text{in}}$, which approaches Xavier initialization for the fan-in-only case.

Rather than implementing a separate function for each activation, we can create a unified initialization function that computes the appropriate variance based on the activation type:

The ReLU initialization has the largest standard deviation because it uses the factor of $\sqrt{2}$ to compensate for variance halving. Tanh uses a gain of approximately 5/3 to account for its non-linearity. SELU uses a standard LeCun initialization with variance $1/n_{\text{in}}$, resulting in the smallest standard deviation.

So far, we've focused on how activations propagate forward. But neural network training depends equally on how gradients propagate backward. A good initialization must preserve gradient magnitudes during backpropagation; otherwise, learning signals won't reach the early layers.

Let's analyze gradient flow and verify that our initializations maintain stable gradients across network depth:

The gradient analysis confirms what we learned from activation analysis: Xavier initialization works well for tanh, maintaining reasonable gradient magnitudes across layers. For ReLU networks, He initialization provides more stable gradient flow.

Modern deep learning frameworks provide built-in initialization functions. Let's see how to use PyTorch's initialization utilities:

All initialization methods produce weights centered at zero, as expected. The He methods have slightly larger standard deviations than Xavier because they use only fan-in rather than averaging fan-in and fan-out. The uniform variants show tighter ranges compared to normal variants because the uniform distribution has bounded support.

Let's compare training dynamics with different initializations on a simple classification task:

The zero initialization fails to learn at all, stuck at the initial loss. Xavier and He initialization both enable learning, with similar convergence curves for this shallow network. The difference between Xavier and He becomes more pronounced in deeper networks with ReLU activations.

While we've focused on weight initialization, biases also need consideration. The common practice is to initialize biases to zero:

Zero biases are almost always a safe choice. Unlike weights, initializing biases to the same value doesn't cause symmetry problems because the gradients for biases depend on the (different) input-weight products.

For ReLU networks, some practitioners use small positive biases (like 0.01) to ensure neurons are active initially. This prevents "dead neurons" that never activate because their inputs are always negative. However, modern techniques like batch normalization and careful learning rate scheduling have reduced the importance of this trick.

The following visualization shows how initialization scale affects the fraction of dead neurons in a ReLU network:

When a network fails to train, initialization is a common culprit. Here are diagnostic techniques:

Key warning signs include:

• Activations near zero or saturated: Indicates vanishing activations
• Large gradient variance: May cause unstable training
• High percentage of dead neurons: Common in ReLU networks with poor initialization
• Extreme activation values: Risk of numerical overflow

While Xavier and He initialization remain the defaults, modern deep learning has introduced techniques that reduce sensitivity to initialization:

Batch normalization normalizes activations during training, which provides some robustness to initialization:

Common in transformers, layer normalization similarly helps with training stability:

For networks with skip connections, careful initialization of residual branches is important:

The scaling factor (often $1/\sqrt{n}$ where $n$ is the number of residual blocks) prevents the variance from growing linearly with depth.

Let's walk through initializing a small transformer-style language model:

The GPT-2 initialization uses a small standard deviation (0.02) for all linear layers. This relatively uniform approach works because layer normalization handles the variance normalization during training.

Weight initialization, while important, is not a magic solution to all training difficulties.

Initialization is just the starting point. A good initialization sets up favorable conditions for learning, but it cannot compensate for basic problems like inappropriate architectures, poor data quality, or incorrect hyperparameters. A network with He initialization will still fail if the learning rate is orders of magnitude too large.

Normalization techniques reduce sensitivity. Batch normalization and layer normalization have reduced the importance of precise initialization. Networks with these components can often train successfully with a wider range of initialization scales. This is one reason modern transformer architectures use a simple standard deviation of 0.02 rather than layer-specific calculations.

Very deep networks remain challenging. For networks with hundreds of layers (like some vision transformers), even careful initialization may not prevent training instabilities. Additional techniques like learning rate warmup, gradient clipping, and careful residual scaling become necessary.

Despite these limitations, understanding initialization principles remains valuable. When a network fails to train, checking initialization is often a productive first step. The mathematical framework we developed, analyzing how variance propagates through layers, provides insight into network behavior that extends beyond just initialization.

The following parameters control weight initialization behavior in PyTorch:

For nn.init.kaiming_normal_ and nn.init.kaiming_uniform_:

• Use mode='fan_in' (default) to preserve forward pass variance
• Use mode='fan_out' to preserve backward pass variance
• Set nonlinearity='relu' for standard ReLU or specify 'leaky_relu' with the a parameter

For nn.init.xavier_normal_ and nn.init.xavier_uniform_:

• These use fan_avg mode internally (no mode parameter)
• The gain parameter multiplies the standard deviation (default 1.0)
• Use gain=nn.init.calculate_gain('tanh') for tanh activations

Weight initialization determines whether a neural network can learn effectively from the start. The core principle is variance preservation: weights should be scaled so that activations and gradients maintain reasonable magnitudes as they propagate through layers.

Key takeaways:

• Zero initialization causes symmetry: All neurons learn the same thing, wasting network capacity
• Xavier initialization uses variance $\frac{2}{n_{\text{in}} + n_{\text{out}}}$, designed for tanh and sigmoid activations
• He initialization uses variance $\frac{2}{n_{\text{in}}}$, accounting for ReLU's variance-halving effect
• Biases are typically initialized to zero, which doesn't cause symmetry problems
• Modern architectures with batch or layer normalization are more robust to initialization, but still benefit from reasonable starting values
• Diagnostic tools can identify initialization problems by examining activation and gradient statistics

The next chapter covers batch normalization, a technique that normalizes activations during training and further reduces sensitivity to initialization while enabling training of even deeper networks.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about weight initialization in neural networks.