Weight Initialization: Xavier, He & Variance Preservation

Training a neural network means finding good values for its weights, but where do you start? The answer matters more than you might expect. If you initialize weights poorly, activations can explode to infinity or vanish to zero before a single gradient update is applied, 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, see how modern architectures like BERT and GPT approach initialization, and compare empirically how each method performs in practice. By the end, you'll have the conceptual foundations to reason about initialization decisions for any architecture you encounter, and you'll understand precisely when the choice of initialization scheme makes the difference between a network that trains and one that doesn't.

Before principled initialization schemes existed, training deep neural networks was largely a matter of craft and intuition. Practitioners in the early 2000s knew that random initialization was necessary for symmetry breaking, but the question of how much randomness and what scale remained largely unanswered by theory. Researchers would manually tune initialization hyperparameters, sometimes spending significant compute searching for settings that produced reasonable gradient magnitudes in the first few iterations.

This difficulty was one of the primary reasons deep networks were considered impractical for much of the 1990s and early 2000s. Shallow networks with one or two hidden layers could be trained reliably, but adding more layers introduced instabilities that were difficult to diagnose. When a 10-layer network failed to learn, the culprit could be any combination of bad initialization, inappropriate learning rates, the wrong activation function, or all three simultaneously.

The field began to change with a pair of foundational papers. In 2010, Xavier Glorot and Yoshua Bengio published "Understanding the Difficulty of Training Deep Feedforward Neural Networks," which provided the first rigorous analysis of how activation variance propagates through layers and derived the initialization scheme now known as Xavier or Glorot initialization. Five years later, Kaiming He and colleagues extended this analysis to ReLU networks in "Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification," introducing He initialization and demonstrating that it enabled training of networks much deeper than had previously been practical. Together, these papers transformed initialization from an empirical art into a theoretically grounded engineering discipline.

Understanding what these authors figured out, and why it took the field so long to get there, gives the derivations that follow their proper context. The math is not difficult, but the framing is everything.

The most intuitive initialization strategy is to set all weights to the same constant value. Zero seems like a natural, neutral starting point. But this approach is fundamentally broken.

Consider what happens when all weights in a layer are zero. Every neuron receives the same input (zero dot product) and produces the same output. During backpropagation, every neuron also receives identical gradients and applies identical updates. No matter how long you train, every neuron in the layer remains identical to every other. A 512-neuron layer with zero initialization behaves exactly like a 1-neuron layer.

This is the symmetry problem: neurons that start identical stay identical. The network's representational capacity collapses. Any initialization scheme that makes all weights in a layer the same value, whether zero, one, or any constant, triggers this failure mode.

The intuition here is worth dwelling on. Imagine a network designed to recognize different types of objects in images. You want early neurons to detect different low-level features: edges at different orientations, color contrasts, texture patterns. If all neurons start the same and receive the same gradient, they can only ever develop the same feature detector, no matter how many neurons you use or how long you train. The entire point of a wide network is to represent many features in parallel, and this capacity is completely destroyed by symmetric initialization.

The fix is randomness. If we initialize each weight independently from a random distribution, neurons start with different values, receive different gradient signals, and diverge to learn different features. But random initialization introduces a new challenge: what scale should those random values be?

Suppose we fix the symmetry problem by using random initialization. We draw weights from a standard normal distribution $\mathcal{N}(0, 1)$. At first this seems reasonable, but a deep network reveals a severe problem.

In a layer with $n_{\text{in}}$ inputs, each pre-activation sums $n_{\text{in}}$ weighted inputs. If each weight has variance 1 and each input has variance 1, then the pre-activation has variance $n_{\text{in}}$. Its standard deviation is $\sqrt{n_{\text{in}}}$. For a layer with 256 inputs, every layer amplifies the signal by a factor of $\sqrt{256} = 16$.

After 10 layers, the signal has grown by a factor of $16^{10} \approx 10^{12}$. This is exploding activations: values become astronomically large, causing numerical overflow or saturating activation functions.

The opposite problem occurs with too-small initialization. If weights are drawn from $\mathcal{N}(0, 0.01^2)$, each layer reduces the signal by a factor of roughly $0.01 \cdot \sqrt{n_{\text{in}}}$. For 256 inputs this is $0.01 \times 16 = 0.16$. After 10 layers, the signal shrinks to $0.16^{10} \approx 10^{-8}$. This is vanishing activations: the network loses the ability to distinguish inputs, and gradients vanish with the activations.

