RMSNorm Deep Dive — Math + Kernels

The Jacobian is a projector, and its null space is the input

Write ms = (1/d) Σ_j x_j², r = √(ms + ε), x̂ = x / r. Since ∂r/∂x_j = x_j / (d·r), the normalize step’s Jacobian is:

∂x̂_i/∂x_j = δ_ij/r − x_i·x_j / (d·r³)
J          = (1/r) · ( I − x̂ x̂ᵀ / d )     J: [d, d]

That rank-one correction is exactly the piece that makes J singular. Note x̂·x̂ = d·ms/(ms+ε), which is d when ε = 0. Then J·x = (1/r)(x − x̂·(d·r)/d) = 0: the input direction is annihilated. J is (1/r) times an orthogonal projector onto the complement of x, of rank d − 1. LayerNorm’s Jacobian is the same shape with one more subtraction, (1/d)·11ᵀ, and therefore rank d − 2: it kills the constant direction too. Dropping the mean gives the layer back one dimension of expressiveness.

Why the input gradient is exactly orthogonal to the input

Let g = ∂L/∂y and u = γ ⊙ g (the upstream gradient pulled through the gain). Applying Jᵀ:

∂L/∂x = (1/r) · [ u − x̂ · (u·x̂)/d ]
∂L/∂γ_i = Σ_tokens  g_i · x̂_i

One reduction — the scalar u·x̂ — and that is the whole backward pass. LayerNorm needs two (one over x̂, one over u⊙x̂), which is why a fused RMSNorm backward is the simpler kernel.

Now dot the input gradient with x. With ε = 0, x̂·x̂ = d, so the two terms cancel and (∂L/∂x)·x = 0 exactly, for every input and every upstream gradient. This is not an approximation; it is the derivative form of RMSNorm(c·x) = RMSNorm(x). With a real epsilon the residual is (u·x̂)·ε/(ms+ε) — a tiny, quantifiable leak of scale sensitivity.

A worked backward pass, by hand

Take d = 4, x = [2, 2, −2, 2], γ = [1, 2, 1, 0.5], ε ≈ 0. Then Σx² = 16, ms = 4, r = 2, and x̂ = [1, 1, −1, 1] (check: x̂·x̂ = 4 = d). Forward output y = γ⊙x̂ = [1, 2, −1, 0.5].

Let the upstream gradient be g = [1, 0, −1, 2], so u = γ⊙g = [1, 0, −1, 1] and u·x̂ = 1 + 0 + 1 + 1 = 3.

x̂·(3/4) = [0.75, 0.75, −0.75, 0.75]
u − that   = [0.25, −0.75, −0.25, 0.25]
∂L/∂x = (1/2)·that = [0.125, −0.375, −0.125, 0.125]
∂L/∂γ = g⊙x̂ = [1, 0, 1, 2]

Verify the orthogonality claim: 0.125(2) − 0.375(2) − 0.125(−2) + 0.125(2) = 0.25 − 0.75 + 0.25 + 0.25 = 0. Scale the input by ten (r = 20) and x̂, y, and ∂L/∂γ come out identical.

What scale invariance does to your optimizer

Invariance is not free; it relocates a knob. Suppose a weight matrix W produces a vector that feeds a norm directly. Because L(c·W) = L(W), the gradient is orthogonal to W and ∇L(c·W) = (1/c)·∇L(W), so gradient magnitude falls as 1/||W||.

An orthogonal update can only grow the norm: ||W + ηg||² = ||W||² + η²||g||². Since only the direction of W affects the loss, the meaningful step size is the angular one, which scales as η·||g|| / ||W|| ∝ η / ||W||². So the network self-anneals: as weights grow, the effective learning rate decays — and weight decay, by shrinking ||W||, raises it. Your decay coefficient is partly a learning-rate schedule in disguise. This is exact only where a weight’s output enters a norm unmediated; in a pre-norm stack the sublayer output joins the stream first, so the invariance is approximate — but the coupling survives.

No mean, no cancellation — by construction

A large share of LayerNorm kernel folklore — the E[x²] − E[x]² trap, catastrophic cancellation, two-pass sweeps, Welford updates and their parallel merges — exists to solve one problem: a variance requires subtracting two large, nearly equal quantities.

RMSNorm does not have that problem, structurally rather than luckily. Σx_j² is a sum of non-negative terms, so nothing ever cancels; there is no intermediate statistic to wait on, so a single sweep is natural rather than a risky optimization; and the result can never come out negative, so √ can never return NaN from a sign error. A plain fp32 accumulation is the numerically stable algorithm. Deleting the mean does not merely save a pass — it deletes a category of bug, and with it the branch of the kernel that used to defend against it.

Epsilon in squared units, and what actually overflows

