RMSNorm Scaling

Worked Example

Consider calculating RMSNorm and LayerNorm for a small vector. Let the input vector be x = [1, 2, 3, 4]. We use gain = 1 and epsilon = 1e-6.

1. Calculate Mean Square: The mean of the squares of the elements is (1^2 + 2^2 + 3^2 + 4^2) / 4 = (1 + 4 + 9 + 16) / 4 = 30 / 4 = 7.5.
2. Calculate RMS Denominator: sqrt(7.5 + 1e-6) ≈ 2.7386129701.
3. Calculate RMSNorm Output (Coordinate 0): For the first element x_0 = 1, the output is 1 * 1 / sqrt(7.5 + 1e-6) ≈ 0.3651483473268884.
4. Calculate Mean: The mean of the elements is (1 + 2 + 3 + 4) / 4 = 2.5.
5. Calculate Variance: The population variance is ((1-2.5)^2 + (2-2.5)^2 + (3-2.5)^2 + (4-2.5)^2) / 4 = 1.25.
6. Calculate LayerNorm Output (Coordinate 0): The output is 1 * (1 - 2.5) / sqrt(1.25 + 1e-6) ≈ -1.3416402498438813.

Experiment

Use the RMSNorm Scaling module to observe how RMSNorm and LayerNorm behave under different input vectors and parameter settings.

1. Set the Vector control to Mixed, Epsilon to 1e-06, Gain to 1, and Coordinate to 0. Observe the calculated outputs match the worked example.
2. Change the Vector control to Offset (which adds 10 to each element of the Mixed vector). Notice how the LayerNorm output remains identical because it recenters the data, while the RMSNorm output changes significantly.
3. Change the Gain control to 2. Observe that all outputs scale proportionally.
4. Change the Vector control to Zero and Epsilon to 0. Note that both normalization techniques result in undefined behavior (division by zero).
5. Change the Vector control to Constant and keep Epsilon at 0. Observe that RMSNorm is well-defined and outputs 1, while LayerNorm remains undefined because the variance is zero.

Limitations

RMSNorm assumes that the mean of the activations is naturally close to zero and does not require explicit centering. If the input distribution shifts significantly away from zero, RMSNorm might perform differently than LayerNorm. Additionally, the inclusion of a non-zero epsilon term slightly breaks exact scale invariance, which becomes noticeable for very small input values. Finally, when the input vector is identically zero and epsilon is zero, the operation is mathematically undefined and must be handled carefully in implementations.

References

The following primary source discusses the detailed mechanics and motivations for Root Mean Square Normalization:

• Zhang, B., & Sennrich, R. (2019). Root Mean Square Layer Normalization. Explains the motivation behind RMSNorm and its comparison to Layer Normalization.