Weight Initialization

Why you cannot just start at zero

The tempting first idea — set every weight to zero — fails immediately, and the reason is instructive. If every weight in a layer is identical, then every neuron in that layer computes the exact same function of the input, receives the exact same gradient, and updates by the exact same amount. They stay identical forever. This is the symmetry problem: a layer of 512 neurons initialized to a constant has the representational capacity of a single neuron, permanently. You must break symmetry, and the only way to do that is randomness — each weight drawn independently so that neurons start life differentiated.

So the weights must be random. But which random? Once you accept that weights are samples from some distribution, the design question becomes a question about that distribution’s statistics — specifically its mean and variance. We almost always center it at mean zero (no systematic bias in any direction), which leaves the variance as the single knob that matters. Setting that variance correctly is the entire game, and the wrong value — even by a modest factor, compounded across many layers — is the difference between a network that trains and one that produces NaN on the first batch.

The signal is a random variable with a scale

To reason about initialization we treat the activations flowing through the network as random variables and track one number about them: their variance, a proxy for the typical magnitude of the signal. Call the input to a layer x and the output y. If the layer systematically shrinks the variance — Var(y) < Var(x) — then after many layers the signal is crushed toward a constant and the network can no longer distinguish inputs (vanishing activations). If it systematically inflates it — Var(y) > Var(x) — the signal explodes toward saturation or overflow (exploding activations).

Neither is a small effect, because layers compound multiplicatively. If each layer multiplies the variance by a factor γ, then after L layers the variance is scaled by γ^L. A modest γ = 1.5 over 50 layers is 1.5^50 ≈ 6×10^8; a modest γ = 0.7 is 0.7^50 ≈ 2×10^-8. Depth turns a small per-layer bias into an astronomical one. The design goal writes itself: choose the initialization so that γ ≈ 1 — the layer preserves variance — and the signal keeps a sane scale no matter how deep the stack.

The forward-pass variance derivation

Take a single linear layer with no activation yet: y = Wx, where W is [n_out, n_in] and x is a length-n_in input. Each output coordinate is a sum over the fan-in:

y_i = Σ_{j=1..n_in} W_ij * x_j

Assume the weights W_ij are drawn i.i.d. with mean 0 and variance Var(W), that the inputs x_j are i.i.d. and independent of the weights, and (crucially) that the weights are zero-mean. For a product of two independent variables where one is zero-mean, the variance of the product is Var(W) · E[x^2] — note it is the second moment E[x^2], not the variance of x, that carries through (they coincide only when x is itself zero-mean). Summing n_in independent such terms:

Var(y_i) = n_in · Var(W) · E[x^2]

For the first layer the input is zero-mean, so E[x^2] = Var(x) and the condition to preserve variance — Var(y) = Var(x) — is simply n_in · Var(W) = 1, i.e. Var(W) = 1/n_in. That is the fan-in rule. Writing the propagation through E[x^2] rather than Var(x) is what will let the ReLU case fall out cleanly later.

The backward pass has its own condition

Preserving the forward signal is only half the job; the gradients must survive the trip backward too, or the network stops learning in its early layers. The gradient with respect to the layer input flows through the transpose of the weight matrix:

∂L/∂x = W^T (∂L/∂y)

By the identical argument — a sum, now over the n_out outputs that each input feeds — the variance of the incoming gradient relates to the outgoing gradient by Var(∂L/∂x) = n_out · Var(W) · Var(∂L/∂y). Preserving gradient variance therefore demands Var(W) = 1/n_out — the fan-out rule.

Now the tension is explicit. The forward pass wants Var(W) = 1/n_in; the backward pass wants Var(W) = 1/n_out. Unless the layer is square (n_in = n_out) you cannot satisfy both exactly with a single scalar variance. Every classical initialization scheme is, at heart, a particular way of resolving this one conflict — either splitting the difference between the two fans, or deliberately committing to one of them.

Xavier / Glorot: averaging the two fans

Glorot and Bengio’s answer (2010), universally called Xavier initialization, is a compromise. Rather than obey one condition and violate the other, it averages the two fans so that the single variance is as close as possible to satisfying both. The two ideal conditions are n_in · Var(W) = 1 and n_out · Var(W) = 1; replacing the fan by its average (n_in + n_out)/2 gives the condition ((n_in + n_out)/2) · Var(W) = 1, hence:

Var(W) = 2 / (n_in + n_out)

Be careful here: this is not the average of the two target variances 1/n_in and 1/n_out — averaging the variances would give a different (and wrong) number. It is the average of the fans that goes in the denominator. In the square case n_in = n_out = n it collapses back to 1/n, satisfying both conditions exactly, as it should. Xavier was derived assuming the activation is roughly linear around zero — it is the right choice for tanh and other symmetric, zero-centered nonlinearities, and it is the sensible default whenever a layer is not followed by a ReLU.

Normal or uniform? Same variance, two dresses

A variance target does not by itself pin down a distribution — any zero-mean distribution with the right variance works, and in practice we use either a Gaussian or a uniform. For the Gaussian form you set the standard deviation to the square root of the target variance directly: σ = √Var(W). For the uniform form you exploit that a uniform distribution on [-a, a] has variance a^2/3, so to hit a target variance V you solve a^2/3 = V, giving a = √(3V).

