LayerNorm and RMSNorm

$\mu$ and $\sigma^2$ are the mean and variance of one token’s $d$ features (Part 4: over dim=-1). $\epsilon$ is a tiny constant against dividing by zero. $\gamma$ and $\beta$ are learned, one per feature.

The mean of $[2, 4, 6, 8]$ is 5. Subtracting it leaves values that sum to 0.

Now the features have mean 0 and variance 1 (ignoring $\epsilon$). Then $\gamma$ and $\beta$ rescale and shift each feature, so the model can undo any of it that it doesn’t want.

Centering removes the all-ones direction (Part 3’s projections). Dividing by $\sigma$ then puts every token on a sphere of radius $\sqrt d$: here $\sqrt 4 = 2$, and indeed $2(1.342^2) + 2(0.447^2) \approx 4$. Then $\gamma$ stretches each axis (a diagonal matrix, Part 2).

No mean, no $\beta$: just divide by the root mean square and apply the gain. It’s cheaper and works about as well (Zhang & Sennrich, 2019). Llama uses RMSNorm.

Take $\gamma = 1$, $\beta = 0$ and ignore $\epsilon$. Write $\operatorname{LN}(x)$ and $\operatorname{RMSNorm}(x)$.

The mean is already 0, so centering does nothing; $\sigma = \operatorname{rms} = 2$. LayerNorm and RMSNorm only differ when the mean isn’t 0.

nn.LayerNorm(d)      # (x - mean) / sqrt(var + eps) * weight + bias
nn.RMSNorm(d)        # x / sqrt(mean(x**2) + eps) * weight

Both normalize over the last axis: one token at a time, never across the batch.

• LayerNorm: subtract the mean, divide by the standard deviation, then $\odot\,\gamma + \beta$.
• RMSNorm: divide by the root mean square, then $\odot\,\gamma$.
• Per token, over its $d$ features.
• Geometry: remove the mean direction, land on a sphere of radius $\sqrt d$, rescale each axis.