RMSNorm

LayerNorm without the mean subtraction, on the observation that the re-scaling does the work and the re-centring mostly does not.

[normalisation][core]

Theory

RMSNorm divides an activation vector by its root mean square and applies a learned per-channel gain. Dropping the mean subtraction removes one reduction pass and one term from the backward pass, at no measured cost in convergence for transformer language models.

\mathrm{RMSNorm}(x) = g \odot \frac{x}{\sqrt{\tfrac{1}{d}\sum_i x_i^2 + \varepsilon}}

The $\varepsilon$ sits inside the square root, not outside it — a detail that changes the gradient near zero and differs between implementations often enough to be worth checking when porting weights.

Implementation

python · torch ≥ 2.1

import torch
from torch import Tensor, nn


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x: Tensor) -> Tensor:
        dtype = x.dtype
        x = x.float()                              # accumulate in fp32
        rms = x.pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(rms + self.eps)
        return (x.to(dtype) * self.weight)

The upcast is not optional. In bf16 the sum of squares over a 4096-wide vector loses enough mantissa that the norm drifts by a fraction of a percent per layer, and the drift compounds with depth.

scaled-dot-product-attention gradient-checkpointing

References

[1]Zhang & Sennrich — Root Mean Square Layer Normalization (2019)arXiv:1910.07467

[2]Xiong et al. — On Layer Normalization in the Transformer Architecture (2020)arXiv:2002.04745