Chapter 5 · Section 5.2 · 50 min

LayerNorm and RMSNorm

Section 5.1 left us with a problem. The residual stream is a running sum, and sums drift. After a few dozen layers a token's vector can be several times longer than it started, and different tokens can drift by different amounts. Every sublayer then reads its input at an unpredictable scale. That matters because the sublayers are scale-sensitive. Attention logits grow with the square of the input scale and push the softmax into saturation (the failure we analyzed in §4.3), and an activation function like GELU behaves like a linear map near zero and like a ReLU far from it.

The fix is to normalize what each sublayer reads. That sounds like a detail. It is not. Where you put the normalization decides whether a 48-layer Transformer trains at all, whether it needs a learning-rate warmup, and how its gradients are distributed over depth. This section derives LayerNorm and RMSNorm from scratch, forward and backward, then works through the argument for why modern models normalize before each sublayer rather than after.

Why not BatchNorm

You met BatchNorm in §0.7. It normalizes each feature using the mean and variance over the batch. For language models that is a poor fit, for three reasons.

It leaks across sequences and positions. Statistics over a batch of sequences mix tokens from different documents and, worse, from future positions of the same document. A causal language model is supposed to predict token $t+1$ from tokens up to $t$ , but a BatchNorm over all positions lets information from later tokens flow into earlier ones through the shared mean.

It behaves differently at test time. BatchNorm uses running averages at inference, so a model's output for a sequence depends on whether it is in training mode. Generation one token at a time would be normalized with statistics of batches it never sees.

Batches are small and ragged. Per-device batches in large-model training can be a handful of sequences, padded to different lengths, which makes batch statistics noisy and padding-dependent.

LayerNorm (Ba, Kiros & Hinton, 2016) sidesteps all three. It normalizes each token using statistics over that token's own features. No other token, no other sequence, and no running averages are involved, so training and inference compute exactly the same function.

LayerNorm normalizes each token by itself

Take one token's vector $\mathbf{x} \in \R^d$ (a row of the stream $\mathbf{X} \in \R^{T\times d}$ ). LayerNorm computes its mean and variance over the $d$ features,

\mu = \frac{1}{d}\sum_{j=1}^{d} x_j, \qquad \sigma^2 = \frac{1}{d}\sum_{j=1}^{d} (x_j - \mu)^2, \qquad s = \sqrt{\sigma^2 + \epsilon},

then standardizes and applies a learned per-feature gain $\boldsymbol{\gamma} \in \R^d$ and shift $\boldsymbol{\beta} \in \R^d$ ,

\hat{\mathbf{x}} = \frac{\mathbf{x} - \mu\mathbf{1}}{s}, \qquad \operatorname{LN}(\mathbf{x}) = \boldsymbol{\gamma} \odot \hat{\mathbf{x}} + \boldsymbol{\beta}.

Here $\mathbf{1}$ is the all-ones vector, $\odot$ is the elementwise product, and $\epsilon$ (typically $10^{-5}$ ) keeps the division safe for a constant vector. For a whole sequence, each row gets its own $\mu$ and $s$ , while $\boldsymbol{\gamma}$ and $\boldsymbol{\beta}$ are shared by every token. A LayerNorm has $2d$ parameters.

In plain words, shift the token's features to mean zero, scale them to unit variance, then let the model choose a new scale and offset for each feature.

A worked example. Let $\mathbf{x} = (2, 4, 6, 8)$ , so $d = 4$ . The mean is $\mu = 5$ . The deviations are $(-3, -1, 1, 3)$ , so $\sigma^2 = (9 + 1 + 1 + 9)/4 = 5$ and, ignoring $\epsilon$ , $s = \sqrt 5 \approx 2.236$ . Then $\hat{\mathbf{x}} \approx (-1.342, -0.447, 0.447, 1.342)$ . Notice two facts you can check by hand. The entries of $\hat{\mathbf{x}}$ sum to zero, and their squares sum to $4 = d$ .

The geometry, a projection and then a sphere

Those two facts are the whole geometry of LayerNorm. Subtracting the mean is a linear projection,

\mathbf{x} - \mu\mathbf{1} = \Big(\mathbf{I} - \tfrac{1}{d}\mathbf{1}\mathbf{1}^\top\Big)\mathbf{x},

which removes the component of $\mathbf{x}$ along $\mathbf{1}$ and lands in the $(d-1)$ -dimensional hyperplane of vectors whose entries sum to zero. Dividing by $s$ then rescales the result so that $\lVert\hat{\mathbf{x}}\rVert^2 = d\,\sigma^2/(\sigma^2 + \epsilon) \approx d$ . So before the gain and shift, every token is mapped onto a sphere of radius $\sqrt d$ inside that hyperplane. The gain stretches the sphere into an ellipsoid and the shift moves it.

