Pre-Norm vs Post-Norm: Why It Matters
# Post-norm (Vaswani 2017): norm after the add
def post_norm(x, attn, ffn, ln1, ln2):
x = ln1(x + attn(x))
x = ln2(x + ffn(x))
return x
# Pre-norm (modern default): norm before the sublayer
def pre_norm(x, attn, ffn, ln1, ln2):
x = x + attn(ln1(x))
x = x + ffn(ln2(x))
return x
Why it matters: In pre-norm the residual stream is never normalized, so every block has a clean identity gradient path back to the input — deep stacks (≥12 layers) train without learning-rate warmup and tolerate worse init. Post-norm, by contrast, needs careful warmup schedules to avoid divergence (Xiong et al. 2020, On Layer Normalization in the Transformer Architecture). Every modern LLM is pre-norm for exactly this reason.