Normalization — keeping activations in a trainable range

A transformer block multiplies its input by dense weight matrices at every sublayer — two projections inside attention, three inside the FFN. After 32 or 96 such blocks, the magnitude of the hidden states either explodes (values reach 1e8+, causing NaN in float16) or collapses (values underflow to zero, killing all gradient signal). Normalization layers reset the scale of activations at fixed points in the network, bounding the dynamic range that each sublayer must handle.

LayerNorm

Layer normalization (Ba, Kiros & Hinton 2016) operates on a single token's feature vector independently. Given a vector x of dimension d_model, it computes the mean and variance across that vector's features, normalizes to zero mean and unit variance, then applies a learned affine transformation:

where μ = mean(x) and σ² = var(x) are computed over the d_model dimension for each token independently. γ (scale) and β (bias) are learnable parameters of shape (d_model,), initialized to ones and zeros respectively. ε is a small constant (typically 1e-5 or 1e-6) that prevents division by zero when variance is extremely small.

The parameter count is minimal: 2 × d_model per norm layer. For a model with $d_{\text{model}} = 4096$, that's 8,192 parameters per norm — negligible compared to the millions in each attention or FFN sublayer. A 32-layer transformer with two norms per layer has 64 norm layers, contributing 64 × 8,192 = 524,288 parameters total (0.006% of an 8B model).

The operation normalizes each token independently — token at position 3 is normalized using only its own d_model feature values, regardless of what other tokens in the sequence look like. This is distinct from BatchNorm (Ioffe & Szegedy 2015), which normalizes across the batch dimension. BatchNorm is unsuitable for autoregressive models because sequences have different lengths, causal masking makes batch statistics ill-defined during inference, and the batch dimension in LLM training is often small (batch size depends on available GPU memory, not statistical requirements).

The mean-and-variance computation requires two reductions over the feature dimension: one for the mean, one for the variance (or equivalently, one for the sum and one for the sum of squares). On GPU hardware, reductions are expensive relative to pointwise operations because they require communication across thread blocks. A naive implementation of LayerNorm thus bottlenecks on two sequential reductions per token per norm layer.

import torch
import torch.nn as nn

class LayerNorm(nn.Module):
    def __init__(self, d_model, eps=1e-5):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(d_model))
        self.beta = nn.Parameter(torch.zeros(d_model))
        self.eps = eps

    def forward(self, x):
        # x shape: (batch, seq_len, d_model)
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        x_norm = (x - mean) / torch.sqrt(var + self.eps)
        return self.gamma * x_norm + self.beta

Post-norm vs. pre-norm placement

The placement of the normalization layer relative to the sublayer changes training dynamics profoundly.

Post-norm (original transformer)

Apply LayerNorm after adding the residual:

x = LayerNorm(x + Sublayer(x))

This was the arrangement in Vaswani et al. 2017. In this formulation, the gradient from the loss must flow through the normalization at every layer during backpropagation. The norm's Jacobian rescales gradients based on the current activation statistics — it has a normalizing effect on the gradient's magnitude, which helps in shallow networks (the original transformer used 6 layers) but creates problems in deep ones.

The issue becomes clear at 32+ layers. The gradient of LayerNorm with respect to its input depends on the current batch's statistics (the mean and variance of the activations flowing through that layer at that training step). As training progresses, these statistics shift, making the effective learning rate for each layer fluctuate unpredictably. Xiong et al. 2020 ("On Layer Normalization in the Transformer Architecture") showed formally that post-norm transformers require carefully tuned learning rate warmup to avoid divergence — and that even with warmup, the last layers receive well-behaved gradients while early layers see noisy, poorly-scaled ones. Training a 48-layer post-norm transformer requires 10,000+ warmup steps; training a 96-layer one is impractical.

Pre-norm (modern standard)

Apply normalization before the sublayer, and let the residual bypass the norm entirely:

x = x + Sublayer(Norm(x))

Now the residual connection creates a direct additive path from any layer's output to the input of any later layer. During backpropagation, the gradient of the loss with respect to an early layer has a component that flows directly through the chain of residual additions — it never passes through any normalization layer on this path. The norm only affects the gradient flowing through the sublayer (attention or FFN), not the gradient flowing along the residual highway.

The practical consequence: pre-norm transformers train stably without learning rate warmup at depths up to 100+ layers. This is the same principle that makes ResNets trainable to 1000+ layers (He et al. 2016). All modern LLMs — GPT-3 (Brown et al. 2020), PaLM (Chowdhery et al. 2022), Llama (Touvron et al. 2023), Mistral (Jiang et al. 2023), Qwen (Bai et al. 2023) — use pre-norm.

One subtlety: post-norm models, when they converge, sometimes achieve slightly better final quality on small-scale experiments. The hypothesis is that normalization after the sublayer output constrains the representational scale more tightly, acting as a form of regularization. But this advantage vanishes at the scale and depth of modern LLMs, where training stability dominates. No one is willing to spend weeks tuning warmup schedules to save 0.1% on loss when pre-norm trains reliably out of the box.

RMSNorm

Root mean square normalization (Zhang & Sennrich 2019) simplifies LayerNorm by dropping the mean-centering step entirely. It normalizes by the root-mean-square of the vector only:

where $$\text{mean}(x^{2}) = (\frac{1}{d_{\text{model}}}) \cdot \sum x_{i}^{2}$$. There is no subtraction of the mean and no β parameter — just γ (scale), reducing the learnable parameters to d_model per norm layer (one parameter per feature instead of two).

