Pre-Norm vs Post-Norm › Summary

The placement of layer normalization relative to residual connections determines whether a transformer uses pre-norm or post-norm architecture. This choice affects gradient behavior during training and the amount of tuning required.

Key takeaways from this chapter:

• Post-norm normalizes after the residual addition: The original transformer design applies normalization to the sum of the sublayer output and residual. This ensures each layer produces normalized outputs but can create gradient flow challenges in deep networks.

• Pre-norm normalizes before the sublayer: Moving normalization inside the residual branch provides a direct gradient path through the residual connection. This enables stable training of very deep networks without complex warmup schedules.

• Gradient highways explain stability: The $+1$ term from the residual connection in pre-norm ensures gradients can flow backward without attenuation, regardless of sublayer behavior. This mathematical property is why pre-norm trains more stably.

• Activation accumulation is the trade-off: Pre-norm outputs are not normalized, leading to growing activation magnitudes. A final layer normalization before the output projection addresses this for most purposes.

• Modern architectures favor pre-norm: Most large language models, including GPT and LLaMA, use pre-norm due to its stability benefits. Post-norm persists in some encoder models like BERT where network depth is more modest.

• Learning rate tolerance differs: Pre-norm allows training with larger learning rates and simpler schedules. Post-norm requires careful warmup and learning rate selection to avoid divergence.

In the next chapter, we'll explore feed-forward networks, the other major component of transformer blocks alongside attention. Understanding how these two components work together completes the picture of transformer block architecture.