Initialization Scaling

The one goal: keep variance constant across layers

Every initialization rule is a means to one end: hold the variance of activations roughly constant as data flows forward, and the variance of gradients roughly constant as errors flow backward. Consider a single linear layer y = Wx with input x of dimension n_in, weights drawn i.i.d. with mean 0 and variance Var(W), and inputs also zero-mean with variance Var(x).

Each output is a sum of n_in independent products, so variances add:

Var(y_j) = Σ_i Var(W_ji · x_i)
         = n_in · Var(W) · Var(x)

For the output to carry the same variance as the input we need n_in · Var(W) = 1, i.e. Var(W) = 1/n_in. That equation — scale the weight variance by the inverse of fan-in — is the seed from which every practical rule below grows.

The fan-in rule and where it comes from

The relation Var(W) = 1/n_in is the fan-in rule. It says a neuron that sums many inputs must draw each weight from a tighter distribution, because a wide layer accumulates more terms and would otherwise inflate the output. A layer with n_in = 4096 initializes weights with standard deviation 1/√4096 = 0.0156; a narrow n_in = 64 layer uses 0.125 — eight times larger.

In code this is a fan_in lookup, not a fixed number. PyTorch’s defaults, framework helpers, and hand-rolled inits all compute the layer’s fan-in and divide by it. The subtlety is which dimension is fan-in: for a weight matrix W: [n_out, n_in] it is n_in on the forward pass, but the backward pass sums over n_out instead — which is exactly why a single fan-in rule cannot perfectly balance both directions, motivating the compromise in the next section.

Xavier / Glorot: balancing forward and backward

The forward pass wants Var(W) = 1/n_in; the backward pass, by the same argument applied to gradients, wants Var(W) = 1/n_out. You cannot satisfy both unless the layer is square, so Xavier (Glorot) initialization takes their harmonic-style compromise — the average of the two denominators:

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

Glorot derived this assuming a linear or symmetric, unit-derivative activation around zero (identity, tanh near the origin). Under that assumption each layer’s Jacobian passes variance through unchanged, so a stack of layers neither amplifies nor damps signal. Xavier was the default that made deep tanh and sigmoid networks trainable at all, and it is still the sensible choice for the linear projections inside attention and for any layer whose activation is roughly linear at initialization.

Kaiming / He: correcting for ReLU-family activations

Xavier’s assumption breaks for ReLU and its relatives, because ReLU zeroes out the negative half of its inputs. If pre-activations are symmetric about zero, ReLU discards roughly half the variance, so signal shrinks by a factor of two per layer — over many layers, it vanishes. Kaiming (He) initialization compensates by doubling the target variance:

Var(W) = 2 / n_in        (the 2 offsets ReLU’s halving)
std(W) = √(2 / n_in)

The general form is Var(W) = gain² / n_in, where gain depends on the nonlinearity: gain = √2 for ReLU, 1 for identity, and intermediate values for leaky-ReLU, GELU, or SiLU. Modern transformer MLP blocks use GELU/SiLU gates whose gain sits near ReLU’s, so a He-style √(2/n_in) is the right starting point for the up-projection, while attention’s linear W_Q, W_K, W_V lean Xavier.

A worked number

Take one MLP layer of a small transformer: hidden width d = 768, expanding to 4d = 3072 with a GELU. The up-projection is W: [3072, 768], so n_in = 768.

He init sets std = √(2/768) = √0.0026 ≈ 0.051. Feed it a unit-variance input x: the pre-activation variance is n_in · Var(W) · Var(x) = 768 · 0.0026 · 1 ≈ 2. After GELU roughly halves it, the output variance lands near 1 — preserved, exactly as intended.

Now suppose a rule left each layer shrinking variance by just 5% instead of preserving it. Over 48 layers that compounds to 0.95^48 ≈ 0.08: the signal at the top is a twelfth of its input, and gradients are starved. Small per-layer errors are lethal at depth.

Depth changes everything: the residual stream

The worked example exposes the real problem in transformers: depth. Even a perfectly variance-preserving layer, repeated, can still accumulate variance through the residual connections. A residual block computes x ← x + f(x). If x and the sublayer output f(x) each have variance 1 and are roughly independent, their sum has variance 2. After N such blocks the residual stream’s variance grows like 1 + N — it drifts linearly with depth.

