LayerNorm vs RMSNorm architecture — math, numerics, fusion, placement

The two formulas, just enough to build them

Both norms operate on a single row x of length d — the hidden vector for one token — and both end with a learned per-channel scale. LayerNorm centers and scales:

μ   = (1/d) Σ_i x_i
var  = (1/d) Σ_i (x_i - μ)^2
y_i  = γ_i * (x_i - μ) / sqrt(var + ε) + β_i

RMSNorm drops the mean and the bias, keeping only the scale by root-mean-square:

ms   = (1/d) Σ_i x_i^2
y_i  = γ_i * x_i / sqrt(ms + ε)

From an implementation standpoint the difference is stark: LayerNorm needs the mean before it can compute the variance (a data dependency), plus a bias add. RMSNorm needs only Σ x_i^2, a single reduction with no prior mean to wait on. Everything downstream — numerics, fusion, cost — follows from that one structural gap.

Where epsilon actually goes

The ε exists to stop a divide-by-zero when a row is constant (LayerNorm) or all-zero (RMSNorm), and to bound the gradient when the variance is tiny. Placement is not cosmetic. The correct form is sqrt(var + ε) — epsilon inside the radical — which keeps the denominator well-conditioned for the backward pass. A stray sqrt(var) + ε is a different function: it barely regularizes small variances and shifts the whole normalizer, and it is a genuine source of train/inference mismatch when two code paths disagree.

Magnitude matters too. LayerNorm typically uses ε = 1e-5; RMSNorm implementations commonly use 1e-6. If activations are stored in fp16, an epsilon of 1e-8 is below the representable resolution near the values it guards and effectively vanishes — another reason the reduction and the epsilon add both belong in fp32.

Mixed precision: compute in fp32, store in bf16

Modern models keep activations in fp16 or bf16 to halve memory traffic, but the normalization statistics must be computed in fp32. The reason is range and resolution. Summing d = 4096 squared activations in fp16 risks overflow — fp16 tops out near 65504, and a handful of activations with magnitude 10 already push Σ x_i^2 past that. bf16 shares fp32’s exponent range so it will not overflow, but its 8-bit mantissa loses so much precision in a long accumulation that the variance is visibly wrong.

The standard recipe: load bf16 x, upcast to fp32, accumulate the sum (and sum-of-squares) in fp32, compute 1/sqrt(·) in fp32, then multiply and downcast the result back to bf16 on store. The gain γ is usually held in fp32 as well. This costs nothing in bandwidth — the tensor is still bf16 in memory — and it is the single most common fix for a norm that trains fine at small batch and diverges at large.

Catastrophic cancellation and the two-pass problem

The textbook variance identity var = E[x^2] - E[x]^2 is a trap in finite precision. When the mean is large relative to the spread, you are subtracting two big, nearly equal numbers to recover a small one, and the leading significant digits cancel — catastrophic cancellation — leaving noise or even a negative ‘variance’ that makes sqrt return NaN.

The numerically safe route is the two-pass method the LayerNorm formula implies literally: first pass computes μ, second pass computes Σ(x_i - μ)^2 from the already-centered values, which never cancels. The cost is reading the row twice. For a row that fits in cache that is cheap; for streaming or when the row is large relative to cache, the second read is real bandwidth, which is exactly what one-pass and Welford methods try to avoid.

One-pass and Welford variance

A single pass can still be stable if you update the mean and the centered sum of squares together. Welford’s algorithm does this incrementally:

for i in 1..d:
    δ   = x_i - mean
    mean += δ / i
    M2  += δ * (x_i - mean)   # uses the *updated* mean
var = M2 / d

Because M2 only ever accumulates products of deviations, it never subtracts large like-magnitude quantities, so it stays accurate in one sweep. The catch on wide SIMD hardware is that the scalar recurrence is serial. In practice GPU and CPU kernels use a parallel variant: each lane keeps its own (count, mean, M2), and a tree reduction merges partial triples with Chan’s parallel formula. RMSNorm sidesteps all of this — it needs only Σ x_i^2, and centering is never involved, so a plain fp32 sum is already stable.

Why the op is memory-bound

Count the arithmetic: LayerNorm is a few multiply-adds per element plus one rsqrt per row — on the order of 5d to 10d flops for a row of width d. Count the memory: you read d activations, read d gain (and bias) parameters, and write d outputs. The arithmetic intensity is roughly one flop per byte moved, far below the hundreds of flops per byte a modern GPU or CPU can sustain. Normalization is therefore memory-bandwidth bound: its runtime is set by how fast you can stream the tensor, not by the math.

