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.

Initializing embeddings

Embeddings do not fit the fan-in story because a token embedding is a lookup, not a matrix-vector product — each row is selected whole, so there is no sum over a fan-in whose variance we need to preserve. The relevant constraint is instead that each embedding vector should start at a sensible magnitude relative to what the first layer expects. GPT-2 simply draws token and positional embeddings from N(0, 0.02^2) — the same small fixed std as the rest of the model — and that works well in practice.

Two wrinkles are worth knowing. First, the original ‘Attention Is All You Need’ transformer multiplies its embeddings by √d_model before adding positional encodings, deliberately scaling them up so the learned embeddings are not swamped by the fixed sinusoidal positions — a reminder that the right embedding scale depends on what it is added to. Second, many models tie the input embedding matrix to the output (un-embedding) projection, sharing one weight tensor; then the initialization has to be reasonable for both roles at once, which is another reason the small, neutral 0.02 std is a safe default rather than a fan-based formula.

Biases, LayerNorm, and the interaction with normalization

Not every parameter follows the variance rules. Biases are almost always initialized to zero: they carry no fan-in sum to preserve, and symmetry is already broken by the random weights, so zero is the neutral, no-op starting point. LayerNorm (and its RMSNorm cousin) initialize their gain γ to 1 and shift β to 0 — the identity transform — so that at step zero normalization neither scales nor shifts the signal.

Normalization also changes how much the matrix init has to carry. A LayerNorm sitting in front of each sub-layer re-standardizes the activations to unit variance regardless of what the previous layer did, which forgives a lot of variance drift and is part of why transformers are less brittle to init than plain deep stacks. But it forgives, it does not eliminate: LayerNorm normalizes the input to each sub-layer, yet the residual stream it adds back into is exactly the unnormalized accumulator whose variance still grows with depth — which is why the 1/√(2N) residual scaling is still needed even in a fully normalized network. Init and normalization are complementary, not redundant.

Practical implications for small CPU models and common pitfalls

For the small language models this series cares about — the kind you might train or fine-tune on a CPU — getting init right is cheap insurance against expensive failure. The rules of thumb: use He/Kaiming (fan-in) for any layer followed by ReLU or GELU, Xavier for tanh/linear, zero the biases, set LayerNorm to identity, and if you are building a transformer, apply the 1/√(2N) scaling to the residual projections. Frameworks give you these for free (kaiming_normal_, xavier_uniform_), but only if you call the right one — the default init on a bare Linear layer is not always what your activation wants.

The pitfalls are predictable. Using the framework default (often a Xavier-style rule) under a ReLU quietly halves your signal every layer — survivable when shallow, fatal when deep. Forgetting the residual scaling gives a transformer that trains but is oddly learning-rate-sensitive and unstable at depth. Initializing with a hand-picked σ like 0.1 ‘because it seems small’ ignores fan-in entirely and, as the worked example showed, can be catastrophically wrong for a wide layer. And the classic silent bug: a network that looks fine on the forward pass but never improves, because the gradient variance vanished on the way back. When a deep model will not train, initialization is one of the first things to check — and one of the easiest to fix.