Batch Norm, Layer Norm, and RMSNorm — the normalization tour

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 21 / 80 · Module M04 · ~3 hours

0 · Normalization is axes plus state

Normalization is an operation plus axes, affine parameters and state policy. BatchNorm and LayerNorm share a centering formula but not the same statistical dependency. RMSNorm does not center. The distinction matters when batch size changes, a backbone is frozen, a padded sequence is processed or a checkpoint omits buffers.

The history is useful without a mythical depth wall or invented ablation. BatchNorm (2015) introduced minibatch normalization; LayerNorm (2016) emphasized per-case statistics; RMSNorm (2019) studied removing recentering. These developments do not prove that all transformers use one norm or that a fixed H100 speedup applies everywhere. Read a model's implementation rather than a family-wide slogan.

Prerequisites

Know tensor reductions, registered buffers, .train()/.eval() and the dropout expectation. The exact programs run on PyTorch 2.9 CPU with matching epsilon and affine settings. GPU throughput and paper-scale ablations are not executed.

1 · The core operation and its limits

For a chosen group of entries S,

With positive epsilon, pre-affine variance is $v/(v+\varepsilon)$, not exactly one; constant inputs become zero before affine shift. After affine parameters, mean/variance need not be zero/one at all. Affine parameters are initialized to one/zero in these standard examples, but their sharing shape differs across modules.

Learned gamma/beta can restore a fixed scale/offset, not generally reconstruct the removed input-dependent mean and variance. Two vectors differing by a constant offset have identical centered normalization, so one fixed beta cannot invert that many-to-one operation. Norms change representation, not just decorate a reversible identity.

The spreadsheet analogy is precise for a (N,C) matrix: BN uses each column's batch statistics; LN(C) uses each row's feature statistics; RMSNorm(C) uses each row's mean square without centering. For higher-rank tensors, name the actual axes rather than repeating 'rows versus columns'.

2 · BatchNorm — batch and spatial axes

For (N,C), statistics are per channel over N. For (N,C,H,W), BatchNorm2d reduces over (N,H,W) and shares one gamma/beta per channel. A 64×224×224 tensor has 3,211,264 entries per channel, but spatial correlation means these are not that many independent samples. Effective sample size and feature-map size matter; normalization later in a CNN may see far fewer spatial positions.

Two examples traveling together affect each other's training-mode outputs. This coupling can inject noise and sometimes regularize, but it is not a universal explanation for every large-batch performance gap. For independent scalar samples, variance of a sample mean is $\sigma^2/N$; correlated spatial/sequence samples violate the simple effective-N interpretation.

2.1 Training, running estimates and evaluation

By default, training uses current-batch statistics and updates running estimates. PyTorch uses biased variance (unbiased=False) for the training normalization and an unbiased variance estimate for the running-variance update. With momentum m,

where M is the number of values per channel and exceeds one. The momentum is the weight of the new observation, unlike a common optimizer-momentum convention. Evaluation uses running values when tracking is enabled. With track_running_stats=False, batch statistics are used in eval as well.

Batch size one is not automatically invalid for BatchNorm2d: spatial dimensions can provide multiple values per channel. PyTorch's training BatchNorm1d on (1,C) rejects the single-value case; it does not silently return all zeros. Evaluating with valid running statistics supports that input.

2.2 Small batches and correlated examples

Small per-device batches can give noisy or unrepresentative statistics; gradient accumulation does not enlarge the BatchNorm batch. SyncBatchNorm combines the participating ranks' statistics for a synchronized forward, not all sequential microbatches across time. GroupNorm and LayerNorm remove batch coupling by choosing different axes, not by estimating the same batch statistic more accurately.

2.3 Calibration and frozen backbones

Setting requires_grad=False freezes parameter gradients, not running buffers. A backbone still in train mode can change BN running statistics. To recalibrate, place the overall model in eval mode (disabling dropout), then put only BatchNorm modules in train mode and run representative training/calibration, not test, data under no-grad. Reset/average running statistics according to the intended estimator; restore modes afterward. An eval-only pass will not update standard tracked BN buffers.

