Normalization

Key idea

LayerNorm = (x − mean) ÷ std × γ + β; RMSNorm = x ÷ rms × γ; pre-norm keeps the residual path clean

Real world

Grading on a curve. A raw mark of 72 means little until you know the class average was 60 and the spread was 8. Converting every mark to “1.5 spreads above average” makes marks from a hard exam and an easy exam comparable.

In an LLM

GPT-2’s residual vector for “ is” has length 5 at the bottom and 430 at the top. Each layer reads it through a LayerNorm, so every layer sees numbers of the same scale, and the final norm shrinks it back to about 255 before the prediction.

Why normalise

The residual stream keeps growing: every layer adds to it. If each layer read the raw vector, the early layers would see vectors of length about 55 and the late ones vectors several times longer (in GPT-2 small, from about 55 to about 430: see the chart on the previous page), and every layer would need weights tuned to its own scale. Worse, during training the scale shifts as earlier layers change, so later layers would be chasing a moving target.

The fix is to rescale each token’s vector to a standard size before a layer reads it. The layer always sees numbers of roughly the same size, however long the stream has grown.

LayerNorm

LayerNorm takes one token’s vector $x$ with $d$ numbers and does three things:

1. Subtract the mean. $\mu = \frac{1}{d}\sum_i x_i$. After this the numbers average to zero.
2. Divide by the standard deviation. $\sigma = \sqrt{\frac{1}{d}\sum_i (x_i - \mu)^2}$. After this they have spread 1.
3. Scale and shift. Multiply each number by a learned gain $\gamma_i$ and add a learned bias $\beta_i$.

The tiny $\epsilon$ (usually $10^{-5}$) stops a division by zero. The gain and bias let the model put back any scale it actually needs, per dimension. Each token is normalised on its own: no information moves between tokens.

RMSNorm is the cheaper version that most models now use. It skips the mean and the bias:

Drag the bars and watch both at work:

Two things to try. Move the scale slider: the input bars grow or shrink, the output does not change at all. That is the point of normalising: the size of the input no longer matters, only its shape. Then move the shift slider: LayerNorm ignores it (it subtracts the mean), RMSNorm does not (it never removes the mean). Zhang and Sennrich, 2019 showed that the re-centring step barely matters in practice, so RMSNorm gets the same quality with fewer operations. Llama, Mistral, Qwen and Gemma all use it.

Pre-norm or post-norm

There are two places to put the norm. The original Transformer (Vaswani et al., 2017) normalised after adding each sub-layer’s output: post-norm, $x \leftarrow \operatorname{LN}(x + f(x))$. GPT-2 moved it to the start of each branch: Pre-norm, $x \leftarrow x + f(\operatorname{LN}(x))$.

The difference looks small but matters a lot for deep models. In pre-norm the main path from the input to the output has nothing on it but additions, so the gradient flows straight down it (the identity path from the residual stream). In post-norm the gradient must pass through a LayerNorm at every sub-layer, and each one rescales it. Xiong et al., 2020 showed that post-norm gradients near the output are large at the start of training, which is why post-norm models need a slow learning-rate warm-up and still often diverge when they are deep.

Pre-norm has one side effect: since nothing on the main path normalises it, the stream grows with depth (you saw this in GPT-2). So pre-norm models add one final norm after the last block, before the output layer.

QK-norm

Some recent models (Gemma 3, OLMo 2, Qwen 3) add one more norm inside attention: they apply RMSNorm to every query and key vector before the dot product. This caps how large the attention scores can get, which stops a failure in large-scale training where a few scores grow without limit, softmax becomes one-hot, and the loss spikes. It is a small, cheap change: two extra norms per layer.

def layer_norm(x, gamma, beta, eps=1e-5):          # x: (..., d)
    mu = x.mean(-1, keepdim=True)
    var = ((x - mu) ** 2).mean(-1, keepdim=True)    # biased variance
    return (x - mu) / torch.sqrt(var + eps) * gamma + beta
 
def rms_norm(x, gamma, eps=1e-5):
    return x / torch.sqrt((x ** 2).mean(-1, keepdim=True) + eps) * gamma