RMSNorm

Normalization layers keep activation scale under control so that deep stacks train at all. RMSNorm does the job with one reduction instead of two: it rescales a vector by its root mean square and learns a single per-channel gain. It is the normalization used in most modern transformer blocks.

The idea

For a vector $x \in \mathbb{R}^d$ and a learned gain $\gamma$:

Two differences from LayerNorm matter:

• No mean subtraction. LayerNorm re-centers the vector to zero mean; RMSNorm leaves the mean where it is and only controls magnitude.
• No bias. LayerNorm carries 2d learned values (weight and bias), RMSNorm carries d — for d = 4096 that is 4096 saved parameters per norm per block.

The scaling property is the point. Multiplying the input by any positive constant cancels against the denominator, so the output is invariant to the input's overall scale. That keeps a residual stream from drifting into tiny or huge magnitudes.

The ε inside the square root is doing real work: the denominator can never fall below √ε, so the gradient of the sum of outputs is bounded by 1/√ε.

Worked example

Take x = (1, 2, 3, 4) with γ = 1 and ε = 1e-5.

• Mean square: (1 + 4 + 9 + 16) / 4 = 30/4 = 7.5
• Root mean square: √7.5 ≈ 2.7386
• Output: (0.3651, 0.7303, 1.0954, 1.4606)

LayerNorm on the same input gives (−1.3416, −0.4472, 0.4472, 1.3416): zero mean, unit variance. RMSNorm's output has mean 0.9129, which is fine — the block downstream can learn to shift it.

Scale invariance is easy to see. Feed in 100·x and the output is identical to four decimal places.

In code

import torch

x = torch.tensor([1.0, 2.0, 3.0, 4.0])
gamma = torch.ones(4)

def rmsnorm(x, gamma, eps=1e-5):
    return x / torch.sqrt(x.pow(2).mean(-1, keepdim=True) + eps) * gamma

print([round(v, 4) for v in rmsnorm(x, gamma).tolist()])
print([round(v, 4) for v in rmsnorm(100 * x, gamma).tolist()])
print([round(v, 4) for v in torch.nn.functional.layer_norm(x, (4,)).tolist()])
print(sum(p.numel() for p in torch.nn.RMSNorm(4).parameters()))    # 4

tiny = (1e-8 * x).requires_grad_(True)
rmsnorm(tiny, gamma).sum().backward()
print(tiny.grad[0].item())          # 316.22778 == 1/sqrt(1e-5), bounded by eps

Check yourself

1. You multiply an activation vector by 1000 before feeding it to RMSNorm. What changes in the output?
2. RMSNorm never subtracts the mean, so why does it still need ε?
3. Why can RMSNorm get away with dropping the bias term that LayerNorm keeps?

Key takeaways

• RMSNorm divides by the root mean square and applies one learned gain per channel.
• Outputs depend on the direction of x, not on its overall magnitude.
• One reduction and no mean subtraction make it cheaper than LayerNorm at nearly the same quality.