Both failure modes prevent learning. The goal is to find an initialization scale that keeps activations in a reasonable range throughout the network, regardless of depth.

The gradient side of the picture is equally important and often more immediately damaging to training. When activations vanish, the gradients that flow backward during backpropagation vanish with them. This happens because the gradient of the loss with respect to a weight in an early layer requires multiplying many terms together through the chain rule, and if any activation is near zero, its contribution to the gradient chain is near zero. This is the vanishing gradient problem, and it is what made networks with more than a few layers so notoriously difficult to train before principled initialization and techniques like batch normalization were introduced.

When activations explode, the gradients can explode too, driving weight updates to extreme values and destabilizing the optimization. In practice, activation saturation with sigmoid or tanh is another failure mode: once activations reach the flat regions of these functions, their derivatives are essentially zero, and gradients stop flowing regardless of whether the pre-activations have exploded or not.

The violin plots make the failure modes visible. Small initialization compresses all activations toward zero layer by layer. Large initialization immediately saturates them into two spikes near $\pm 1$. We need a principled way to find the scale that keeps activations in the useful middle range.

The experiments reveal the symptoms, but to find the cure we need to understand the cause mathematically. The key question is: how does the variance of activations change as they pass through each layer?

If we can answer this question, we can work backward to determine what weight variance preserves activation variance across layers. This is the central insight behind principled initialization: treat variance preservation as a design constraint, then solve for the weight scale that satisfies it. The beauty of this approach is that it converts a vague intuition ("weights shouldn't be too big or too small") into a concrete formula with a clear derivation.

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

where:

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

The pre-activation $z$ is a sum of $n_{\text{in}}$ random terms. We want to understand how its variance relates to the variances of the weights and inputs.

To derive how variance propagates, we make two assumptions that hold approximately during the early stages of training: weights $w_i$ and inputs $x_i$ are independent random variables with zero mean.

Step 1: Variance is additive for independent random variables. Because the $w_i x_i$ terms are independent of each other (each uses a different weight-input pair):

Step 2: Variance of a product of two independent zero-mean variables. For any two independent, zero-mean random variables $A$ and $B$:

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

where:

• $n_{\text{in}}$: the number of input neurons (fan-in)
• $\text{Var}(w)$: the variance of each weight, assumed identical across all weights in the layer
• $\text{Var}(x)$: the variance of each input activation, assumed identical across all inputs

This formula immediately explains the explosion problem. With $n_{\text{in}} = 256$ and unit-variance weights ($\text{Var}(w) = 1$), the output variance is $256 \times \text{Var}(x)$. Each layer multiplies variance by $n_{\text{in}}$, producing exponential growth.

The assumptions we made (independence, zero mean) are strictly true only at initialization, before any training has occurred. As training proceeds, weights and activations become correlated. However, the variance relationships remain approximately valid throughout early training, which is precisely the period when initialization matters most. By the time the network has trained for many epochs, the optimization landscape has been significantly reshaped and initialization conditions matter much less.

To maintain stable variance across layers, we want $\text{Var}(z) = \text{Var}(x)$. Starting from:

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

Solving for the weight variance:

The weight variance should scale inversely with the number of input connections. Summing $n_{\text{in}}$ independent random terms amplifies variance by $n_{\text{in}}$, so each weight's variance must be $1/n_{\text{in}}$ to counteract this amplification.

This result is sometimes called LeCun initialization, after Yann LeCun who described it in his 1998 work on convolutional networks. It is the optimal forward-pass initialization for networks with linear or nearly-linear activation functions, and it forms the theoretical foundation for both Xavier and He initialization.

Our variance analysis derived the optimal weight variance for forward propagation. But training a neural network involves both forward and backward propagation, and gradients face the same variance amplification problem in reverse.

Xavier Glorot and Yoshua Bengio, in their 2010 paper "Understanding the Difficulty of Training Deep Feedforward Neural Networks," recognized that we must consider both signal directions simultaneously. If we optimize only for forward variance, gradients might explode or vanish during backpropagation. The key insight was that a good initialization must care about the health of the gradient signal just as much as the health of the forward activations, because training depends on both flowing cleanly through the entire network.

The forward pass constraint from our derivation above gives:

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 the same reasoning applied to the gradient equations, maintaining gradient variance during the backward pass requires:

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

To see why the backward pass requirement involves $n_{\text{out}}$, recall the chain rule during backpropagation. The gradient of the loss with respect to an input activation $x_i$ is:

