Whether the signal and the gradient survive many layers depends on the scale of the initial weights and on where the LayerNorms sit. Pre-LN against post-LN through depth and in training, RMSNorm, and the maximal update parameterisation that lets one learning rate serve every width.
This deck follows chapter 7, initialisation, normalisation and µP. Measured numbers: results.md §7a–7e. The Decoder’s design: LayerNorm and residuals.
Before the first step every weight is a random number, and how large those numbers are decides whether the signal and the gradient survive the trip through many layers. A matrix whose entries have standard deviation \(\sigma\) multiplies a vector of \(d\) numbers into one about \(\sigma\sqrt d\) times as large.
Keeps the size the same layer after layer. For \(d = 16\): 0.25.
The Decoder’s choice (Radford et al. 2019), far smaller here: every sublayer starts out adding almost nothing. GPT-2 also scales residual layers’ weights by \(1/\sqrt N\) for \(N\) residual layers, because every block adds to the same stream.
$$x_{l+1} = x_l + F_l(\operatorname{LN}(x_l))$$ The sublayer reads a normalised copy; the stream itself is only ever added to.
$$x_{l+1} = \operatorname{LN}(x_l + F_l(x_l))$$ Every block normalises the sum, so every gradient must pass through the LayerNorms above it.
The experiment: an 8-block version of the Decoder’s model at initialisation, on “alice was”, with weights at std 0.02 and at std 0.25: the RMS of the residual stream entering each block, and the norm of each block’s weight gradients (results.md §7a).
With std 0.02 every sublayer adds almost nothing: the pre-LN stream is essentially the embedding plus the positional encoding, RMS 0.70 at every depth. With std 0.25 the pre-LN stream grows with depth, from 0.72 entering block 1 to 4.30 leaving block 8, because each block adds to it. Post-LN normalises after every addition, so its stream stays at RMS 1.
Post-LN pays for its tidy stream in the gradients: at initialisation the blocks nearest the output get the largest (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 (deck 03’s highway).
A 6-block model, AdamW at a peak of 0.01 with a cosine, 1,500 steps, three seeds (results.md §7b):
| Norm placement | Warmup steps | Final validation loss (mean ± sd) | Seeds 1, 2, 3 |
|---|---|---|---|
| pre-LN | 0 | 2.378 ± 0.128 | 2.302, 2.306, 2.526 |
| pre-LN | 200 | 2.310 ± 0.074 | 2.272, 2.395, 2.263 |
| post-LN | 0 | 2.949 ± 0.002 | 2.950, 2.950, 2.947 |
| post-LN | 200 | 2.877 ± 0.124 | 2.735, 2.950, 2.947 |
Pre-LN trains with or without warmup. Post-LN without warmup is still on the plateau in every seed; with warmup, some runs start learning, not all. This is why GPT-2 and nearly every large model since put the norm first.
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 \frac{x}{\sqrt{\tfrac1d \sum_i x_i^2 + \epsilon}}$$| Norm | Parameters | Final validation loss (mean ± sd) | Seeds 1, 2, 3 |
|---|---|---|---|
| LayerNorm | 5,376 | 2.277 ± 0.090 | 2.299, 2.354, 2.178 |
| RMSNorm | 5,296 | 2.140 ± 0.008 | 2.148, 2.132, 2.141 |
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 with fewer parameters, but LayerNorm’s spread over the seeds is wide: read it as “at least as good” (results.md §7c).
Tuning the learning rate of a large model directly costs several large runs. One would like to tune it on a narrow model and reuse it on a wide one. Under the standard parameterisation (SP) that does not work.
With Adam each weight moves by about the learning rate per step (deck 04). A wider matrix sums more of those moves into each output, so the same learning rate is a bigger effective step for a wider model, and the best rate shifts down as the width grows.
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 at any width. Then the best rate found on a narrow model carries over (“µTransfer”).
What µP changes for Adam, relative to a base width \(d_0\) (here 16, the Decoder’s) with \(m = d/d_0\): hidden matrices start at std \(0.02/\sqrt m\) and their learning rate is divided by \(m\); attention scores are scaled by \(1/d_k\) instead of \(1/\sqrt{d_k}\); the logits are divided by \(m\) (Yang et al., Table 3).
The coordinate check: 10 AdamW steps at learning rate 0.01 for widths 16 to 256, and how far the logits move (RMS change, on 16 fixed windows; results.md §7d). Under SP they move 1.43 at width 16 and 4.48 at width 256; the last residual stream moves 124. Under µP the logits move 1.43 and 1.42: width-independent, which is the whole point.
Widths 16 and 64, three learning rates, three seeds, 1,500 steps (results.md §7e; at width 16 the two parameterisations are the same model):
| Parameterisation | Width | Parameters | Peak learning rate | Final validation loss |
|---|---|---|---|---|
| standard | 16 | 5,376 | 0.01 | 2.420 ± 0.101 |
| standard | 16 | 5,376 | 0.03 | 2.457 ± 0.213 |
| standard | 16 | 5,376 | 0.1 | 2.656 ± 0.025 |
| standard | 64 | 70,656 | 0.01 | 2.440 ± 0.193 |
| standard | 64 | 70,656 | 0.03 | 2.531 ± 0.205 |
| standard | 64 | 70,656 | 0.1 | 2.811 ± 0.083 |
| µP | 64 | 70,656 | 0.01 | 2.314 ± 0.073 |
| µP | 64 | 70,656 | 0.03 | 2.310 ± 0.048 |
| µP | 64 | 70,656 | 0.1 | 2.682 ± 0.127 |
With SP the wider model is no better than the narrow one at any rate tried. With µP the same width-64 model improves, best at the same rates as width 16. The grid is coarse and the spreads wide (± 0.03 to 0.2: the plateau ends at a different step in every run), so the toy shows the direction rather than an optimum. The site’s notebook 07 runs widths 128–1,024 on a GPU: GPU run pending, so no number from it appears here.
Animated: chapter 7, initialisation, normalisation and µP · next: deck 07, stability