In plain words, LayerNorm throws away exactly two numbers per token (the mean and the overall scale) and keeps the direction of what is left.

This explains its invariances. For any $a > 0$ and any constant $c$ ,

\operatorname{LN}(a\,\mathbf{x} + c\,\mathbf{1}) = \operatorname{LN}(\mathbf{x}) \quad (\text{ignoring } \epsilon).

A sublayer that adds the same constant to every coordinate of the stream, or rescales the whole stream, is invisible to the next LayerNorm. Play with this in the widget, where you can drag the bars of a 6-dimensional vector.

You wrote LayerNorm forward and backward, for any batch shape, in Exercise 0.7e of §0.7. Two details from that exercise matter in every implementation. The variance is the biased one, dividing by $d$ rather than $d - 1$ (NumPy's default, while PyTorch's torch.var defaults to the unbiased one, a classic source of mismatches), and $\epsilon$ goes inside the square root. There we borrowed the backward formula from BatchNorm. Here we derive it from scratch, because its structure is what explains the pre-LN argument later in this section.

The backward pass, derived

Interviewers ask for this one constantly, because it tests whether you can differentiate through a function whose statistics depend on every input. Let $\mathbf{y} = \operatorname{LN}(\mathbf{x})$ for one token, and suppose the loss gradient $\partial\mathcal{L}/\partial\mathbf{y}$ arrives from above. The gain multiplies elementwise, so define $\mathbf{g} = \boldsymbol{\gamma} \odot \partial\mathcal{L}/\partial\mathbf{y}$ , the gradient with respect to $\hat{\mathbf{x}}$ . We need $\partial\mathcal{L}/\partial\mathbf{x} = \mathbf{J}^\top\mathbf{g}$ where $\mathbf{J} = \partial\hat{\mathbf{x}}/\partial\mathbf{x}$ is the Jacobian. The result is

\frac{\partial \hat{\mathbf{x}}}{\partial \mathbf{x}} = \frac{1}{s}\Big(\mathbf{I} - \tfrac{1}{d}\mathbf{1}\mathbf{1}^\top - \tfrac{1}{d}\hat{\mathbf{x}}\hat{\mathbf{x}}^\top\Big), \qquad \frac{\partial \mathcal{L}}{\partial \mathbf{x}} = \frac{1}{s}\Big(\mathbf{g} - \operatorname{mean}(\mathbf{g})\,\mathbf{1} - \operatorname{mean}(\mathbf{g}\odot\hat{\mathbf{x}})\,\hat{\mathbf{x}}\Big).

∂ The derivation, one partial derivative at a time

Start with the pieces. Since $\mu = \frac1d\sum_k x_k$ , we have $\partial\mu/\partial x_j = 1/d$ .

For the variance, $\sigma^2 = \frac1d\sum_k (x_k - \mu)^2$ , so

\frac{\partial\sigma^2}{\partial x_j} = \frac{2}{d}\sum_k (x_k - \mu)\Big(\delta_{kj} - \frac1d\Big) = \frac{2}{d}(x_j - \mu) - \frac{2}{d^2}\sum_k (x_k - \mu) = \frac{2}{d}(x_j - \mu),

because the deviations sum to zero. Then $s = \sqrt{\sigma^2 + \epsilon}$ gives

\frac{\partial s}{\partial x_j} = \frac{1}{2s}\cdot\frac{2}{d}(x_j - \mu) = \frac{x_j - \mu}{d\,s} = \frac{\hat{x}_j}{d}.

Now differentiate $\hat{x}_i = (x_i - \mu)/s$ with the quotient rule,

\frac{\partial \hat{x}_i}{\partial x_j} = \frac{1}{s}\Big(\delta_{ij} - \frac1d\Big) - \frac{x_i - \mu}{s^2}\cdot\frac{\hat{x}_j}{d} = \frac{1}{s}\Big(\delta_{ij} - \frac1d - \frac{\hat{x}_i\hat{x}_j}{d}\Big).

That is the Jacobian in matrix form. It is symmetric, so $\mathbf{J}^\top\mathbf{g} = \mathbf{J}\mathbf{g}$ , and

(\mathbf{J}\mathbf{g})_j = \frac1s\Big(g_j - \frac1d\sum_i g_i - \frac{\hat{x}_j}{d}\sum_i g_i\hat{x}_i\Big),

