Pre-norm vs post-norm architecture, in depth: where the norm sits, how it changes training, and how modern LLMs place it

The two placements in one picture

Write a sublayer (attention or MLP) as F and the normalization as N. Post-norm computes the next residual stream as N(x + F(x)). Pre-norm computes it as x + F(N(x)). The difference is whether N sits on the skip path, which every later layer's gradient must cross, or only on the branch that feeds F.

Two consequences follow directly from the picture. In post-norm, the output of every block is normalized, so the residual stream has a controlled scale everywhere. In pre-norm, nothing ever normalizes the stream itself: each block adds its contribution to a running sum whose magnitude tends to grow with depth. That is why pre-norm models end with one extra norm, applied after the last block and before the unembedding. Forget it and the logits inherit whatever scale the stream has drifted to.

Both blocks in code

The minimal PyTorch versions differ by two lines per sublayer. The attention and MLP modules are placeholders for whatever you use; the normalization can be LayerNorm or RMSNorm in either layout.

import torch
import torch.nn as nn

class PostNormBlock(nn.Module):
    def __init__(self, d, attn, mlp):
        super().__init__()
        self.attn, self.mlp = attn, mlp
        self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d)

    def forward(self, x):
        x = self.n1(x + self.attn(x))      # norm on the residual path
        x = self.n2(x + self.mlp(x))
        return x

class PreNormBlock(nn.Module):
    def __init__(self, d, attn, mlp):
        super().__init__()
        self.attn, self.mlp = attn, mlp
        self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d)

    def forward(self, x):
        x = x + self.attn(self.n1(x))      # norm inside the branch
        x = x + self.mlp(self.n2(x))
        return x

class PreNormLM(nn.Module):
    def __init__(self, blocks, d, vocab):
        super().__init__()
        self.blocks = nn.ModuleList(blocks)
        self.final_norm = nn.LayerNorm(d)  # required for pre-norm
        self.head = nn.Linear(d, vocab, bias=False)

    def forward(self, h):
        for b in self.blocks:
            h = b(h)
        return self.head(self.final_norm(h))

A post-norm model does not need the final norm because the last block already ends in one. Embeddings are a separate question: original post-norm models such as BERT normalize the embedding sum before the first block, while pre-norm models feed raw embeddings and let the first block's branch norm handle scale.

What the evidence says about training

Two 2020 papers frame most of what practitioners believe. Xiong et al., On Layer Normalization in the Transformer Architecture, analysed gradients at initialization and found that in post-norm the gradients of parameters near the output are large, which is why post-norm needs a learning-rate warmup to avoid early divergence. With pre-norm the gradients were well behaved, and they trained BERT-style and translation models without warmup and with faster convergence. The warmup article covers schedules in detail; the point here is that warmup in post-norm is structural, not a tuning nicety.

Liu et al., Understanding the Difficulty of Training Transformers, added the other half. They argued that the main problem with post-norm is amplification: each block's output depends heavily on its residual branch, so small parameter updates are amplified through depth and destabilize training. Pre-norm reduces that dependence and is therefore more stable, but the authors also reported that when post-norm does train successfully on their translation tasks, it often reaches better final quality than pre-norm. Their remedy, Admin, initializes the skip path so post-norm starts in a stable regime.

The pre-norm cost is easy to see in the residual stream. Because the stream is a sum that is never renormalized, its norm grows with depth, and a late block's branch output, which comes from a normalized input, becomes a smaller and smaller fraction of the stream. Later layers then change the representation less, which several analyses describe as deep pre-norm layers being underused. That is the trade in one line: pre-norm buys trainability with some loss of per-layer effectiveness at depth.

Worked example: measure it on your own model

Rather than trust the folklore, record two curves for any model you train: the norm of the residual stream after each block, and the gradient norm of each block's parameters. The hooks below work for either layout and cost one reduction per block.

def attach_diagnostics(model):
    stats = {"stream": {}, "grad": {}}
    for i, block in enumerate(model.blocks):
        def fwd_hook(mod, inp, out, i=i):
            # mean L2 norm of the residual stream per token, after block i
            stats["stream"][i] = out.detach().float().norm(dim=-1).mean().item()
        block.register_forward_hook(fwd_hook)
    return stats

def grad_norms(model, stats):
    for i, block in enumerate(model.blocks):
        sq = sum(p.grad.float().pow(2).sum() for p in block.parameters()
                 if p.grad is not None)
        stats["grad"][i] = float(sq) ** 0.5

# in the training loop, after loss.backward():
#   grad_norms(model, stats); log stats every N steps

On a small pre-norm model, expect the stream norm to rise roughly steadily from the first block to the last and the per-block gradient norms to be fairly flat across depth. On a post-norm model the stream norm is pinned near the scale set by the norm's gain, and early in training the gradient norms of blocks near the output are much larger than those near the input, which is the pattern Xiong et al. predicted. If a post-norm run diverges in the first few hundred steps, this plot usually shows the top blocks' gradients exploding first.