One epoch is neither always sufficient nor always insufficient: at momentum .1, the initial estimate's weight after k updates is $.9^k$. That is about .349 after ten updates and .00515 after fifty, but a changing/correlated data distribution still matters.

3 · LayerNorm — specified trailing dimensions

For (N,T,D), LayerNorm(D) normalizes each token over D independently of other tokens and batch members. LayerNorm((T,D)) instead pools both trailing dimensions per example; it is a different operation. LayerNorm(T) on (N,T,D) normally errors if T differs from D—it does not magically transpose and normalize time. Equal sizes can conceal a mistaken axis assumption, so test shapes and dependence directly.

LayerNorm has no running mean/variance and uses the same formula in train and eval. This says nothing about other layers in the containing model. Padding tokens normalized separately do not contaminate other tokens' normalization statistics, but attention masks and loss masks still matter. LN is not a substitute for them.

4 · A worked example — same input, three norms

The two current exact programs are retained. Row one has mean 2.5, biased variance 1.25 and mean square 7.5. Row two has mean 6.5, variance 1.25 and mean square 43.5. Column means are (3,4,5,6) and biased variances all 4.

import torch
torch.manual_seed(0)
 
x = torch.tensor([[1.0, 2.0, 3.0, 4.0],
                  [5.0, 6.0, 7.0, 8.0]])       # (N=2, C=4)
eps = 1e-5
 
# --- BatchNorm: per feature, over batch (dim=0) ---
bn_mean = x.mean(dim=0)                          # [3, 4, 5, 6]
bn_var  = x.var(dim=0, unbiased=False)           # [4, 4, 4, 4]
bn_out  = (x - bn_mean) / torch.sqrt(bn_var + eps)
# [[-1, -1, -1, -1],
#  [+1, +1, +1, +1]]                             # each COLUMN is normalized
 
# --- LayerNorm: per example, over features (dim=1) ---
ln_mean = x.mean(dim=1, keepdim=True)            # [[2.5], [6.5]]
ln_var  = x.var(dim=1, unbiased=False, keepdim=True)  # [[1.25], [1.25]]
ln_out  = (x - ln_mean) / torch.sqrt(ln_var + eps)
# [[-1.34, -0.45, 0.45, 1.34],
#  [-1.34, -0.45, 0.45, 1.34]]                   # each ROW is normalized
 
# --- RMSNorm: per example, over features, no mean subtraction ---
ms      = (x ** 2).mean(dim=1, keepdim=True)     # [[7.5], [43.5]]
rms_out = x / torch.sqrt(ms + eps)
# [[0.365, 0.730, 1.095, 1.461],                  # row 1: not zero mean!
#  [0.758, 0.910, 1.061, 1.213]]                  # row 2: not zero mean!

import torch.nn as nn
 
bn = nn.BatchNorm1d(4, eps=eps, affine=False, track_running_stats=False)
bn.train()
assert torch.allclose(bn(x), bn_out, atol=1e-5)
 
ln = nn.LayerNorm(4, eps=eps, elementwise_affine=False)
assert torch.allclose(ln(x), ln_out, atol=1e-5)
 
rms = nn.RMSNorm(4, eps=eps, elementwise_affine=False)  # PyTorch 2.4+
assert torch.allclose(rms(x), rms_out, atol=1e-5)
 
print("all three norms match to 1e-5")

The printed rounded comments ignore tiny epsilon corrections. BN's pre-affine columns are zero mean with variance slightly below one; LN's rows behave analogously. RMSNorm rows have mean square slightly below one and generally nonzero mean. Compare exact formulas, not rounded comments, at tolerance 1e-5. Different epsilon, affine parameters, train/eval state or variance convention can fail a comparison even when the axis is right.

5 · Why per-token norms suit transformers