which is the vector formula above. Note that $\epsilon$ never had to be dropped. The formula is exact with $\epsilon$ included, as long as $\hat{\mathbf{x}}$ and $s$ are the values used in the forward pass.

The parameter gradients are simpler. Since $y_j = \gamma_j\hat{x}_j + \beta_j$ , summing over every token $t$ that shares the parameters gives $\partial\mathcal{L}/\partial\gamma_j = \sum_t (\partial\mathcal{L}/\partial y_{t,j})\,\hat{x}_{t,j}$ and $\partial\mathcal{L}/\partial\beta_j = \sum_t \partial\mathcal{L}/\partial y_{t,j}$ .

In plain words, the input gradient is the incoming gradient with two components removed, its average and its projection onto $\hat{\mathbf{x}}$ , all divided by the token's scale $s$ .

There is a cleaner way to see the whole formula. Let $\mathbf{u} = \mathbf{1}/\sqrt d$ and $\mathbf{w} = \hat{\mathbf{x}}/\sqrt d$ . Both are unit vectors (ignoring $\epsilon$ , since $\lVert\hat{\mathbf{x}}\rVert^2 = d$ ), and they are orthogonal, because the entries of $\hat{\mathbf{x}}$ sum to zero. The Jacobian is then

s\,\frac{\partial\hat{\mathbf{x}}}{\partial\mathbf{x}} = \mathbf{I} - \mathbf{u}\mathbf{u}^\top - \mathbf{w}\mathbf{w}^\top,

the orthogonal projection onto the $(d-2)$ -dimensional subspace perpendicular to both. In plain words, LayerNorm's backward pass projects the incoming gradient onto the directions LayerNorm does not throw away, then divides by the token's scale. The three properties below are this one statement, read three ways.

Each piece has a meaning worth saying out loud in an interview.

The gradient is orthogonal to $\mathbf{1}$ . Summing the entries of $\partial\mathcal{L}/\partial\mathbf{x}$ gives $\frac1s(\sum g_j - d\operatorname{mean}(\mathbf{g}) - \operatorname{mean}(\mathbf{g}\odot\hat{\mathbf{x}})\sum_j\hat{x}_j) = 0$ , because $\sum_j\hat{x}_j = 0$ . That is necessary. Moving $\mathbf{x}$ along $\mathbf{1}$ cannot change the output, so the loss cannot have a slope in that direction.

The gradient is orthogonal to $\hat{\mathbf{x}}$ , up to $\epsilon$ . The dot product with $\hat{\mathbf{x}}$ is $\frac1s\operatorname{mean}(\mathbf{g}\odot\hat{\mathbf{x}})(d - \lVert\hat{\mathbf{x}}\rVert^2)$ , and $\lVert\hat{\mathbf{x}}\rVert^2 = d\sigma^2/(\sigma^2+\epsilon)$ is $d$ up to $\epsilon$ . Rescaling $\mathbf{x}$ does not change the output either, so no gradient flows along it.

The gradient shrinks with the input's scale. Everything is divided by $s$ . A token whose stream has grown four times larger receives a gradient four times smaller through the norm. Keep this in mind, because it is the mechanism behind the pre-LN gradient argument below.

Core

LayerNorm's Jacobian is a scaled projection

Write layer_norm_jacobian(x, eps=1e-5) for one token x of shape (d,), returning the (d, d) Jacobian $\partial\hat{\mathbf{x}}/\partial\mathbf{x}$ of LayerNorm without gain and shift. Then write layer_norm_vjp(x, g, eps=1e-5) returning $\mathbf{J}^\top\mathbf{g}$ with the fast formula, without building the (d, d) matrix. The tests check the Jacobian against finite differences, check that $s\mathbf{J}$ is an orthogonal projection of rank $d - 2$ when $\epsilon = 0$ , and check that the two functions agree.

import numpy as np

def layer_norm_jacobian(x, eps=1e-5):
    raise NotImplementedError

def layer_norm_vjp(x, g, eps=1e-5):
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

Here is a compact implementation we will reuse for the rest of the section, together with a check of the three properties.

Python · cell 1

🔬 Scale invariance and the growth of weights