This is a weighted sum of $n_{\text{out}}$ gradient terms, structurally identical to the forward pass equation but with fan-out replacing fan-in. The same variance analysis applies, giving the backward requirement $\text{Var}(w) = 1/n_{\text{out}}$.

We face a dilemma. The forward pass requires $\text{Var}(w) = 1/n_{\text{in}}$, but the backward pass requires $\text{Var}(w) = 1/n_{\text{out}}$. These two constraints conflict whenever $n_{\text{in}} \neq n_{\text{out}}$, which is the common case for layers that change dimensionality.

Glorot and Bengio proposed taking the harmonic mean of the two fan values:

where:

• $n_{\text{in}}$: fan-in (input dimension of the weight matrix)
• $n_{\text{out}}$: fan-out (output dimension of the weight matrix)
• $\frac{2}{n_{\text{in}} + n_{\text{out}}}$: a compromise variance that partially satisfies both the forward and backward variance preservation constraints

Notice that when $n_{\text{in}} = n_{\text{out}}$, this formula reduces to $1/n_{\text{in}}$, recovering the LeCun result. For layers where the input and output dimensions differ substantially, the formula takes a middle path between the two competing requirements. This is an elegant engineering compromise: it is not optimal for either direction alone, but it avoids catastrophically failing in either direction.

This is Xavier initialization (also called Glorot initialization). To sample weights, we convert this target variance to the parameters of either a normal or uniform distribution.

For a normal distribution $\mathcal{N}(0, \sigma^2)$, the variance is $\sigma^2$. Setting $\sigma^2 = 2/(n_{\text{in}} + n_{\text{out}})$ and taking the square root gives the standard deviation:

Weights are drawn from:

For a uniform distribution $U[-a, a]$, the variance is $a^2/3$. Setting $a^2/3 = 2/(n_{\text{in}} + n_{\text{out}})$ and solving for $a$:

Weights are drawn from:

In practice, the uniform variant is the more commonly referenced form, particularly in the original Glorot and Bengio paper. Both versions achieve similar results in practice because the variance formula is the same; only the shape of the distribution differs. The uniform variant is bounded, so it avoids the (very small) probability of drawing a very large outlier weight from the tails of a normal distribution.

The derivation assumes the activation function does not distort the variance of the signal. Tanh and sigmoid satisfy this assumption near zero, where they are both approximately linear. Specifically:

• $\tanh(z) \approx z$ for small $|z|$
• $\sigma(z) \approx 0.25 z + 0.5$ for small $|z|$, with derivative $\approx 0.25$

For tanh, the derivative near zero is approximately 1, meaning the function barely affects variance. For sigmoid, the derivative near zero is about 0.25, which does systematically reduce variance. This is why PyTorch's implementation includes an empirical correction factor of $5/3$ for tanh (accounting for the actual saturation properties) and why some implementations of Xavier for sigmoid include additional adjustment. The original Glorot and Bengio paper noted that sigmoid is more problematic than tanh, because its non-zero mean output shifts the input distribution seen by subsequent layers.

Xavier initialization was a breakthrough for networks with tanh or sigmoid activations. But by the mid-2010s, ReLU (Rectified Linear Unit) had become the dominant activation function due to its simplicity and its resistance to the vanishing gradient problem. Unfortunately, Xavier initialization performs poorly in ReLU networks.

The problem is fundamental. Xavier's derivation assumes the activation function is approximately linear near zero, so it does not systematically distort variance. Tanh and sigmoid satisfy this: near zero, both behave approximately like the identity function. ReLU does not. It zeros all negative inputs, keeping only the positive half of the distribution. This asymmetry means ReLU is far from linear around zero, and the linear approximation underlying Xavier's derivation breaks down completely.

The practical consequence is severe. A deep ReLU network initialized with Xavier will see its activations slowly collapse toward zero as the depth increases, because each ReLU layer systematically discards half the information by zeroing all negative values. The network can still learn, but it starts from a much worse position than necessary.

Consider a pre-activation $z$ drawn from a symmetric distribution centered at zero with variance $\sigma^2$. ReLU applies $\text{ReLU}(z) = \max(0, z)$, zeroing all negative values while keeping positive values unchanged. Since the distribution is symmetric around zero, approximately half the values are negative and become zero.

The variance of the ReLU output is:

where:

• $\text{ReLU}(z) = \max(0, z)$: the rectified linear unit function that passes positive inputs and zeros negative ones
• $\sigma^2$: the variance of the input pre-activation $z$ before applying ReLU
• $\frac{1}{2}$: the variance reduction factor, because ReLU eliminates roughly half the distribution (all values below zero)