The most common misreading of ε is dimensional. It is added to ms, a mean of squares, not to an activation, so ε = 1e−6 guards rows whose RMS is around √(1e−6) = 1e−3, not 1e−6. Compare epsilons by their square roots or you will be off by three orders of magnitude when porting a config. Representability compounds it: in fp16 the smallest normal is 6.1e−5 and the smallest subnormal 6e−8, so 1e−6 is subnormal (and gone under flush-to-zero) and 1e−8 rounds to exactly zero. The guard you wrote silently does not exist.

Overflow gets misread the same way. Two numbers, two different stories:

fp16 max        ≈ 65504
single element  overflows when |x_i| > √65504 ≈ 256   (rare)
accumulator     d = 4096, |x_i| ≈ 4  →  Σx² ≈ 65536  (overflows!)

An individual activation almost never reaches 256; the sum reaches the fp16 ceiling with entirely ordinary activations, purely because d is large. bf16 carries fp32’s exponent range and so never overflows, but with only 8 significand bits it drops any addend more than 2⁸ = 256× below the running sum — at Σx² ≈ 65536 with terms near 16, the tail of the row contributes nothing. Silent stagnation, not an inf. Hold ms, the ε add, and the rsqrt in fp32 whatever the activation dtype; the tensor stays 16-bit in memory, so it costs no bandwidth.

Folding the gain into the next matmul

At inference γ can disappear entirely. The norm output feeds a linear layer, so for W: [d_out, d_in]:

z = W · (γ ⊙ x̂)  =  (W · diag(γ)) · x̂
i.e. scale column j of W by γ_j, once, offline

Per token, per norm, per layer you delete a d-wide elementwise multiply and a d-element load of γ — and on a memory-bound op the vanished load is the larger win. In pre-norm attention the norm feeds W_Q, W_K, and W_V, so fold into all three. Quantization cuts both ways: if W is quantized per input channel the per-column scales absorb γ for free, but per output channel the fold changes the range each row scale was calibrated on, so you must re-quantize. RMSNorm folds cleanly precisely because there is no β bias and no mean subtraction to leave behind an affine remainder.

One sweep, two statistics: fusing norm and int8 quantization

On a CPU SLM the tensor leaving a norm usually heads into an int8 matmul, which needs a per-token scale — that is, absmax(y). The naive kernel takes an extra sweep over the row to find it. It does not have to, because r is a positive scalar that factors straight out:

absmax(y) = max_i |γ_i x_i / r| = (1/r) · max_i |γ_i x_i|
pass 1: accumulate Σx²  AND  max|γ_i x_i|   (both from x alone)
then:   r = √(ms+ε);  scale = 127·r / max|γ⊙x|

Both statistics come out of the same load of x, in one SIMD loop, and the second pass writes int8 directly. On the earlier example γ⊙x = [2, 4, −2, 1], so absmax = 4 and absmax(y) = 4/2 = 2 — matching y = [1, 2, −1, 0.5]. LayerNorm cannot do this: |γ_i(x_i − μ)| depends on μ, which is not known until the pass has ended.

Residual-stream growth and what it does above it

In a pre-norm stack, x_ℓ₊₁ = x_ℓ + F_ℓ(RMSNorm(x_ℓ)), and the stream’s norm grows with depth — roughly like √L if block outputs are treated as uncorrelated increments, faster in practice. That variance bookkeeping belongs to the signal-propagation and variance-maintenance articles in this series; what matters here is how it interacts with the norm.

RMSNorm makes each block’s input scale-stable regardless, which is the point. But the block’s relative edit to the stream shrinks like 1/√ℓ: deep layers write into an increasingly loud residual and are progressively quieter contributors. Two visible consequences: learned γ often drifts upward in later layers to claw back amplitude, and r itself grows monotonically with depth, so a per-layer fp16 norm that was safe at layer 4 can overflow at layer 40. That depth dependence is the argument for an fp32 accumulator everywhere rather than only where you measured a problem.

Pitfalls specific to the gain

The γ-only parameterization concentrates the failure modes into one vector. Weight decay on γ is the classic: decay pulls it toward zero, which is signal death, not regularization. Exclude norm parameters from decay, or use Gemma’s reparameterization y = x̂ · (1 + w) with w initialized to zero, where decay pulls toward the identity.

Three more. Keep γ and its gradient in fp32: ∂L/∂γ accumulates over every token in the batch, so a bf16 accumulator stagnates for exactly the reason the forward reduction does. Save 1/r from the forward pass rather than recomputing it in the backward, so both use a bit-identical denominator. And check the reduction axis: RMSNorm reduces over the hidden dimension only, and a kernel that accidentally folds in the sequence axis still trains — badly, and with a batch-dependent inference path that is miserable to find later.