📝 Layer Normalization

Description

< What is it? >

Layer normalization (layer norm) normalizes a layer's activations across the features of a single example, then rescales the result with learnable $\gamma$ and $\beta$. It is the same idea as batch normalization with the axis swapped: batch norm looks along the batch at one feature, layer norm looks along the features of one example.

                 features ──▶
              ┌───────────────┐
   example 1  │ ▓ ▓ ▓ ▓ ▓ ▓ ▓ │  ← layer norm: one μ, σ² per row
   example 2  │ ▓ ▓ ▓ ▓ ▓ ▓ ▓ │
   example 3  │ ▓ ▓ ▓ ▓ ▓ ▓ ▓ │
              └───────────────┘
                 ▲
                 batch norm: one μ, σ² per column

That one change removes every batch dependency: each example is normalized entirely on its own, so the result does not depend on batch size, on which examples were batched together, or on whether the model is training or serving.

< The computation >

For one example with $H$ features:

$$\mu = \frac{1}{H}\sum_{i=1}^{H} x_i, \qquad \sigma^2 = \frac{1}{H}\sum_{i=1}^{H}\left(x_i - \mu\right)^2$$ $$y_i = \gamma_i \cdot \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta_i$$

Identical in form to batch norm — see Reading the symbols for the Greek letters. The only difference is what the sums range over.

Key points

< Layer norm vs. batch norm >

	Batch norm	Layer norm
Normalizes over	The batch, per feature	The features, per example
Batch size	Needs a reasonable one; undefined at 1	Irrelevant — works at batch size 1
Train vs. inference	Different: batch statistics, then running averages	Identical; no running statistics to keep
Variable-length inputs	Awkward — padding pollutes the statistics	Fine, each token normalizes itself
Typical home	CNNs, vision	Transformers, RNNs, anything sequential

< Why transformers use it >

• Sequences vary in length. Batch statistics computed across padded positions are statistics of the padding as much as the data.
• Inference is often batch size 1. Generating one token at a time gives batch norm nothing to average over, while layer norm behaves exactly as it did in training.
• No running statistics means nothing to synchronize across devices in distributed training, and no train/eval discrepancy to forget about.

Note the trade-off: layer norm gives up batch norm's mild regularizing noise, because there is no longer any cross-example randomness in the normalization.

< Where it goes in a block >

• Post-LN — x + Sublayer(x), then normalize. The original transformer. It needs learning-rate warmup to train stably at depth.
• Pre-LN — normalize first, then x + Sublayer(LN(x)). Leaves a clean residual path from input to output, trains without warmup, and is what essentially every modern model uses.

< In PyTorch >

import torch.nn as nn

layer_norm = nn.LayerNorm(512)          # normalize over the last dim of (N, L, 512)
x = layer_norm(x)                       # same behaviour in train() and eval()

The argument is the shape to normalize over, not a channel count — nn.LayerNorm(512) normalizes the last dimension, nn.LayerNorm([C, H, W]) the last three. Set elementwise_affine=False to drop $\gamma$ and $\beta$.

< RMSNorm >

Modern LLMs mostly use RMSNorm, which skips the mean subtraction and divides by the root-mean-square alone:

$$y_i = \gamma_i \cdot \frac{x_i}{\sqrt{\frac{1}{H}\sum_j x_j^2 + \epsilon}}$$

Cheaper — one pass instead of two, and no $\beta$ — and in practice it matches layer norm's quality, which suggests the rescaling was doing the real work all along.

< Related ideas >

• Batch Normalization is the batch-axis counterpart, with the train/inference split layer norm avoids.
• Data Normalization compares all three in one table.
• Vanishing & Exploding Gradients is what keeping activation scales stable is ultimately for.

Reference

• Layer Normalization — Ba, Kiros & Hinton (2016)
• On Layer Normalization in the Transformer Architecture — Xiong et al. (2020) — the Pre-LN vs. Post-LN analysis
• Root Mean Square Layer Normalization — Zhang & Sennrich (2019)
• PyTorch nn.LayerNorm documentation