36. RMSNorm From Scratch (and How It Differs From LayerNorm)

Problem

RMSNorm is LayerNorm with the mean-centering step removed: instead of normalizing by the standard deviation of (x - mean), it normalizes purely by the root-mean-square of x itself, and it drops the additive bias term entirely. The claim behind it is that re-centering isn't what makes normalization help — re-scaling is — so you can skip a reduction (the mean) and a parameter (the bias) at little to no quality cost, for a cheaper op. RMSNorm is now the default normalization layer in most modern efficient architectures.

Implement RMSNorm's forward pass, then be ready to state precisely how it relates to LayerNorm: which term does it skip, and under what condition does it become numerically identical to LayerNorm?

Source: src/36_rms_norm_from_scratch.py

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

Examples:
>>> x = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
>>> weight = torch.ones(4)
>>> rms_norm(x, weight).shape
torch.Size([1, 4])

Step-by-Step Approach

1. Skip the mean entirely — RMSNorm never computes x.mean(dim=-1). This is the single structural difference from LayerNorm.
2. Compute the mean of squares over the last dim: x.pow(2).mean(dim=-1, keepdim=True) — this is the "energy" of the feature vector.
3. Take the square root with eps added inside, for the same divide-by-zero protection as LayerNorm: rms = sqrt(mean_sq + eps).
4. Divide x by rms, then multiply by the learned weight — there is no bias/beta term to add afterward.
5. Convince yourself of the relationship to LayerNorm algebraically: if a row of x already has mean 0, then var(x) = mean((x - 0)^2) = mean(x^2), so RMSNorm's denominator becomes identical to LayerNorm's — with unit gamma and zero beta, the two layers produce the same output.
6. Note the practical payoff: RMSNorm skips one reduction (the mean) and one learned parameter (beta) per normalization call, which adds up across every layer of a deep, long-sequence model.

The key insight is that RMSNorm isn't a different idea from LayerNorm, it's LayerNorm's re-scaling term with the re-centering term deleted -- understanding that relationship (not just the two formulas separately) is what the "compare to LayerNorm" half of this problem is testing.

Reference solution

def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
    # No mean subtraction at all -- this is the entire difference from
    # LayerNorm. We only measure the "energy" (root-mean-square) of x.
    ms = x.pow(2).mean(dim=-1, keepdim=True)      # shape (..., 1)
    rms = torch.sqrt(ms + eps)                     # eps inside sqrt again, same reason as LayerNorm
    x_norm = x / rms                                # shape (..., d)
    return x_norm * weight                          # no beta/bias term to add

Key Functions & Tricks

• x.pow(2).mean(dim=-1, keepdim=True) — the mean-square ("energy") that replaces LayerNorm's variance.
• torch.sqrt(ms + eps) — eps inside the root, identical placement convention to LayerNorm.
• No x.mean(dim=-1) call anywhere — its absence is the entire point of the layer.
• F.layer_norm(...) — used in the test suite as an independent cross-check: on a mean-zero row with unit weight/zero bias, it must equal rms_norm's output.
• Broadcasting * weight — single-parameter affine (no additive term), shape (d,) broadcasting over leading dims.

How to Recognize This Pattern