RMSNorm RMSNorm, a normalization technique proposed by Zhang and Sennrich in 2019, rescales neural network activations by dividing by the root mean square of each row, skipping the mean-centering step of LayerNorm. PyTorch 2.4 and later ships torch.nn.RMSNorm, which performs the same computation. The technique cuts layer normalization's runtime by 7-64% depending on the model, with minimal impact on quality, and is used in modern language models like Llama and Mistral. RMSNorm 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 python 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=