Most modern language models, like Llama and Mistral, normalize activations with Root Mean Square Normalization (RMSNorm). Here is what it does, and a PyTorch module you can drop into a model.
The problem it solves
Inside a neural network, a layer’s output feeds the next layer as input. Those values can become too big or too small as they pass through layer after layer. A value that doubles at each of 40 layers is unusable by the end. Normalization rescales the values back to a steady range before they move on, so training stays stable.
What RMSNorm does
$$ \mathrm{RMS}(x) = \sqrt{\frac{1}{n}\sum_{i=1}^{n} x_i^2 + \epsilon} $$
$$ y_i = \frac{x_i}{\mathrm{RMS}(x)} \gamma_i $$
Take a vector $x$ of $n$ values. Square each value, sum them, average the squares, take the square root. That gives the result of $\mathrm{RMS}(x)$, one number for the row’s typical size.
$\epsilon$ is a tiny constant so $x_i$ never divides by zero. Divide every value in the row by $\mathrm{RMS}(x)$, and the row’s own RMS becomes about 1. Multiply by $\gamma$, a learned weight for each position, to get the output $y$.
Example: row $x = (2, -4, 6, -8)$. $\mathrm{RMS}(x) = \sqrt{30} \approx 5.48$. With weights $\gamma = (1.5, 0.5, 1.0, 2.0)$, the output is $y = (0.55, -0.37, 1.10, -2.92)$.
The code
import torch
import torch.nn as nn
class RMSNorm(nn.Module):
def __init__(self, dim: int, eps: float = 1e-6):
super().__init__()
self.eps = eps
self.weight = nn.Parameter(torch.ones(dim))
def forward(self, x: torch.Tensor) -> torch.Tensor:
rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
return (x / rms) * self.weight
dim
is the length of the row being normalized, the model’s hidden size in a transformer. self.weight
starts at all 1s, so a freshly initialized RMSNorm
passes its input through unchanged (aside from the rescaling), and training moves each weight away from 1 as needed.
forward
runs as the same two formulas above.
The RMS division and the weight multiplication both broadcast: rms
has shape (..., 1)
and self.weight
has shape (dim,)
, so each applies to every row of x
without a loop.
Running it on a random batch:
torch.manual_seed(0)
x = torch.randn(2, 8) * 10
norm = RMSNorm(8)
y = norm(x)
print(y.pow(2).mean(dim=-1).sqrt())
tensor([1.0000, 1.0000], grad_fn=<SqrtBackward0>)
Two rows, each starting on a different scale, both land at an RMS of 1 after the call. torch.nn.RMSNorm, shipped in PyTorch 2.4 and later, does the same computation.
RMSNorm vs. LayerNorm
LayerNorm does one more step: before dividing, it subtracts the row’s average from every value. This centers the row on 0. RMSNorm skips that. It rescales, but never re-centers.
Zhang and Sennrich, who proposed RMSNorm in 2019, measured how much the centering step in LayerNorm actually mattered. Across machine translation and language modeling, dropping it barely moved model quality, while removing the mean and variance calculation cut layer normalization’s runtime by 7-64%, depending on the model.
Full code: gist.github.com/soasme/rms_norm.py (module in rms_norm.py
, the demo above in demo.py
).