Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

LLaMA Components › Pre-Norm with RMSNorm

Michael BrenndoerferUpdated August 2, 20256 min read

Language AI HandbookAll chapters

Part 29 of 67Modern Decoder Models

Chapter 208LLaMA Components › Pre-Norm with RMSNorm

Reading section 2 of 12 · Pre-Norm with RMSNorm · 6 minSection 2 of 12 · 6 min

Section 2 of 12 · 6 min read

LLaMA Components › Pre-Norm with RMSNormOne section · Sections in this chapter2 / 12 ↓

Sections in this chapter12

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

LLaMA places normalization before each sublayer rather than after. This pre-norm configuration, combined with the simpler RMSNorm instead of LayerNorm, defines the normalization strategy throughout the architecture.

Why Pre-Norm?Link Copied

The original transformer used post-norm: normalize after adding the residual connection. This creates a sequence where the residual and sublayer output are summed first, then normalized together. While mathematically elegant, post-norm creates training instabilities in deep networks. Gradients must pass through the normalization layer before reaching the residual path, which can attenuate or distort the gradient signal.

Pre-norm flips this order: normalize the input before the sublayer, then add the unnormalized input via the residual connection. The key benefit is that gradients can flow directly through the residual path without encountering normalization. For a stack of $L$ blocks, this means gradients from the output reach early layers with minimal degradation.

The pre-norm equations for a single transformer block are:

\begin{aligned} \mathbf{h} &= \mathbf{x} + \text{Attention}(\text{Norm}(\mathbf{x})) \\ \mathbf{y} &= \mathbf{h} + \text{FFN}(\text{Norm}(\mathbf{h})) \end{aligned}

where:

$\mathbf{x} \in \mathbb{R}^{n \times d}$ : input to the block with $n$ tokens and dimension $d$
$\text{Norm}(\cdot)$ : normalization function (RMSNorm in LLaMA)
$\text{Attention}(\cdot)$ : multi-head self-attention
$\text{FFN}(\cdot)$ : feed-forward network with SwiGLU
$\mathbf{h}$ : intermediate representation after the attention sublayer
$\mathbf{y}$ : block output

Notice that $\mathbf{x}$ is added directly to the attention output without normalization, and similarly $\mathbf{h}$ is added directly to the FFN output. This creates an unimpeded gradient highway through the model.

Out[4]:

Visualization

Line plot comparing gradient magnitude vs layer depth for pre-norm and post-norm configurations, showing pre-norm maintains more stable gradients in deep networks. — Simulated gradient magnitude through transformer layers. Pre-norm maintains more consistent gradients across depth, while post-norm shows gradual attenuation. The residual connection in pre-norm creates a direct path for gradients to reach early layers.

Why RMSNorm?Link Copied

To understand RMSNorm, let's first consider what LayerNorm does. LayerNorm performs two operations on each token's representation: centering (subtracting the mean to shift values toward zero) and scaling (dividing by the standard deviation to make the spread consistent). Both operations seem intuitively important for stabilizing neural network training.

RMSNorm challenges this intuition by asking: do we really need both? The answer, surprisingly, is no. Empirical studies found that centering contributes little to training stability, while scaling does the heavy lifting. RMSNorm removes the centering step entirely, normalizing only by the root mean square:

\text{RMSNorm}(\mathbf{x}) = \boldsymbol{\gamma} \odot \frac{\mathbf{x}}{\sqrt{\frac{1}{d}\sum_{i=1}^{d} x_i^2 + \epsilon}}

where:

$\mathbf{x} \in \mathbb{R}^d$ : input vector (one token's representation)
$\boldsymbol{\gamma} \in \mathbb{R}^d$ : learnable scale parameters
$\epsilon$ : small constant for numerical stability (typically $10^{-6}$ )
$\odot$ : element-wise multiplication

Why does removing centering not hurt performance? The key insight is that the learnable scale parameter $\boldsymbol{\gamma}$ gives the network flexibility to adjust the distribution as needed. If centering were truly needed, the model could learn to approximate it through subsequent layers. In practice, models converge just as well without explicit centering.

The practical benefit is speed. RMSNorm saves one reduction operation (computing the mean) and one element-wise operation (subtracting the mean) compared to LayerNorm. These savings seem minor for a single forward pass, but normalization runs at every layer, for every token, at every training step. Over billions of operations, the efficiency gains compound significantly.

Let's implement both normalizations and compare their behavior:

In[5]:

Code

import numpy as np


class RMSNorm:
    """
    Root Mean Square Layer Normalization as used in LLaMA.

    Normalizes by RMS without mean centering, saving computation
    while maintaining training stability.
    """

    def __init__(self, dim, eps=1e-6):
        self.eps = eps
        self.weight = np.ones(dim)  # Learnable scale

    def __call__(self, x):
        # Compute RMS: sqrt(mean(x^2))
        rms = np.sqrt(np.mean(x**2, axis=-1, keepdims=True) + self.eps)
        # Normalize and apply scale
        return self.weight * (x / rms)


class LayerNorm:
    """Standard LayerNorm for comparison."""

    def __init__(self, dim, eps=1e-6):
        self.eps = eps
        self.weight = np.ones(dim)
        self.bias = np.zeros(dim)

    def __call__(self, x):
        mean = np.mean(x, axis=-1, keepdims=True)
        var = np.var(x, axis=-1, keepdims=True)
        return self.weight * (x - mean) / np.sqrt(var + self.eps) + self.bias

Out[6]:

Visualization

Histogram showing RMSNorm normalized activation distribution centered near zero. — RMSNorm output distribution. RMSNorm produces outputs centered near zero but not exactly, since the mean wasn't explicitly subtracted.

Histogram showing LayerNorm normalized activation distribution precisely centered at zero. — LayerNorm output distribution. LayerNorm produces zero-centered outputs by construction through explicit mean subtraction.

LayerNorm outputs are precisely centered at zero by construction. RMSNorm outputs cluster near zero but aren't exactly centered, since the mean wasn't explicitly subtracted. In practice, this difference rarely matters: the learnable scale parameter $\boldsymbol{\gamma}$ can adjust the distribution as needed, and the model learns to work with either normalization scheme.