Layer Normalization & RMSNorm

Why It Matters for Modern Models

LayerNorm and RMSNorm normalize each activation vector with its own statistics rather than statistics over the batch, so they do exactly the same computation at training and test time (Ba, Kiros & Hinton, 2016).
The original Transformer applied LayerNorm after each residual addition (post-norm); GPT-3 and LLaMA instead normalize the input of each sublayer (pre-norm), and LLaMA uses RMSNorm (Brown et al., 2020; Touvron et al., 2023).
Where the normalization sits changes how training behaves: Xiong et al. (2020) showed that post-norm Transformers have large gradients near the output at initialization and need a learning-rate warmup, while pre-norm Transformers can train without one.
RMSNorm skips the mean; Zhang & Sennrich (2019) found comparable quality to LayerNorm with running time reduced by 7% to 64% across their models, and the actual saving depends on the implementation.

Questions to Think About

What do the learned γ and β restore after normalization, and why can they not recover the mean and scale that each input vector had before it was normalized?
How does moving the normalization from after the residual addition (post-norm) to the start of each sublayer (pre-norm) change the path that gradients take through a deep stack of layers?

Interactive Demo

The interactive demo for this concept is in the newer notebook.

Open the notebook demo

Core Math

If you prefer intuition first, try the demo in the newer notebook, then come back for the math.

Key Equation

\begin{aligned}\mu(x)&=\frac{1}{d}\sum_{i=1}^d x_i,\\ \sigma_\epsilon(x)&=\sqrt{\frac{1}{d}\sum_{i=1}^d\bigl(x_i-\mu(x)\bigr)^2+\epsilon},\\ \mathrm{LN}_\epsilon(x)&=\gamma\odot\frac{x-\mu(x)}{\sigma_\epsilon(x)}+\beta,\\ \mathrm{RMS}_\epsilon(x)&=\sqrt{\frac{1}{d}\sum_{i=1}^d x_i^2+\epsilon},\\ \mathrm{RMSNorm}_\epsilon(x)&=\gamma\odot\frac{x}{\mathrm{RMS}_\epsilon(x)}\end{aligned}

For one activation vector $x\in\mathbb R^d$ , with integer $d\ge 1$ , let $\gamma,\beta\in\mathbb R^d$ and set $\epsilon=10^{-5}$ for this page. Write $x_i$ for coordinate $i$ of $x$ , and let $\odot$ denote elementwise multiplication.

LayerNorm centers the $d$ features and then applies an elementwise affine map:

\begin{aligned} \mu(x) &= \frac{1}{d}\sum_{i=1}^d x_i,\\ \sigma_\epsilon(x) &= \sqrt{\frac{1}{d}\sum_{i=1}^d\bigl(x_i-\mu(x)\bigr)^2+\epsilon},\\ \mathrm{LN}_\epsilon(x) &= \gamma\odot\frac{x-\mu(x)}{\sigma_\epsilon(x)}+\beta. \end{aligned}

RMSNorm omits mean centering and uses the stabilized root mean square:

\begin{aligned} \mathrm{RMS}_\epsilon(x) &= \sqrt{\frac{1}{d}\sum_{i=1}^d x_i^2+\epsilon},\\ \mathrm{RMSNorm}_\epsilon(x) &= \gamma\odot\frac{x}{\mathrm{RMS}_\epsilon(x)}. \end{aligned}

The demo in the linked notebook uses one scalar $\gamma$ and one scalar $\beta$ for all $d$ features; the equations above allow a separate value for each feature. RMSNorm has no $\beta$ term under this page's convention.

LayerNorm centres and rescales each vector (Ba, Kiros & Hinton, 2016); RMSNorm only rescales it (Zhang & Sennrich, 2019). This page and the linked notebook demo place $\epsilon=10^{-5}$ inside each square root as an explicit numerical convention; library defaults, normalized axes, and affine variants may differ.

Key Papers