Weight initialization and variance propagation

The problem: variance drifts layer to layer

One linear layer, $y=Wx$, $x\in\mathbb{R}^{n_{\text{in}}}$, $W\in\mathbb{R}^{n_{\text{out}}\times n_{\text{in}}}$ with i.i.d. entries, mean $0$, variance $\sigma_W^2$, independent of $x$. For independent, mean-$0$ $A,B$: $\operatorname{Var}(AB)=\operatorname{Var}(A)\operatorname{Var}(B)$ exactly (both cross terms in $\mathbb{E}[A^2B^2]-(\mathbb{E}[AB])^2$ factor cleanly). So, for $x_i$ i.i.d. with variance $\sigma_x^2$:

$$ \operatorname{Var}(y_j) = \operatorname{Var}\Big(\sum_{i=1}^{n_{\text{in}}}W_{ji}x_i\Big) = \sum_{i=1}^{n_{\text{in}}}\operatorname{Var}(W_{ji})\operatorname{Var}(x_i) = n_{\text{in}}\,\sigma_W^2\,\sigma_x^2. $$

Each layer multiplies variance by a fixed factor, $n_{\text{in}}\sigma_W^2$. Across $L$ layers that compounds to the $L$-th power of that factor.

Intuition: this is Dynamics in the eigenbasis again, one step removed, repeated multiplication by a factor above $1$ blows up, below $1$ collapses to $0$; here the factor is $n_{\text{in}}\sigma_W^2$ and depth plays the role of the exponent.

The fix: force the per-layer factor to 1

Require $\operatorname{Var}(y_j) = \operatorname{Var}(x_i)$, i.e. the factor equals $1$:

$$ n_{\text{in}}\,\sigma_W^2 = 1 \quad\Longrightarrow\quad \sigma_W^2 = \frac1{n_{\text{in}}}. $$

This is LeCun initialization: sample each weight from a distribution with variance $1/n_{\text{in}}$ (e.g. $\mathcal{N}(0,1/n_{\text{in}})$), and the forward-pass variance neither grows nor shrinks with depth.

Intuition: wider layers (larger $n_{\text{in}}$) sum more terms, so each individual weight must be proportionally smaller to keep the total sum's variance the same.

Xavier/Glorot: balancing forward and backward

The backward pass has its own variance-propagation equation. Backprop sends a gradient through $W^T$ (Backpropagation via Jacobian-vector products), so by the same computation with the roles of $n_{\text{in}}$ and $n_{\text{out}}$ swapped:

$$ \operatorname{Var}\big(\nabla_x\mathcal{L}\big) = n_{\text{out}}\,\sigma_W^2\,\operatorname{Var}\big(\nabla_y\mathcal{L}\big). $$

Matching the forward factor ($n_{\text{in}}\sigma_W^2=1$) and the backward factor ($n_{\text{out}}\sigma_W^2=1$) at once is only possible if $n_{\text{in}}=n_{\text{out}}$. Xavier/Glorot initialization compromises with their average:

$$ \sigma_W^2 = \frac{2}{n_{\text{in}}+n_{\text{out}}}. $$

Intuition: you cannot satisfy two different constraints on one shared parameter $\sigma_W^2$ unless they happen to agree; split the difference instead of favoring one direction and letting the other drift.

He initialization: correcting for ReLU

$\text{ReLU}(z)=\max(z,0)$ zeros out roughly half of a zero-mean, symmetric pre-activation. For $z$ symmetric around $0$:

$$ \mathbb{E}\big[\text{ReLU}(z)^2\big] = \mathbb{E}\big[z^2 \mid z>0\big]\,P(z>0) = \tfrac12\,\mathbb{E}[z^2], $$

since the positive half is left untouched and contributes exactly half of $\mathbb{E}[z^2]$ by symmetry. So ReLU halves the variance that would otherwise pass through. Doubling $\sigma_W^2$ exactly compensates:

$$ \sigma_W^2 = \frac{2}{n_{\text{in}}}. $$

Intuition: He initialization is LeCun initialization with the ReLU tax priced in up front, since the earlier derivation implicitly assumed the nonlinearity preserves variance, which only $\sigma_W^2=1/n_{\text{in}}$ needs, and ReLU does not.

Why this matters: compounding over depth

With the per-layer factor off by even a small multiplicative error $1+\epsilon$, variance after $L$ layers scales by $(1+\epsilon)^L$: exponential in depth. A network with hundreds of layers makes even a slightly wrong $\sigma_W^2$ catastrophic, pushing activations (and, by the same mechanism, backpropagated gradients) to $0$ or $\infty$ long before training can do anything useful.

Intuition: initialization is not a detail, it is the boundary condition on a multiplicative recurrence; get the per-layer factor to exactly $1$, or depth itself amplifies the error for you, in either direction.