RMSNorm

The problem

RMSNorm (root mean square layer normalization) is a simpler LayerNorm (layer normalization). LayerNorm shifts each vector to mean 0 and rescales it; RMSNorm only rescales it, so that its typical size is 1, then applies a learned gain $\gamma$ of $D$ numbers. There is no bias. LLaMA (Large Language Model Meta AI) applies it before every attention and MLP (multilayer perceptron) block.

How it works

The typical size of a vector is measured by its RMS (root mean square): square every entry, take the mean, take the square root.

\mathrm{RMS}(x) = \sqrt{\frac{1}{D}\sum_j x_j^2 + \epsilon}

y = \gamma \odot \frac{x}{\mathrm{RMS}(x)}

where:

$x$ is one token's vector of $D$ numbers, and $x_j$ its entry $j$ ;
$\epsilon$ (epsilon) is eps, which avoids dividing by zero for an all-zero vector;
$\gamma$ (gamma) is the learned gain, and $\odot$ is elementwise multiplication.

Example. $x = (1, 2, 3, 4)$ , $\gamma = (1, 1, 1, 1)$ .

Squares: $(1, 4, 9, 16)$ . Mean: $30 / 4 = 7.5$ . Square root: $\mathrm{RMS}(x) \approx 2.74$ .

$y = x / 2.74 \approx (0.37, 0.73, 1.10, 1.46)$ .

LayerNorm on the same vector gives $(-1.34, -0.45, 0.45, 1.34)$ : it removes the mean first, so half of the entries become negative. RMSNorm keeps them all positive.

What to implement

rmsnorm(x, gamma, eps=1e-6): x has shape $(\dots, D)$ , with any number of leading axes, and is normalized over the last axis. Return $y$ , same shape as x.

Requirements

Do not subtract the mean: that is the difference from LayerNorm.
eps goes inside the square root, as in PyTorch's torch.nn.RMSNorm.
No Python loops.

Hint

Take the mean of x**2 over the last axis with keepdims=True, so that the result has shape $(\dots, 1)$ and divides every entry of its own vector.