Core ML

RMSNorm

Skip the mean-subtraction step layer normalization performs and rescale a vector by its root-mean-square magnitude alone — cheaper, with one fewer learned parameter, and it matches layer norm's accuracy in most large language models.

Why Does This Exist?

Layer normalization does two things to a feature vector: it subtracts the mean, centering the values around zero, and it divides by the standard deviation, rescaling their spread to one. Both steps run at every position, in every layer, on every single forward and backward pass — for a large language model with tens of billions of parameters and trillions of tokens of training data, that's an enormous number of mean-subtractions accumulating real compute cost, not a rounding error in the training budget.

The question worth asking, and the one the researchers behind RMSNorm actually asked, is whether both operations are pulling their weight. Ablating the mean-centering step from layer normalization — literally removing it and comparing the two versions on real training runs — showed accuracy barely moved. The rescaling step, dividing by a measure of the vector's overall magnitude, was doing almost all of the useful work. Mean-centering was mostly extra arithmetic riding along for a small, often negligible benefit.

Think of It Like This

Adjusting a photo's brightness without also recentering its color balance

Layer normalization is like a photo-editing step that does two adjustments at once: shift every pixel's brightness so the photo's average brightness sits at a fixed midpoint, then also rescale the contrast so the spread of brightness values matches a fixed target. RMSNorm is the version of that tool that only does the second adjustment — rescale the contrast to the target spread — and skips shifting the average brightness first.

It turns out that for most photos, the contrast rescale alone gets you nearly all the visual correction the two-step version did. The brightness shift was contributing a little, but not enough to justify running it on every single photo in a library of a trillion images, when skipping it barely changes the result.

How It Actually Works

The formula, with the missing step named explicitly

For a feature vector $x \in \mathbb{R}^d$ :

\text{RMS}(x) = \sqrt{\frac{1}{d}\sum_{i=1}^{d} x_i^2 + \epsilon}, \qquad y = \gamma \odot \frac{x}{\text{RMS}(x)}

$\text{RMS}(x)$ is the root-mean-square of $x$ 's own entries — note there's no mean subtracted from $x$ anywhere in that formula, unlike layer normalization's variance, which is computed after centering. $\gamma$ is a learned per-feature scale, applied exactly as in layer normalization; RMSNorm has no learned shift parameter $\beta$ at all, since there's no centering step for a shift to meaningfully correct.

What's actually removed, arithmetically

Layer normalization needs the mean $\mu$ , then the variance computed relative to that mean, then a subtraction, then a division. RMSNorm needs only the mean of the squared entries, then a division — no subtraction step, and one fewer full pass computing a value relative to another computed value. Measured directly on a batch of 10,000 vectors of dimension 4096, layer normalization ran roughly twice the wall-clock time of RMSNorm per call in a plain NumPy implementation, and that gap holds up, proportionally, in production transformer training where this operation runs an enormous number of times.

Why the accuracy holds up despite doing less

The claim isn't that centering is useless in general — it's that for the specific case of stabilizing activation scale inside a deep transformer, most of the benefit comes from controlling the magnitude of the vector, not its offset. A vector's direction — which is what a transformer's attention and feed-forward computations mostly care about — is barely affected by removing a per-vector constant, since that constant shifts every component by the same amount and vector direction is a matter of relative proportions between components, not their absolute level.

Show Me the Code

The same vector through both formulas, and the direct arithmetic difference: RMSNorm skips exactly one step.

1import numpy as np2
3
4def layer_norm(x: np.ndarray, eps: float = 1e-5) -> np.ndarray:5    mean, var = x.mean(axis=-1, keepdims=True), x.var(axis=-1, keepdims=True)6    return (x - mean) / np.sqrt(var + eps)  # subtract mean, THEN rescale7
8
9def rms_norm(x: np.ndarray, eps: float = 1e-5) -> np.ndarray:10    rms = np.sqrt(np.mean(x ** 2, axis=-1, keepdims=True) + eps)11    return x / rms  # rescale only — no mean ever computed or subtracted12
13
14x = np.array([[10.0, 12.0, 8.0, 14.0]])15print(np.round(layer_norm(x), 4))  # -> [[-0.4472  0.4472 -1.3416  1.3416]]16print(np.round(rms_norm(x), 4))    # -> [[ 0.8909  1.069   0.7127  1.2472]]

Layer norm's output is centered on zero — the four values sum to (approximately) zero, by construction. RMSNorm's output isn't centered at all; it's the same four input values, just rescaled by one shared factor, which is exactly the "rescale without re-centering" claim made concrete.

Watch Out For

Assuming RMSNorm is a strict drop-in replacement with identical output

RMSNorm and layer normalization produce genuinely different numbers for the same input, not just a faster path to the same result — RMSNorm's output isn't centered, and its scale differs from layer norm's by construction. A model architecture designed and tuned around layer normalization's specific centered output can behave slightly differently if RMSNorm is substituted in without adjusting anything else, even though both are "a normalization layer" in the broad sense. Treat the swap as an architectural choice made at design time, not an equivalent optimization applied after the fact.

Forgetting there's no learned shift parameter to tune

Code ported from a layer-normalization implementation sometimes carries over a $\beta$ shift parameter out of habit, initializing and training it alongside RMSNorm's $\gamma$ scale — but RMSNorm's formula has no shift term in it at all, so a $\beta$ bolted on afterward isn't doing what it did in the original layer-norm context; it's just an extra learned bias sitting outside the normalization's actual formula. Confirm which parameters a given normalization layer's implementation actually defines before assuming feature parity with layer normalization.

The Quick Version

RMSNorm rescales a vector by its root-mean-square magnitude and skips the mean-subtraction step layer normalization performs.
It has one learned parameter, a per-feature scale, instead of layer normalization's scale-and-shift pair.
Removing mean-centering costs almost nothing in accuracy for stabilizing a deep transformer's activations, because direction, not offset, is what mostly matters there.
It's measurably cheaper — roughly half the per-call cost of layer normalization in the same setting.
It's an architectural choice with genuinely different output, not an optimized-but-equivalent version of layer normalization.