BatchNorm cut ImageNet training time in half overnight in 2015. The trick is embarrassingly simple — re-center and re-scale the activations at every layer so the next layer always sees a clean, standardized input. LayerNorm (transformers), GroupNorm (small batches), RMSNorm (LLaMA) are all variations on "subtract the mean, divide by the std." Pick the wrong axis and your model silently breaks at inference.
Learning Objectives
After this lesson, you will be able to:
Explain why deep networks need normalization layers — activations drift across the layer stack and the optimizer ends up chasing a moving target instead of fitting your data
Pick between BatchNorm, LayerNorm, GroupNorm, InstanceNorm, and RMSNorm based on which dimension you reduce over and what your batch size and architecture look like
Use BN correctly in production: switch to model.eval() at inference so you use running statistics, not batch statistics from a single test sample
Read the per-token normalization in a transformer block, recognize pre-norm vs post-norm placement, and explain why LayerNorm (or RMSNorm) — never BatchNorm — became the transformer default
Don't worry if normalization layers feel mysterious — they all do roughly the same thing (subtract a mean, divide by a standard deviation, then learnable scale and shift). The only real question is which slice of the activation tensor you compute the mean and variance over. Once you see that, the four norms become four flavors of the same idea.
The original BN paper called this problem internal covariate shift — the input distribution of each layer keeps shifting as earlier layers learn. Later work (Santurkar et al. 2018, "How Does Batch Normalization Help Optimization?") argued the real win is that BN smooths the loss landscape: bigger steps remain safe because the loss gradient is less erratic when activations are normalized. Either way, the empirical fact stands — normalization makes deep networks dramatically more trainable.
The only difference between BatchNorm, LayerNorm, GroupNorm, InstanceNorm, and RMSNorm is which axes go into the reduction S. Get that one detail right and everything else falls out.
A typical CNN activation has shape (N, C, H, W): N samples in the batch, C channels, height H, width W. A typical transformer activation has shape (N, T, D): N sequences, T tokens, D feature dimensions.
For a CNN activation of shape (N, C, H, W), BatchNorm computes C separate (μ, σ²) pairs — one per channel — averaging across the batch dimension and spatial dimensions for each channel.
The train-vs-eval split is the key BN gotcha. During training, BN uses the current batch's statistics. During inference (model.eval()), it uses running averages of mean and variance accumulated during training (with exponential decay controlled by momentum, typically 0.1). This is necessary because at inference you might serve one sample at a time — you cannot compute a meaningful batch statistic from N=1.
Try it! Open the Python REPL and verify the train-vs-eval split: import torch; bn = torch.nn.BatchNorm2d(3); x = torch.randn(8, 3, 4, 4); bn.train(); print(bn(x).mean().item()); bn.eval(); print(bn(x).mean().item()) — the train output is exactly 0 (batch mean was subtracted), the eval output is small but nonzero (running mean was used).
#LayerNorm: reduce over the feature dimension per sample
For a transformer activation of shape (N, T, D), LayerNorm computes N · T separate (μ, σ²) pairs — one per token — averaging across the Dfeature dimension for each individual token. The batch and time dimensions are not part of the reduction.
LayerNorm has no train/eval mode difference and no running statistics — every forward pass computes fresh statistics from the input. That is one less bug surface than BN, and it is exactly why generative inference (one token at a time, batch=1) works seamlessly with LN.
#GroupNorm: reduce over (H, W) within a channel group
For a CNN activation of shape (N, C, H, W), GroupNorm splits the C channels into G groups (typically 32) and computes N · G separate (μ, σ²) pairs — averaging across the spatial dimensions and the channels within each group for each sample.
GroupNorm sits between BatchNorm (one stat per channel, mixing the batch) and InstanceNorm (one stat per channel per sample, no mixing). When G = C, GroupNorm becomes InstanceNorm. When G = 1, GroupNorm becomes LayerNorm (almost — applied over channel and spatial axes). This makes it a flexible knob: high G if your batch is small, low G if you have plenty of channels per group.
GroupNorm is the default fix for object detection (Mask R-CNN trains with batch size 1-2 per GPU) and medical imaging (3D scans where memory limits batches to 1-4).
InstanceNorm reduces over (H, W) per sample per channel — one statistic per (N, C) slot. Used in style transfer (Ulyanov 2016) where you want to wipe out per-image contrast and color bias, leaving only structural content.
RMSNorm drops the mean-centering step entirely and only rescales by the root-mean-square:
Your CNN trains beautifully on ImageNet with batch size 32 per GPU. You move to a higher-resolution dataset that only fits batch size 2 per GPU, and your validation accuracy collapses while training loss looks fine. Most likely cause and fix?
The right answer is the second one. With batch size 2, BN's per-batch mean and variance are essentially noise — the channel statistics swing wildly run-to-run, the running averages converge to garbage, and inference accuracy drops. GroupNorm (Wu & He 2018) was created exactly to fix this: it ignores the batch dimension entirely.
The transformer story is the same logic taken to the limit. A transformer at inference might be generating one token at a time, autoregressively, with effective batch size 1. There is no batch over which to compute meaningful statistics. LayerNorm just works: it normalizes across the feature dimension of a single token, regardless of how many sequences are in the batch and regardless of train vs eval mode. Layer Norm fits the architecture; BatchNorm fights it.
The original transformer (Vaswani 2017) put LayerNorm after the residual addition (LayerNorm(x + Attention(x))). This is post-norm. Modern transformers (GPT-2 onward) put LayerNorm before the sublayer: x + Attention(LayerNorm(x)). This is pre-norm.
Post-norm has slightly higher final accuracy on small models but blows up at large depth without careful warmup. Pre-norm is more stable during training (gradients flow cleanly through the residual connection) and won out as transformers scaled past 24 layers. All large LLMs use pre-norm: GPT-2, GPT-3, GPT-4, LLaMA, Mistral.
Tests · Manual BN should match torch.BatchNorm2d in train mode. Manual LN should match torch.LayerNorm. The train-vs-eval demo should show ≈0 train output and ≈10 eval output for a single-sample input.
BatchNorm Distribution TrackerInteractive
Watch how BatchNorm reshapes the per-channel activation distribution across training steps, and what happens to the running mean and variance.
Loading visualization...
BatchNorm vs LayerNorm Side by SideInteractive
Same activation tensor, two different reduction axes. Toggle between BN and LN to see which slice of the tensor each norm computes statistics over.
All four normalizations are the same recipe with a different reduction axis. Subtract a mean, divide by a standard deviation, learnable scale and shift; the only choice is which slice of the activation tensor counts as "the same group"
BatchNorm fights you at small batch sizes. Below batch 8 it loses accuracy, below batch 4 it stops working, and at batch 1 it produces zeros; switch to GroupNorm or LayerNorm in those regimes
LayerNorm is the transformer default for a reason. It has no batch dependence and no train/eval mode split, so it works identically at training time and at single-token autoregressive inference
RMSNorm is LayerNorm minus the mean-centering. Saves 7-30% compute at LLM scale with no measurable accuracy loss; LLaMA, PaLM, Mistral, and most modern open LLMs use it
Forgetting model.eval() before inference is the silent BN bug that costs production teams weeks — always switch modes explicitly, and prefer frameworks (Lightning, Hugging Face) that do it for you
Why do transformers use LayerNorm instead of BatchNorm?
Now that activations are stabilized, the next question is everything that turns an unstable training run into a reliable one — warmup, gradient clipping, mixed precision, and the OneCycle schedule. Up next: Practical Training Tricks.