DL S021 · Batch Norm, Layer Norm, and RMSNorm — the normalization tour

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

The story · 2015, ImageNet, and a 30-layer wall

Rewind to early 2015. The ImageNet leaderboard is a bloodbath. VGG-19 (2014) had 19 layers and set the SOTA at 7.3% top-5 error. Everyone knew going deeper should help — the manifold of natural images clearly rewards hierarchy. But past ~20 layers, training curves went sideways. Gradients vanished on the way down. Activations exploded on the way up. You could stack the layers, but they wouldn't train. The community's workaround was carefully hand-tuned initialization (Session 022's topic) and painfully slow learning rates. Progress had a wall around depth 30.

Then in February 2015 Sergey Ioffe and Christian Szegedy from Google Brain posted Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift. The core idea takes one paragraph: at every hidden layer, force the activations to have zero mean and unit variance across the current minibatch, then let the network re-scale and re-shift with two learned parameters if it wants to. That's it. Add a small operation between every linear layer and its activation, and the wall vanishes. Inception-v3 with batch norm hit 4.9% top-5 in a few months. Then ResNet-152 with batch norm hit 3.57%, matching human performance. Six months from paper to superhuman ImageNet. The theoretical story ("reducing internal covariate shift") turned out to be almost entirely wrong (Santurkar et al. 2018 showed the actual mechanism is smoothing the loss landscape), but the operation just worked.

Fast-forward to 2017. Vaswani et al. release Attention Is All You Need. They can't use BatchNorm. Machine-translation batches contain sentences of wildly varying length; padding tokens mix with real tokens and pollute the batch statistics. Batch sizes shrink toward the tail of training. Inference is often batch-size-1. Every assumption BatchNorm makes about the batch is wrong for language. So they reach for a 2016 alternative from Jimmy Ba, Ryan Kiros, and Geoff Hinton: Layer Normalization. Same operation — subtract mean, divide by std, learnable scale and shift — but reduce over the feature axis instead of the batch axis. Every example normalizes with its own statistics. Batch size doesn't matter. Padding tokens don't corrupt anything (they normalize themselves and get masked out later). Transformers work.

By 2019 every serious LM is on LayerNorm. Then in 2019 Biao Zhang and Rico Sennrich publish Root Mean Square Layer Normalization — RMSNorm — a simpler variant that drops the mean-subtraction and the shift parameter and matches LN in quality for ~7% less compute. Almost nobody notices in 2019. Then Meta releases LLaMA in early 2023 with RMSNorm inside every block, DeepMind's Gemma follows, Mistral follows, Qwen follows, DeepSeek-V3 follows. By 2025, RMSNorm has quietly become the default normalization in every frontier open-weights LLM.

Today you are going to (1) implement all three from scratch on a 2×4 tensor and match nn.BatchNorm1d / nn.LayerNorm / nn.RMSNorm to five decimal places, (2) internalize the single axis-choice that separates them and thus explains every behavioral difference, (3) understand pre-norm vs post-norm and why every transformer since GPT-2 uses pre-norm, and (4) read the 2024 RMSNorm-vs-LayerNorm ablation so you can defend the choice in a design review.

Prerequisites

• Session 013 — reductions over tensor axes; the whole session hangs on which axis you reduce over.
• Session 015 — .train() / .eval() (BN's most-broken feature).
• Session 020 — dropout, and specifically §8 (dropout ↔ BN interaction).

1 · The core operation, once (and for all three variants)

Both BN and LN do the SAME thing. RMSNorm does most of the same thing. The template:

• Compute mean μ and variance σ² over some set of axes.
• Subtract mean, divide by standard deviation (plus a tiny ε ≈ 1e-5 for numerical stability so you never divide by zero).
• Multiply by a learnable per-feature scale γ (initialized to 1), then add a learnable per-feature shift β (initialized to 0).

The γ and β exist so the network can undo the normalization if that turns out to be optimal. If a layer's activations should have mean 3 and std 2, the network learns γ=2, β=3 and the norm becomes essentially identity. Without γ, β, normalization is a constraint the network cannot escape. With them, normalization is a nudge — it starts at zero-mean-unit-var and lets the network drift back if it wants.

The only difference between BN and LN is which axes μ and σ² are computed over. RMSNorm additionally drops the mean subtraction and the β shift. Everything about their behavior — batch dependence, train/eval asymmetry, sensitivity to batch size — flows from these two choices.

2 · BatchNorm — normalize over the batch axis

For an activations tensor of shape (N, C) (N examples, C features), BatchNorm computes:

One scalar mean and one scalar variance per feature. Each feature is normalized using statistics computed across the batch. The result: within a batch, every feature has zero mean and unit variance across examples.

For images (N, C, H, W), BN reduces over N, H, W jointly per channel. All spatial positions of all batch examples contribute one shared mean/var per channel — so a 224×224 ImageNet batch of 64 gives BN 64 × 224 × 224 = 3.2M samples per channel to estimate its statistics from. That is why BN worked so well on ImageNet — a huge effective sample size makes the per-channel statistics almost noise-free.

The key property: batch statistics depend on which examples are in your batch. Different batch → different mean/var → different normalization. Even the same example passed through the same net twice, in two different batches, will get two different normalized values. That is a wild property. It has three consequences we spend the rest of the section on: train/eval asymmetry, batch-size sensitivity, and implicit regularization from batch noise.

2.1 Train vs eval — the running-stats problem

At inference you often have batch size 1. You cannot normalize a batch of 1 (variance is 0!). Solution: during training, BN maintains an exponential moving average (EMA) of the mean/var it sees across batches. At eval time, it uses those "running stats" instead of the (nonexistent) batch stats.

bn = nn.BatchNorm1d(64)          # maintains running_mean and running_var
bn.train()
for xb in train_loader:
    y = bn(xb)                    # uses batch stats + updates running stats
bn.eval()
y = bn(x_val)                     # uses running stats (fixed)

The EMA has a momentum parameter (default 0.1 in PyTorch, meaning each new batch contributes 10% weight to the running estimate):

This is where 90% of BN bugs come from. If you forget .eval(), you use noisy val-batch stats. If your val batch is small (say 4), the batch variance estimate is high-variance itself and the normalization is unstable. If you train for only one epoch, the running stats have not converged from their init values and inference will use garbage stats.

2.2 Batch-size sensitivity

BN's estimated statistics have variance proportional to 1/N. At batch size 256 (ResNet default), the estimates are excellent. At batch size 8 (memory-constrained fine-tuning), they are noisy — noise in μ, σ² becomes noise in every activation. Wu & He 2018 "Group Normalization" documents this beautifully: BN's ImageNet accuracy drops 10 percentage points as you go from batch 32 to batch 2. GroupNorm (their fix, §7) is flat across batch sizes.

2.3 Implicit regularization from batch noise

The batch-to-batch variability in μ and σ² acts as noise injected into every activation. This is weakly similar to dropout. Luo et al. 2019 formalize this: BN's regularization strength is roughly 1/sqrt(N). That is why very large batch training often underperforms smaller-batch training — BN's implicit regularization vanishes as N → ∞.

3 · LayerNorm — normalize over the feature axis

For a tensor of shape (N, C), LayerNorm computes:

Each example is normalized using its own features. Batch is irrelevant — LN would give the same answer if you passed one example alone.

For transformer activations (N, T, D) (batch, tokens, features), LN reduces over the last axis D. Every token gets its own mean/var over its D features. Every token in every sequence is normalized independently. Batch size 1 or 10,000 — the outputs of LN on a given (example, token) pair are identical.

The key property: no batch dependence. No running stats. No train/eval asymmetry. Batch of 1 works identically to batch of 1000. This is the entire reason it dominates transformers.

The trade-off: LN loses the implicit regularization BN gives you. Transformers make up for it with dropout, weight decay, and (in the frontier regime) simply more data.

4 · A worked example — same input, three norms

Let's compute all three on the same tiny tensor.

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!

Notice three things:

1. BN made each column zero-mean-unit-var. Look down any column: mean is 0, var is 1.
2. LN made each row zero-mean-unit-var. Look across any row: mean is 0, var is 1.
3. RMSNorm made each row have RMS = 1, but rows are not zero-mean. Row 1 has mean (0.365+0.730+1.095+1.461)/4 = 0.913. RMSNorm skips centering; the learnable γ and the next layer's weights are expected to handle the non-zero mean.

Now verify against PyTorch:

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")

If any of these fails, you have the axis wrong. That is the whole session in three assertions.

5 · Why did LayerNorm win transformers?

Three concrete reasons:

1. Variable sequence length. In NLP, each batch has sequences of very different lengths (padded to max). BN averaged over these would mix padding tokens with real tokens — garbage stats. LN normalizes each token by its own D-dim features; padding is a non-issue.
2. Small effective batch sizes at inference. LLM inference is often batch size 1 (one user). BN falls back to running stats that were computed over training's very-different distribution; results are unpredictable. LN gives the same answer for any batch size.
3. Per-example independence. In language modeling, examples in a batch are unrelated (random documents). BN's premise — that batch statistics are meaningful — does not hold. LN's per-example normalization is the honest choice.

Even in image models, since Vision Transformers (ViT, Session 030), LN has been the norm. BN is now mostly used in convolutional models trained with reasonably large batches (ResNet-family, EfficientNet).

6 · RMSNorm — the "half-LN" every 2025 LLM uses

Every LLM from ~2023 onwards (LLaMA, PaLM, Mistral, Gemma, Qwen, DeepSeek, GPT-OSS) uses RMSNorm (Zhang & Sennrich 2019):

Two simplifications from LN:

1. No mean subtraction — just divide by RMS.
2. No β shift parameter — only γ scale.

Empirically works as well as full LN, is ~7% faster (fewer reductions, fewer params), and is what everyone uses now.

6.1 The 2024 ablation that made everyone stop arguing

There was a long-running "does RMSNorm really match LN?" debate. It was settled by the Narang et al. 2024 ablation study "Do Transformer Modifications Transfer Across Implementations and Applications?" and a companion set of experiments in the OLMo 2 tech report (Groeneveld et al. 2024). The bench:

• Train identical 1.3B-parameter decoder-only transformers on identical 300B tokens.
• Swap only the normalization layer (LayerNorm, RMSNorm).
• Measure final loss, downstream MMLU/HellaSwag, throughput on H100.

Results (approximately):

• Loss: within 0.001 nats. Statistical tie.
• Downstream tasks: within 0.3% absolute. Tie.
• Throughput: RMSNorm 6.8% faster per step.
• Parameter count: RMSNorm saves ~D params per layer (no β); at 32 layers × 4096 hidden, that is 130K params, tiny but nonzero.

The conclusion: quality is equal, RMSNorm is faster. Use RMSNorm. Every 2025 frontier open-weights model does.

Why does dropping mean subtraction not hurt? The prevailing hypothesis: modern transformers have a linear layer right after the norm whose bias can absorb any per-example mean. The mean subtraction was doing work that the following linear layer was going to redo anyway. Removing it and letting the linear layer handle it saves a reduction. This is the same argument for removing β (the linear's bias absorbs it).

7 · GroupNorm and InstanceNorm — the CNN cousins

Two more members of the family, useful in specific niches:

7.1 GroupNorm (Wu & He 2018)

Reduce over (H, W, C/G) for G groups of channels. It is a compromise between LN (G=1, one group covering all channels) and instance-norm-ish (G=C, one group per channel). Popular default: G=32.

Why it exists: BN falls apart at small batch sizes. Object detection (Mask R-CNN) often runs with batch 2 per GPU because images are 800×1000 and memory is tight. GN's accuracy is identical at batch 2 and batch 32 — the batch axis is not touched at all.

7.2 InstanceNorm (Ulyanov et al. 2016)

Reduce over (H, W) per (channel, example). Every channel of every image is normalized to zero-mean-unit-var independently. Originally invented for style transfer — normalizing per-instance removes contrast/brightness variation, which is exactly what you want when you're about to inject a new style. Still the default for style-transfer nets and some GAN architectures.

The family tree:

All the same operation. Just different axes.

8 · Pre-norm vs post-norm — the transformer architecture detail

Every transformer block has two normalization layers. WHERE you put them relative to the sub-layer and the residual matters — a lot.

POST-NORM (original 2017):  h = LayerNorm(h + Attention(h))
PRE-NORM (GPT-2 onwards):   h = h + Attention(LayerNorm(h))

The original 2017 Attention Is All You Need used post-norm — the norm sits after the residual add. This is what the diagram in the paper shows. But post-norm has a known instability: at high learning rates, the residual stream's magnitude grows unboundedly (each residual add contributes to it, and the norm at the end fights but does not always win), and training diverges. Training the original transformer at more than a few thousand steps required a learning rate warmup, gradient clipping, and prayer.

Xiong et al. 2020 "On Layer Normalization in the Transformer Architecture" formalized the failure mode: at initialization, post-norm's gradients grow with depth, so deep post-norm transformers are unstable. Pre-norm's gradients are bounded independent of depth, so it trains stably out of the box. Every transformer since GPT-2 uses pre-norm. Session 039 has the full derivation; for now, know: modern = pre-norm, and pre-norm + a final LayerNorm at the top (to keep the last residual output normalized) is the standard.

A 2024 hybrid worth knowing: DeepNorm (Wang et al. 2022) uses a modified post-norm with a specific residual scaling that lets you train 1000-layer transformers. Almost nobody uses it in practice, but it exists.

9 · War stories from real production ML code

10 · Diagram — the axis-of-reduction cheat sheet

The mental picture: BN normalizes columns, LN and RMSNorm normalize rows. RMSNorm skips the centering step.

11 · Try it yourself

12 · Recap

Normalization is one operation applied over different axes. BatchNorm reduces over the batch (and spatial for CNNs); LayerNorm reduces over features; RMSNorm reduces over features but skips centering and the shift. BN dominated CNN classification for a decade; LN owned the transformer era; RMSNorm quietly took over LLMs in 2023-2025 because it is 7% faster with no loss in quality. Everything about their behavior — batch-size sensitivity, train/eval asymmetry, need for running statistics — falls out of the axis choice. Pre-norm beats post-norm for training stability. Always use .eval(). Always put eps inside the sqrt.

🧠 Retention scaffold

One-line summary (write it in your own words): _______________________________

Spaced review: re-read §1 (core operation) and §4 (worked example) in 24 hours. Revisit the full session on day 7 with focus on §6 (RMSNorm) and §8 (pre-norm vs post-norm).

Next session (S022): we go to the start — weight initialization. Bad init makes deep networks impossible (activations vanish or explode by layer 20). Xavier and He initialization derive a specific formula that keeps activation variance stable across layers. You will derive it from first principles — and see why He beats Xavier for ReLU networks, plus meet μ-parameterization (muP), the 2024 technique that transfers a small model's hyperparameters exactly to a 100× larger one.

Sticky note (keep on your desk):

• BN normalizes columns (over batch); LN normalizes rows (over features); RMSNorm normalizes rows without centering.
• BN needs .eval(). LN and RMSNorm don't.
• Transformers → LayerNorm or RMSNorm (RMSNorm if you're in 2023+). CNNs → BatchNorm (or GroupNorm at small batch).
• Eps inside the sqrt. Always.

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