Variance Preservation in Weight Initialization

A feedforward network with fifty layers trained from a naive initialization does not train badly; it does not train at all. The usual small random numbers drawn from a narrow uniform or normal distribution carry a scale that is wrong by an amount that grows multiplicatively with depth. By layer ten the activations may be a thousandth of their original magnitude, or a thousand times it, depending on whether the scale is slightly too small or slightly too large. The derivation that fixes this is short, and it explains most of the empirical advice about initializers that followed.

The cleanest way to see the problem is to follow a single scalar statistic through the network: the variance of the activations and of the backpropagated gradients, layer by layer, ignoring any structure beyond independence assumptions.

The variance recursion for one layer

Consider a layer with $n_{\text{in}}$ inputs $x_j$, weights $w_{ij}$, and no bias for the moment. The pre-activation is $z_i = \sum_j w_{ij} x_j$. Assume the weights are independent of each other and of the inputs, are drawn with zero mean and variance $\sigma_w^2$, and that the inputs are also zero-mean and mutually independent with variance $q$. Then the variance of $z_i$ is

$$\operatorname{Var}(z_i) = \sum_j \operatorname{Var}(w_{ij} x_j) = n_{\text{in}} , \sigma_w^2 , q$$

because the variance of a product of independent zero-mean variables is the product of their variances. The factor $n_{\text{in}}$ is the entire problem. A layer that would preserve the input variance needs $\sigma_w^2 = 1/n_{\text{in}}$. Draw weights with a fixed variance such as $0.01$ and a layer with a thousand inputs multiplies the variance by ten at the first layer and by $10^L$ after $L$ layers.

The same recursion holds for gradients in reverse. If the forward pass preserves variance, the gradient with respect to the layer inputs has variance scaled by $n_{\text{out}} \sigma_w^2$, where $n_{\text{out}}$ is the number of units in the layer. A weight matrix that is a linear map from $\mathbb{R}^{n_{\text{in}}}$ to $\mathbb{R}^{n_{\text{out}}}$ therefore wants $\sigma_w^2 = 1/n_{\text{in}}$ going forward and $\sigma_w^2 = 1/n_{\text{out}}$ going backward. When the layer is not square the two conditions cannot both hold; the standard compromise sets the variance to the harmonic mean of the two reciprocals, which is $2/(n_{\text{in}} + n_{\text{out}})$.

That last constant is the Glorot (Xavier) initialization. For a uniform distribution on $[-a, a]$, the variance is $a^2/3$, so preserving $\sigma_w^2 = 2/(n_{\text{in}} + n_{\text{out}})$ requires $a = \sqrt{6/(n_{\text{in}}+n_{\text{out}})}$. The familiar bound follows from a one-line computation, not from experiment.

What the derivation assumed

Every step above used linearity of the layer and independence of the weights, inputs, and each other’s contributions. Neither is strictly true. The first assumption is the important one: the recursion computed the variance of the pre-activation $z$, and a nonlinearity $\phi$ placed after it changes the variance by a factor that depends on the shape of $\phi$. For $\tanh$ near zero the factor is close to one, which is why the Glorot argument was stated in the context of sigmoid and tanh networks. For the rectifier $\phi(z) = \max(0, z)$ the factor is not one, and the conclusion changes.

Why the rectifier needs a different constant

Take $z$ symmetric about zero with variance $q$. After a ReLU, half the values are zero and the other half are unchanged, so the output has mean $\sqrt{q/2\pi}$ and variance $q(1/2 - 1/2\pi) \approx 0.34 q$. The activation halves the variance. Applied $L$ times, that factor compounds to $(0.34)^{L/2}$ in standard deviation, so a network of twenty rectified layers loses roughly nine orders of magnitude of signal scale if the weights are drawn at Glorot scale.

The fix is to compensate for the halving at each layer, which means doubling the weight variance: $\sigma_w^2 = 2/n_{\text{in}}$. Written as a normal distribution this is the He initializer.

import math

def he_normal(fan_in, rng):
    # std chosen so that Var(Wx) is preserved through a rectifier,
    # because the rectifier contributes a factor of 1/2.
    return rng.gauss(0.0, math.sqrt(2.0 / fan_in))