Three architectural advantages explain the usual choice without claiming BN is impossible for sequences:

1. Variable length: reducing only over a token's features avoids pooling padding into shared normalization statistics. Sequence-aware BN variants can mask or alter statistics, but that is extra design.
2. Autoregressive decoding: one-token inference can use the same per-token norm formula as training, without a frozen estimate of a pooled training distribution.
3. Batch independence: moving the same token vector between batches leaves its LN/RMSNorm output unchanged. BN with frozen eval statistics is also batch-independent, but its training formula and estimated state differ.

Unrelated examples do not make a population mean 'a quantity that does not exist'; mixed distributions have means. The question is whether pooling is useful and consistent with the training/inference objective. Per-token normalization is a convenient contract, not the only conceivable architecture.

6 · RMSNorm — scale without recentering

For the explicit PyTorch epsilon convention,

It omits mean subtraction and, in PyTorch's standard module, affine bias. A scale parameter remains when elementwise_affine=True. RMSNorm is not LN with a slightly faster implementation: it preserves information about an input's offset that LN removes. The default epsilon is dtype-sensitive in this API; set it explicitly for cross-implementation comparisons.

An epsilon outside the root defines a different function. Neither formula returns infinity on an exactly zero vector with positive epsilon, so 'test zeros' cannot distinguish the placement. Test small nonzero inputs and compare gradients. An input-dependent mean cannot generally be absorbed by the next layer's fixed bias, especially in architectures whose Linear layers have no bias.

6.1 Study claims versus your own benchmark

The RMSNorm paper studies particular architectures and reports measured efficiency; it does not establish an invariant 7% or 6.8% H100 gain at equal loss. The old 'Narang 2024 ablation at arXiv:2410.06264' and associated exact model/token/loss/MMLU figures are not accepted evidence. That identifier actually names Think While You Generate: Discrete Diffusion with Planned Denoising, not a normalization comparison. Likewise, an OLMo report is not interchangeable with a controlled one-change ablation unless its actual experiment says so.

If removing one beta vector per norm of width D, savings are D parameters per norm. With two such norms in each of 32 blocks at D=4096, savings are 262,144, not 131,072; 131,072 assumes one per block. Final norms and implementation choices alter the total.

A valid benchmark matches epsilon, affine settings, dtype, layout, shape and warmup; times synchronized forward/backward separately and together; checks outputs/gradients against references; and measures trained quality with repetitions. Fewer reductions can help, but fusion, memory bandwidth and kernel dispatch can dominate. GPU (64,4096) timing and transformer swapping remain unexecuted protocols, not prefilled 5–10% wins.

7 · GroupNorm and InstanceNorm

For (N,C,H,W), GroupNorm(G,C) splits channels into G groups and reduces over (C/G,H,W) for each example/group. C must be divisible by G. Statistics do not depend on other batch members. GroupNorm(1,C) shares the full (C,H,W) reduction with LayerNorm((C,H,W)), but standard affine parameter shapes differ: GN has one scale/bias per channel; LN may have one per element. Equality requires matching affine behavior.

InstanceNorm2d reduces over (H,W) separately per example/channel. Defaults differ from BN: affine and running-stat tracking are normally disabled, so train/eval use input statistics. Optional tracked state changes that. GroupNorm(C,C) and InstanceNorm2d have matching reduction sets under compatible epsilon/affine/state settings. Style-transfer applications motivate removing instance contrast; this does not make it the universal choice for every GAN.

GroupNorm's batch independence does not imply identical trained accuracy for all physical batch sizes: optimizer noise/data order/update counts can still change training.

8 · Pre-norm, post-norm and the residual Jacobian

For a sublayer F and normalization N,

Their local Jacobians are $I+J_FJ_N$ and $J_N(I+J_F)$ respectively. Pre-norm leaves an explicit identity summand; post-norm puts normalization on that route too. If F is initially zero, pre-norm gives identity, but post-norm still removes the constant-shift direction. That is a concrete gradient-path distinction, not a guarantee that all pre-norm gradients are depth-independent or training is always stable.

