Normalization: BatchNorm, LayerNorm, and RMSNorm

Take [1, 3], temporarily ignoring epsilon and learned parameters. LayerNorm subtracts the mean, 2, and divides by the standard deviation, 1, giving [-1, 1]. RMSNorm instead divides by the root mean square, √5, without centering. Both change scale, but they are different operations. Next we place them inside a batch of sequences and identify the reduction axes.

The core difference: normalization axes

BatchNorm takes its statistics across a batch of examples. LayerNorm takes them across one example's own features.

Start with the axes, then check running statistics and train / eval modes. Together they determine behavior with variable lengths, batch size 1, and autoregressive decoding. The axis alone is not the entire implementation contract.

Each row in the figure is one feature vector. For a sequence shaped [B,T,d], token-wise normalization means LayerNorm(d) / RMSNorm(d). Passing (T,d) includes time as well. Check normalized_shape, not just the layer's name. Both LayerNorm and RMSNorm allow normalization over multiple trailing dimensions.

Three formulas

BatchNorm (for each feature $j$, statistics over the batch dimension):

$$\mu_j = \frac{1}{N}\sum_{i=1}^{N} x_{ij}, \qquad \sigma_j^2 = \frac{1}{N}\sum_{i=1}^{N}(x_{ij}-\mu_j)^2$$

$$\hat{x}_{ij} = \frac{x_{ij}-\mu_j}{\sqrt{\sigma_j^2+\epsilon}}, \qquad y_{ij} = \gamma_j \hat x_{ij} + \beta_j$$

LayerNorm (for each example $i$, statistics over the feature dimension):

$$\mu_i = \frac{1}{d}\sum_{j=1}^{d} x_{ij}, \qquad \sigma_i^2 = \frac{1}{d}\sum_{j=1}^{d}(x_{ij}-\mu_i)^2, \qquad y_{ij} = \gamma_j\,\frac{x_{ij}-\mu_i}{\sqrt{\sigma_i^2+\epsilon}} + \beta_j$$

RMSNorm (root-mean-square scaling, commonly with learned gain $\gamma$ only):

$$y_{ij} = \gamma_j\,\frac{x_{ij}}{\sqrt{\frac{1}{d}\sum_{k} x_{ik}^2+\epsilon}}$$

The learned $\gamma$, and $\beta$ where present, run per feature, with length $d$. BatchNorm and LayerNorm differ in their reduction axes. LayerNorm and RMSNorm use the same feature axis but different statistics: variance versus the raw second moment. RMSNorm is not LayerNorm assuming that the mean is already zero.

Add 10 to the opening example: [1, 3] becomes [11, 13]. Ignoring epsilon, LayerNorm still returns [-1, 1]; the RMSNorm denominator changes from $\sqrt{5}$ to $\sqrt{145}$, so its output changes too. It does not remove the shift. Pre-Norm versus Post-Norm is a separate choice about placement, not the normalization formula. The RMSNorm paper investigates scale normalization without re-centering.

Why can train and eval differ on the same batch?

The BatchNorm formula above uses a two-dimensional input. On [N,C,L], BatchNorm1d reduces over N,L per channel. By default, training uses current-input statistics and eval uses accumulated running statistics. With track_running_stats=False, eval uses the current input too. These are statistical state, not optimizer-trained parameters such as gamma / beta. See PyTorch 2.14 (checked October 2026).

A distribution shift is not the only possible explanation. Even the same batch [1,3] can produce different outputs. Set eps=0, disable affine parameters, and use momentum=1 so the running statistics come directly from this batch:

There are two denominators: training uses $v_b=\sum(x-\mu)^2/m$, while the running-variance update uses $v_u=\sum(x-\mu)^2/(m-1)$. Here $m$ counts the values reduced per channel. With momentum=a, the update is $v_{\rm run}\leftarrow(1-a)v_{\rm run}+av_u$; the mean uses the same weighting rule. This is not optimizer momentum. The original BN paper, §3.1, distinguishes training statistics from inference estimates too.

import torch

def batchnorm_train_eval_example():
    inputs = torch.tensor([[1.0], [3.0]])
    norm = torch.nn.BatchNorm1d(1, eps=0, momentum=1, affine=False)
    training = norm(inputs)
    stored_variance = norm.running_var.clone()
    norm.eval()
    evaluation = norm(inputs)
    return training, evaluation, stored_variance

training, evaluation, stored_variance = batchnorm_train_eval_example()
print(training.flatten().tolist())
print(evaluation.flatten().tolist(), stored_variance.item())

Gradient accumulation does not enlarge BN's batch

A forward over [0,2,10,12] has mean 6. Two forwards over [0,2] and [10,12] instead use means 1 and 11. You can accumulate gradients before updating weights, but that does not combine two completed normalizations into one. Running statistics also update on each forward, not at optimizer.step().

Starting with running mean 0 and variance 1, with momentum=0.1:

The first row uses sample variance $104/3$. Each microbatch in the second row has sample variance 2, but updates the state twice. Normalized activations, subsequent losses, and gradients may differ too. Matching the effective batch size therefore does not establish equivalent training with BN.

One more distinction: torch.no_grad() disables gradient recording; it does not switch BN modes. A training-mode module still updates running statistics. Validation commonly uses both .eval() and no_grad() because they control different behaviors. Autograd documentation

Four concepts: answer first, then go deeper

For input-feature min–max and Z-score scaling, see preprocessing versus network normalization.

Post-LN (2017 original)          Pre-LN (common today)
x = LN(x + Attn(x))             x = x + Attn(LN(x))
x = LN(x + FFN(x))              x = x + FFN(LN(x))
                                 ...
                                 x = LN(x)   ← final LN

Where modern LLMs add more norms: QK-norm and pre + post norm

Pre-LN controls the $x$ that enters each sublayer, but attention has another quantity that can grow on its own: the logit $q\cdot k/\sqrt{d_k}$. The input $x$ is normalized, yet the lengths of $q=xW_Q$ and $k=xW_K$ still follow the projection weights.

Take a small example. Only score gaps matter here, so leave out the constant $\sqrt{d_k}$. With $q=[2,1]$ and keys $[1,0]$ and $[0,1]$, the scores are 2 and 1 and the softmax is about [0.731, 0.269]. If training scales $W_Q$ and $W_K$ each by 3, then $q=[6,3]$, the keys become $[3,0]$ and $[0,3]$, and the scores are 18 and 9. The gap grows from 1 to 9 and the softmax is about [0.99988, 0.00012]: nearly one-hot, so the gradient back through the softmax to the scores becomes small (derivation in multi-head attention). ViT-22B reported exactly this: attention logits became extremely large, attention entropy fell to near zero, and the loss diverged after a few thousand steps.

QK-norm applies RMSNorm to each head's $q$ and $k$ before the dot product, each with its own learnable gain:

$$ \hat q=\operatorname{RMSNorm}(q),\qquad \hat k=\operatorname{RMSNorm}(k),\qquad s=\frac{\hat q\cdot\hat k}{\sqrt{d_k}}. $$

As shown above, RMSNorm ignores an overall scale when epsilon is ignored: $[6,3]$ and $[2,1]$ normalize to the same vector, so scaling the weights by 3 no longer changes the score. With unit gain, $\|\hat q\|=\|\hat k\|=\sqrt{d_k}$, so $|s|\le\sqrt{d_k}$; for $d_k=128$ the bound is about 11.3. Normalization removes growth caused by uniformly scaling the projection norms. The directions of q and k can still become more aligned and raise a score with fixed gains; learned gains can change the bound as well. Neither effect is guaranteed to happen slowly. RoPE is a rotation and preserves length, so the bound does not depend on whether the norm comes before or after RoPE; Gemma 3's Transformers implementation applies it before.

The Qwen3 technical report says it removed Qwen2's QKV bias and introduced QK-Norm for stable training; the Gemma 3 technical report §2 replaced Gemma 2's soft-capping with QK-norm. Soft-capping squeezes a score into $(-c,c)$ with $c\tanh(s/c)$ after the fact; QK-norm limits the lengths of $q$ and $k$ at the source.

Norm before and after. Gemma 3 puts an RMSNorm both before and after each sublayer:

$$ h=x+\operatorname{Norm}_{\text{post}}\big(\operatorname{Attn}(\operatorname{Norm}_{\text{pre}}(x))\big), $$

and the FFN does the same. This post-norm acts only on the branch output; nothing normalizes the sum after it is added back, so the trunk keeps the Pre-LN identity path above. It shares a name with the 2017 Post-LN, which normalizes the whole residual sum, but sits in a different place. One way to read it is as a limit on how large an update each sublayer writes into the residual stream.

How to choose a normalization method

One rule: whenever "the same input must produce the same output in different batches" is a hard requirement, do not let normalization use batch statistics.

Modern LLMs often use RMSNorm. Removing mean centering can simplify computation, but actual speed and synchronization depend on kernel fusion and tensor partitioning, not just the number of sums in the formula.

Verify it

../code/norm_compare.py compares reduction axes and three distinct cases: BatchNorm training on [1,C] raises an error; evaluation with running statistics works; training on [1,C,L] also works when L>1, but its statistics span time. These are not all the same “small-batch problem.”

The figure is generated by ../code/make_norm_figures.py.

Self-check

Next

Normalization keeps each layer's input scale under control; trainable depth also needs residual connections. Once you have read both, continue to FFN and SwiGLU, then assemble the pieces in the vanilla Transformer.

The 30-second mental model

Normalization is not about keeping every representation permanently at mean zero and variance one. It is about controlling the numerical scale fed into each sublayer and making a deep network easier to optimize.