Whenever a normalization directly follows a linear map, $\operatorname{LN}(\mathbf{x}\mathbf{W})$ is unchanged when $\mathbf{W}$ is multiplied by any $a > 0$ . The same argument as above then shows that the gradient with respect to $\mathbf{W}$ is orthogonal to $\mathbf{W}$ itself (as a vector in $\R^{d \times d}$ ). A gradient step orthogonal to $\mathbf{W}$ can only increase $\lVert\mathbf{W}\rVert$ , since $\lVert\mathbf{W} - \eta\mathbf{G}\rVert^2 = \lVert\mathbf{W}\rVert^2 + \eta^2\lVert\mathbf{G}\rVert^2$ . And because the gradient scales like $1/\lVert\mathbf{W}\rVert$ , the effective step size on the direction of $\mathbf{W}$ scales like $\eta/\lVert\mathbf{W}\rVert^2$ . Weight norms grow and the effective learning rate decays on its own, unless weight decay pulls the norm back (van Laarhoven, 2017). In a Transformer this applies exactly to matrices followed by a norm with no residual add in between, such as the query and key projections under the QK-norm of §8.8, and only approximately elsewhere.

RMSNorm drops the mean

Zhang & Sennrich (2019) asked which half of LayerNorm does the work, the re-centering or the re-scaling. Their hypothesis was that re-scaling is what matters, and they proposed RMSNorm, which skips the mean entirely,

r = \sqrt{\frac{1}{d}\sum_{j=1}^{d} x_j^2 + \epsilon}, \qquad \operatorname{RMSNorm}(\mathbf{x}) = \boldsymbol{\gamma} \odot \frac{\mathbf{x}}{r}.

In plain words, divide the token's vector by its root mean square, so that it lands on the sphere of radius $\sqrt d$ , then apply a learned per-feature gain. There is no shift parameter, so RMSNorm has $d$ parameters.

The geometry is the sphere without the hyperplane. The mean direction $\mathbf{1}$ survives, and RMSNorm is invariant to $\mathbf{x} \mapsto a\mathbf{x}$ for $a > 0$ but not to adding a constant. In the residual stream, $\mathbf{1}$ is one direction among $d$ with no special meaning, so losing invariance to it costs little. In their experiments RMSNorm matched LayerNorm's quality while cutting running time by 7% to 64% depending on the model. It needs one reduction over the features instead of two, and no subtraction. T5, LLaMA, Mistral, Gemma, Qwen, and most recent open models use it.

The backward pass is the LayerNorm derivation with the mean terms deleted. With $\hat{\mathbf{x}} = \mathbf{x}/r$ and $\mathbf{g} = \boldsymbol{\gamma}\odot\partial\mathcal{L}/\partial\mathbf{y}$ ,

\frac{\partial r}{\partial x_j} = \frac{x_j}{d\,r}, \qquad \frac{\partial \hat{x}_i}{\partial x_j} = \frac{1}{r}\Big(\delta_{ij} - \frac{\hat{x}_i\hat{x}_j}{d}\Big), \qquad \frac{\partial\mathcal{L}}{\partial\mathbf{x}} = \frac{1}{r}\Big(\mathbf{g} - \operatorname{mean}(\mathbf{g}\odot\hat{\mathbf{x}})\,\hat{\mathbf{x}}\Big).

In plain words, remove from the incoming gradient its component along the normalized input and divide by the scale. No component along $\mathbf{1}$ is removed, because RMSNorm does not ignore shifts.

Interview

RMSNorm forward and backward

Write rms_norm(x, gamma, eps=1e-6) and rms_norm_backward(dy, x, gamma, eps=1e-6) for inputs of shape (..., d). The backward returns (dx, dgamma), with dgamma of shape (d,).

import numpy as np

def rms_norm(x, gamma, eps=1e-6):
    raise NotImplementedError

def rms_norm_backward(dy, x, gamma, eps=1e-6):
    """Returns dx (like x), dgamma (d,)."""
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

Details that matter in real code

Compute the statistics in float32. bfloat16 carries only 8 bits of precision (§8.5), about two to three significant decimal digits per number. A running sum of 4,096 squares accumulated in it is far worse, because once the total is about 256 times larger than the next term, adding that term changes nothing and the sum stalls. LLaMA's reference RMSNorm upcasts to float32, normalizes, and casts back. Fused kernels do the same internally.

Know your parameterization. Gemma's RMSNorm multiplies by $(1 + \boldsymbol{\gamma})$ with $\boldsymbol{\gamma}$ initialized at zero, which is the same function with a different starting point for weight decay. OLMo's first release used a LayerNorm with no learned gain or shift at all. PaLM removed every bias from its norms and dense layers and reported better stability at scale.

Check $\epsilon$ . Values of $10^{-5}$ and $10^{-6}$ are both common, and putting it outside the square root ( $\sigma + \epsilon$ ) is a different function. When you port weights between frameworks, mismatched $\epsilon$ produces small but real differences in the logits.

Where does the norm go, pre-LN or post-LN

