Layer Normalization, and Why Transformers Need It

The problem normalization solves

Deep networks are hard to train because the distribution of activations at each layer keeps shifting as the parameters below it change. Gradients computed for one layer’s inputs can become stale by the time a weight update actually lands, especially with dozens of stacked layers. Normalization layers fight this by keeping each layer’s input distribution roughly fixed in scale and location, which in turn keeps the loss surface better conditioned for gradient descent.

Batch normalization was the first widely used answer, but it normalizes across the batch dimension, which creates two problems for sequence models: it couples examples within a minibatch (bad for varying batch sizes and for RNNs unrolled over time), and it behaves differently at train versus inference time because it relies on running statistics. Layer normalization sidesteps both issues by normalizing across the feature dimension of a single example, independent of batch size.

The mechanics

For a single token’s activation vector $x \in \mathbb{R}^d$, layer norm computes the mean and variance across the $d$ features, then rescales:

where $\mu = \frac{1}{d}\sum_i x_i$, $\sigma^2 = \frac{1}{d}\sum_i (x_i - \mu)^2$, and $\gamma, \beta \in \mathbb{R}^d$ are learned scale and shift parameters. The epsilon term just avoids division by zero when variance collapses.

Because the statistics are computed per example rather than per batch, layer norm gives identical behavior at training and inference. There is no running mean to track, no discrepancy between train-mode and eval-mode forward passes, and no dependency on batch composition.

import torch

def layer_norm(x, gamma, beta, eps=1e-5):
    # x: (..., d)
    mu = x.mean(dim=-1, keepdim=True)
    var = x.var(dim=-1, keepdim=True, unbiased=False)
    x_hat = (x - mu) / torch.sqrt(var + eps)
    return gamma * x_hat + beta

Pre-norm versus post-norm

The original transformer paper applied layer norm after the residual addition (post-norm): x = LayerNorm(x + Sublayer(x)). This works, but it means the residual stream itself is never normalized, and gradients flowing back through many post-norm blocks can still explode or vanish without careful learning rate warmup.

Most modern architectures (GPT-2 onward, LLaMA, etc.) use pre-norm instead: x = x + Sublayer(LayerNorm(x)). Here the residual path stays a clean, unnormalized sum of contributions, which makes it act like a highway for gradients, and each sublayer sees a normalized input regardless of how large the residual stream has grown. The tradeoff is that pre-norm networks can suffer from representation collapse in very deep stacks, since each layer’s contribution is normalized before being added, so later layers may need extra tricks (like a final normalization layer before the output projection) to compensate.

In practice, pre-norm transformers train more reliably without a learning rate warmup schedule, which is a big part of why they became the default for large language models, but the community still debates whether pre-norm’s flatter loss landscape trades away some of the expressivity that post-norm’s stricter normalization provides.

RMSNorm as a simplification

A common variant, RMSNorm, drops the mean-centering step entirely and only rescales by the root mean square:

This removes the $\beta$ shift parameter and the mean subtraction, which saves a reduction pass and slightly speeds up training, with negligible quality loss in most large-scale language model recipes. LLaMA, PaLM, and several other recent architectures use RMSNorm for exactly this reason: it is cheaper per call and, over billions of forward passes, that adds up.

Practical takeaways

When implementing or debugging a transformer stack, layer norm placement is one of the first things worth checking if training diverges early. A few rules of thumb I keep coming back to: pre-norm plus a small warmup is a safe default for anything more than a dozen layers, RMSNorm is a reasonable default when compute is the binding constraint, and post-norm is worth revisiting only if you are chasing the exact behavior of an older checkpoint or paper you are trying to reproduce faithfully.