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 whether the first is needed. LLaMA, Qwen and most current decoders answer 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 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 large its input is. Re-centering is the part it can drop.

What happened?

• "Double every entry" leaves both outputs unchanged. Both norms are invariant to scale.
• "Add 1 to every entry" leaves LayerNorm unchanged and moves RMSNorm. RMSNorm is not 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 nearly coincide.

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 has a gain only, and octlm follows LLaMA. The note records that the bias variant was not 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, with 3.3 million parameters, 400 steps, three seeds per variant, on a Colab T4. Reading the grid explains the grid and its noise.

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 far inside the noise.
• It saves 2,304 parameters, one bias vector of 256 in each of 9 norms.
• It is 8.7 percent slower per step.

The slowdown is the opposite of the paper's result, and the note explains it. PyTorch's LayerNorm is one fused GPU kernel. octlm's RMSNorm is a chain of separate tensor operations (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. The paper counts arithmetic saved. The measurement counts which kernels the framework provides. Both are true, and the second one 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 did not 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 decides nothing.
• The bias variant. Recorded as untested.