Initialisation, normalisation and µP

The scale of the initial weights. Each weight matrix multiplies a vector of $d$ numbers; if its entries have standard deviation $\sigma$, the output's size is about $\sigma\sqrt d$ times the input's. Choosing $\sigma \approx 1/\sqrt{d}$ (Xavier initialisation; He initialisation adds a factor $\sqrt 2$ for ReLU) keeps it the same, layer after layer. The Decoder's model uses $\sigma = 0.02$, GPT-2's choice (Radford et al. 2019), which for $d = 16$ is far smaller than $1/\sqrt{16} = 0.25$: every sublayer starts out adding almost nothing, the residual stream is just the embedding plus the positional encoding (Decoder, chapter 5), and the two designs look alike (pick std 0.02 in the animation). GPT-2 also scales the residual layers' weights by $1/\sqrt N$ for $N$ residual layers, because every block adds to the same stream and their sum grows with depth.

Pre-LN vs post-LN. With std 0.25 the difference shows. In the pre-LN model each block adds its output to the stream, so the stream grows with depth: its RMS is 0.72 entering block 1 and 4.30 leaving block 8. Post-LN normalises after every addition, so the stream stays at RMS 1, but the gradient pays for it: the gradient reaching a block has passed through every LayerNorm above it, and at initialisation the blocks nearest the output get the largest gradients (post-LN: 1.29 at block 1, 2.29 at block 8). Xiong et al. (2020) proved this for post-LN at initialisation, and showed it is why post-LN transformers need a learning-rate warmup: a full-size step on those large gradients makes training unstable. In pre-LN the residual path carries the gradient back unchanged (chapter 4's "highway"), so the gradients are of similar size at every depth.

The chart shows what that means for training: a 6-block model at a peak learning rate of 0.01, three seeds each. Pre-LN trains with or without warmup (2.378 ± 0.128 without, 2.310 ± 0.074 with 200 warmup steps). Post-LN without warmup ends at 2.949 ± 0.002: every seed is still on chapter 1's plateau after 1,500 steps. With warmup it ends at 2.877 ± 0.124 (seeds: 2.735, 2.950, 2.947): warmup lets some runs start learning, not all. This is why GPT-2 and nearly every large model since put the norm first.

RMSNorm. LayerNorm subtracts the mean and divides by the standard deviation, then scales and shifts. RMSNorm (Zhang & Sennrich 2019) skips the mean and the shift: $y = \gamma \odot x / \operatorname{RMS}(x)$. It is cheaper and, in practice, as good; LLaMA uses it, as pre-normalisation (Touvron et al. 2023, §2.2). On the Decoder's model (the default run, three seeds), RMSNorm even ends lower: 2.140 ± 0.008 against LayerNorm's 2.277 ± 0.090, with 5,296 parameters instead of 5,376; LayerNorm's spread over the seeds is wide, so read that as "at least as good".

µP: one learning rate for every width. Tuning the learning rate of a large model directly costs several large runs. The maximal update parameterisation (µP; Yang et al. 2022) rescales the initialisation, the learning rate and a few multipliers with the width so that each step changes every layer's activations by about the same amount whatever the width. Then the best learning rate found on a narrow model carries over to a wide one ("µTransfer"). Under the standard parameterisation it does not: with Adam, each weight moves by about the learning rate per step, and a wider matrix sums more of those moves, so the same learning rate is a bigger effective step for a wider model.

The coordinate check below measures exactly that: ten AdamW steps at the same learning rate for widths 16 to 256, and how far the logits move. Under the standard parameterisation they move 1.43 at width 16 and 4.48 at width 256; under µP 1.43 and 1.42.

What µP changes, relative to a base width $d_0$ (here 16, the Decoder's, so at width 16 the two are identical) with $m = d/d_0$: hidden matrices start at std $0.02/\sqrt m$ and Adam's learning rate for them is divided by $m$; the attention scores are scaled by $1/d_k$ instead of $1/\sqrt{d_k}$; the logits are divided by $m$. The second chart trains widths 16 and 64 at three learning rates, three seeds each, for 1,500 steps. With the standard parameterisation the wider model is no better than the narrow one at any rate tried (width 64: 2.440, 2.531, 2.811 at 0.01, 0.03, 0.1; width 16's best 2.420). With µP the same width-64 model reaches 2.314, 2.310 and 2.682, best at the same rates as width 16. The grid is coarse and the spreads are wide (± 0.03 to 0.2: the plateau of chapter 1 ends at a different step in every run), so the toy shows the direction rather than an optimum. A GPU notebook runs the real version, widths 128 to 1,024 (below).

Run it yourself on a GPU. Notebook 07 (open in Colab) sweeps widths 128–1,024 and seven learning rates under both parameterisations on a GPU (about an hour on an A100). The notebook is one of the planned GPU replications listed under Runs to come; every number on this page comes from the in-browser model.