This single fact drives every serving optimization that follows. If the op is bound by moving bytes, the way to make it faster is to move fewer bytes — which means never writing the normalized tensor to memory just to read it back into the next layer. That is the case for fusion.

Fusing the norm into the residual and the next matmul

In a transformer block the norm never stands alone. Pre-norm layout computes h = x + sublayer(norm(x)), so a naive implementation writes norm(x) to HBM, reads it back for the attention or FFN matmul, then reads x again for the residual add — three tensor-sized memory trips around a one-flop-per-byte op. A fused kernel keeps the normalized row in registers or shared memory and hands it straight to the matmul epilogue, and it folds the residual add into the same launch.

Collapsing those trips removes redundant reads/writes and one kernel-launch overhead per norm, per layer, per token — the kind of saving that compounds into a low-double-digit percent of end-to-end throughput. RMSNorm is easier to fuse precisely because it has no mean dependency and no bias: one reduction, one scale, done, which is part of why inference-oriented architectures adopted it.

A worked example, LN and RMS side by side

Take a width-8 row x = [2, 4, 4, 4, 5, 5, 7, 9]. The mean is 40 / 8 = 5. Deviations are [-3, -1, -1, -1, 0, 0, 2, 4]; their squares sum to 32, so var = 32 / 8 = 4 and σ = 2. With γ = 1, β = 0, ε ≈ 0, LayerNorm gives (x - 5) / 2 = [-1.5, -0.5, -0.5, -0.5, 0, 0, 1, 2] — a zero-mean, unit-variance row.

RMSNorm instead uses ms = (1/8) Σ x_i^2. The squares sum to 232, so ms = 29 and sqrt(ms) ≈ 5.385. The output is x / 5.385 ≈ [0.37, 0.74, 0.74, 0.74, 0.93, 0.93, 1.30, 1.67] — note it is not zero-mean, because RMSNorm never subtracted the mean. That preserved offset is exactly the representational difference the gain and the rest of the block learn to absorb.

A CPU SIMD kernel

On a CPU serving an SLM, the norm is a tight vectorized loop. For RMSNorm over a bf16 row you widen to fp32, accumulate the sum of squares into several SIMD accumulators, then reduce. Multiple accumulators matter: a single one serializes on the multiply-add latency, whereas 4–8 independent lanes keep the pipeline full and are summed only at the end.

acc0..acc3 = 0                      # fp32 SIMD registers
for i in steps of 8:               # AVX2: 8 fp32 lanes
    v = fp32(load_bf16(x + i))
    acc_k = fma(v, v, acc_k)        # rotate across k
ss  = horizontal_add(acc0+..+acc3)  # one scalar reduction
inv = rsqrt(ss / d + eps)
for i in steps of 8:
    store_bf16(y + i, fp32(load(x+i)) * inv * gamma_i)

The horizontal reduction across lanes is the one serial step and wants to happen exactly once. LayerNorm needs the same shape twice — a reduction for the mean, then a second for the centered variance — before the scale loop, which is the concrete extra cost of centering.

How much cheaper is RMSNorm, really

The honest answer: not much in raw flops, but meaningfully in the things that gate a memory-bound kernel. RMSNorm skips one full reduction (no mean), skips the centering subtraction on every element, and skips reading and adding the bias β. That trims a reduction pass and a per-element op, and it removes the mean-then-variance data dependency that limits how tightly LayerNorm can pipeline.

On a memory-bound op the practical win is modest per call — often in the single-digit to low-double-digit percent range on the norm itself — but it recurs twice per block across dozens of layers for every token, and the simpler kernel fuses more readily, so the compounded serving effect is larger than the microbenchmark suggests. The parameter saving (dropping β) is negligible for memory but real for a fused register budget.

Pitfalls worth a checklist

A short list of the mistakes that actually ship. Epsilon dtype: add epsilon in fp32; adding it after downcasting can null it out. Epsilon placement: inside the sqrt, always — the two forms are different functions and silently diverge between frameworks. Gain init: initialize γ to 1 (not 0), or the first forward pass zeros the signal. Reduction axis: normalize over the hidden dimension only; accidentally including the sequence or batch axis turns LayerNorm into a different, wrong statistic.

Two more bite in the backward pass. The LayerNorm gradient contains two reduction terms (over the normalized values and over their product with the upstream gradient), so a fused backward must carry both — RMSNorm has one. And keep the saved 1/sqrt(var + ε) from the forward pass rather than recomputing it, so forward and backward use bit-identical denominators.