RMSNorm › Worked Example: Step-by-Step Numerical Trace

Before moving to code, let's trace through a complete numerical example of both LayerNorm and RMSNorm on a small vector. This gives you a concrete feel for the computations and reinforces the mathematical relationship we derived.

Suppose we have a small input vector with $d = 4$ elements:

$$\mathbf{x} = [2.0, -1.0, 0.5, 1.5]$$

We'll normalize this vector with both methods and compare the results step by step.

First we compute the mean $\mu$ and the root mean square $\text{RMS}(\mathbf{x})$:

$$\mu = \frac{1}{4}(2.0 + (-1.0) + 0.5 + 1.5) = \frac{3.0}{4} = 0.75$$

For the RMS, we first square each element: $[4.0, 1.0, 0.25, 2.25]$. The mean of these squares is:

$$\frac{1}{4}(4.0 + 1.0 + 0.25 + 2.25) = \frac{7.5}{4} = 1.875$$

So $\text{RMS}(\mathbf{x}) = \sqrt{1.875} \approx 1.369$.

Let's verify the Pythagorean relationship. The variance is:

$$\sigma^2 = \frac{1}{4}\left[(2.0 - 0.75)^2 + (-1.0 - 0.75)^2 + (0.5 - 0.75)^2 + (1.5 - 0.75)^2\right]$$ $$= \frac{1}{4}\left[1.5625 + 3.0625 + 0.0625 + 0.5625\right] = \frac{5.25}{4} = 1.3125$$

So $\sigma = \sqrt{1.3125} \approx 1.146$. Now let's check: $\sigma^2 + \mu^2 = 1.3125 + 0.5625 = 1.875 = \text{RMS}^2$. The identity holds exactly.

LayerNorm subtracts the mean, then divides by the standard deviation. Using $\epsilon = 10^{-6}$ (negligible here):

$$\mathbf{x} - \mu = [2.0 - 0.75, -1.0 - 0.75, 0.5 - 0.75, 1.5 - 0.75] = [1.25, -1.75, -0.25, 0.75]$$

Dividing by $\sigma \approx 1.146$:

$$\frac{\mathbf{x} - \mu}{\sigma} \approx [1.090, -1.527, -0.218, 0.654]$$

With $\gamma = \mathbf{1}$ and $\beta = \mathbf{0}$, the LayerNorm output is $[1.090, -1.527, -0.218, 0.654]$. The output has mean exactly 0 and standard deviation exactly 1 (up to floating point precision).

RMSNorm simply divides by the RMS, with no mean subtraction. Using $\text{RMS} \approx 1.369$:

$$\frac{\mathbf{x}}{\text{RMS}} \approx \left[\frac{2.0}{1.369}, \frac{-1.0}{1.369}, \frac{0.5}{1.369}, \frac{1.5}{1.369}\right] \approx [1.461, -0.730, 0.365, 1.096]$$

With $\gamma = \mathbf{1}$, the RMSNorm output is $[1.461, -0.730, 0.365, 1.096]$.

The two outputs differ because the mean $\mu = 0.75$ is non-trivial relative to the standard deviation $\sigma \approx 1.146$. Their ratio is $\mu / \sigma \approx 0.65$, which is substantial. Using the geometric picture, this corresponds to a triangle with a large vertical leg relative to the horizontal leg, so the hypotenuse (RMS) differs noticeably from the horizontal leg ($\sigma$).

In a real transformer, the mean of activations is typically much smaller relative to the standard deviation. If we had $\mu = 0.05$ and $\sigma = 1.0$, then $\text{RMS} = \sqrt{1.0 + 0.0025} \approx 1.001$, which is indistinguishable from $\sigma = 1.0$ for practical purposes. This is the regime that justifies RMSNorm's approximation in practice.

We can verify that the RMSNorm output has RMS equal to 1 (when $\gamma = \mathbf{1}$):

$$\text{RMS}([1.461, -0.730, 0.365, 1.096]) = \sqrt{\frac{1.461^2 + 0.730^2 + 0.365^2 + 1.096^2}{4}} = \sqrt{\frac{2.134 + 0.533 + 0.133 + 1.201}{4}} = \sqrt{\frac{4.001}{4}} \approx 1.0$$

The output RMS equals 1.0 exactly (up to rounding), confirming that RMSNorm correctly normalizes the magnitude of the input vector to 1.