On Layer Normalization in the Transformer Architecture

By Ruibin Xiong, Yunchang Yang, Di He, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, Tie-Yan Liu. 2020

1,522 citations recorded by the source corpus.

The Transformer is widely used in natural language processing tasks. To train a Transformer however, one usually needs a carefully designed learning rate warm-up stage, which is shown to be crucial to the final performance but will slow down the optimization and bring more hyperparameter tunings. In this paper, we first study theoretically why the learning rate warm-up stage is essential and show that the location of layer normalization matters. Specifically, we prove with mean field theory that at initialization, for the original-designed Post-LN Transformer, which places the layer normalization between the residual blocks, the expected gradients of the parameters near the output layer are large. Therefore, using a large learning rate on those gradients makes the training unstable. The warm-up stage is practically helpful for avoiding this problem. On the other hand, our theory also shows that if the layer normalization is put inside the residual blocks (recently proposed as Pre-LN Transformer), the gradients are well-behaved at initialization. This motivates us to remove the warm-up stage for the training of Pre-LN Transformers. We show in our experiments that Pre-LN Transformers without the warm-up stage can reach comparable results with baselines while requiring significantly less training time and hyper-parameter tuning on a wide range of applications.

Subfields: Electrical and Electronic Engineering, Electronic, Optical and Magnetic Materials, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Electrical and Electronic Engineering

In Electronic, Optical and Magnetic Materials

In Model Architectures

Citations in the corpus

PanGu-α: Large-scale Autoregressive Pretrained Chinese Language Models with Auto-parallel Computation (2021, methodology): “For both MHA and FFN, we take the pre-layer normalization scheme, which can make the training of Transformer model easier and faster [14].”
Longformer: The Long-Document Transformer (2020, background): “…and learned position embeddings, dropout values of [0.1, 0.2] (small model) and [0.1, 0.4] (large model), pre-layernorm and post-layernorm (Xiong et al., 2020), learning rate (LR) of phase1 of values [2.5e-5, 5e-4, 1e-4] constant and cosine LR schedules, and different conﬁgurations for…”
CogView: Mastering Text-to-Image Generation via Transformers (2021, background): “Pre-LN [50] is proven to converge faster and more stable than the original Post-LN, and becomes the default structure of Transformer layers in recent works.”
GLM-130B: An Open Bilingual Pre-trained Model (2022, methodology): “We experiment with existing practices, e.g., Pre-LN (Xiong et al., 2020), Post-LN (Ba et al., 2016), Sandwich-LN (Ding et al., 2021), which are unfortunately incapable of stabilizing our GLM-130B test runs (Cf. Figure 3 (a) and Appendix B.2 for details).”
Scaling Laws for Neural Machine Translation (2021, methodology): “We train a series of pre-layer norm Transformer networks with varying sizes (Xiong et al., 2020).”
FLAVA: A Foundational Language And Vision Alignment Model (2021, positive): “In addition, the ViT transformer architecture (which applies layer norm [3] before the multi-head attention rather than after [112]) provides more robust learning for the text encoder under large learning rate than the BERT [28] transformer architecture.”
Primer: Searching for Efficient Transformers for Language Modeling (2021, background): “• Pre and Post Normalization: The standard practice for Transformers has become putting normalization before both the self-attention and feed forward transformations [44, 45].”
ReZero is All You Need: Fast Convergence at Large Depth (2020, methodology): “Since the introduction of Transformers [21], there have been several competing placements of the LayerNorm within the Transformer to achieve better convergence [4, 26].”
Do Transformer Modifications Transfer Across Implementations and Applications? (2021, methodology): “This small change has been unanimously adopted by all current Transformer implementations because it leads to more eﬀective training (Baevski and Auli, 2019; Xiong et al., 2020).”
BASE Layers: Simplifying Training of Large, Sparse Models (2021, methodology): “, 2016) is applied to the inputs of each residual block (Xiong et al., 2020) and to the outputs of the transformer.”
Group Equivariant Stand-Alone Self-Attention For Vision (2020, background): “These are in fact open questions for self-attention in general (Xiong et al., 2020; Liu et al., 2020; Zhao et al., 2020) and developments in this direction are of utmost importance.”
On the Stability of Fine-tuning BERT: Misconceptions, Explanations, and Strong Baselines (2020, background): “Another line of work investigates optimization difﬁculties when pre-training transformer-based language models (Xiong et al., 2020; Liu et al., 2020).”
Swin Transformer V2: Scaling Up Capacity and Resolution (2021, background): “Postnormalization harms training stability [60], and sandwich normalization sacrifices representation power due to too many normalization layers.”
Luna: Linear Unified Nested Attention (2021, background): “Technically, the position-wise feedforward layer operates on each position independently and layer normalization plays a crucial role in controlling the gradient scales (Xiong et al., 2020).”
General-purpose, long-context autoregressive modeling with Perceiver AR (2022, methodology): “In all experiments except where mentioned, we use pre-layernorm attention modules (Xiong et al., 2020) and squared-ReLU nonlinearities (So et al.”
Do Transformers Really Perform Badly for Graph Representation? (2021, methodology): “In addition, we apply the layer normalization (LN) before the multi-head self-attention (MHA) and the feed-forward blocks (FFN) instead of after [48].”
A Practical Survey on Faster and Lighter Transformers (2021, background): “However, adaptive learning rates have a problematically large variance in the early stages of optimization, resulting in convergence issues [72]; and (ii) the magnitude of the error signal propagated through LayerNorm is inversely proportional to the magnitude of the input [136].”