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.