The original 2017 Transformer used post-norm. Xiong et al. (2020) analyze initialization gradient behavior and warmup implications under their assumptions. Pre-norm is common, often with a final norm, but modern models also use post/hybrid arrangements. DeepNorm (2022, not a 2024 invention) modifies residual scaling and initialization to study very deep transformers. Post-norm's outputs are normalized; the simplistic claim that its residual output simply grows without bound despite the norm is not the mechanism.

9 · Complete custom modules and failure controls

The CPU program implements the requested BN/LN/RMS modules, including learnable affine parameters and BN running buffers. This BN is intentionally bounded to 2-D (N,C), fixed numeric momentum, tracking enabled, and at least two training samples—not every framework option. The comparison checks forward, input gradients, buffers and eval, not only one output tensor.

import copy
import torch
from torch import nn
 
class MyBatchNorm1d(nn.Module):
    def __init__(self,c,eps=1e-5,momentum=.1):
        super().__init__(); self.eps=eps; self.momentum=momentum
        self.weight=nn.Parameter(torch.ones(c)); self.bias=nn.Parameter(torch.zeros(c))
        self.register_buffer("running_mean",torch.zeros(c))
        self.register_buffer("running_var",torch.ones(c))
        self.register_buffer("num_batches_tracked",torch.tensor(0,dtype=torch.long))
    def forward(self,x):
        if x.ndim != 2 or x.shape[1] != len(self.weight): raise ValueError("shape")
        if self.training:
            if len(x) < 2: raise ValueError("need two values per feature")
            mean=x.mean(0); var=x.var(0,unbiased=False)
            with torch.no_grad():
                self.running_mean.lerp_(mean,self.momentum)
                self.running_var.lerp_(var*len(x)/(len(x)-1),self.momentum)
                self.num_batches_tracked.add_(1)
        else: mean,var=self.running_mean,self.running_var
        return (x-mean)*torch.rsqrt(var+self.eps)*self.weight+self.bias
 
class MyLayerNorm(nn.Module):
    def __init__(self,d,eps=1e-5):
        super().__init__(); self.eps=eps
        self.weight=nn.Parameter(torch.ones(d)); self.bias=nn.Parameter(torch.zeros(d))
    def forward(self,x):
        return (x-x.mean(-1,keepdim=True))*torch.rsqrt(x.var(-1,unbiased=False,keepdim=True)+self.eps)*self.weight+self.bias
 
class MyRMSNorm(nn.Module):
    def __init__(self,d,eps=1e-5):
        super().__init__(); self.eps=eps; self.weight=nn.Parameter(torch.ones(d))
    def forward(self,x):
        return x*torch.rsqrt(x.square().mean(-1,keepdim=True)+self.eps)*self.weight
 
torch.manual_seed(21)
for custom,ref in [(MyBatchNorm1d(4),nn.BatchNorm1d(4)),
                   (MyLayerNorm(4),nn.LayerNorm(4)),
                   (MyRMSNorm(4),nn.RMSNorm(4,eps=1e-5))]:
    custom.double(); ref.double(); ref.load_state_dict(custom.state_dict())
    a=torch.randn(7,4,dtype=torch.float64,requires_grad=True)
    b=a.detach().clone().requires_grad_(); probe=torch.randn_like(a)
    u,v=custom(a),ref(b)
    assert torch.allclose(u,v,atol=1e-10,rtol=1e-9)
    (u*probe).sum().backward(); (v*probe).sum().backward()
    assert torch.allclose(a.grad,b.grad,atol=1e-9,rtol=1e-8)
    for name,state in custom.state_dict().items():
        assert torch.allclose(state,ref.state_dict()[name],atol=1e-10)
    custom.eval();ref.eval()
    assert torch.allclose(custom(a),ref(b),atol=1e-10)