To see this more carefully, consider the expected squared output of ReLU applied to $z \sim \mathcal{N}(0, \sigma^2)$:

The indicator $\mathbf{1}_{z > 0}$ is 1 when $z$ is positive and 0 when negative. Since $z$ has zero mean, $\mathbb{E}[\text{ReLU}(z)] = \mathbb{E}[z \cdot \mathbf{1}_{z > 0}] = \sigma / \sqrt{2\pi}$, which is nonzero. The variance therefore equals $\mathbb{E}[\text{ReLU}(z)^2] - (\mathbb{E}[\text{ReLU}(z)])^2$. For large networks where the number of neurons is large, the mean contribution to variance is much smaller than $\sigma^2/2$, and the approximation $\text{Var}(\text{ReLU}(z)) \approx \sigma^2/2$ holds well.

This variance halving has severe consequences for deep networks. After $L$ ReLU layers, variance shrinks by a factor of $2^L$. A 10-layer ReLU network reduces activation variance by $2^{10} = 1024$. A 20-layer network reduces it by over a million. Xavier initialization, calibrated assuming no such reduction, produces vanishing activations in deep ReLU networks.

Kaiming He and colleagues addressed this in their 2015 paper "Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification." Their derivation modifies the Xavier variance equation to include ReLU's variance reduction factor.

Starting from the variance propagation equation with the ReLU correction:

where:

• $\text{Var}(a)$: variance of the activation output, that is, the distribution after applying ReLU
• $\frac{1}{2}$: the variance reduction factor from ReLU zeroing all negative pre-activation values
• $n_{\text{in}}$: fan-in (number of input connections)
• $\text{Var}(w)$: variance of the weights in this layer
• $\text{Var}(x)$: variance of the input activations from the previous layer

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

Dividing both sides by $\text{Var}(x)$:

Solving for the weight variance:

This is He initialization (also called Kaiming initialization). The factor of 2 compensates precisely for ReLU's variance-halving effect. Notice that He uses only $n_{\text{in}}$ rather than averaging $n_{\text{in}}$ and $n_{\text{out}}$ as Xavier does. He et al. found that matching forward pass variance was the more critical constraint for deep ReLU networks in practice, and empirical results confirmed that omitting the fan-out term worked better for the very deep convolutional networks they were studying.

Converting to distribution parameters:

Normal distribution:

where $\sqrt{2/n_{\text{in}}}$ is the standard deviation (the square root of the target variance).

Uniform distribution: Using the same derivation as for Xavier, where the variance of $U[-a, a]$ equals $a^2/3$, we solve $a^2/3 = 2/n_{\text{in}}$ to obtain $a = \sqrt{6/n_{\text{in}}}$:

We've seen Xavier for symmetric activations and He for ReLU. The deep learning toolkit includes many more activations, each with its own variance characteristics. The general principle remains: analyze how the activation distorts variance and adjust the initialization accordingly. This generalization is what PyTorch's nn.init.calculate_gain() function systematizes. This provides gain factors for each activation that can be plugged into the LeCun base formula.

The following table summarizes the recommended initialization for common activations:

The SELU activation (Scaled Exponential Linear Unit) is a special case that deserves its own mention. SELU is designed to be self-normalizing: under the right conditions, activations in a SELU network converge to a fixed-point distribution with mean zero and unit variance regardless of network depth. The activation function itself applies a specific scaling and shift to maintain this fixed-point property.

For SELU to work correctly, the weights must be initialized with LeCun initialization, $\text{Var}(w) = 1/n_{\text{in}}$. This is the simplest of all the initialization schemes and the theoretical precursor to Xavier. The SELU paper by Klambauer et al. (2017) provides rigorous proofs that the combination of LeCun initialization, SELU activations, and specific rescaling factors produces self-normalizing behavior, eliminating the need for batch normalization in networks designed around this activation function.

For Leaky ReLU with negative slope $\alpha$ (typically 0.01 or 0.2), negative inputs are scaled by $\alpha$ rather than zeroed. This preserves some variance from the negative half of the distribution. The corrected variance formula is:

where:

• $\alpha$: the negative slope of Leaky ReLU (the multiplier applied to negative pre-activations)
• $(1 + \alpha^2)$: a correction factor accounting for variance from both positive inputs (coefficient 1, so contribution is $1^2 = 1$) and negative inputs (coefficient $\alpha$, so contribution is $\alpha^2$)
• $n_{\text{in}}$: fan-in (number of input connections)

