Home / Fundamentals / Transformer Architecture
Why do modern LLMs use pre-norm and RMSNorm? What are the trade-offs?
30-second answer
Pre-norm (x + f(Norm(x))) keeps an identity residual path from input to output, so gradients flow directly and deep stacks train stably with little warmup; post-norm normalizes after each residual add and is prone to divergence at depth. The cost is that the residual stream grows with depth, so later layers contribute relatively less; scaled initialization and a final norm help. RMSNorm drops LayerNorm's mean-centering and bias (x / RMS(x) · γ), which is cheaper and empirically as good. QK-norm is often added to stop attention logits from growing and causing loss spikes.
Try answering out loud before reading on. The walkthrough below is paced for a 40-minute interview round.
0–5 min Clarify and frame
Framing: "Normalization placement and type are mostly about optimization stability at depth. I'll compare post-norm and pre-norm, LayerNorm and RMSNorm, then cover what goes wrong at scale and the fixes."
5–15 min Core explanation
Post-norm: x_{l+1} = Norm(x_l + F(x_l))
Pre-norm: x_{l+1} = x_l + F(Norm(x_l)) (final Norm before the LM head)
LayerNorm(x) = γ ⊙ (x − mean(x)) / √(var(x) + ε) + β
RMSNorm(x) = γ ⊙ x / √(mean(x²) + ε)
Pre-norm leaves an identity path through the residual stream: x_L = x_0 + Σ F_l(…). The gradient to any layer includes a direct term through that sum, so very deep stacks train stably, usually with modest warmup.
Post-norm puts a normalization on the residual path after every block. At initialization, gradients to early layers can be poorly scaled, which is why the original Transformer needed careful learning-rate warmup and deep post-norm models often diverge.
RMSNorm removes mean-centering and the bias. It's cheaper (one reduction instead of two), and in practice the re-centering adds little; most recent LLMs use it.
15–25 min Deep dive
The pre-norm trade-off. Because blocks only add to the stream, its norm tends to grow with depth (roughly like √L if contributions are similar in size). Each later block's contribution is then relatively smaller, so later layers can be under-used, sometimes described as representation collapse or reduced effective depth. Mitigations:
- Scaled initialization: initialize each block's output projection with std scaled by
1/√(2L)so the stream doesn't blow up early. - DeepNorm and related schemes: post-norm with up-weighted residuals and down-scaled initialization, enabling very deep post-norm models.
- Extra norms: normalize sublayer outputs as well as inputs ("sandwich" norm).
Attention logit growth. In large models, query and key norms can grow during training, making attention entropy collapse toward one-hot and causing loss spikes. QK-norm (RMSNorm on q and k per head before the dot product) bounds the logits. Some models also soft-cap logits (e.g. c · tanh(s / c)).
Output logit drift. Large output logits can destabilize training; an auxiliary z-loss, α · (log Σ_j e^{z_j})², keeps the softmax normalizer near zero.
Where norms sit in a modern block:
h = x + Attn(RMSNorm(x))
y = h + SwiGLU(RMSNorm(h))
...
logits = LMHead(RMSNorm(y_L))
25–32 min Trade-offs
| Choice | Pros | Cons |
|---|---|---|
| Post-norm | Can give slightly better final quality when it trains | Unstable at depth; heavy warmup needs |
| Pre-norm | Stable, simple | Stream growth; later layers contribute less |
| LayerNorm | Re-centers; standard | Slightly more compute; bias terms |
| RMSNorm | Cheaper, works as well | No re-centering (rarely matters) |
| QK-norm | Prevents logit blow-up | Small compute cost; changes attention temperature dynamics |
Norm weights (γ) are typically excluded from weight decay.
32–37 min Production example
A team scaling a model from 1B to 13B parameters saw periodic loss spikes beginning mid-training. Diagnostics showed growing attention logit magnitudes in a few heads. Adding QK-norm and a small z-loss, together with a slightly lower peak learning rate, removed the spikes in a restarted run; training throughput was unchanged within noise.
37–40 min Likely follow-ups
- Why not BatchNorm? Statistics across the batch are ill-suited to variable-length sequences and autoregressive inference, and couple examples together.
- Do you apply weight decay to γ? Usually not; decaying normalization gains toward zero hurts.
- Can you remove normalization entirely? With careful initialization and scaling it's possible, but it's fragile at scale.
What interviewers listen for
- The identity-path argument for pre-norm, and its cost.
- Correct LayerNorm and RMSNorm formulas.
- Practical stability tools: scaled init, QK-norm, z-loss.
Red flags
- "Pre-norm is just better" with no trade-offs.
- Confusing LayerNorm with BatchNorm.