Log both curves every few hundred steps through training, not only at initialization. A pre-norm stream norm that keeps climbing late in training, or a sudden spike in one block's gradient norm, is an early warning of the loss spikes that the variants below were designed to suppress.

The variant family

Since 2020, most new designs have kept pre-norm's clean skip path and added normalization somewhere else to control the stream or the branch output. The table lists the placements you will meet in papers and model configs.

DeepNorm is the one post-norm descendant worth knowing in detail. It multiplies the skip path by a constant alpha before the norm and scales down the initialization of selected weights by beta. For a decoder-only model with N layers the DeepNet paper uses alpha = (2N)^(1/4) and beta = (8N)^(-1/4); encoder-decoder models have different constants. The effect is to keep each update's contribution bounded as depth grows, which let the authors train a 1,000-layer model.

The OLMo 2 report describes moving the norm from the input of attention and MLP to their outputs, h = x + RMSNorm(Attention(x)), alongside QK-norm, and reports that the combination reduced the growth and spikiness of the gradient norm. Gemma 2 uses norms on both the input and the output of each sublayer. Both keep the skip path free of normalization, which is the property that matters most for trainability.

Choosing a placement for a new model

1. Default to pre-norm with RMSNorm and a final norm. It is the best-tested layout, every serving stack supports it, and it trains without heroic warmup.
2. If you see loss spikes at scale, add branch-output control before changing anything else. QK-norm, a post-branch norm (sandwich or the OLMo 2 output norm) or both are cheap, and published reports attribute stability gains to them.
3. Consider post-norm only for shallow models or with a stabilizer. For encoders of a dozen layers with a proper warmup, post-norm is fine and may score slightly higher. For anything deep, use DeepNorm or Admin-style initialization rather than plain post-norm.
4. Decide before pretraining, not after. The placement changes what every weight learned; it is not a hyperparameter you can revisit at fine-tuning time.
5. Pair placement with initialization and depth scaling. Pre-norm models often scale down the output projections of each branch by a depth-dependent factor; see depth scaling for why the two choices interact.

Inference and kernels

Placement also shapes the kernels you can fuse. In a pre-norm block, the residual add that ends one sublayer is immediately followed by the norm that starts the next, so serving engines commonly fuse them into a single add-plus-RMSNorm kernel that reads the stream once, writes the updated stream and the normalized branch input together, and saves a full round trip to memory per sublayer. Post-norm has the same add-then-norm adjacency in the opposite order and can be fused too. Hybrids with a norm on the branch output add one more normalization per sublayer, which costs a little bandwidth but fuses into the same add.

Numerics differ as well. Because the pre-norm stream grows with depth, many implementations keep the residual stream in FP32 or BF16 even when the matmuls run in lower precision, and quantization methods treat the stream and the normalized branch inputs differently. If you quantize activations, the inputs to attention and MLP are the normalized branch inputs in pre-norm, and the norm's gain vector is where scaling tricks such as SmoothQuant fold their factors.

Working with pretrained checkpoints

• You cannot switch placement on a trained model. Moving the norm changes the function the weights compute. Converting a post-norm BERT into pre-norm without retraining produces garbage, not a slightly different model.
• Read the config, not the paper. Families evolve between versions; check where each norm sits in the reference modelling code before you write a converter or a custom kernel.
• Do not drop the final norm in conversions. A pre-norm checkpoint loaded into code that skips the final norm runs, produces plausible-looking logits, and scores badly. Compare outputs layer by layer against the reference on a fixed input.
• Fine-tuning post-norm models needs warmup too. The structural reason for warmup does not disappear after pretraining; a short warmup is cheap insurance.

Failure modes

• Early divergence in post-norm. Loss becomes NaN within the first few hundred steps. Add or lengthen warmup, lower the peak rate, or switch to DeepNorm-style scaling.
• Late loss spikes in deep pre-norm. The stream norm and attention logits grow until a step blows up. Add QK-norm or a branch-output norm, clip gradients, and check the learning rate for the largest layers.
• Underused deep layers. In very deep pre-norm stacks, later blocks change the representation little; check by measuring the ratio of each block's branch output norm to the stream norm.
• Mismatched epsilon or gain. Reimplementations that change the norm's epsilon, or load a gain into the wrong layer, show up as small but stubborn quality drops.

Trade-offs at a glance

For the step-by-step algebra behind these rows, the residual connections math article shows how the identity path keeps gradients alive and why anything placed on it matters.

What to do next

1. Open the modelling code for a model you use and label where each norm sits: input of attention, input of MLP, branch outputs and the final norm.
2. Implement the two blocks above in a small model of 12 to 24 layers and train both with and without warmup; record whether post-norm diverges without it.
3. Attach the diagnostics hooks and plot stream norm and per-block gradient norm across depth at initialization and halfway through training.
4. Add QK-norm and a branch-output norm to the pre-norm model, rerun, and compare loss-spike frequency and the stream-norm curve.
5. For any checkpoint conversion, write a layer-by-layer equivalence test against the reference implementation before using it.
6. Record the chosen placement, norm type, epsilon and initialization scaling in your model card so later fine-tuners and kernel authors do not have to guess.