RMSNorm

Implement RMSNorm, the normalisation layer used in LLaMA and most recent large language models.

LayerNorm subtracts the mean of each feature vector, divides by its standard deviation, then applies a scale and a shift. RMSNorm drops the mean and the shift. It only divides by the root mean square of the vector and applies a learned scale. It is cheaper and works as well in practice.

For each vector x along the last axis:

rms(x) = sqrt(mean(x ** 2) + eps)
out    = x / rms(x) * weight

Function signature

def rms_norm(x, weight, eps=1e-6):

Example

Input

rms_norm(np.array([[1.0, 2.0, 3.0, 4.0]]), weight=np.ones(4))

Output

array([[0.36514835, 0.73029669, 1.09544504, 1.46059339]])

Reasoning

The mean of the squares is (1 + 4 + 9 + 16) / 4 = 7.5, and its square root is 2.7386. Each value is divided by that. The output has a root mean square of 1. Unlike LayerNorm, the mean of the output is not 0.

Constraints and notes

• x has shape (..., d) and weight has shape (d,).
• Normalise over the last axis only.
• Add eps inside the square root.