When $\alpha = 0$ (standard ReLU), the factor $(1 + 0) = 1$ recovers He initialization. When $\alpha = 1$ (linear activation), the factor $(1 + 1) = 2$ gives $\text{Var}(w) = 1/n_{\text{in}}$, approaching LeCun initialization for the forward-only case. For typical Leaky ReLU with $\alpha = 0.01$, the correction factor is barely distinguishable from He initialization, which is why many practitioners simply use He for Leaky ReLU without the explicit correction.

GELU (Gaussian Error Linear Unit), used in BERT and GPT, applies a smooth, probabilistic gating that scales each input by the probability that a standard normal random variable is less than the input. Like ReLU, it zeros or nearly zeros a fraction of inputs, reducing variance by roughly half. He initialization works well for GELU networks in practice, as the variance reduction is similar in magnitude to standard ReLU.

Swish, another smooth activation used in some architectures, has a similar profile to GELU. Both can be approximated as having a gain factor of $\sqrt{2}$, identical to ReLU, making He initialization a reasonable default. The key insight is that any activation function that zeros or heavily suppresses roughly half its inputs requires the $2/n_{\text{in}}$ variance formula to compensate, while activations that are closer to linear around zero (tanh, sigmoid) do not need this correction factor.

ReLU and GELU initializations have the largest standard deviation because they use the $\sqrt{2}$ gain factor to compensate for variance halving. Tanh uses a gain of approximately 5/3. SELU uses LeCun initialization with a gain of 1, giving variance $1/n_{\text{in}}$.

Xavier and He initialization focus on setting the right variance. Orthogonal initialization takes a different approach: it constructs weight matrices whose columns are perpendicular unit vectors, a property called orthogonality.

An orthogonal matrix $W$ satisfies $W^T W = I$ (the identity matrix). This means the transformation preserves the length of vectors: $\|Wx\|_2 = \|x\|_2$ for any input vector $x$. In terms of variance, orthogonal initialization guarantees exact variance preservation during the forward pass, with no dependence on the input distribution or layer size.

The construction uses singular value decomposition (SVD). Starting from a random matrix $A$ with entries drawn from $\mathcal{N}(0, 1)$, we compute the SVD $A = U \Sigma V^T$ and use $U$ (or $V^T$) as the initial weight matrix. Both $U$ and $V$ are unitary matrices satisfying the orthogonality condition.

Beyond variance preservation, orthogonal matrices have an important property: their eigenvalues all have magnitude exactly 1. For a general matrix $W$, the eigenvalues can be larger or smaller than 1 in magnitude. When a matrix with eigenvalues much larger than 1 is applied repeatedly (as in a recurrent network processing a long sequence), the signal grows exponentially. When eigenvalues are much smaller than 1, the signal decays exponentially. This is the spectral radius perspective on exploding and vanishing gradients.

An orthogonal matrix, with eigenvalues of magnitude exactly 1, neither amplifies nor attenuates the signal regardless of how many times it is applied. This makes orthogonal initialization particularly well-suited for recurrent neural networks, where the same hidden-to-hidden weight matrix is multiplied at every time step. Even for very long sequences (hundreds or thousands of steps), an orthogonally initialized recurrent matrix maintains bounded signal propagation.

Orthogonal initialization achieves near-perfect variance preservation in linear networks. This makes it particularly valuable for recurrent neural networks, where the hidden-to-hidden weight matrix is applied at every time step. We'll explore this in more depth when we cover RNNs.

For feedforward networks, the practical advantages of orthogonal initialization over He are less clear-cut, because activation functions break the orthogonality property immediately after the first layer. Xavier and He initialization work well for typical feedforward architectures, while orthogonal initialization is the preferred choice for RNN hidden-to-hidden matrices.

Modern large language models take a more pragmatic approach to initialization than the theoretically derived schemes above. Both BERT and GPT use a simple, uniform strategy: normal distribution with a fixed small standard deviation, regardless of layer size. This works because layer normalization, applied throughout these architectures, normalizes activations during training and reduces the network.s sensitivity to the exact initialization scale.

The choice of a fixed standard deviation rather than a layer-size-dependent formula is a deliberate engineering decision. In transformer architectures, every attention layer and feed-forward layer is followed by layer normalization, which rescales activations to have unit variance at every step. This means the network continuously self-corrects its activation distribution, making the initialization less critical for preventing vanishing or exploding activations. The initialization instead needs to be "reasonable enough" that the network can start making productive gradient updates from the first iteration.

