RMS Norm

View from layer norm

Given a token's embedding or feature vector $\mathbf{x}\equiv [x_1, \cdots, x_d]^T$, the layer normalization first computes the mean and std among its components \begin{equation}\mu = \frac{1}{d}\sum_{i=1}^dx_i\,\,\,\,\sigma =\sqrt{\frac{1}{d}\sum_{i=1}^d (x_i-\mu)^2}\,, \end{equation} and then normalize \begin{equation}\text{LayerNorm}(\mathbf{x})=\frac{\mathbf{x}-\mu}{\sigma}\odot \boldsymbol{\gamma}  +\boldsymbol{\beta}\,,\tag{1}\end{equation} where $\boldsymbol{\gamma}$ and $\boldsymbol{\beta}$ are learnable parameters.

A common view of RMS norm is to think it as a special case of the layer norm (1) in which we ignore all terms related to bias, i.e. $\mu=0$ and $\boldsymbol{\beta}=0$: \begin{equation}\text{RMSNorm}(\mathbf{x})=\frac{\mathbf{x}}{\sqrt{\displaystyle\frac{1}{d}\sum_{i=1}^d x_i^2}}\odot \boldsymbol{\gamma}\equiv\frac{\mathbf{x}}{\sqrt{\displaystyle\frac{1}{d}\left|\left|\mathbf{x}\right|\right|^2_2}}\odot \boldsymbol{\gamma}\,.\tag{2}\end{equation} In this view, roughly speaking, the RMS norm is essentially normalizes each component's magnitude such that \begin{equation}x_i \sim \mathcal{O}(1)\,,\,\,\,  \left|\left|\mathbf{x}\right|\right|_2 \sim \mathcal{O}(\sqrt{d})\,.\tag{3} \end{equation}

View from RoPE

To me, a more natural view of RMS norm is from RoPE. Modern LLM uses RoPE to encode positions as rotations applied to token features $\mathbf{x}$. When considering rotations, it is natural to think about unit vectors originating at the origin: \begin{equation}\text{UnitNorm}(\mathbf{x}) \equiv \frac{\mathbf{x}}{\left|\left|\mathbf{x}\right|\right|_2}\odot \boldsymbol{\gamma}\,,\tag{4}\end{equation} which then relates to the RMS norm (2) as \begin{equation}\text{RMSNorm}(\mathbf{x})=\sqrt{d}\,\text{UnitNorm}(\mathbf{x})\,.\tag{5}\end{equation} Both RMS norm and Unit norm are compatible with rotations as we only requires vector having fixed length when considering rotations, not necessarily unit length.

Remarks:

• The view from RoPE provides a natural justification of the RMS norm (2). 
• Unlike RMS norm (3), the unit norm essentially normalize the vector length rather each component magnitude: \begin{equation}x_i \sim \mathcal{O}(1/\sqrt{d})\,,\,\,\,  \left|\left|\mathbf{x}\right|\right|_2 \sim \mathcal{O}(1)\,.\tag{6} \end{equation}

Normalization Consistency

The RMS norm (2) and unit norm (4) are the same up to a constant $\sqrt{d}$. We can use either of them, but need to keep normalization constant consistent.

For example, when using RMS norm (2) on $\mathbf{q}$, $\mathbf{k}$ vectors, their scaled dot attention is of the form \begin{equation}\frac{\mathbf{q}\cdot\mathbf{k}}{\sqrt{d}}\,,\tag{7}\end{equation} in which the denominator constant ${\sqrt{d}}$ is to make the result (7) of the unit variance.

In another example, consider the forward linear pass $\mathbf{y}=\mathbf{W}\cdot\mathbf{x}$, if we apply RMS norm (2) to $\mathbf{x}$ and also requires the output $\mathbf{y}$ preserving the component magnitude $y_i \sim \mathcal{O}(1)$, we then need to initialize $\mathbf{W}\sim \mathcal{O}(1/\sqrt{d})$, which is exactly the Xavier initialization when $d_{\text{out}}=d_{\text{in}}=d$.

In contrast, if we use the unit norm (6) instead, then we need to define scaled dot attention as \begin{equation}\sqrt{d}\mathbf{q}\cdot\mathbf{k}\end{equation} for the unit variance.