The theoretical justification: Zhang & Sennrich 2019 show empirically that the re-centering operation (mean subtraction) in LayerNorm contributes negligibly to the model's final performance. The normalization of scale (dividing by a magnitude estimate) is what provides the training stability benefit. Since $\text{RMS}(x) = \sqrt{\text{mean}(x^{2})}$ and $$\text{std}(x) = \sqrt{\text{mean}(x^{2}) - \text{mean}(x)^{2}}$$, RMSNorm and LayerNorm produce identical outputs when the mean of x is zero — which is approximately true in practice after the first few training steps because the learned representations tend to center themselves through the bias-free linear layers.

The computational saving comes from eliminating one reduction operation. RMSNorm computes a single reduction (the mean of squares), while LayerNorm requires two (mean and variance, or equivalently, sum and sum-of-squares). On GPU, each reduction requires synchronization across the feature dimension. For $d_{\text{model}} = 4096$, this synchronization is the dominant cost of the norm operation — the pointwise multiply and divide are negligible by comparison. Dropping one reduction saves roughly 5–10% of the norm's total latency.

At scale — billions of tokens, thousands of GPUs, weeks of training — a 5–10% speedup on an operation that occurs 64+ times per forward pass (2 norms × 32 layers + final norm) accumulates into meaningful wall-clock savings. If training Llama 3 405B took 30.84M GPU-hours (Meta, 2024), a 5% speedup on normalization (which accounts for ~3% of total compute) saves ~46,000 GPU-hours.

class RMSNorm(nn.Module):
    def __init__(self, d_model, eps=1e-6):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(d_model))
        self.eps = eps

    def forward(self, x):
        # x shape: (batch, seq_len, d_model)
        rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return self.gamma * (x / rms)

Adoption

RMSNorm is used by: Llama 1/2/3 (Meta), PaLM and Gemma (Google), Qwen (Alibaba), Mistral and Mixtral (Mistral AI), GPT-NeoX (EleutherAI), DeepSeek (DeepSeek AI). The original GPT-2 and GPT-3 used LayerNorm with β; the shift to RMSNorm happened around 2022–2023 as the field converged on the understanding that mean-centering is unnecessary overhead.

No published benchmark shows a meaningful quality difference between LayerNorm and RMSNorm at scale. The Llama technical report (Touvron et al. 2023) states the choice was made purely for training efficiency. Gemma (Google, 2024) confirms the same finding. The quality difference is in the noise floor of evaluation benchmarks.

Numerical precision and the ε constant

The ε constant matters more than it appears. In float16 (the dominant training precision before bfloat16), the smallest representable positive normal value is ~6e-8, and the smallest subnormal is ~6e-8. If var(x) + ε or mean(x²) + ε evaluates to zero after rounding, the division produces infinity, which propagates NaN through the rest of the network irreversibly.

Common choices across model families:

• PyTorch's built-in nn.LayerNorm default: ε = 1e-5
• Llama RMSNorm: ε = 1e-6 (safe because Llama trains in bfloat16)
• Gemma RMSNorm: ε = 1e-6
• GPT-2 LayerNorm: ε = 1e-5

With bfloat16 (which has the same exponent range as float32 — 8 exponent bits — but only 7 bits of mantissa vs. float32's 23), the risk of underflow in the denominator is negligible for any ε ≥ 1e-7. The shift from float16 to bfloat16 training (standard since ~2020) removed one of the practical headaches of normalization implementation.

Comparing outputs and performance

import torch
import time

d_model = 4096
batch_size = 32
seq_len = 2048

layer_norm = LayerNorm(d_model).cuda()
rms_norm = RMSNorm(d_model).cuda()

x = torch.randn(batch_size, seq_len, d_model, device="cuda", dtype=torch.bfloat16)

# Verify outputs are similar
with torch.no_grad():
    ln_out = layer_norm(x.float()).bfloat16()
    rms_out = rms_norm(x.float()).bfloat16()

    cos_sim = torch.nn.functional.cosine_similarity(
        ln_out.view(-1, d_model), rms_out.view(-1, d_model), dim=-1
    )
    print(f"Mean cosine similarity between LayerNorm and RMSNorm: {cos_sim.mean():.6f}")
    # Typically 0.98-0.99 — very similar but not identical

# Benchmark
torch.cuda.synchronize()
start = time.perf_counter()
for _ in range(1000):
    _ = layer_norm(x)
torch.cuda.synchronize()
ln_time = time.perf_counter() - start

torch.cuda.synchronize()
start = time.perf_counter()
for _ in range(1000):
    _ = rms_norm(x)
torch.cuda.synchronize()
rms_time = time.perf_counter() - start

print(f"LayerNorm: {ln_time*1000:.1f} ms for 1000 iters")
print(f"RMSNorm:   {rms_time*1000:.1f} ms for 1000 iters")
print(f"Speedup:   {ln_time/rms_time:.2f}x")
# Typical result on A100: ~1.06-1.10x speedup for RMSNorm

Where normalization sits in the modern transformer

In a pre-norm decoder-only transformer (Llama architecture), normalization appears at three points:

• Before the attention sublayer: x = x + MHA(RMSNorm(x))
• Before the FFN sublayer: h = h + FFN(RMSNorm(h))
• After the final transformer block, before the language model head: logits = LM_Head(RMSNorm(final_hidden))

That final norm is critical. Without it, the hidden states entering the LM head would have unconstrained magnitude — their scale depends on how many layers' residual contributions accumulated (which varies by position and by input). The final norm ensures the LM head receives inputs with a consistent scale regardless of which layers contributed most to the final representation.

Llama 3 8B has 32 blocks × 2 norms per block + 1 final norm = 65 RMSNorm layers, each with 4,096 learnable γ parameters. Total norm parameters: 65 × 4,096 = 266,240 — about 0.003% of the model's 8 billion parameters, yet removing any single one would make training collapse within hundreds of steps as activations escape the representable range.