Pre-Norm and Post-Norm

Two lines, differing in where one function call goes.

In eq. 2 there is an unbroken identity path from the embedding to the output: $x$ passes through additions and nothing else. In eq. 1 every layer’s output goes through a normalisation, so the path from layer 1 to layer 50 traverses fifty non-linear rescalings.

Why post-norm needs warmup

Xiong et al. worked out the gradient at initialisation. For post-norm, the expected gradient norm at the parameters of the final layer is $O(d\sqrt{\ln d})$ and does not shrink with depth — the early steps of training take enormous, badly-conditioned steps and the model diverges. For pre-norm the same quantity carries a $1/\sqrt{L}$ factor, and depth reduces the gradient rather than leaving it unbounded.

Warmup costs more than the steps it consumes. It couples the schedule to the architecture, it interacts with the batch size, and it is one more thing to tune per run. Removing it is most of why pre-norm won.

What pre-norm gives up

The identity path is also the problem. Each block adds its output to a stream that is never rescaled, so the variance of $x_\ell$ grows roughly linearly with depth. Later blocks see an input whose norm is dominated by the accumulated sum, and their own contribution — normalised before $F$, and therefore of bounded size — is a progressively smaller fraction of it.

The observable consequence is that deep pre-norm models behave like shallower ones: the last layers change the representation less than the first, and at matched parameter count a post-norm model that does converge reaches slightly better loss. This is sometimes called representation collapse, and it is the motivation for every variant below.

from torch import Tensor, nn


class PreNormBlock(nn.Module):
    """The modern arrangement. Note the norms are inside the branches."""

    def __init__(self, dim: int, **kw):
        super().__init__()
        self.n1, self.n2 = nn.RMSNorm(dim), nn.RMSNorm(dim)
        self.attn, self.ffn = Attention(dim, **kw), FeedForward(dim)

    def forward(self, x: Tensor) -> Tensor:
        x = x + self.attn(self.n1(x))
        return x + self.ffn(self.n2(x))


class Stack(nn.Module):
    def __init__(self, dim: int, layers: int, **kw):
        super().__init__()
        self.blocks = nn.ModuleList(PreNormBlock(dim, **kw) for _ in range(layers))
        # Required: the residual stream itself is never normalised on its path.
        self.final = nn.RMSNorm(dim)

    def forward(self, x: Tensor) -> Tensor:
        for block in self.blocks:
            x = block(x)
        return self.final(x)

Sandwich norm normalises both entering and leaving the branch — $x + \mathrm{LN}_2(F(\mathrm{LN}_1(x)))$ — which bounds each block’s contribution without touching the identity path. Gemma 2 ships this, and it costs one extra norm per branch.

Residual scaling multiplies each branch output by a constant, commonly $1/\sqrt{2L}$, so the accumulated variance stays $O(1)$ instead of $O(L)$. GPT-2 did a version of this at initialisation rather than as a persistent factor.

DeepNorm goes the other way: keep post-norm, and pick the residual weight and the initialisation scale so that the update to each layer is bounded by construction. It is the only one of the three that recovers post-norm’s quality rather than patching pre-norm’s deficit, and it is the least used, because it requires getting two constants right as a function of depth.