RMS Normalization

Root Mean Square Layer Normalization (RMSNorm).

Implements RMSNorm and its partial variant pRMSNorm, both introduced by Zhang & Sennrich (2019), arXiv:1910.07467.

RMSNorm simplifies LayerNorm by removing the mean-subtraction (re-centering) step and rescaling activations solely by their root mean square magnitude. This preserves the re-scaling invariance property while reducing computation.

pRMSNorm further reduces overhead by estimating the RMS statistic from only the first *p*% of the hidden dimensions, exploiting the assumption that neurons within a layer are approximately i.i.d.

class src.model.norm.rms.RMSNorm(*args: Any, **kwargs: Any)[source]

Bases: Module

Root Mean Square Layer Normalization.

Applies RMS normalization to the last dimension of the input:

y_i = g_i * x_i / RMS(x),   RMS(x) = sqrt( (1/n) * sum_j x_j^2 + eps )

where g is a learnable per-dimension scale (initialized to 1) and no bias is used, following the original formulation.

When partial_ratio is greater than 0, the RMS is estimated from only the first k = ceil(dim * partial_ratio) elements (pRMSNorm). This is the partial variant from Section 5 of the paper; the recommended default ratio is 6.25%.

Parameters:

  • dim – Number of features in the input (normalized dimension).

  • eps – Small constant for numerical stability. Defaults to 1e-6.

  • partial_ratio – Fraction of dimensions used for RMS estimation. 0.0 (default) uses all dimensions (standard RMSNorm). Values in (0, 1] activate pRMSNorm. The number of dimensions used is max(1, ceil(dim * partial_ratio)).

weight

Learnable per-dimension scale of shape (dim,), initialized to 1.

eps

Epsilon value for numerical stability.

partial_ratio

Fraction of dimensions used for RMS estimation.

k

Number of dimensions used for RMS estimation (dim when partial_ratio is 0, otherwise ceil(dim * partial_ratio)).

References

Zhang & Sennrich (2019). “Root Mean Square Layer Normalization.” arXiv:1910.07467.

forward(x: torch.Tensor) → torch.Tensor[source]

Apply (partial) RMS normalization.

Parameters:

x – Input tensor of shape (..., dim).

Returns:

Tensor of same shape as x, normalized by RMS and scaled by the learnable weight.