Pre-layer norm vs post-layer norm: what is the difference, where are they used and why?

Pre-layer norm and post-layer norm are two options for the placement of normalization layers (usually LayerNorm) in transformer architectures.

Post-layer norm is the classic approach, where normalization is applied after adding the layer outputs and the input (residual connection). The scheme is: x + Sublayer(x) -> LayerNorm. It was used in the original Vaswani et al. (2017) paper.

Pre-layer norm applies normalization before the main layer, with the scheme: LayerNorm(x) -> Sublayer -> x + output.

Differences and consequences

• Training stability: Pre-layer norm generally provides more stable and faster training, especially in deep transformers. This is because gradients propagate better through the normalized input.
• Gradient flow: In post-layer norm, gradients can vanish or explode, making training of deep models more difficult.
• Usage: Modern large models (GPT-3, T5, PaLM) more often use pre-layer norm due to its advantages in stability and training efficiency.

Summary

• If building a deep transformer model, it is recommended to use pre-layer norm.
• Post-layer norm is simpler and was the first approach, but may require additional tricks for stability.

# Example of pre-layer norm in pseudocode
def transformer_block(x):
    norm_x = LayerNorm(x)
    out = Sublayer(norm_x)
    return x + out

# Example of post-layer norm
def transformer_block(x):
    out = Sublayer(x)
    return LayerNorm(x + out)