Home / Fundamentals / Transformer Architecture

Transformer ArchitectureMidQuestion 6 of 9

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.