25. RMSNorm From Scratch

Problem

RMSNorm is a close variant of LayerNorm built on the observation that the mean-centering step (subtracting the mean) contributes little to LayerNorm's benefit — most of the value comes from rescaling by the activation magnitude. Dropping mean-centering removes one reduction and one learned bias term per layer, which adds up across a deep network's worth of norm layers. DeepMind's own large language models adopted exactly this simplification: both Gopher and Chinchilla use RMSNorm instead of LayerNorm.

Implement the forward pass: rms = sqrt(mean(x^2, dim=-1) + eps), then y = (x / rms) * weight. Note there's no mean subtraction and no additive beta/bias term — only a single learned per-feature scale.

Source: src/25_rmsnorm_from_scratch.py

def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: ...

>>> x = torch.randn(2, 5, 8)
>>> weight = torch.ones(8)
>>> rms_norm(x, weight).shape
torch.Size([2, 5, 8])
>>> rms_norm(torch.zeros(1, 4), torch.ones(4))
tensor([[0., 0., 0., 0.]])

Step-by-Step Approach

1. Square every element and take the mean over the last dimension with keepdim=True — this is the “mean square,” with no mean-subtraction anywhere in the computation.
2. Add eps inside the square root before taking it: sqrt(mean_sq + eps). This is what keeps an all-zero input row finite (sqrt(eps), not sqrt(0)) instead of producing a divide-by-zero.
3. Divide the input by the RMS to get a unit-scale-ish tensor, then multiply elementwise by the learned per-feature weight — there is no additive term to apply afterward.
4. Confirm the zero-row edge case produces a finite output (exactly zero, since 0 / sqrt(eps) == 0), not NaN or Inf.

RMSNorm is LayerNorm with exactly one step removed (mean-subtraction) and one parameter removed (the additive bias) — if you can already implement LayerNorm, the entire delta is deleting two lines, not learning a new algorithm.

Reference solution

import torch


def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
    # mean of squares over the feature dim, no mean-subtraction anywhere --
    # that's the entire distinction from LayerNorm
    mean_sq = x.pow(2).mean(dim=-1, keepdim=True)
    # eps inside the sqrt (not just added to a variance) is what keeps this
    # finite for an all-zero row: sqrt(0 + eps) is a small positive number,
    # never zero, so the division below never produces NaN/Inf
    rms = torch.sqrt(mean_sq + eps)
    return (x / rms) * weight

Key Functions & Tricks

• Tensor.pow(2).mean(dim=-1, keepdim=True) — the mean-square statistic RMSNorm normalizes by, computed with no prior mean-subtraction
• torch.sqrt(mean_sq + eps) — eps placed inside the square root is the numerical-stability trick that prevents NaN on an all-zero input row
• x / rms — the core rescaling step; unlike LayerNorm there is no x - mean term feeding into it
• elementwise * weight — the sole learned parameter — a per-feature scale with no corresponding additive bias, unlike LayerNorm's gamma/beta pair

How to Recognize This Pattern