There are two ways to combine a normalization with a residual sublayer $F$ . The original Transformer (Vaswani et al., 2017) normalized after the residual add, which is called post-LN,

\text{post-LN:}\quad \mathbf{x} \leftarrow \operatorname{LN}\big(\mathbf{x} + F(\mathbf{x})\big).

GPT-2 (Radford et al., 2019) moved the norm to the input of each sublayer and added one final norm before the unembedding, which is called pre-LN,

\text{pre-LN:}\quad \mathbf{x} \leftarrow \mathbf{x} + F\big(\operatorname{LN}(\mathbf{x})\big), \qquad \text{logits} = \operatorname{LN}_f\big(\mathbf{x}^{(S)}\big)\,\mathbf{W}_U.

Here $\mathbf{x}^{(S)}$ is the stream after all $S$ sublayers and $\mathbf{W}_U \in \R^{d \times V}$ is the unembedding. The difference looks cosmetic. It is not. In post-LN the residual stream itself is normalized at every layer, so the "identity path" of §5.1 passes through a LayerNorm Jacobian at every step. In pre-LN the identity path is clean all the way from the loss to the embeddings, and only the branches see normalized inputs.

The gradient-scale argument

Xiong et al. (2020) analyzed both architectures at initialization and found the answer. Here is the argument in a form you can reproduce at a whiteboard, using two facts we already have. Fact one, from the derivation above, is that the norm's Jacobian scales like $\sqrt{d}/\lVert\mathbf{x}\rVert$ (it divides by $s = \lVert\mathbf{x} - \mu\mathbf{1}\rVert/\sqrt d$ ). Fact two, from §5.1, is how the stream's norm grows.

Pre-LN. Each branch reads a normalized input, so each write has roughly the same expected size, a fixed multiple of $d$ in squared norm. The writes are close to uncorrelated at initialization, so their squared norms add,

\mathbb{E}\,\lVert\mathbf{x}^{(\ell)}\rVert^2 \approx \lVert\mathbf{x}^{(0)}\rVert^2 + \ell\,c\,d,

linear growth in the layer index $\ell$ for some constant $c$ (Xiong et al. prove bounds of exactly this form, between $(1 + \ell/2)d$ and $(1 + 3\ell/2)d$ in their setup). The final norm divides by $\lVert\mathbf{x}^{(S)}\rVert/\sqrt d \sim \sqrt{S}$ , so the gradient that enters the stream from the loss is scaled down by about $1/\sqrt{S}$ . Every sublayer receives that gradient through the clean identity path. A branch's output weight gradient is its (order one, because its input was normalized) activation times that stream gradient, so every layer's weight gradient is of order $1/\sqrt{S}$ .

Post-LN. The stream is renormalized after every sublayer, so its norm stays near $\sqrt d$ forever. The last sublayer's gradient therefore does not shrink with depth at all. Going down one layer, the gradient passes through $\partial\operatorname{LN}(\mathbf{z})/\partial\mathbf{z}$ with $\mathbf{z} = \mathbf{x} + F(\mathbf{x})$ . Since $\mathbf{z}$ is the sum of two vectors of comparable size, $\lVert\mathbf{z}\rVert$ is larger than $\sqrt d$ and the Jacobian shrinks the gradient by a constant factor on every layer. The result is gradients that grow toward the output, with the top layers getting the largest ones.

Xiong et al. state the conclusion as a theorem about the last FFN layer's second weight matrix $\mathbf{W}^{2,L}$ in a model with $L$ layers. With high probability at initialization,

\text{post-LN:}\ \ \Big\lVert\frac{\partial\mathcal{L}}{\partial\mathbf{W}^{2,L}}\Big\rVert_F \le \mathcal{O}\big(d\sqrt{\ln d}\big), \qquad \text{pre-LN:}\ \ \Big\lVert\frac{\partial\mathcal{L}}{\partial\mathbf{W}^{2,L}}\Big\rVert_F \le \mathcal{O}\Big(d\sqrt{\tfrac{\ln d}{L}}\Big).

In plain words, post-LN hands its top layers large gradients no matter how deep the model is, while pre-LN's gradients shrink like $1/\sqrt{L}$ and are spread evenly across layers.

Why should you care? Large gradients at the top combined with a large learning rate make the early updates of the top layers too aggressive, and warmup is what keeps them in check. The argument is cleanest for SGD, where the step is proportional to the gradient. With Adam, which normalizes each coordinate's step (§5.7), the story also involves Adam's unreliable second-moment estimates in the first steps (Liu et al., 2019), but the conclusion is the same. Xiong et al. showed the effect directly. On IWSLT14 German to English, their post-LN Transformer trained with Adam reached a BLEU score of about 34 with warmup and only 8.45 without it, while a pre-LN Transformer trained without warmup reached comparable results in less training time. Warmup was a patch for post-LN's gradient profile.

