34  Normalization: BatchNorm to RMSNorm

TipTL;DR

Normalization rescales activations to keep them in a well-conditioned range at every layer: BatchNorm (Ioffe & Szegedy 2015) → LayerNorm (Ba 2016, the Transformer default) → RMSNorm (Zhang 2019, LLaMA’s choice) → GroupNorm (vision models with small batches). All four solve the same problem; they differ only in which axis they average over.

Depends on: Gradient Flow & Vanishing Gradients

34.1 Why this matters

Gradient Flow and Optimization both showed that training is sensitive to scale: an ill-conditioned loss surface makes gradient descent zig-zag, and saturated activations kill gradients outright. Normalization attacks this from a different angle. Rather than relying only on careful initialization and a well-tuned learning rate, rescale the activations themselves, at every layer, so they land in a stable range no matter how deep the stack.

34.2 The mechanism

Every method below computes the same two numbers, a mean \(\mu\) and a variance \(\sigma^2\), and applies the same rescaling,

\[ \hat{x} = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}}, \qquad y = \gamma\,\hat{x} + \beta, \]

with \(\gamma, \beta\) learned so the network can undo the normalization if that helps. What changes across methods is only which values get averaged together to compute \(\mu\) and \(\sigma^2\).

BatchNorm averages down each column ↕ one feature, across the batch batch (examples) LayerNorm / RMSNorm averages across each row ↔ all features, one example

34.2.1 BatchNorm: average across the batch

BatchNorm (Ioffe & Szegedy, 2015) computes \(\mu_B, \sigma_B^2\) for one feature by averaging over every example currently in the mini-batch (the highlighted column above). Two consequences follow directly from that choice:

  • A train/test mismatch. At test time there may be no batch at all (e.g. one example at a time), so BatchNorm keeps a running average of \(\mu_B, \sigma_B^2\) collected during training and uses that instead.
  • A batch-size dependence. With a very small batch, \(\mu_B,\sigma_B^2\) are noisy estimates, and BatchNorm’s benefit degrades. This, plus the train/test mismatch, motivated every variant below.

34.2.2 LayerNorm: average across features, per example

LayerNorm (Ba et al., 2016) computes \(\mu_L, \sigma_L^2\) for one example by averaging over that example’s own features instead (the highlighted row above). Because the statistics come from a single example, LayerNorm has no train/test mismatch and no batch-size dependence. This is exactly what a Transformer needs: variable batch sizes, variable sequence lengths, and autoregressive decoding one token at a time, where a “batch” of one token has no useful batch statistics at all. This is the LayerNorm used throughout The Transformer.

34.2.3 RMSNorm: drop the mean, keep the scale

RMSNorm (Zhang & Sennrich, 2019) observes that LayerNorm’s re-centering (subtracting \(\mu_L\)) contributes little of its benefit; most of the stabilization comes from rescaling by magnitude alone. It normalizes with only the root-mean-square, no mean subtraction and no separate shift parameter:

\[ \hat{x}_i = \frac{x_i}{\sqrt{\frac{1}{n}\sum_j x_j^2 + \epsilon}}, \qquad y_i = \gamma\,\hat{x}_i. \]

Cheaper to compute and empirically just as effective, RMSNorm is the normalization used in LLaMA and most modern open LLMs.

34.2.4 GroupNorm: average within channel groups

GroupNorm (Wu & He, 2018) splits channels into groups and normalizes within each group across the spatial dimensions, independent of batch size entirely. It targets vision architectures (detection, segmentation) where memory constraints force small batches, exactly the regime where BatchNorm’s batch-statistics estimate becomes unreliable.

import numpy as np

def batchnorm(x, eps=1e-5):        # average down the batch axis (axis 0)
    mu, var = x.mean(0), x.var(0)
    return (x - mu) / np.sqrt(var + eps)

def layernorm(x, eps=1e-5):        # average across features (axis -1), per example
    mu, var = x.mean(-1, keepdims=True), x.var(-1, keepdims=True)
    return (x - mu) / np.sqrt(var + eps)

def rmsnorm(x, eps=1e-5):          # rescale only, no re-centering
    rms = np.sqrt((x**2).mean(-1, keepdims=True) + eps)
    return x / rms

x = np.array([[1.0, 5.0, 3.0], [2.0, 2.0, 8.0]])
print("BatchNorm:\n", np.round(batchnorm(x), 2))
print("LayerNorm:\n", np.round(layernorm(x), 2))
print("RMSNorm:\n",  np.round(rmsnorm(x), 2))
BatchNorm:
 [[-1.  1. -1.]
 [ 1. -1.  1.]]
LayerNorm:
 [[-1.22  1.22  0.  ]
 [-0.71 -0.71  1.41]]
RMSNorm:
 [[0.29 1.46 0.88]
 [0.41 0.41 1.63]]

34.2.5 What to observe

  • All four methods share one formula; only the averaging axis changes. BatchNorm averages down the batch, LayerNorm and RMSNorm average across features within one example, GroupNorm averages within a channel group.
  • RMSNorm’s simplification (drop re-centering) is a deliberate finding, not an oversight: the mean turned out not to matter much.
  • Every variant keeps a learned scale (\(\gamma\), and sometimes \(\beta\)) so the network can recover the original distribution if normalization is not helpful at some layer.

34.3 Application & impact

Concept here What it became Where you see it today
BatchNorm The default for CNNs, 2015–2020 Vision backbones with large batch training
LayerNorm The Transformer default Every sublayer in The Transformer
RMSNorm The efficient LLM default LLaMA, Mistral, and most modern open LLMs
GroupNorm Small-batch vision alternative Detection and segmentation backbones
NoteKey takeaway

Normalization is a second lever on training stability, alongside careful initialization and optimizer choice: rescale activations at every layer so depth does not compound scale problems. The lineage is one of simplification, not reinvention. BatchNorm’s per-batch statistics became LayerNorm’s per-example statistics, and RMSNorm dropped the one part (re-centering) that turned out not to matter much.

← Back to: The Transformer