Layer normalization

Layer normalization

Let $\mathbf{x} \in \mathbb{R}^{n}$ be an activation vector.

We first compute the mean $\mu_{\mathbf{x}}$ and variance $\sigma_{\mathbf{x}}^2$ of the activation vector as follows:

$$\begin{aligned} \mu_{\mathbf{x}} &= \frac{1}{n} \sum_{i=1}^{n} x_i, \quad \sigma_{\mathbf{x}}^2 = \frac{1}{n} \sum_{i=1}^{n} (x_i - \mu_{\mathbf{x}})^2. \end{aligned}$$

The $\text{LayerNorm}$ operation[1]Layer normalization [link]
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. arXiv preprint. 2016. is then defined as follows:

$$\begin{aligned} \text{LayerNorm}(\mathbf{x}) &= \mathbf{g} \odot \frac{\mathbf{x} - \mu_{\mathbf{x}}}{\sqrt{\sigma_{\mathbf{x}}^2 + \epsilon}} + \mathbf{b} \end{aligned}$$

where $\mathbf{g}, \mathbf{b} \in \mathbb{R}^{n}$ are learnable parameters, $\epsilon$ is a small constant, and $\odot$ denotes the element-wise multiplication.

Geometric interpretation

We can break down the $\text{LayerNorm}$ operation into three distinct steps.

1. Mean-centering

We first subtract the mean $\mu_{\mathbf{x}}$ from the activation vector $\mathbf{x}$:

$$\begin{aligned} \mathbf{x}^{(1)} = \mathbf{x} - \mu_{\mathbf{x}}. \end{aligned}$$

This can equivalently be thought of as projecting the activation vector onto the all-ones direction $\hat{\mathbf{1}}$, and subtracting this component away:

$$\begin{aligned} \mathbf{x}^{(1)} &= \mathbf{x} - \mu_{\mathbf{x}}\\ &= \mathbf{x} - \mu_{\mathbf{x}} \mathbf{1}\\ &= \mathbf{x} - \left(\frac{1}{n} \sum_{i=1}^{n} x_i \right) \mathbf{1}\\ &= \mathbf{x} - \left( \frac{1}{n} \mathbf{x} \cdot \mathbf{1} \right) \mathbf{1}\\ &= \mathbf{x} - \left( \mathbf{x} \cdot \frac{\mathbf{1}}{\sqrt{n}} \right) \frac{\mathbf{1}}{\sqrt{n}}\\ &= \mathbf{x} - \underbrace{\left( \mathbf{x} \cdot \hat{\mathbf{1}} \right) \hat{\mathbf{1}}}_{\text{proj}_{\hat{\mathbf{1}}}(\mathbf{x})}. \end{aligned}$$

Thus, mean-centering can be thought of as projecting the $n$-dimensional activation onto an $(n-1)$-dimensional subspace, namely the orthogonal complement of $\hat{\mathbf{1}}$.

2. Variance normalization

After mean-centering, we normalize the variance of the activation vector:

$$\begin{aligned} \mathbf{x}^{(2)} = \frac{\mathbf{x}^{(1)}}{\sqrt{\sigma_{\mathbf{x}}^2 + \epsilon}}. \end{aligned}$$

Note that $\epsilon$ is a small constant ($\approx 10^{-5}$) to prevent division by zero, and to improve numerical stability.

Variance normalization can equivalently be thought of as projecting the activation onto the $\sqrt{n}$-radius sphere.

To see this, first note that the variance of the original activation vector $\mathbf{x}$ is equal to the variance of the mean-centered activation vector $\mathbf{x}^{(1)}$, since variance is invariant under translation:

$$\begin{aligned} \sigma_{\mathbf{x}}^2 = \sigma_{\mathbf{x}^{(1)}}^2. \end{aligned}$$

Next, notice that the variance of a mean-centered vector can be expressed in terms of its squared norm:

$$\begin{aligned} \sigma_{\mathbf{x}^{(1)}}^2 &= \frac{1}{n} \sum_{i=1}^{n} \left(x_i^{(1)} - \mu_{\mathbf{x}^{(1)}}\right)^2\\ &= \frac{1}{n} \sum_{i=1}^{n} \left(x_i^{(1)}\right)^2\\ &= \frac{1}{n} \| \mathbf{x}^{(1)} \|^2. \end{aligned}$$

Putting things together:

$$\begin{aligned} \mathbf{x}^{(2)} &= \frac{\mathbf{x}^{(1)}}{\sqrt{\sigma_{\mathbf{x}}^2 + \epsilon}}\\ &\approx \frac{\mathbf{x}^{(1)}}{\sqrt{\sigma_{\mathbf{x}}^2}}\\ &= \frac{\mathbf{x}^{(1)}}{\sqrt{\sigma_{\mathbf{x}^{(1)}}^2}}\\ &= \frac{\mathbf{x}^{(1)}}{\sqrt{\frac{1}{n} \| \mathbf{x}^{(1)} \|^2}}\\ &= \sqrt{n} \frac{\mathbf{x}^{(1)}}{\| \mathbf{x}^{(1)} \|}. \end{aligned}$$

3. Affine transformation

Finally, we apply an affine transformation, scaling by a learned gain $\mathbf{g} \in \mathbb{R}^{n}$, and shifting by a learned bias vector $\mathbf{b} \in \mathbb{R}^{n}$:

$$\begin{aligned} \mathbf{x}^{(3)} = \mathbf{g} \odot \mathbf{x}^{(2)} + \mathbf{b}. \end{aligned}$$

Root mean squared layer normalization

Root Mean Square Layer Normalization[2]Root mean square layer normalization [link]
Biao Zhang and Rico Sennrich. Advances in Neural Information Processing Systems. 2019. is a simplification of Layer Normalization that skips mean-centering, and normalizes by the root mean square of the activations.

We first compute the root mean square of the activation vector as follows:

$$\begin{aligned} \text{RMS}(\mathbf{x}) &= \sqrt{\frac{1}{n} \sum\nolimits_{i=1}^{n} x_i^2}. \end{aligned}$$

The $\text{RMSNorm}$ operation is defined as follows:

$$\begin{aligned} \text{RMSNorm}(\mathbf{x}) &= \mathbf{g} \odot \frac{\mathbf{x}}{\text{RMS}(\mathbf{x})} + \mathbf{b}. \end{aligned}$$

References

References cited in the text are listed first, in order of citation; additional references follow, ordered alphabetically.

1. Layer normalization [link]
   Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. arXiv preprint. 2016.
2. Root mean square layer normalization [link]
   Biao Zhang and Rico Sennrich. Advances in Neural Information Processing Systems. 2019.
3. What is LayerNorm folding? [link]
   Neel Nanda. 2022.
4. Geometry and dynamics of LayerNorm [link]
   Paul M. Riechers. arXiv preprint. 2024.