BERT uses $\mathcal{N}(0, 0.02)$ for all weight matrices and zeros for all bias vectors. The 0.02 standard deviation was chosen empirically by the original BERT authors based on what produced stable training for their specific architecture. The layer normalization parameters (scale $\gamma$ and bias $\beta$) are initialized to 1 and 0, respectively, establishing an initial identity transformation.

The choice of 0.02 is smaller than what Xavier or He would prescribe for typical transformer layer dimensions. For a layer with $n_{\text{in}} = 768$ (the hidden dimension of BERT-base), He initialization gives $\sigma = \sqrt{2/768} \approx 0.051$, and Xavier gives $\sigma = \sqrt{2/(768+768)} \approx 0.036$. The BERT value of 0.02 is smaller. This conservatism is intentional: smaller initialization means smaller initial activations, and with layer normalization continuously rescaling, the network quickly adapts to useful activation scales regardless of the initialization.

GPT-2 uses the same $\mathcal{N}(0, 0.02)$ base distribution but introduces a residual scaling factor for the projection matrices at the end of each attention block and each feed-forward block. These weights are scaled by $1/\sqrt{n_{\text{layers}}}$:

where:

• $n_{\text{layers}}$: the total number of transformer layers in the model
• $0.02$: the base standard deviation shared with BERT
• $\frac{1}{\sqrt{n_{\text{layers}}}}$: a residual scaling factor that prevents variance from accumulating across layers

This scaling addresses a specific problem with residual connections. In a residual network, each layer computes $x_{\ell+1} = x_\ell + F(x_\ell)$, adding the residual branch output $F(x_\ell)$ to the input. If each residual branch contributes variance $\sigma^2$, then after $L$ layers the total variance grows to $L \sigma^2$. Scaling the projection weights by $1/\sqrt{L}$ reduces each branch's variance contribution to $\sigma^2/L$, keeping the total variance constant at $\sigma^2$.

This residual scaling trick has been adopted by many subsequent transformer architectures. Larger models like GPT-3 and various open-source LLMs use the same basic approach, sometimes with slightly different base standard deviations tuned to their specific architectures. The key insight transfers directly: when residual branches accumulate, scale the contributing weights to keep variance bounded.

Theory tells us which initialization should work best for a given activation function, but practice sometimes differs from theory. Let's run a controlled empirical comparison to see how the different methods perform during training.

The comparison shows a clear hierarchy. Zero initialization fails completely. Xavier performs well but converges slightly more slowly than He on this ReLU network, because it was designed for tanh. He and orthogonal initialization both perform well and converge at similar rates. The BERT-style fixed standard deviation (0.02) underperforms because it is not calibrated to these layer dimensions.

Modern deep learning frameworks provide built-in initialization functions. PyTorch offers the full suite of initialization methods through the torch.nn.init module. Understanding how to use these functions and when to apply them is an essential practical skill.

The key PyTorch parameters for Kaiming (He) initialization are:

• mode ('fan_in' or 'fan_out'): Controls whether to use the input dimension or output dimension for the variance calculation. Use 'fan_in' (the default) to preserve forward pass variance, or 'fan_out' to preserve backward pass variance.
• nonlinearity ('relu', 'leaky_relu', 'tanh', etc.): Specifies the activation function so PyTorch computes the correct gain factor automatically.
• a (for Leaky ReLU): The negative slope parameter. PyTorch computes the gain as $\sqrt{2 / (1 + a^2)}$, matching our derived formula.

For Xavier initialization, the gain parameter multiplies the standard deviation. PyTorch provides nn.init.calculate_gain('tanh') and similar functions to look up the recommended gain for a given activation.

In practice, you'll often want to apply initialization to an entire model at once. The cleanest pattern is to write an _init_weights method and call self.apply() to walk through all submodules.

The _init_weights pattern scales cleanly from simple MLPs to complex architectures. When building transformer models, you'd typically call self.apply(self._init_weights) to recursively initialize all submodules, then optionally scale specific weights (like GPT-2's projection weights) with a second pass.

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

Unlike weights, initializing all biases to the same value (zero) does not cause a symmetry problem. The gradient of the loss with respect to a bias depends on the upstream gradient, not on other neurons' biases. Each neuron receives a distinct upstream gradient because its incoming weights are different (due to random weight initialization), so biases update independently from the start.