Concretely, Xavier-uniform sets a = √(6/(n_in + n_out)) (plug V = 2/(n_in+n_out) into √(3V) — the 3 and the 2 multiply to the 6 you see in every framework’s source). Xavier-normal sets σ = √(2/(n_in + n_out)). The two give statistically equivalent scale; the uniform version bounds the extreme weights (nothing beyond ±a) while the Gaussian has thin infinite tails. The choice is minor in practice — what matters is the variance, and both hit the same one. Knowing the √(3V) conversion lets you read any of these formulas back into the single variance target they encode.

He / Kaiming: the factor of 2 that ReLU demands

ReLU changes the accounting, and the derivation through E[x^2] shows exactly how. A ReLU zeroes every negative preactivation, so it does not propagate a zero-mean signal — it propagates a rectified one. For a preactivation z that is symmetric about zero, half its mass is negative and gets clipped to 0, so the second moment of the activation x = relu(z) is exactly half the second moment of z:

E[x^2] = E[relu(z)^2] = ½ · E[z^2] = ½ · Var(z)

(Note this is a statement about E[relu(z)^2], the quantity that actually propagates — not about Var(relu(z)), which is smaller still because the ReLU output is not zero-mean.) Feed that halved second moment into the layer relation Var(y) = n_in · Var(W) · E[x^2] and the preserve-variance condition becomes n_in · Var(W) · ½ = 1, i.e. Var(W) = 2/n_in. That is He (Kaiming) initialization: exactly Xavier’s fan-in rule with a factor of 2 bolted on to compensate for the energy ReLU throws away. He commits to the fan-in condition (forward preservation) rather than averaging, which is the right call for the deep ReLU stacks it was designed for.

Worked example: variance through a 50-layer stack

Numbers make the stakes concrete. Take a plain ReLU network, 50 layers, each n_in = 256. Per layer the variance factor is γ = ½ · n_in · Var(W) (the ½ from ReLU). Compare three initializations:

The middle row uses a fixed σ = 0.1 that someone might pick by feel; γ = 0.5 · 256 · 0.01 = 1.28, and 1.28^50 ≈ 160,000 — the activations are five orders of magnitude too large before the first gradient step. The bottom row shows that even Xavier’s fan-in variance 1/n_in, which is correct for tanh, halves the signal every ReLU layer and annihilates it over depth (0.5^50 ≈ 10^-15). Only He’s factor-of-2 lands on γ = 1. This is the whole argument for why the right init is not a nicety: with fan-in 256 the He standard deviation is √(2/256) = 0.088, and being off it by a factor that looks harmless compounds into overflow or underflow.

A small map of variance through depth

The three regimes above are worth seeing at a glance. Preserve variance and the signal rides flat through the stack; miss high and it curves off the top; miss low and it decays into the floor. The vertical axis is (log) signal variance, the horizontal axis is layer depth.

The gradients follow the mirror image of this picture on the way back, which is why the backward condition matters as much as the forward one. A network can have healthy-looking forward activations and still fail to train because the gradient signal vanished before it reached the early layers — the reason Xavier bothers to average the two fans at all.

Transformers break the variance assumption

The clean derivation above assumes a simple feed-forward stack. Transformers violate one of its premises structurally, and that forces an extra correction. The culprit is the residual stream: every sub-layer (each attention block and each feed-forward block) does not replace its input but adds to it — x ← x + Sublayer(x). Variance of a sum of roughly independent terms adds, so if each of N sub-layers contributes a term of variance v, the residual stream’s variance grows toward N · v as you go up the stack.

That is a slow, linear-in-depth inflation rather than the exponential blow-up of a bad matrix init, but for a 48- or 96-layer model it is very real: the activations entering the final layers can be an order of magnitude larger than those entering the first, which destabilizes training and makes the model sensitive to learning rate. A transformer has N = 2 × n_layers residual additions (two per block), so a 24-layer model already has 48 contributions piling into the stream. Plain He or Xavier init on the projection matrices does nothing to counter this, because the growth comes from the addition structure, not from any single matrix.

The GPT-2 rule: scale residual projections by 1/√(2N)

The fix, popularized by GPT-2, is to shrink each sub-layer’s contribution so the accumulated variance stays bounded. If you want N terms to sum to a controlled total, scale each one’s standard deviation by 1/√N — equivalently its variance by 1/N — so that N copies of variance v/N sum back to v. GPT-2 applies exactly this to the weights that write into the residual stream: the attention output projection and the second feed-forward projection (the c_proj layers). Their initialization std is scaled:

σ_residual = 0.02 / √(2 · n_layers)

where 0.02 is GPT-2’s base std for every other weight and 2 · n_layers = N is the residual-path count. For a 12-layer GPT-2, σ = 0.02/√24 = 0.02/4.90 ≈ 0.0041; for a 48-layer model, 0.02/√96 ≈ 0.0020 — the deeper the model, the more each writer is damped. Only the residual-writing projections get this treatment; the query/key/value and first-FFN matrices keep the base std. It is a targeted patch for the one place the standard variance argument does not apply, and it is why GPT-2-lineage models train stably at depth without exotic tricks.