Normalization layers

Deep dive

How it works

In: activations [N, D] (or [N, C, H, W]), learned γ and β · Out: same shape, standardized per the chosen axes · State: γ, β; BatchNorm also keeps running mean and variance · Cost: O(N·D), one read and one write when fused

For one row x ∈ ℝᴰ (LayerNorm) or one channel across the batch (BatchNorm):

μ = mean(x)          σ² = mean((x − μ)²)          rstd = 1/√(σ² + ε)
x̂ = (x − μ)·rstd     y = γ ⊙ x̂ + β
RMSNorm:  r = 1/√(mean(x²) + ε)    y = γ ⊙ x · r

Backward for LayerNorm, with g = dy ⊙ γ:

dx = rstd · (g − mean(g) − x̂ · mean(g ⊙ x̂))
dγ = Σ_rows dy ⊙ x̂        dβ = Σ_rows dy

For RMSNorm, with g = dy ⊙ γ: dx = r·g − r³·x·mean(g ⊙ x).

The two projections in the LayerNorm backward remove the components of g along the mean direction and along x̂. Scaling or shifting the input does not change the output, so the gradient has no component in those directions.

Code

// LayerNorm and RMSNorm forward/backward over rows of length d (row-major [rows, d]).
pub struct Saved { pub mean: Vec<f32>, pub rstd: Vec<f32> }

pub fn layernorm_fwd(x: &[f32], g: &[f32], b: &[f32], d: usize, eps: f32, y: &mut [f32]) -> Saved {
    let rows = x.len() / d;
    let (mut mean, mut rstd) = (vec![0.0; rows], vec![0.0; rows]);
    for r in 0..rows {
        let xr = &x[r * d..(r + 1) * d];
        let mu = xr.iter().map(|&v| v as f64).sum::<f64>() / d as f64;    // accumulate in f64
        let var = xr.iter().map(|&v| (v as f64 - mu).powi(2)).sum::<f64>() / d as f64; // two-pass: no cancellation
        let rs = 1.0 / (var + eps as f64).sqrt();
        for i in 0..d { y[r * d + i] = ((xr[i] as f64 - mu) * rs) as f32 * g[i] + b[i]; }
        mean[r] = mu as f32; rstd[r] = rs as f32;
    }
    Saved { mean, rstd }
}

pub fn layernorm_bwd(x: &[f32], g: &[f32], s: &Saved, dy: &[f32], d: usize,
                     dx: &mut [f32], dg: &mut [f32], db: &mut [f32]) {
    for r in 0..x.len() / d {
        let (mu, rs) = (s.mean[r], s.rstd[r]);
        let (mut m1, mut m2) = (0.0f32, 0.0f32);
        for i in 0..d {
            let xh = (x[r * d + i] - mu) * rs;
            let gi = dy[r * d + i] * g[i];
            m1 += gi; m2 += gi * xh;
            dg[i] += dy[r * d + i] * xh; db[i] += dy[r * d + i];
        }
        let (m1, m2) = (m1 / d as f32, m2 / d as f32);
        for i in 0..d {
            let xh = (x[r * d + i] - mu) * rs;
            dx[r * d + i] = rs * (dy[r * d + i] * g[i] - m1 - xh * m2);
        }
    }
}

pub fn rmsnorm_fwd(x: &[f32], g: &[f32], d: usize, eps: f32, y: &mut [f32]) -> Vec<f32> {
    (0..x.len() / d).map(|r| {
        let xr = &x[r * d..(r + 1) * d];
        let ms = xr.iter().map(|&v| v * v).sum::<f32>() / d as f32;
        let rr = 1.0 / (ms + eps).sqrt();
        for i in 0..d { y[r * d + i] = xr[i] * rr * g[i]; }
        rr
    }).collect()
}

pub fn rmsnorm_bwd(x: &[f32], g: &[f32], rinv: &[f32], dy: &[f32], d: usize, dx: &mut [f32]) {
    for r in 0..x.len() / d {
        let rr = rinv[r];
        let dot = (0..d).map(|i| dy[r * d + i] * g[i] * x[r * d + i]).sum::<f32>() / d as f32;
        for i in 0..d { dx[r * d + i] = rr * dy[r * d + i] * g[i] - rr * rr * rr * x[r * d + i] * dot; }
    }
}

/// Fold inference-mode BatchNorm into the preceding conv/linear: w' = w·k, b' = (b − μ)·k + β.
pub fn fold_bn(w: &mut [f32], b: &mut [f32], per_out: usize, gamma: &[f32], beta: &[f32],
               mean: &[f32], var: &[f32], eps: f32) {
    for c in 0..b.len() {
        let k = gamma[c] / (var[c] + eps).sqrt();
        w[c * per_out..(c + 1) * per_out].iter_mut().for_each(|v| *v *= k);
        b[c] = (b[c] - mean[c]) * k + beta[c];
    }
}

