Xiong et al. (2020), "On Layer Normalization in the Transformer Architecture," compare two ways of placing layer normalization relative to each sublayer's residual connection: Post-LN, used in the original Transformer, and Pre-LN. What did they find about training stability, and what practical consequence follows for Post-LN training?
- Pre-LN and Post-LN produce identical gradient magnitudes at initialization, so the choice between them has no measurable effect on how training proceeds
- Post-LN, which applies layer normalization after the residual addition, produces expected gradients near the output layer that grow large at initialization, which the authors show makes a learning-rate warm-up stage necessary; placing layer normalization inside the residual block instead (Pre-LN) keeps gradients well-behaved at initialization without requiring warm-up
- Pre-LN requires a longer warm-up stage than Post-LN because normalizing before each sublayer slows down how quickly gradient magnitudes stabilize during the first training steps
- The placement of layer normalization only affects inference-time computation cost and has no bearing on gradient behavior or the need for a warm-up stage during training
Why B? And why not the others?
Correct answer: B. Post-LN, which applies layer normalization after the residual addition, produces expected gradients near the output layer that grow large at initialization, which the authors show makes a learning-rate warm-up stage necessary; placing layer normalization inside the residual block instead (Pre-LN) keeps gradients well-behaved at initialization without requiring warm-up
Using mean field theory, Xiong et al. show that in the Post-LN Transformer -- where layer normalization is applied after the residual addition, as in the original architecture -- the expected gradients near the output layer are large at initialization, and they demonstrate that this is why a learning-rate warm-up stage has empirically been necessary to train Post-LN Transformers successfully. When layer normalization is instead placed inside each residual block, before the sublayer's transformation (Pre-LN), the authors show gradients are well-behaved at initialization, and Pre-LN Transformers can be trained without the warm-up stage while reaching comparable results in less training time. The option claiming identical gradient magnitudes for both placements is wrong because the paper's central contribution is precisely that the two placements behave very differently at initialization. The option claiming Pre-LN needs a longer warm-up is wrong because the paper's finding is the reverse: Pre-LN is the placement that removes the need for warm-up, not the one that requires more of it. The option claiming the placement only affects inference cost is wrong because the paper's analysis and experiments concern training-time gradient behavior and optimization stability, not inference-time computation.
Source: Xiong et al., "On Layer Normalization in the Transformer Architecture" (2020), arXiv:2002.04745