Skip to content

← Modern Techniques step 1 of 10

Easy Research

RMSNorm

RMSNorm rescales a vector by its root mean square (RMS). Square its entries, average those squares, add epsilon, then take a square root. Divide the original entries by that value and apply a learned scale gamma.

Unlike LayerNorm, RMSNorm does not subtract the mean. For [1, 1] with gamma [1, 1] and epsilon 0, it returns [1, 1]; centering that row would instead produce zeros. Its output need not have mean zero.

Your task

Implement rmsnorm(x, gamma, eps=1e-6):

  • x is a floating-point tensor of shape (N, D), with positive dimensions.
  • gamma is a floating-point scale tensor of shape (D,).
  • eps is nonnegative. If it is zero, every row has nonzero mean square.
  • Return a tensor of shape (N, D).

For each row, compute denominator = sqrt(mean(x ** 2) + eps) over the D features, then return (x / denominator) * gamma. Preserve the reduced axis so each row's denominator broadcasts over its features. Do not subtract the mean, add a bias, or normalize across the batch.

For x = [[0.001, -0.001]], gamma = [1, 2], and eps = 1e-6, the mean square is 1e-6 and the denominator is sqrt(2e-6). The output is approximately [[0.707107, -1.414214]]. Epsilon has a visible effect at this scale: adding it after the square root gives a different answer. With positive epsilon, an all-zero row returns zeros.

The diagram shows which entries share statistics. This exercise has one feature per column and no spatial axes. With RMSNorm selected, each row shares a denominator; the diagram does not display the resulting values.

Loading visualization…

Reference: Root Mean Square Layer Normalization.