# Exact running variance convention: batch biased var 4, unbiased 8.
bn=nn.BatchNorm1d(1,momentum=.1); bn(torch.tensor([[1.],[5.]]))
assert torch.allclose(bn.running_mean,torch.tensor([.3]))
assert torch.allclose(bn.running_var,torch.tensor([1.7]))
# Frozen weights do not freeze running buffers; eval does.
bn.requires_grad_(False); before=bn.running_mean.clone()
with torch.no_grad(): bn(torch.tensor([[3.],[7.]]))
assert not torch.equal(before,bn.running_mean)
bn.eval();before=bn.running_mean.clone();bn(torch.tensor([[10.]]))
assert torch.equal(before,bn.running_mean)
try: nn.BatchNorm1d(4)(torch.ones(1,4))
except ValueError: pass
else: raise AssertionError("singleton training BN accepted")
assert torch.isfinite(nn.BatchNorm2d(4)(torch.randn(1,4,2,2))).all()
# Different batch companions: BN changes, LN does not.
a=torch.tensor([[1.,2.,3.,4.]])
same=torch.cat([a,a]);other=torch.cat([a,torch.tensor([[9.,0.,8.,-3.]])])
bn=nn.BatchNorm1d(4,track_running_stats=False);ln=nn.LayerNorm(4)
assert not torch.allclose(bn(same)[0],bn(other)[0])
assert torch.equal(ln(same)[0],ln(other)[0])
ln.train();u=ln(a);ln.eval();assert torch.equal(u,ln(a)) and not list(ln.buffers())
# Offset invariance and epsilon-placement negative controls.
tiny=torch.full((1,4),1e-7,dtype=torch.float64);eps=1e-5
inside=tiny/torch.sqrt(tiny.square().mean(-1,keepdim=True)+eps)
outside=tiny/(torch.sqrt(tiny.square().mean(-1,keepdim=True))+eps)
assert not torch.allclose(inside,outside,atol=1e-6)
assert torch.allclose(ln(a),ln(a+10),atol=1e-6)
rms=nn.RMSNorm(4,eps=eps)
assert not torch.allclose(rms(a),rms(a+10))
z=torch.randn(2,4,3,3)
assert torch.allclose(nn.GroupNorm(1,4,affine=False)(z),
                      nn.LayerNorm((4,3,3),elementwise_affine=False)(z),atol=1e-6)
assert torch.allclose(nn.GroupNorm(4,4,affine=False)(z),
                      nn.InstanceNorm2d(4,affine=False)(z),atol=1e-6)
# Pre-norm with zero branch has identity Jacobian; post-norm does not.
t=torch.randn(4,dtype=torch.float64,requires_grad=True)
jpre=torch.autograd.functional.jacobian(lambda a:a+0*ln.double()(a),t)
jpost=torch.autograd.functional.jacobian(lambda a:ln.double()(a+0*a),t)
assert torch.equal(jpre,torch.eye(4,dtype=torch.float64))
assert torch.allclose(jpost@torch.ones(4,dtype=torch.float64),torch.zeros(4,dtype=torch.float64),atol=1e-12)
print("custom norms, gradients, buffers, axes and negative controls passed")

11 · Try it yourself — calibration and fair model comparison

A calibration helper must selectively update BN while keeping dropout disabled and restore all original mode flags even on error. The helper below restores modes and momenta, but it is not transactional: if a later batch fails, earlier running-buffer updates remain. Empty input leaves freshly reset buffers. Require nonempty, representative data and a successful full pass before accepting the calibrated checkpoint; otherwise reload the prior buffers and investigate. For equal-size batches, momentum=None forms a cumulative average of batch estimates; unequal sizes require care because equal weighting of batch means is not a pooled-example mean. Even with equal batch sizes, the average within-batch variance is not the pooled sample variance: pooling also includes differences between batch means. For batches (0,2) and (10,12), BN averages unbiased batch variances to 2, whereas all four values have unbiased variance 104/3. Both estimates have a role; do not label this helper an exact pooled-statistics estimator. It updates tracked BN buffers only; track_running_stats=False leaves no running estimates to calibrate.

