Weight Initialization › He Initialization for ReLU Networks

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:

$$\text{Var}(\text{ReLU}(z)) \approx \frac{\sigma^2}{2}$$

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)$:

$$\mathbb{E}[\text{ReLU}(z)^2] = \mathbb{E}[z^2 \cdot \mathbf{1}_{z > 0}] = \frac{1}{2}\mathbb{E}[z^2] = \frac{\sigma^2}{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:

$$\text{Var}(a) = \frac{1}{2} \cdot n_{\text{in}} \cdot \text{Var}(w) \cdot \text{Var}(x)$$

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:

$$\text{Var}(x) = \frac{1}{2} \cdot n_{\text{in}} \cdot \text{Var}(w) \cdot \text{Var}(x)$$

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

$$1 = \frac{1}{2} \cdot n_{\text{in}} \cdot \text{Var}(w)$$

Solving for the weight variance:

$$\text{Var}(w) = \frac{2}{n_{\text{in}}}$$

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:

$$w \sim \mathcal{N}\!\left(0, \sqrt{\frac{2}{n_{\text{in}}}}\right)$$

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}}}$:

$$w \sim U\!\left[-\sqrt{\frac{6}{n_{\text{in}}}}, \sqrt{\frac{6}{n_{\text{in}}}}\right]$$