38. RMSNorm From Scratch

Problem

RMSNorm is LayerNorm with the mean-centering step dropped: it rescales a vector by its root-mean-square magnitude instead of its standard deviation, and applies a learned per-feature scale but no learned bias. Modern open-weight LLM families (Llama, Mistral) use RMSNorm instead of LayerNorm specifically because skipping mean subtraction removes one reduction per normalization call, which adds up across every layer of a deep model.

Implement RMSNorm.forward: divide the input by its root-mean-square over the last dimension (with eps added inside the square root for numerical stability when a vector is near zero), then scale by the learned per-feature weight. __init__ already creates self.weight — your job is only the normalization math in forward.

Source: src/38_rms_norm_from_scratch.py

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6): ...
    def forward(self, x: torch.Tensor) -> torch.Tensor: ...  # normalize over last dim

>>> torch.manual_seed(0)
>>> norm = RMSNorm(dim=4)
>>> x = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
>>> norm(x).shape
torch.Size([1, 4])

Step-by-Step Approach

1. Square every element: x.pow(2).
2. Take the mean over the last dimension, keeping it for broadcasting: .mean(dim=-1, keepdim=True).
3. Add eps inside the square root (not outside): torch.sqrt(mean_sq + eps) — this is what protects against division by zero for an all-zero row, not just a cosmetic offset.
4. Divide x by that RMS value — this is the rescaling step; note there is no mean subtraction anywhere, unlike LayerNorm.
5. Multiply elementwise by self.weight, the learned per-feature scale.

The key insight is the one-line difference from LayerNorm: LayerNorm normalizes by std(x - mean(x)) and has both a learned scale and bias; RMSNorm normalizes by sqrt(mean(x²)) directly and has only a learned scale. When a row is already zero-mean, the two produce identical results, since var(x) == mean(x²) in that case.

Reference solution

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:
        # root-mean-square over the last dim, no mean-centering; eps sits
        # *inside* the sqrt so it protects against div-by-zero even when x
        # is exactly all-zero
        rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return x / rms * self.weight

Key Functions & Tricks

• tensor.pow(2) / tensor.mean(dim=-1, keepdim=True) — the reduction that replaces LayerNorm's variance computation; keepdim=True keeps broadcasting against the unreduced x working without an explicit unsqueeze.
• eps inside the square root, not added after — puts the stabilizer where it actually prevents a zero denominator.
• nn.Parameter(torch.ones(dim)) — initializing the scale to 1 so a freshly-constructed RMSNorm starts as a near-identity transform (up to the RMS rescaling itself).
• nn.RMSNorm (PyTorch ≥ 2.4) — the built-in equivalent, useful as an oracle to check a from-scratch implementation bit-for-bit.

How to Recognize This Pattern