Normalize Hidden States with RMSNorm

A hidden-state vector can grow or shrink in magnitude as it passes through a neural network. RMSNorm controls that scale by dividing the vector by its root mean square magnitude, then applying a trainable gain. Unlike LayerNorm, it does not subtract the vector mean before rescaling.

That missing centering operation is the defining distinction. RMSNorm constrains scale while leaving a uniform shift across coordinates present in the normalized representation.

RMSNorm uses the second raw moment

For a hidden vector x with d coordinates, its root mean square is:

rms(x) = sqrt((1 / d) * sum(x_i^2))

A practical implementation includes a small positive epsilon in the denominator computation:

y_i = g_i * x_i / sqrt((1 / d) * sum(x_j^2) + epsilon)

Here, g_i is a trainable gain for coordinate i. The exact placement and numeric value of epsilon are implementation details that should match the model definition used to create a checkpoint.

The statistic is based on squared coordinate values rather than deviations from their mean. As a result, RMSNorm does not need to compute x_i - mean(x) before normalization.

Consider two vectors:

a = [1, -1]
b = [3, 1]

Both have the same centered values after subtracting their respective means: [1, -1]. LayerNorm therefore removes the common offset between them before scaling. RMSNorm sees different raw squared magnitudes and retains information associated with that offset.

Scale changes cancel before the gain

Ignoring epsilon, multiply every coordinate by a positive scalar c:

rms(c * x) = c * rms(x)

The normalized ratio becomes:

(c * x_i) / (c * rms(x)) = x_i / rms(x)

So a positive uniform rescaling of the input does not change the normalized direction. A negative scalar adds a sign reversal because the RMS denominator is nonnegative.

This property concerns input scale, not arbitrary changes to individual coordinates. Multiplying one feature by a large factor changes both that coordinate and the shared RMS statistic, which can alter every normalized coordinate.

The trainable gain also matters. RMS normalization produces a scale-controlled intermediate vector, but g can assign different output scales to different coordinates. The module is therefore not equivalent to forcing every output vector to a fixed Euclidean norm after all trainable parameters are applied.

RMSNorm and LayerNorm preserve different information

LayerNorm typically computes a mean and variance across the normalized feature dimension:

mu = mean(x)
variance = mean((x - mu)^2)
y = g * (x - mu) / sqrt(variance + epsilon) + b

RMSNorm instead uses the raw second moment:

mean_square = mean(x^2)
y = g * x / sqrt(mean_square + epsilon)

Since:

mean(x^2) = variance(x) + mean(x)^2

the two denominators coincide only when the feature mean is zero, apart from details such as epsilon. Even then, parameterization can differ if one module has an additive bias and the other does not.

A constant offset makes the distinction direct. Add scalar k to every coordinate. LayerNorm removes that common shift during centering. RMSNorm does not. Its numerator changes, and its RMS statistic changes as well.

This means RMSNorm is not simply a cheaper spelling of LayerNorm. Replacing one with the other changes the transformation represented by the network. A checkpoint trained with one normalization rule cannot generally be converted by swapping the module name while retaining identical activations.

The normalized dimension must match the architecture

In transformer blocks, normalization commonly operates across the hidden dimension independently for each token position. If a tensor has conceptual shape:

[batch, sequence, hidden]

the RMS statistic is commonly taken across hidden, producing one scale statistic for each [batch, sequence] position.

Normalizing across the sequence dimension instead couples token positions to one another and implements a different operation. The same issue appears with extra spatial, channel, or expert dimensions: the formula alone does not specify which axes participate.

This axis choice also determines the shape of the trainable gain. A per-hidden-coordinate gain has hidden elements and broadcasts across batch and sequence positions. Code that accidentally reduces across another dimension can still produce shape-compatible tensors while changing model behavior.

Numeric precision affects the reduction

The sum of squares is a reduction, so its numeric behavior depends on the dtype used for accumulation. Low-precision inputs can have a narrower representable range and less precision than a wider accumulator.

Implementations may therefore compute the mean square in a wider dtype and cast the normalized result back afterward. That is an implementation choice rather than a property implied by the mathematical definition. Compatibility work should inspect the actual kernel or framework operation used by the target model instead of assuming identical accumulation rules across runtimes.

epsilon also prevents division by zero for an all-zero vector and limits the denominator near zero. Once epsilon is significant relative to the mean square, exact scale cancellation no longer holds. For ordinary magnitudes where the mean square dominates epsilon, the formula approaches the scale-invariant form above.

Placement changes the residual computation

RMSNorm specifies a normalization operation, not its location inside a residual block. A pre-normalized transformer can apply it before attention or a feed-forward sublayer:

h = x + attention(rmsnorm(x))

Another architecture can place normalization after a residual addition:

h = rmsnorm(x + attention(x))

These expressions are not interchangeable. In the first, the residual path carries x directly while the sublayer receives a rescaled representation. In the second, normalization transforms the combined residual and sublayer output.

For model implementation and checkpoint compatibility, the relevant contract therefore includes both the RMSNorm formula and its position in the computation graph.

RMSNorm is most useful to reason about as a specific invariance choice: it suppresses uniform magnitude changes without erasing the feature mean. That boundary explains both its compact computation and the cases where substituting it for a centering normalization changes the model rather than merely changing an implementation detail.