That drift is usually tamed by LayerNorm, which renormalizes the stream at every block. But relying on normalization to absorb an ever-growing residual is fragile: at initialization the sublayer contributions swamp the identity path, and the first few hundred steps are the ones most likely to diverge. The cleaner fix is to scale the init so each block adds less.

The 1/sqrt(2N) residual scaling

The standard remedy, popularized by GPT-2, is to shrink the output projection of each residual sublayer — the attention output W_O and the MLP down-projection — by a factor that depends on the number of residual layers N:

std(W_out) = base_std / √(2N)

The 2N counts the residual paths: each transformer block contributes two (one from attention, one from the MLP), so a model with N blocks has 2N additive contributions to the stream. Dividing each output std by √(2N) makes each sublayer’s variance contribution 1/(2N), so summing all 2N of them adds up to a bounded O(1) instead of growing with depth. GPT-2 implements this by taking a base std = 0.02 and rescaling residual-output weights by 1/√(2N) — the single most important trick for training deep transformers without divergence.

Why 0.02? The empirical LLM default

Production language models rarely quote Xavier or He by name; they set a small fixed std, almost always 0.02, from a truncated normal. Where does it come from? For the widths GPT-style models used, √(2/n_in) lands near 0.02–0.05, so a flat 0.02 is a slightly conservative, width-agnostic stand-in that errs on the safe (small) side.

It is not a law of nature. The value interacts with hidden width: as models grew wider, a fixed 0.02 became too large relative to the fan-in rule, contributing to instability. The principled alternative is to make init explicitly width-dependent — which is exactly what muP does, and why the next section ties init to the learning rate rather than treating them as separate knobs.

Init and learning rate must scale together (muP)

Initialization does not act alone. A weight’s effect on the network is set by both its initial value and how far the optimizer moves it per step, so init variance and learning rate are coupled — tuning one in isolation is fighting with one hand. Maximal Update Parametrization (muP) makes the coupling explicit by asking that, as width d grows, every activation and every update stays O(1) in size.

Satisfying that constraint forces init variance and learning rate to scale with width in opposite directions: hidden-layer init variance scales as 1/d (the fan-in rule again), while the effective per-parameter learning rate for those layers scales as 1/d under Adam. The payoff is hyperparameter transfer: tune the learning rate on a small proxy model, then widen the network and reuse it. The lesson for initialization scaling is that the ‘right’ variance is the one that keeps the update size sane, not the one that only looks good at step zero.

Special cases: embeddings, LayerNorm, and biases

The fan-in logic applies to matrix multiplies, but a transformer has parameters that are not standard linear layers, and each takes a different rule.

Embeddings are lookups, not sums, so there is no fan-in to divide by; they are typically initialized directly to the target std (e.g. 0.02), sometimes scaled by √d at the input so the embedding enters the residual stream at unit variance. LayerNorm gains start at 1 and biases at 0, so the layer is an identity at init and adds no variance surprise. Ordinary biases start at 0: a nonzero bias would inject a mean into a carefully zero-mean signal. These boundary parameters sit at the entry and exit of the residual stream, where a mistake is not averaged away by depth but injected straight into every downstream layer.

CPU-SLM implications

For a small language model you train yourself on a CPU, initialization scaling is disproportionately valuable, because you cannot afford the compute to recover from a bad start. A CPU run measured in hours or days has no slack for a diverged attempt, so getting init right on the first try — He-style √(2/n_in) on MLP up-projections, Xavier on attention projections, 1/√(2N) on residual outputs, zeros on biases — directly buys wall-clock time.

Correct scaling also lets you shorten learning-rate warmup, which exists largely to survive the unstable early steps that bad init creates. And because SLMs are narrow (small n_in), the fan-in rule yields relatively large stds — so the flat 0.02 borrowed from big models is often too small here, and computing √(2/n_in) per layer is the better default.

Common pitfalls

The failure modes are consistent. Forgetting the ReLU/GELU factor of 2 uses Xavier where He is needed, so activations shrink by half per layer and deep gradients vanish. Skipping the residual 1/√(2N) scaling lets the residual stream variance grow linearly with depth, producing early-step divergence that no amount of learning-rate tuning fully cures.

Copying 0.02 blindly ignores width: it is too large for very wide models and too small for narrow SLMs. Initializing biases or LayerNorm gains randomly injects a mean or a scale the variance analysis never accounted for. And tuning init and learning rate separately wastes runs chasing a stability that only exists at the right product of the two. The through-line: initialization is a variance budget spread across width, activation, and depth — spend it deliberately.