This is a subtle but important distinction. The symmetry problem for weights arises because two neurons receiving the same input with the same weights produce the same output and therefore receive the same gradient. Biases do not have this issue because the gradient for a bias in layer $\ell$ depends on the product of all upstream activations and weights from layer $\ell$ to the output. Since those upstream weights differ between neurons, the bias gradients differ too.

For ReLU networks, some practitioners use small positive biases (0.01) to ensure neurons start in an active state. This prevents "dead neurons" where all inputs are negative from the start. However, batch normalization and careful learning rate selection have largely eliminated the need for this adjustment in modern architectures.

There are special cases where bias initialization does matter. The output layer bias can be initialized to the log-odds of the class priors (for classification) or the mean target value (for regression), giving the network a sensible starting point that reduces the number of epochs needed to converge. Gate biases in LSTM cells are sometimes initialized to positive values (commonly 1 or 2) to keep the forget gate initially open, allowing gradients to flow more easily through the gate during early training.

One of the most important developments in deep learning was the realization that initialization and normalization are complementary tools addressing the same underlying problem: keeping activation distributions healthy throughout a deep network. Understanding their relationship clarifies when each is necessary and when one can substitute for the other.

Batch normalization, introduced by Ioffe and Szegedy in 2015, directly normalizes pre-activation values to have zero mean and unit variance at each layer, then scales and shifts them by learned parameters. This operation continuously corrects the activation distribution during training, reducing the network.s sensitivity to poor initialization. As a result, networks with batch normalization can converge from a much wider range of initializations than those without.

Layer normalization, the variant used in transformers, operates along the feature dimension rather than the batch dimension. It has the same effect of continuously normalizing activations, but it does so per-sample rather than per-batch. This makes it better suited for variable-length sequences and small batch sizes.

The existence of these normalization layers explains a fact that puzzles many practitioners: why does BERT work so well with a simple fixed initialization that theory suggests is suboptimal? The answer is that layer normalization renders the precise initialization scale much less important. Whatever the initial activation distribution, the first layer normalization immediately rescales it. The initialization just needs to avoid catastrophic failure modes (complete collapse or extreme saturation) in the handful of forward passes before the first gradient update.

Without normalization, initialization becomes much more critical. If you design a network that omits batch or layer normalization for any reason, the theoretically derived initializations (Xavier or He, depending on your activation function) are essential for achieving reliable training. This was the regime that the Glorot/Bengio and He et al. papers were addressing, before normalization was ubiquitous.

The practical implication is: if you add batch normalization or layer normalization to your network, you have more flexibility in initialization; if you remove normalization layers for efficiency or other architectural reasons, you should apply principled initialization more carefully.

Proper initialization has measurable effects on training stability beyond just convergence speed. The connection is to how well-conditioned the early gradient signal is: well-initialized networks start in regions of the loss surface where gradients are informative and directions are well-scaled, enabling optimizers to take productive steps from the first iteration.

Poor initialization forces the optimizer to spend many epochs correcting the initial state before it can begin learning the actual task. In the worst cases, networks with bad initialization never converge at all, because early gradient signals are too weak (vanishing) or too noisy (exploding) to provide useful update directions. The optimizer takes large steps in meaningless directions, or takes steps so small that any progress is undetectable.

The condition number of the weight matrices is another lens through which to understand initialization quality. When weights are initialized appropriately, the early layers of the network apply transformations with condition numbers close to 1, meaning all directions are scaled similarly. This makes the optimization landscape well-conditioned and gradient descent more efficient. Poor initialization can produce extremely ill-conditioned transformations where some directions are amplified by orders of magnitude more than others, making optimization very difficult.

The interactions between initialization and other training components are also important. Batch normalization dramatically reduces sensitivity to initialization by normalizing activations at each layer during training. Layer normalization, used in transformers, plays a similar role. This is why BERT and GPT can use a simple fixed standard deviation rather than layer-size-dependent formulas: the normalization layers continuously correct the activation distribution during training, absorbing initialization errors.

When normalization layers are absent (some architectures avoid them for specific reasons), initialization quality becomes more critical. In these settings, using the theoretically derived initialization for your specific activation function is the safest strategy.

When a network fails to train, determining whether initialization is the culprit requires specific diagnostic tools. Simply observing poor performance after many training epochs does not distinguish initialization failure from other problems like inappropriate learning rates, overfitting, or data issues.

