Weight Initialisation in Transformers

A transformer with 80 stacked blocks adds 80 independent contributions to its residual stream before training has adjusted a single weight. Get the initial scale of those contributions even slightly wrong and the consequence is not subtle: too large, and the residual stream's variance compounds across depth until early training is numerically unstable; too small, and gradient signal barely reaches the earliest layers. Initialisation in a deep transformer is not boilerplate, it is a load-bearing hyperparameter, and it needed its own fix distinct from what worked for shallower networks.

Why generic initialisation is not enough

Xavier/Glorot initialisation (Glorot and Bengio, 2010), which scales a layer's initial weight variance by its fan-in and fan-out, was designed to keep activation variance roughly constant through a single forward pass of a moderately deep feedforward network. It says nothing about what happens when the same-scale contribution is added, via a residual connection, dozens of times in a row. Each residual add is x = x + f(x); if f's output variance at initialisation is comparable to x's, the residual stream's variance grows roughly with depth, block after block, purely as an artifact of initialisation, before the model has learned anything.

The GPT-2 fix: scale the output projections by depth

Radford et al., 2019 addressed this directly. Most weights are initialised from a normal distribution with a small fixed standard deviation (about 0.02), but the output projection of each attention sublayer (W_O) and each FFN sublayer (the down-projection) is additionally scaled by an extra factor of 1 / sqrt(2 * n_layers). These are specifically the matrices whose output is added straight into the residual stream, so shrinking them by a factor that grows with depth keeps the expected variance each block contributes small enough that, summed across all blocks, the residual stream's variance at initialisation stays controlled rather than growing linearly with the number of layers. This is a targeted correction for exactly the failure mode generic Xavier init does not account for.

muP: initialisation as part of a width-consistent system

Yang et al., 2022 ("Tensor Programs V") generalised the problem: instead of asking only what initial variance keeps activations stable, ask what combination of initial variance, learning rate, and per-layer multiplier makes the size of the parameter update at each training step independent of model width. This is maximal update parametrization, muP. Its payoff is not just stability, it is transfer: hyperparameters tuned on a small proxy model at one width carry over, without re-tuning, to a much larger model of the same architectural "shape." A search that would otherwise require expensive sweeps at the target scale can instead be run cheaply on a small model and applied directly to the large one.

Why pre-norm makes this more forgiving, not irrelevant

The pre-norm architecture (see layernorm-residual-connections) provides a clean, un-normalised residual highway from input to output regardless of what any individual sublayer's initial output looks like, which is part of why pre-norm models train without the careful learning-rate warmup that post-norm required. But "more forgiving" is not "init-independent": pre-norm reduces the sensitivity to a poorly scaled sublayer output, it does not remove the need to choose that scale sensibly in the first place, especially at the depths and widths frontier models now use.

When it falls down

• Initialisation only governs step zero. A well-chosen initial variance controls behaviour before any gradient step; if the learning rate or data distribution is mismatched to it, the balance it establishes is disrupted within the first few hundred steps regardless, so init reduces but does not remove the need for warmup and careful learning-rate scheduling.
• The GPT-2 heuristic is recipe-specific. The 1/sqrt(2N) scaling was derived for the standard sequential pre-norm block; it does not automatically transfer to a parallel attention-and-FFN block (see parallel-attention-and-ffn) or to an unusual depth-width ratio (see depth-width-aspect-ratio) without re-deriving the appropriate constant for that architecture.
• muP guarantees hyperparameter transfer, not everything else. It says the optimal learning rate and init scale found on a small proxy model apply to the large one; it does not guarantee that data mixture decisions, architectural ablations, or anything else tuned on the small model will transfer the same way, so teams still validate the full recipe at target scale.

Further reading

• Radford et al., 2019, Language Models are Unsupervised Multitask Learners (GPT-2) - the 1/sqrt(N)-scaled output-projection initialisation.
• Glorot and Bengio, 2010, Understanding the Difficulty of Training Deep Feedforward Neural Networks - Xavier initialisation, the pre-transformer baseline.
• Yang et al., 2022, Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer, arXiv:2203.03466 - muP.