RMSNorm › Key Parameters

When implementing or using RMSNorm, several parameters control its behavior. Understanding each one helps you configure it correctly for your use case and debug issues that arise.

The dim parameter specifies the feature dimension to normalize over. This should match the hidden dimension of your model. For transformer models, this is typically the embedding dimension, which ranges from 768 for smaller models like BERT-base up to 8192 or more for large models. Normalizing over a larger dimension gives a more accurate estimate of the RMS because you're averaging over more terms, which reduces the variance of the RMS estimate. This is one reason why normalization works better for large hidden dimensions than small ones.

The eps parameter (default: $10^{-6}$ to $10^{-5}$) is a small constant added inside the square root for numerical stability. It prevents division by zero when all input values are near zero. LLaMA uses $10^{-5}$, while some implementations use $10^{-6}$. The choice involves a trade-off: smaller epsilon gives more precise normalization for typical inputs but increases the risk of numerical issues in edge cases. For most applications, either value works well, and the difference in normalized output is negligible for typical activation magnitudes.

The weight ($\gamma$) parameter is the learnable scale vector of dimension dim, initialized to ones. This initialization ensures that at the start of training, RMSNorm is approximately an identity operation: it normalizes the RMS to 1 and then multiplies by 1, leaving the scale unchanged. The network then adjusts $\gamma$ during training to set each feature's scale to the value that works best for the downstream computation. Unlike LayerNorm, RMSNorm has no bias ($\beta$) parameter.

The computation dtype deserves special attention for mixed-precision training. When using float16 or bfloat16 for model weights and activations, the normalization computation should be performed in float32. Float16 has limited range (roughly $\pm 65504$) and limited precision (about 3 decimal digits). Small RMS values close to $\epsilon$ can lead to significant relative errors in float16 arithmetic. Computing the normalization in float32 and then casting back to the working dtype avoids these precision issues while adding only minimal overhead.

The placement convention in modern architectures applies RMSNorm before each attention and feed-forward sub-layer (pre-norm), not after (post-norm). Pre-norm has been empirically validated to provide better gradient flow and training stability compared to post-norm, particularly for deep models with many layers.