The most informative diagnostic is to examine activation statistics in the first few forward passes, before any gradient updates. If activations are nearly zero at the output of every layer, you have a vanishing activation problem that is almost certainly initialization-related. If activations are extremely large or if many neurons have saturated activations (tanh outputs close to $\pm 1$, sigmoid outputs close to 0 or 1), you have an explosion or saturation problem.

Gradient statistics after the first backward pass provide complementary information. Computing the gradient norm at each layer and plotting it across depth reveals whether gradients are vanishing or exploding as they flow backward through the network. A healthy gradient flow shows gradient norms roughly comparable across all layers. Vanishing gradients appear as norms that decay exponentially with distance from the output; exploding gradients appear as norms that grow exponentially.

This diagnostic output immediately shows the difference between initialization schemes. With zero initialization, standard deviations collapse to zero after the first layer and the dead neuron fraction is 100% from the start. With He initialization, standard deviations remain stable across layers and the dead neuron fraction is around 50% (expected for a ReLU network, since the network is applying ReLU to symmetrically distributed pre-activations). With Xavier on a ReLU network, you see a progressive decay in standard deviation, showing the mismatch between Xavier's tanh-oriented design and the ReLU activation.

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

Initialization is only the starting point. A good initialization establishes favorable conditions for learning, but it cannot compensate for fundamental problems like inappropriate architecture choices, insufficient data, or training hyperparameters that are far off. A network with He initialization will still fail if the learning rate is too large by several orders of magnitude.

Normalization techniques reduce sensitivity. Batch normalization and layer normalization have made precise initialization less critical for many architectures. Modern transformers like BERT and GPT demonstrate that a simple, fixed initialization standard deviation works well when combined with layer normalization. The normalization layers continuously adjust the activation distribution during training, absorbing initialization errors.

Very deep networks remain challenging. For networks with hundreds of layers, even careful initialization may not prevent early training instabilities. Additional techniques like learning rate warmup, gradient clipping, and careful residual branch scaling become necessary. The next chapter covers gradient clipping in detail.

The derivations for Xavier and He initialization make assumptions that do not always hold. They assume weights and inputs are independent, that inputs have zero mean, and that the activation function behaves as modeled (linear for Xavier, variance-halving for He). In practice, these assumptions are approximately but not exactly satisfied, particularly after a few training steps. The initializations work well in practice despite these approximations because the variance analysis captures the dominant first-order behavior.

As a practical decision guide:

• Use He initialization for ReLU and GELU networks (the current default for feedforward and convolutional architectures). This is nn.init.kaiming_normal_() or nn.init.kaiming_uniform_() in PyTorch.
• Use Xavier initialization for tanh and sigmoid networks. This is nn.init.xavier_normal_() or nn.init.xavier_uniform_() in PyTorch.
• Use orthogonal initialization for RNN hidden-to-hidden matrices, where the same matrix is applied many times.
• Follow the BERT/GPT-2 approach for transformer architectures with layer normalization: use $\mathcal{N}(0, 0.02)$ for all weights with residual scaling applied to projection layers.
• If you are uncertain about which activation a layer feeds, default to He initialization; it is more tolerant of mismatch than Xavier because the $\sqrt{2}$ factor provides a buffer against variance decay.

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 collapse: all neurons in a layer learn identically, wasting representational capacity. Any constant initialization has this problem.
• Random initialization introduces scale sensitivity: too small causes vanishing activations, too large causes explosion or saturation
• LeCun initialization uses variance $1/n_{\text{in}}$ to preserve forward pass variance in linear networks. It is the theoretical foundation that Xavier and He build upon.
• Xavier initialization uses variance $2/(n_{\text{in}} + n_{\text{out}})$, balancing forward and backward variance for tanh and sigmoid activations
• He initialization uses variance $2/n_{\text{in}}$, adding a factor of 2 to compensate for ReLU's variance-halving effect
• Orthogonal initialization preserves exact variance in linear transformations, making it valuable for RNNs where spectral radius matters for long-range gradient flow
• BERT and GPT use a simple $\mathcal{N}(0, 0.02)$ approach, relying on layer normalization and residual scaling to maintain stability
• Normalization layers reduce sensitivity to initialization and explain why modern large models need less precise initialization schemes
• Diagnostics matter: examining activation statistics and gradient norms at the start of training can pinpoint initialization problems before wasting compute on a doomed training run

The next chapter covers gradient clipping, a complementary technique that handles gradient explosions that can occur during training even when initialization is correct, particularly in recurrent networks processing long sequences.

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