import torch
from torch import nn
 
def recalibrate_bn(model,batches):
    modules=list(model.modules()); modes=[m.training for m in modules]
    bns=[m for m in modules if isinstance(m,nn.modules.batchnorm._BatchNorm)]
    momenta=[m.momentum for m in bns]
    try:
        model.eval()
        for m in bns:
            m.reset_running_stats();m.momentum=None;m.train()
        with torch.no_grad():
            for x in batches: model(x)
    finally:
        for m,v in zip(bns,momenta):m.momentum=v
        for m,v in zip(modules,modes):m.training=v
 
model=nn.Sequential(nn.BatchNorm1d(2),nn.Dropout(.9))
model.eval()
recalibrate_bn(model,[torch.tensor([[1.,3.],[5.,7.]])]*3)
assert torch.allclose(model[0].running_mean,torch.tensor([3.,5.]))
assert model[0].num_batches_tracked.item()==3 and not model.training
assert not model[0].training and not model[1].training
print("selective BN calibration passed")

For the original batch-size-four vision experiment, compare BN, GN and a precisely specified LN axis/affine shape on the same splits, augmentations and compute. Keep both physical and effective batch counts, train/eval losses, BN running-stat drift, validation quality and throughput. BN may suffer noisy statistics; GN/LN can change inductive bias and optimization despite batch independence. A curve alone cannot identify the failure. Repeat seeds, and do not prescribe a guaranteed winning norm.

12 · Recap and worked recall answers

Q1 — Is the only difference an axis? Axis is central, but affine sharing and running-state/train-eval policy also differ. Batch-coupled statistics affect training predictions; default BN eval uses frozen buffers. RMSNorm additionally omits centering.

Q2 — What are the 2×4 outputs? BN columns are approximately (-1,+1); each LN row is approximately (-1.342,-.447,.447,1.342); RMS rows are approximately (.365,.730,1.095,1.461) and (.758,.910,1.061,1.213). Epsilon prevents exact unit variance/RMS, and RMS rows are not centered.

Q3 — Why per-token normalization? It avoids cross-example normalization dependence, does not pool sequence padding into other tokens' statistics, and uses the same local calculation for one-token decoding. Attention/loss masks remain necessary; specialized BN alternatives are not impossible.

Q4 — RMS simplifications and speed? No recentering and no standard beta vector; keep gamma. The old 6.8% H100 ablation number was unsupported. Measure the actual shape/kernel/precision instead of replacing it with another universal figure.

Q5 — Pre versus post? Pre-norm has Jacobian I+J_F J_N; post-norm has J_N(I+J_F). The identity route helps explain different initialization/warmup behavior. It is not a universal no-divergence theorem; DeepNorm and other hybrid designs show the choice is not binary fashion.

Q6 — Norm selection at batch four? Compare actual axes, affine capacity, training/eval state and held-out metrics with matched data/compute/seeds. Include train-only BN calibration and a frozen-buffer check. Do not equate batch independence with identical trained accuracy.

Sources and execution boundary

PyTorch BatchNorm v2.9.0, normalization modules v2.9.0, InstanceNorm source, RMSNorm paper, GroupNorm paper, Santurkar et al., Xiong et al., and DeepNorm support bounded formula, state and architecture claims. CPU exact programs are executed; no H100 benchmark, trained transformer ablation or universal model-adoption claim is made.

Optional pre-read resources

These inherited resource links are optional, not the primary evidence for this review. Video identity, duration and current availability are not certified.

• Batch Normalization (C2W3L04)
• Batch Norm vs Layer Norm — Umar Jamil
• Batch Normalization (Ioffe & Szegedy 2015)
• Layer Normalization (Ba, Kiros, Hinton 2016)

Previous: ← DL S020 · Dropout · Next: DL S022 · Weight Initialization →