Let's measure it. The cell below stacks residual MLP sublayers in both arrangements, backpropagates a random loss, and reports the gradient norm of each layer's output matrix.

Python · cell 2

The numbers tell the story. In post-LN the last layer's gradient stays near 30 whether the model has 8 or 128 layers, while the first layer's falls from about 14 to about 2, a fifteen-fold imbalance at 128 layers. In pre-LN the last layer's gradient falls from about 14 to about 4 as the model grows sixteen times deeper (close to the $1/\sqrt{16}$ of the theorem), and the spread across layers is much smaller. The widget runs a smaller version of the same simulation live, averaged over four random initializations.

Pre-LN is not free

If pre-LN were strictly better, this would be a short section. It has two known costs.

Later layers matter less. The stream's norm grows like $\sqrt{\ell}$ , and each pre-LN write has a fixed size, so a write at layer 40 is a much smaller fraction of the stream than a write at layer 4. Deep layers can only nudge the representation. Sun et al. (2025) connect this to the observation, which we met in §5.1, that many deep layers of open LLMs can be pruned, and propose scaling the output of each layer's norm by $1/\sqrt{\ell}$ to compensate.

Post-LN can end up better, if it trains. Liu et al. (2020) argued that post-LN's instability comes less from gradient imbalance than from an amplification effect, where heavy reliance on the residual branches makes small parameter updates cause large output changes. They also found that the lighter reliance on the branches that pre-LN enforces limits the quality a model can reach, and proposed an initialization (Admin) that keeps post-LN stable early in training while leaving its full capacity for later. DeepNet (Wang et al., 2022) took the same direction to 1,000 layers by scaling the residual before a post-LN, a scheme we return to in §5.7.

Where modern models put their norms

model	year	norm	placement
Transformer (Vaswani et al.)	2017	LayerNorm	post-LN
BERT	2018	LayerNorm	post-LN
GPT-2, GPT-3	2019, 2020	LayerNorm	pre-LN plus a final LN
T5	2019	RMS scaling, no bias	pre-norm plus a final norm
PaLM	2022	LayerNorm, no bias	pre-LN, attention and MLP in parallel from one norm
LLaMA 1 to 3	2023, 2024	RMSNorm	pre-norm plus a final norm
Gemma 2	2024	RMSNorm	before and after each sublayer
OLMo 2	2024	RMSNorm	on each sublayer's output, $\mathbf{x} + \operatorname{Norm}(F(\mathbf{x}))$ , plus QK-norm

Two newer patterns deserve a sentence each. The parallel block of GPT-J and PaLM computes $\mathbf{x} + \operatorname{Attn}(\operatorname{LN}(\mathbf{x})) + \operatorname{MLP}(\operatorname{LN}(\mathbf{x}))$ , so the two input projections can be fused into one matrix multiplication. PaLM reports about 15% faster training at large scale, with a small quality loss at 8B parameters and none at 62B. Normalizing sublayer outputs (Gemma 2's post-norms, OLMo 2's reordered norm) bounds what each sublayer can write while keeping the identity path clean, which addresses the amplification concern while keeping pre-LN's gradient profile.

🔬 QK-norm, a norm inside attention

Large models have one more place where a norm pays off. Attention logits are $\mathbf{q}\cdot\mathbf{k}/\sqrt{d_k}$ , and nothing in pre-LN stops the query and key projections from growing during training until the softmax saturates. QK-norm normalizes each head's query and key vectors before the dot product. Henry et al. (2020) used an L2 normalization with a learned scale, Dehghani et al. (2023) a LayerNorm, and recent LLMs mostly an RMSNorm. Either way, for fixed norm gains the logits are bounded by construction, however large $\mathbf{W}_Q$ and $\mathbf{W}_K$ become. It was essential for training a 22-billion-parameter vision Transformer, and OLMo 2, Gemma 3, and others use it. We study the failure it prevents in §5.7 and the full stability toolkit, QK-norm included, in §8.8.

Core

One residual sublayer, three placements

Write residual_sublayer(x, F, norm, mode) where x has shape (..., d), F is the sublayer function, and norm is a normalization function (both map (..., d) arrays to the same shape). mode is one of "pre" ( $\mathbf{x} + F(\operatorname{norm}(\mathbf{x}))$ ), "post" ( $\operatorname{norm}(\mathbf{x} + F(\mathbf{x}))$ ), or "output" ( $\mathbf{x} + \operatorname{norm}(F(\mathbf{x}))$ , the OLMo 2 ordering). Raise ValueError for anything else.