fn main() {
    let d = 8;
    let x: Vec<f32> = (0..2 * d).map(|i| ((i * 37 % 11) as f32 - 5.0) * 0.7 + 100.0).collect(); // mean far from 0
    let g: Vec<f32> = (0..d).map(|i| 1.0 + 0.1 * i as f32).collect();
    let b = vec![0.5; d];
    let dy: Vec<f32> = (0..2 * d).map(|i| ((i * 13 % 7) as f32 - 3.0) * 0.1).collect();
    let loss = |x: &[f32]| { let mut y = vec![0.0; x.len()]; layernorm_fwd(x, &g, &b, d, 1e-5, &mut y);
                             y.iter().zip(&dy).map(|(a, c)| (a * c) as f64).sum::<f64>() };
    let mut y = vec![0.0; x.len()];
    let s = layernorm_fwd(&x, &g, &b, d, 1e-5, &mut y);
    let (mut dx, mut dg, mut db) = (vec![0.0; x.len()], vec![0.0; d], vec![0.0; d]);
    layernorm_bwd(&x, &g, &s, &dy, d, &mut dx, &mut dg, &mut db);
    for i in [0, 5, 11] {
        let (mut p, mut m) = (x.clone(), x.clone());
        p[i] += 1e-2; m[i] -= 1e-2;
        let num = (loss(&p) - loss(&m)) / 2e-2;
        assert!((num as f32 - dx[i]).abs() < 2e-2, "{} vs {}", num, dx[i]);
    }
    let rinv = rmsnorm_fwd(&x, &g, d, 1e-6, &mut y);
    rmsnorm_bwd(&x, &g, &rinv, &dy, d, &mut dx);
    println!("layernorm grad ok; rmsnorm y[0..3] = {:?}, dx[0] = {:e}", &y[..3], dx[0]);
}

Complexity and constants

• O(N·D) FLOPs, negligible next to matmuls, but memory-bound. An unfused LayerNorm in eager PyTorch reads and writes the tensor several times. A fused kernel reads x once, keeps the row in registers or shared memory, and writes y once.
• For D = 4096 in bf16 a row is 8 KB and fits in one thread block. At about 3 TB/s HBM bandwidth, normalizing 1M tokens × 4096 (8 GB read + 8 GB write) costs order of 5 ms.
• RMSNorm skips the mean reduction: one reduction instead of two, fewer registers. The saved cost is real but small relative to the whole layer.
• Saved for backward: 2 floats per row (μ, rstd), or 1 for RMSNorm, plus the input.

Efficiency

• Cost: O(N·D) FLOPs, memory-bound. For 1M tokens × 4096 in bf16, 8 GB read plus 8 GB written is order of 5 ms at about 3 TB/s HBM. The backward pass keeps 2 floats per row (μ, rstd), 1 for RMSNorm.
• Levers:
  • Fused kernel, one read and one write of x — replaces several eager passes (above) — always on GPU.
  • RMSNorm instead of LayerNorm — one reduction, running time down 7-64% on the models tested (Zhang and Sennrich 2019) — transformers.
  • Fuse the residual add with the norm — saves one tensor read and write (order of) — pre-norm blocks.
  • bf16 activations with fp32 statistics — half the traffic of fp32 — training and serving.
  • Fold BatchNorm into the previous conv or linear at inference — the layer disappears — CNN deployment.
• Naive — separate mean, variance, subtract, divide, scale ops in eager mode · several passes over the tensor
• Standard — fused LayerNorm kernel · 1 read + 1 write per row
• State of the art — fused add + RMSNorm, or folded into neighboring ops · about 1 pass shared with the residual, or zero cost

Implementation notes

• Compute statistics in fp32 even for bf16 activations, then cast the output. Use two-pass or Welford variance: E[x²] − E[x]² cancels catastrophically when |μ| ≫ σ.
• Exclude γ and β (and all norm parameters) from weight decay.
• Fuse the residual add with the following norm (x = x + h; y = norm(x) in one kernel) to save a full read and write.
• BatchNorm inference: fold γ, β and running stats into the preceding conv weights (fold_bn above). Zero runtime cost.
• PyTorch BatchNorm updates running = (1 − momentum)·running + momentum·batch with momentum 0.1, and uses the unbiased variance for the running estimate but the biased one to normalize.
• Distributed BatchNorm: SyncBatchNorm all-reduces sum and sum-of-squares across devices, one collective per layer.
• QK-norm: LayerNorm or RMSNorm on queries and keys per head stops attention logits from growing in large or long-trained transformers (used in ViT-22B).

Pitfalls

• Forgetting model.eval(): BatchNorm keeps using batch statistics at inference and outputs depend on the other examples.
• Fine-tuning with tiny batches and BatchNorm in train mode: running stats drift. Freeze BN statistics.
• Normalizing over padded positions or over the sequence axis in causal models: future tokens leak into the present through the statistics.
• eps too small for the dtype: with fp16 inputs, 1e-12 underflows to 0 and a zero row divides by zero.
• Post-LN transformers deeper than about 12 layers without long warmup: gradients near the output are large and training diverges.

Variants and alternatives

In production

torch.nn.LayerNorm, nn.RMSNorm (PyTorch ≥ 2.4), nn.BatchNorm2d, nn.GroupNorm. Apex and Triton fused LayerNorm/RMSNorm, the RMSNorm kernels in vLLM and llama.cpp. Llama, Mistral, Qwen and Gemma use pre-RMSNorm. BERT and GPT-2 use LayerNorm (post- and pre-LN respectively). ResNets use BatchNorm.

Implementations

Sources

• Ioffe and Szegedy, Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift, ICML 2015. paper
• Ba, Kiros, Hinton, Layer Normalization, 2016. paper
• Zhang and Sennrich, Root Mean Square Layer Normalization, NeurIPS 2019. paper
• Wu and He, Group Normalization, ECCV 2018. paper
• Xiong et al., On Layer Normalization in the Transformer Architecture, ICML 2020. paper
• Wang et al., DeepNet: Scaling Transformers to 1,000 Layers, 2022. paper
• Dehghani et al., Scaling Vision Transformers to 22 Billion Parameters, ICML 2023. paper
• Ioffe, Batch Renormalization, NeurIPS 2017. paper