def glorot_uniform(fan_in, fan_out, rng):
    limit = math.sqrt(6.0 / (fan_in + fan_out))
    return rng.uniform(-limit, limit)

For a uniform draw the bound becomes $\sqrt{6/n_{\text{in}}}$. The factor of two is often described as a “correction for ReLU”; it is the exact variance ratio of a zero-mean symmetric Gaussian passed through a rectifier, and it is exact only under that distributional assumption.

Choosing the fan that enters the constant

The derivation needs $n_{\text{in}}$, and in a convolutional layer “the number of inputs to a unit” is the number of weights that connect to one output element, not the number of input pixels. For a kernel of spatial size $k \times k$ over $C_{\text{in}}$ channels with groups $g$, the effective fan-in is $k^2 C_{\text{in}} / g$. Using $k^2 C_{\text{in}} C_{\text{out}}$ by mistake produces a variance that is too small by a factor of $C_{\text{out}}$, which for a wide layer is a serious mismatch. Framework implementations distinguish these; the distinction is documented in the initializer arguments under names such as fan_in, fan_out, and mode, and the correct choice is the one matching the variance propagation described above.

Bias terms are normally initialised to zero. There is no variance argument for them, and zero keeps the pre-activations at zero mean, which the forward recursion assumed. A small positive bias on a rectified layer is a different device: it pushes units into the active region to avoid dead units, and it is a deliberate departure from the derivation rather than part of it.

Failure modes beyond the constant

The recursion’s conclusion is about scale, not about conditioning. Two distinct failures survive a correctly scaled initializer.

The first is a degenerate fan-in. If a layer’s inputs are themselves all zero, as with an all-zero input vector or a layer whose preceding rectifier has been killed everywhere, then any weight scale preserves zero. The recursion is exact and useless: gradient descent gets no direction from a layer whose weights all receive equal gradient. Randomly breaking symmetry across units within a layer is the purpose of drawing weights at random at all, and it is why initialising a network’s weights to a constant is a special case with no useful behaviour.

The second is eigenvalue spread within a layer. The recursion computes the expected variance under independence, but a particular draw of $W$ has a singular value spectrum. A layer whose largest singular value is greater than one amplifies the dominant direction; one whose smallest is near zero attenuates it. Repeated multiplication of a random matrix with entries of variance $1/n$ has singular values spread over roughly a factor of two or so per layer around the expectation, and over depth the product’s conditioning degrades. This is why orthogonal initialisation, which draws a matrix whose singular values are all exactly one, was proposed: it satisfies the variance condition and additionally keeps every direction at equal scale.

Orthogonal and identity-based initialisation

An orthogonal matrix $Q$ satisfies $Q^\top Q = I$, so it maps every unit vector to another unit vector, and its singular values are all one. For square layers this is the strongest possible version of “preserve variance”. The computational cost is a QR decomposition of a Gaussian matrix, or the construction by successive Householder reflections, both of which cost $O(n^3)$ for an $n \times n$ matrix. For a full weight matrix with $n_{\text{out}} \ne n_{\text{in}}$ only the first $n_{\text{out}}$ or $n_{\text{in}}$ columns of a larger orthogonal factor are used, which gives semi-orthogonal rows or columns.

For recurrent nets the relevant matrix is the recurrent one applied the same number of times as the sequence length, so its spectral radius rather than its variance governs whether the state decays. A matrix with spectral radius slightly below one keeps the state’s expected norm bounded and its gradients from vanishing exponentially, provided the inputs’ contributions do not vanish along with them. Both properties fail if a gain constant is moved away from one.

The normalization layers that appear in most modern architectures weaken these arguments without replacing them. A batch-normalized layer re-standardises its pre-activations, so the precise weight scale matters less for activation magnitude, but it still sets the scale of the update relative to the parameter, and normalisation does not rescue an initialisation whose gradients are zero or whose weights are symmetric. The variance recursion remains the right first question to ask about any initializer, and the residual questions are about conditioning, symmetry, and the interaction with the non-linearity actually used.