import numpy as np

def residual_sublayer(x, F, norm, mode):
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

In PyTorch

PyTorch ships both norms. nn.RMSNorm arrived in version 2.4. Here we check them against our NumPy versions and let autograd confirm our backward formula.

🔥 In PyTorchfor a machine with PyTorch

	BatchNorm	LayerNorm	RMSNorm
statistics over	the batch, per feature	one token's features	one token's features
removes	mean and scale per feature	the token's mean and scale	the token's scale only
train and test identical	no (running averages)	yes	yes
parameters	$2d$ plus running stats	$2d$	$d$
reductions per token	not applicable	two (mean, variance)	one (mean of squares)
typical use	CNNs	BERT, GPT-2, GPT-3	T5, LLaMA, Mistral, Gemma

Research

Mean and variance in one stable pass

Fused LayerNorm kernels compute the statistics in a single pass over the features. The naive one-pass formula $\sigma^2 = \operatorname{mean}(x^2) - \mu^2$ suffers catastrophic cancellation when $|\mu| \gg \sigma$ . Implement Welford's algorithm, welford_mean_var(x), for a 1-D array, processing one element at a time and returning (mean, var) with the biased variance. The tests check accuracy on data with a large offset, where the naive formula fails. This is the reduction that GPU kernels parallelize (§11.4).

import numpy as np

def welford_mean_var(x):
    """One pass over x (1-D). Returns (mean, biased variance)."""
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

🎯 Interview questions on this section

Why do Transformers use LayerNorm rather than BatchNorm?

LayerNorm computes statistics over each token's own features, so it is independent of the batch, of other sequences, and of other positions. BatchNorm's batch statistics mix tokens across sequences and positions (in a causal model that leaks future tokens into earlier positions), behave differently at inference because of running averages, and are noisy for the small, ragged per-device batches of language model training. LayerNorm computes the identical function in training and inference, including one-token-at-a-time generation.

Derive the gradient of LayerNorm with respect to its input. What properties does it have?

With $\hat{\mathbf{x}} = (\mathbf{x} - \mu\mathbf{1})/s$ and $\mathbf{g} = \boldsymbol{\gamma}\odot\partial\mathcal{L}/\partial\mathbf{y}$ , the Jacobian is $\frac1s(\mathbf{I} - \frac1d\mathbf{1}\mathbf{1}^\top - \frac1d\hat{\mathbf{x}}\hat{\mathbf{x}}^\top)$ , using $\partial\mu/\partial x_j = 1/d$ and $\partial s/\partial x_j = \hat{x}_j/d$ . So $\partial\mathcal{L}/\partial\mathbf{x} = \frac1s(\mathbf{g} - \operatorname{mean}(\mathbf{g})\mathbf{1} - \operatorname{mean}(\mathbf{g}\odot\hat{\mathbf{x}})\hat{\mathbf{x}})$ , which is $1/s$ times the orthogonal projection of $\mathbf{g}$ away from $\mathbf{1}$ and $\hat{\mathbf{x}}$ . It sums to zero (shifting $\mathbf{x}$ cannot change the output), it is orthogonal to $\hat{\mathbf{x}}$ up to $\epsilon$ (rescaling cannot either), and it is divided by $s$ , so large inputs receive small gradients. The gain and shift gradients are $\sum_t \partial\mathcal{L}/\partial\mathbf{y}_t\odot\hat{\mathbf{x}}_t$ and $\sum_t\partial\mathcal{L}/\partial\mathbf{y}_t$ .

What is RMSNorm, and why have most recent LLMs switched to it?

RMSNorm divides by the root mean square, $\boldsymbol{\gamma}\odot\mathbf{x}/\sqrt{\operatorname{mean}(\mathbf{x}^2)+\epsilon}$ , with no mean subtraction and no shift. It keeps LayerNorm's scale invariance and drops its shift invariance, which matters little because the mean direction of the residual stream is not special. It needs one reduction instead of two and has half the parameters, and Zhang and Sennrich (2019) found it matched LayerNorm's quality while running 7% to 64% faster depending on the model. In practice the choice is driven by simplicity and kernel speed.

Explain why post-LN Transformers need learning-rate warmup and pre-LN Transformers often do not.

