RMSNorm against LayerNorm

LayerNorm does two things to a vector, it subtracts the mean and it divides by the standard deviation. The RMSNorm paper asked, do we even need the first one ? LLaMA, Qwen and most current decoders say no, and use RMSNorm. EXP-010 swapped it into the Day 1 baseline and changed nothing else.

Divide by the root mean square, keep a gain

RMSNorm divides each vector by its root mean square and multiplies by a learned gain $g$:

Compare with LayerNorm, $\gamma \odot (x - \mu)/\sigma + \beta$. RMSNorm has no $\mu$ and no $\beta$. If the mean happens to be 0, the root mean square equals the standard deviation and the two agree exactly, apart from the bias.

The paper's argument is that LayerNorm's benefit comes from the rescaling, which makes a layer's output independent of how big its input is. The re-centering is the part you can drop.

So what happened ?

• "Double every entry" leaves both outputs unchanged, both norms don't care about scale.
• "Add 1 to every entry" leaves LayerNorm unchanged and moves RMSNorm. RMSNorm isn't invariant to a shift, so its output mean is whatever the input's shape makes it.
• When the input mean is near 0, the two outputs almost overlap.

Eight lines, statistic in float32

class RMSNorm(nn.Module):
    def __init__(self, width: int, epsilon: float = 1e-6) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(width))
        self.epsilon = epsilon

    def forward(self, x: Tensor) -> Tensor:
        scale = torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.epsilon)
        return (x.float() * scale).type_as(x) * self.weight

rsqrt is one over the square root. The statistic is computed in float32 even when the activations are half precision, because squaring small half-precision numbers loses too much. The gain starts at 1, and make_norm picks LayerNorm or RMSNorm from the config's norm switch.

The paper keeps an optional bias term. LLaMA's version only has a gain, and octlm follows LLaMA. The note records that the bias variant wasn't tested.

Same quality, 8.9 percent slower

Hypothesis, RMSNorm matches LayerNorm's quality and is faster. The paper reports speedups from 7 to 64 percent across models and 6.9 to 9.3 percent on a Transformer.

The test is one row of the Day 2 variant grid, 3.3 million parameters, 400 steps, three seeds per variant, on a Colab T4. Reading the grid explains the grid and its noise.

So what happened ?

• On code, RMSNorm is 0.0013 bits per byte from the baseline. The baseline alone moves 0.0434 across three seeds, so the difference is about 30 times smaller than the noise.
• On prose the difference is 0.0006, also way inside the noise.
• It saves 2,304 parameters, one bias vector of 256 in each of 9 norms.
• It's 8.7 percent slower per step.

The slowdown is the opposite of the paper's result, and the note explains why. PyTorch's LayerNorm is one fused GPU kernel. octlm's RMSNorm is a chain of separate tensor ops (square, mean, reciprocal square root, two multiplies), and each one reads and writes the whole activation. At width 256 the norm is a noticeable share of each block, so the extra memory traffic shows up. The paper counts arithmetic saved, the measurement counts which kernels the framework gives you. Both are true, but it's the second one that decides what a config should use.

Decision, no standalone keep. RMSNorm rides inside the modern stack, the combination Day 4 trains, where it matches the Qwen architecture Day 6 has to load. Re-measure its speed at width 512 and above before calling the slowdown general.

What we didn't build, and why

• A fused RMSNorm kernel. Recent PyTorch releases ship torch.nn.RMSNorm, and torch.compile can fuse the plain version. Either would need its own speed measurement, and at this model size the norm's speed doesn't decide anything.
• The bias variant. Recorded as untested.