In post-LN the stream is renormalized after every sublayer, so its norm is fixed and the gradient reaching the last layers does not shrink with depth, while every LayerNorm Jacobian on the way down shrinks it further. Gradients are largest at the top. Xiong et al. (2020) bound the last FFN layer's gradient by $\mathcal{O}(d\sqrt{\ln d})$ for post-LN and $\mathcal{O}(d\sqrt{\ln d/L})$ for pre-LN. In pre-LN the stream grows like $\sqrt{\ell}$ , the final norm divides by roughly $\sqrt{L}$ , and the clean identity path delivers that damped gradient evenly to all layers. With large top-layer gradients and a large learning rate, post-LN diverges early, and warmup keeps the early steps small until the optimizer's statistics and the model settle. Pre-LN trained without warmup in their experiments.

Pre-LN trains more easily. What does it give up?

Because the stream's norm grows with depth while each write has a fixed size, deep layers can only change the representation by a shrinking fraction, so later layers tend to be underused. Pruning studies remove many deep layers of pre-LN LLMs with little loss, and remedies include scaling each layer's norm output by $1/\sqrt{\ell}$ or normalizing sublayer outputs. Post-LN models that do train stably (with careful initialization such as Admin or DeepNorm) have been reported to reach better quality.

Why does a pre-LN Transformer need a final LayerNorm before the unembedding?

In pre-LN the stream itself is never normalized, so its scale after the last block grows with depth and varies by token. Without a final norm the logits would inherit that scale, giving uncontrolled softmax temperature and a loss landscape that depends on depth. The final norm puts every token on the same scale before the unembedding, and it is also where the $1/\sqrt{L}$ damping of the backward pass comes from.

Your small causal LM's training loss is suspiciously low but its samples are garbage. What is one normalization bug that does exactly this?

Normalizing over the time axis (or using BatchNorm over positions) instead of over features. The mean and variance then include future tokens, so every position can see information from later positions through the statistics. The model learns to exploit the leak during teacher-forced training, which makes the loss look great, but at generation time the future tokens do not exist. Check every reduction axis, and test causality by perturbing a future token and confirming earlier outputs do not change.

🧭 What to remember

LayerNorm standardizes each token over its own features, then applies a learned gain and shift. Geometrically it projects out the mean direction and scales onto a sphere of radius $\sqrt d$ .
Its input gradient is $\frac1s(\mathbf{g} - \operatorname{mean}(\mathbf{g})\mathbf{1} - \operatorname{mean}(\mathbf{g}\odot\hat{\mathbf{x}})\hat{\mathbf{x}})$ . The Jacobian is $1/s$ times the projection that removes the two directions LN ignores, $\mathbf{1}$ and $\hat{\mathbf{x}}$ .
RMSNorm drops the mean and the shift, keeps scale invariance, costs one reduction, and is the default in modern LLMs. Compute norm statistics in float32.
Post-LN normalizes the stream and gives large, depth-independent gradients to the top layers, which is why it needs warmup. Pre-LN keeps a clean identity path and damps gradients evenly by about $1/\sqrt{L}$ .
Pre-LN's growing stream makes deep layers less effective, which motivates output norms, depth-scaled norms, and careful initialization.
Next, §5.3 opens the other sublayer in the block, the feed-forward network, which holds two thirds of the weights.

Ba, Kiros & Hinton (2016). Layer Normalization. The original, motivated by recurrent networks.
Zhang & Sennrich (2019). Root Mean Square Layer Normalization. RMSNorm, its invariances, and the speed results.
Xiong et al. (2020). On Layer Normalization in the Transformer Architecture. The gradient-scale theorem and training pre-LN without warmup.
Liu et al. (2020). Understanding the Difficulty of Training Transformers. The amplification effect and the Admin initialization.
Sun et al. (2025). The Curse of Depth in Large Language Models. Why deep pre-LN layers are underused, and LayerNorm scaling.
OLMo Team (2024). 2 OLMo 2 Furious. A careful modern account of norm placement, QK-norm, and stability, with ablations.
Ioffe & Szegedy (2015). Batch Normalization. The predecessor, to see what LayerNorm changed.

Zero to Frontier is made for larger screens

LayerNorm and RMSNorm

Why not BatchNorm

LayerNorm normalizes each token by itself

The geometry, a projection and then a sphere

The backward pass, derived

LayerNorm's Jacobian is a scaled projection

RMSNorm drops the mean

RMSNorm forward and backward

Details that matter in real code

Where does the norm go, pre-LN or post-LN

The gradient-scale argument

Pre-LN is not free

Where modern models put their norms

One residual sublayer, three placements

In PyTorch

Mean and variance in one stable pass

Read more

Watch next