Title: ON LAYER NORMALIZATION IN THE TRANSFORMER ARCHITECTURE Anonymous authors Paper under double-blind review

Abstract: The Transformer architecture is popularly used in natural language processing tasks. To train a Transformer model, a carefully designed learning rate warm-up stage is usually needed: the learning rate has to be set to an extremely small value at the beginning of the optimization and then gradually increases in some given number of iterations. Such a stage is shown to be crucial to the final performance and brings more hyper-parameter tunings. In this paper, we study why the learning rate warm-up stage is important in training the Transformer and theoretically show that the location of layer normalization matters. It can be proved that at the beginning of the optimization, for the original Transformer, which places the layer normalization between the residual blocks, the expected gradients of the parameters near the output layer are large. Then using a large learning rate on those gradients makes the training unstable. The warm-up stage is practically helpful to avoid this problem. Such an analysis motivates us to investigate a slightly modified Transformer architecture which locates the layer normalization inside the residual blocks. We show that the gradients in this Transformer architecture are well-behaved at initialization. Given these findings, we are the first to show that this Transformer variant is easier and faster to train. The learning rate warm-up stage can be safely removed, and the training time can be largely reduced on a wide range of applications.

1. INTRODUCTION:

The Transformer is one of the most commonly used neural network architectures in natural language processing, and layer normalization is one of the key components in the Transformer. The originally designed Transformer places the layer normalization between the residual blocks, which is usually referred to as the Transformer with Post-Layer Normalization (Post-LN) . This architecture has achieved state-of-the-art performance in many tasks including language modeling (Dai et al., 2019; Al-Rfou et al., 2018) and machine translation (Vaswani et al., 2017; Dehghani et al., 2018; Edunov et al., 2018) . Unsupervised pre-trained models based on the Post-LN Transformer architecture also show impressive performance in many downstream tasks (Radford et al., 2019; Devlin et al., 2018; Yang et al., 2019) .

Although it achieves great success, people usually need to deal with the optimization of the Post-LN Transformer more carefully than convolutional networks (He et al., 2016; Tan & Le, 2019) or other sequence-to-sequence models (Sutskever et al., 2014; Gehring et al., 2017) . In particular, to train the model from scratch, any gradient-based optimization approach requires a learning rate warm-up stage: The optimization starts from using an extremely small learning rate, e.g., 1e -7 , and then gradually increases it to a pre-defined maximum value in a pre-defined number of iterations. After that, the learning rate decays similar to the optimization of other architectures. Both previous works (Vaswani et al., 2017; Popel & Bojar, 2018) , as well as our empirical study, show that such a warm-up stage is essential in training the models. Furthermore, the final model performance is quite sensitive to the value of the maximum learning rate and the number of warm-up iterations. Tuning such sensitive hyper-parameters is costly in training large-scale models, e.g., BERT (Devlin et al., 2018) or XLNet (Yang et al., 2019) .

In this paper, we study why the learning rate warm-up stage is essential in the optimization of the Post-LN Transformer and find it is closely related to the position of the layer normalization. As the warm-up stage happens in the first several iterations, we investigate the optimization behavior at initialization of the Post-LN Transformer. According to our theoretical analysis, when putting the layer normalization between the residual blocks, the expected gradients of the parameters near the output layer are large. Therefore, without the warm-up stage, directly using a large learning rate to those parameters may not lead to an improved model and can even make the optimization process unstable. Using a warm-up stage and training the model from small learning rates practically avoid this problem. As the location of the layer normalization plays a crucial role in controlling the gradient scales, we investigate whether there are some other ways of positioning the layer normalization that lead to better-normalized gradients. In particular, we study another variant, the Transformer with Pre-Layer Normalization (Pre-LN) (Klein et al., 2018) . The Pre-LN Transformer puts the layer normalization inside the residual connection and equips with an additional finallayer normalization before prediction (Please see Figure 1 for the differences between the two variants of the Transformer architectures). In this paper, we show that the gradients are wellbehaved without any exploding or vanishing at initialization for the Pre-LN Transformer both theoretically and empirically.

Given the gradients are well-behaved in the Pre-LN Transformer, it is natural to consider removing the learning rate warm-up stage during training. We conduct extensive experiments, including IWSLT14 German-English translation, WMT14 English-German translation, and BERT pre-training tasks. We show that, in all tasks, the learning rate warm-up stage can be safely removed and thus, the number of hyperparameter is reduced. Furthermore, we observe that the loss decays faster for the Pre-LN Transformer model. It can achieve comparable final performances but use much less training time. This is particularly important for training large-scale models on large-scale datasets.

Our contributions are summarized as follows: First, we investigate two Transformer variants, the Post-LN Transformer and the Pre-LN Transformer. By studying the gradients at initialization, we show why the learning rate warm-up stage is essential in training the Post-LN Transformer. Second, we are the first to show that the learning-rate warm-up stage can be removed for the Pre-LN Transformer. By using proper learning rate schedulers, the training time can be largely reduced.

2. RELATED WORK:

Gradient descent-based methods (Kingma & Ba, 2014; Zeiler, 2012; Duchi et al., 2011; Tieleman & Hinton, 2012) are popularly used in optimizing deep neural networks. For convolutional neural networks and recurrent neural networks, a relatively large learning rate is usually set in the beginning, and then decays along with the optimization process (He et al., 2016; 2017; Sutskever et al., 2014; Gehring et al., 2017; He et al., 2019) . The learning rate warm-up stage has been shown to be essential in dealing with some specific problems, e.g., the large-batch training. Goyal et al. (2017) ; He et al. (2019) and You et al. (2018) showed that when training neural networks with extremely large batch sizes (e.g., 8k in ImageNet), optimizing the model with a large learning rate in the beginning usually leads to poor performance. Training with a learning rate warm-up stage can eliminate the performance gap.

However, when optimizing the Post-LN Transformer models, as far as we know, in almost all previous works (Vaswani et al., 2017; Devlin et al., 2018; Dai et al., 2019; Radford et al., 2019; Lu et al., 2019) , the learning rate warm-up stage is essential and critical for training. Popel & Bojar (2018) investigated the influence of different warm-up strategies for the optimization of the Post-LN Transformer model and found that without or with relatively less warm-up iterations (e.g., 12k in Cz-En translation), the optimization diverges.

In a concurrent and independent work (Liu et al., 2019a) , the authors claimed that the benefit of the warm-up stage comes from reducing the variance for the adaptive learning rate in the Adam optimizer (Kingma & Ba, 2014) . They proposed to rectify the variance of adaptive learning rate by a new variant of Adam called RAdam. However, we identify the problem from the parameter initialization. We show that for the Post-LN Transformer, the scales of gradients of some parameters at initialization are large. First-order optimizers take the gradients as input. Using such gradients on these optimizers (not limit to Adam) with a large learning rate may make the optimization unstable and hurt the final performance.

3. OPTIMIZATION FOR THE TRANSFORMER:



3.1. THE TRANSFORMER ARCHITECTURE WITH POST-LAYER NORMALIZATION:

The Transformer architecture usually consists of stacked Transformer layers (Vaswani et al., 2017; Devlin et al., 2018) , each of which takes a sequence of vectors as input and outputs a new sequence of vectors with the same shape. A Transformer layer has two sub-layers: the (multi-head) selfattention sub-layer and the position-wise feed-forward network sub-layer. Residual connection (He et al., 2016) and layer normalization (Lei Ba et al., 2016) are applied for both sub-layers individually. We first introduce each component of the Transformer layer and then present the entire architecture.

Self-attention sub-layer An attention function can be formulated as querying an entry with keyvalue pairs (Vaswani et al., 2017) . The self-attention sub-layer uses scaled dot-product attention, which is defined as:

Attention(Q, K, V ) = softmax( QK T √ d )V

, where d is the dimensionality of the hidden representations, and Q (Query), K (Key), V (Value) are specified as the hidden representations of the previous layer. The multi-head variant of the self-attention sub-layer is popularly used which allows the model to jointly attend to information from different representation sub-spaces, and is defined as

EQUATION

EQUATION

where

W Q k ∈ R d×d K , W K k ∈ R d×d K , W V k ∈ R d×d V

, and W O ∈ R Hd V ×d are project parameter matrices, H is the number of heads. d K and d V are the dimensionalities of Key and Value. Without any confusion, given a sequence of vectors (x 1 , ..., x n ), we use MultiHeadAtt(

x i , [x 1 , x 2 , • • • , x n ])

as the multi-head self-attention mechanism on position i which considers the attention from x i to the entire sequence.

Position-wise FFN sub-layer In addition to the self-attention sub-layer, each Transformer layer contains a fully connected network, which is applied to each position separately and identically. This sub-layer is a two-layer feed-forward network with a ReLU activation function. Given a sequence of vectors h 1 , ..., h n , the computation of a position-wise FFN sub-layer on any h i is defined as:

EQUATION

where W 1 , W 2 , b 1 and b 2 are parameters.

Residual connection and layer normalization Besides the two sub-layers described above, the residual connection and layer normalization are also key components to the Transformer. For any vector v, the layer normalization is computed as LayerNorm(v) = γ v-µ σ + β, in which µ, σ are the mean and standard deviation of the elements in v, i.e., µ = 1 

x post,1 l,i = MultiHeadAtt(x post l,i , [x post l,1 , • • • , x post l,n ]) x pre,1 l,i = LayerNorm(x pre l,i ) x post,2 l,i = x post l,i + x post,1 l,i x pre,2 l,i = MultiHeadAtt(x pre,1 l,i , [x pre,1 l,1 , • • • , x pre,1 l,n ]) x post,3 l,i = LayerNorm(x post,2 l,i ) x pre,3 l,i = x pre l,i + x pre,2 l,i x post,4 l,i = ReLU(x post,3 l,i W 1,l + b 1,l )W 2,l + b 2,l x pre,4 l,i = LayerNorm(x pre,3 l,i ) x post,5 l,i = x post,3 l,i + x post,4 l,i x pre,5 l,i = ReLU(x pre,4 l,i W 1,l + b 1,l )W 2,l + b 2,l x post l+1,i = LayerNorm(x post,5 l,i ) x pre l+1,i = x pre,5 l,i + x pre,3 l,i

Final LayerNorm:

x pre F inal,i ← LayerNorm(x pre L+1,i )

architecture for the Transformer and BERT (Vaswani et al., 2017; Devlin et al., 2018) follows "selfattention (FFN) sub-layer → residual connection → layer normalization", which we call the Transformer with Post-Layer normalization (Post-LN Transformer), as illustrated in Figure 1 .

Post-LN Transformer Denote x l,i as the input of the l-th Transformer layer at position i, where x l,i is a real-valued vector of dimension d, i = 1, 2, ..., n, l = 1, 2, ..., L. n is the length of the sequence and L is the number of layers. For completeness, we define x 0,i as the input embedding at position i which is usually a combination of word embedding and positional embedding. The computations inside the l-th Post-LN Transformer layer are composed of several steps, and we use super-scripts on x to present the input(output) of different steps as in Table 1 (left), where W1,l , W 2,l , b 1,l and b 2,l are parameters of the FFN sub-layer in the l-th layer.

3.2. THE IMPORTANCE OF THE WARM-UP STAGE IN TRAINING THE POST-LN TRANSFORMER:

We are interested in the learning rate warm-up stage in the optimization of the Post-LN Transformer.

Different from the optimization of many other architectures in which the learning rate starts from a relatively large value and then decays (Bahdanau et al., 2017; He et al., 2016; Dauphin et al., 2017) , a learning rate warm-up stage for the Post-LN Transformer is critical. Specifically, denote the learning rate of the t-th iteration as lr(t). Denote the maximum learning rate during training as lr max . Given a predefined time frame T warmup , the learning rate scheduler for the first T warmup iterations is defined as (Vaswani et al., 2018) 

EQUATION

After this warm-up stage, the learning rate will be set by classical learning rate schedulers, such as the linear decay, the inverse square-root decay, or forced decay at particular iterations. As we can see from Eqn (4) , at the beginning of the training, the learning rate starts from zero 1 and then linearly increases to lr max in T warmup iterations. We conduct experiments to show that this learning rate warm-up stage is essential for training Post-LN Transformer models.

Setting We study the optimization process on the IWSLT14 German-to-English (De-En) machine translation task. We mainly investigate two aspects: whether the learning rate warm-up stage is essential and whether the final model performance is sensitive to the value of T warmup . To study the first aspect, we train the model with the Adam optimizer (Kingma & Ba, 2014) and the vanilla SGD optimizer (Ruder, 2016) respectively. For both optimziers, we check whether the warm-up stage can be removed. We follow (Vaswani et al., 2017) to set hyper-parameter β to be (0.9, 0.98) in Adam.

We also test different lr max for both optimizers. For Adam, we set lr max = 5e -4 or 1e -3 , and for SGD, we set lr max = 5e -3 or 1e -3 . When the warm-up stage is used, we set T warmup = 4000 as suggested by the original paper (Vaswani et al., 2017) . To study the second aspect, we set T warmup to be 1/500/4000 ("1" refers to the no warm-up setting) and use lr max = 5e -4 or 1e -3 with Adam.

For all experiments, a same inverse square root learning rate scheduler is used after the warm-up stage. We use both validation loss and BLEU (Papineni et al., 2002) as the evaluation measure of the model performance. All other details can be found in the Appendix. First, we can see that for both optimizers, the learning rate warm-up stage is essential. Without the warm-up stage, the BLEU score of the model trained with Adam optimizer can only achieve 8.45. As a comparison, the model trained using the warm-up stage can achieve around 34 in terms of BLEU score. The same trend can be also observed on the validation loss curves. Although the performance of the model trained with SGD is significantly worse than Adam, we can still see similar phenomena as Adam. The BLEU score is just above zero in 15 epochs without using the warm-up stage.

Second, we can see that the optimization process is sensitive to the value of T warmup , which means T warmup is an important hyper-parameter in training the Post-LN Transformer. For example, when setting T warmup = 500, the learned models with Adam achieve only 31.16 and 2.77 in term of BLEU score for lr max = 5e -4 and 1e -3 respectively2 .

Discussion First, we can see that the learning rate warm-up stage significantly helps the optimization of the Post-LN Transformer and also significantly affects the final performance. Such a warm-up stage brings additional efforts on hyper-parameter tuning which is computationally expensive for large-scale NLP tasks. Second, at the beginning of the training, the loss value is usually large. Standard optimization algorithms usually start with a large learning rate for fast convergence. However, when using the warm-up stage, the learning rate has to gradually increase from zero, which may slow down the optimization process. Liu et al. (2019a) suggests that the warm-up stage plays a role in reducing the undesirably significant variance in Adam in the early stage of model training. Based on this, they design a new variant of Adam optimizer, RAdam. However, according to our results, the warm-up stage also helps the training of SGD. This suggests that the benefit of the warm-up stage may be not for a particular optimizer.

3.3. UNDERSTANDING THE TRANSFORMER AT INITIALIZATION:

We can see that the Post-LN Transformer cannot be trained with a large learning rate from scratch. This motivates us to investigate what happens at the model initialization. We first introduce the parameter initialization setting for our theoretical analysis and then present our theoretical findings.

Notations We denote L(•) as the loss function of one position, L(•) as the loss function of the whole sequence, • 2 and • F as the l 2 norm (spectral norm) and the Frobenius norm, LN(x) as the standard layer normalization with scale γ = 1 and bias β = 0, and J LN (x) = ∂LN(x) ∂x as the Jacobian matrix of LN(x). Let O(•) denote standard Big-O notation that suppress multiplicative constants.

Parameter Initialization There are multiple parameter matrices in each Transformer layer, and most of the parameter matrices are initialized by the Xavier initialization (Glorot & Bengio, 2010) . Given a matrix of size n in ×n out , the Xavier initialization sets the value of each element by independently sampling from Gaussian distribution N (0, 2 nin+nout ). The bias vectors are usually initialized as zero vectors. The scale γ in the layer normalization is set to one.

For theoretical analysis, we study a simpler setting. First, we focus on single-head attention instead of the multi-head variant and for all layers, we set the shape of W Q,l , W K,l , W V,l , W 1,l ,W 2,l to be d × d. Second, we initialize the parameter matrices in the self-attention sub-layer W Q,l and W K,l to be zero matrices. In this setting, the attention is a uniform distribution at initialization and MultiHeadAtt(x

1 l,i , [x 1 l,1 , x 1 l,2 , • • • , x 1 l,n ]

) can be simplified as 1 n n j=1 x l,j W V,l . We test the optimization process in this setting and find that the problems described previously remain. Third, we assume the input vectors are also sampled from the same Gaussian distribution. This is reasonable since the inputs to the Transformer are linear combinations of word embeddings and learnable positional embeddings, both of which are initialized by Gaussian distributions.

Post-LN Transformer v.s. Pre-LN Transformer We compare the Post-LN Transformer with another variant of the Transformer architecture, the Transformer with Pre-Layer Normalization (Pre-LN). The Pre-LN Transformer was implemented in several systems (Vaswani et al., 2018; Klein et al., 2018; Liu et al., 2019b) . Wang et al. (2019) suggested that when stacking more layers, the Pre-LN Transformer is better than its Post-LN counterpart. Different from the Post-LN Transformer that puts the layer normalization between the residual blocks, the Pre-LN Transformer puts the layer normalization inside the residual connection and places it before all other non-linear transformations. Additionally, the Pre-LN Transformer uses a final layer normalization right before the prediction. We provide the mathematical formulations and visualizations of the Post-LN Transformer and the Pre-LN Transformer in Table 1 and Figure 1. For both architectures, each x L,i passes through a softmax layer to produce a distribution over the dictionary V . The loss function is defined on the softmax distribution. For example, in sequence prediction, the loss function is defined as L(x post L+1,i ) = -log(softmax yi (W emb x post L+1,i )) for the Post-LN Transformer and L(x pre F inal,i ) = -log(softmax yi (W emb x pre F inal,i )) for the Pre-LN Transformer, where softmax yi is the probability of ground truth token y i outputted by the softmax distribution and W emb is the word embedding matrix. The loss of the whole sequence is an average of the loss on each position. Without loss of generality, we assume that all the derivatives are bounded. We introduce the following concentration property of random variables which will be further used in the theorem.

Definition 1. A random variable Z ≥ 0 is called ( , δ)-bounded if with probability at least 1 -δ, Z-EZ EZ

≤ , where > 0 and 0 < δ < 1.

Intuitively, if the random variable Z is ( , δ)-bounded, then with a high probability its realization will not get too far away from its expectation. For example, if Y is a d-dimensional standard Gaussian random vector, then Z = Y 2 2 is ( , δ)-bounded with δ = exp(-d 2 /8), 0 < < 1 (see Appendix G for details). As parameter matrices in self-attention sub-layers and FFN sub-layers are initialized by Gaussian distributions, if the norm of the hidden states in the Transformer satisfies the concentrated condition above, we have the following theorem to characterize the scale of the gradients.

Theorem 1 (Gradients of the last layer3 in the Transformer). Assume that x post,5 L,i 2 2 and x pre L+1,i 2 2 are ( , δ)-bounded for all i, where and δ = δ( ) are small numbers. Then with probability at least 0.99 -δ -0.9+ , for the Post-LN Transformer with L layers, the gradient of the parameters of the last layer satisfies

∂ L ∂W 2,L F ≤ O(d √ ln d)

and for the Pre-LN Transformer with L layers,

∂ L ∂W 2,L F ≤ O d ln d L .

From Theorem 1, we can see that for the Post-LN Transformer, the scale of the gradients to the last FFN layer is of order O(d √ ln d) which is independent of L. For the Pre-LN Transformer, the scale of the gradients is much smaller. We first study the forward propagation of the Post-LN Transformer and the Pre-LN Transformer. Lemma 1 will be served as a basic tool to prove the main theorem and other lemmas. instead. As we can see from Lemma 2, the scale of the hidden states in the Post-LN Transformer keeps to be the same in expectation while the scale of the hidden states in the Pre-LN Transformer grows linearly along with the depth. The next lemma shows that the scale of the hidden states highly relates to the scale of the gradient in the architectures using layer normalization.

Lemma 1. If X ∈ R d is a Gaussian vector, X ∼ N (0, σ 2 I d ), then E( ReLU(X) 2 2 ) = 1 2 σ 2 d.

Lemma 3. For x ∈ R d , we have J LN (x) 2 = O( √ d x 2 ) in which J LN (x) = ∂LN(x) ∂x .

The proof of Lemma 1, Lemma 2, Lemma 3, and Theorem 1 can be found in the Appendix. The main idea is that the layer normalization will normalize the gradients. In the Post-LN Transformer, the scale of the inputs to the layer normalization is independent of L, and thus the gradients of parameters in the last layer are independent of L. While in the Pre-LN Transformer, the scale of the input to the final layer normalization is linear in L, and thus the gradients of all parameters will be normalized by

√ L.

Extending to other layers/parameters We have provided a formal proof on the gradients of the last FFN sub-layer as above. In order to fully understand the optimization, we also make some preliminary analysis for other layers and other parameters. Our main result is that the gradient norm in the Post-LN Transformer is large for the parameters near the output and will be likely to decay as the layer index l decreases. On the contrary, the gradient norm in the Pre-Transformer will be likely to stay the same for any layer l. All the preliminary theoretical results are provided in Appendix F.

Empirical study:

We also conduct experiments to study the gradients at initialization for the PostLN/Pre-LN Transformer in real scenarios. The model and training configuration exactly follows Section 3.2. The experiments are repeated ten times using different random seeds. Given an initialized model, we record the hidden states in the Post-LN/Pre-LN Transformer and find that the norm of the hidden states satisfies the concentration property ((0.1,0.125)-bounded). We also record the gradient for each parameter for different mini-batches. For elements in a parameter matrix, we calculate their expected gradients and use the Frobenius norm of those values as the scale of the expected gradient of the matrix. Figure 3 We further study the gradient statistics for the Post-LN Transformer after the warm-up stage with Adam. It can be seen from the figure that the scale of the gradients are very small, and the model can be trained with large learning rates. We believe the gradient scale is one of the reasons that the Post-LN Transformer needs a careful learning rate scheduling in the beginning. Since the gradients are large for some layers, using a large learning rate without warm-up may make the training unstable (see Appendix I). As the gradients are well-behaved for the Pre-LN Transformer, we will show that the learning rate warm-up stage can be removed for this model architecture in the next section.

4. EXPERIMENTS:

The Pre-LN Transformer has been implemented in several systems (Liu et al., 2019b; Baevski & Auli, 2018) , but most of them still follow Vaswani et al. (2017) to use the learning rate warm-up stage. We conduct experiments for the Pre-LN Transformer to test whether the learning rate warmup stage can be removed and how to set learning rate schedulers.

4.1. EXPERIMENT SETTINGS:

Machine Translation We conduct our experiments on two widely used tasks: the IWSLT14 German-to-English (De-En) task and the WMT14 English-to-German (En-De) task. For the IWSLT14 De-En task, we use the same model configuration as in Section 3. For the WMT14 En-De task, we use the Transformer base setting. More details can be found in the Appendix.

For training the Pre-LN Transformer, we remove the learning rate warm-up stage. On the IWSLT14 De-En task, we set the initial learning rate to be 5e -4 and decay the learning rate at the 8-th epoch by 0.1. On the WMT14 En-De task, we run two experiments in which the initial learning rates are set to be 7e -4 /1.5e -3 respectively. Both learning rates are decayed at the 6-th epoch followed by the inverse square root learning rate scheduler.

We train the Post-LN Transformer using the learning rate warm-up stage as the baseline. In both IWSLT14 De-En task and WMT14 En-De task, we set the number of the warm-up stage to be 4000 2017) and then use the inverse square root learning rate scheduler. For all experiments above, we use the Adam optimizer and set the hyper-parameter β to be (0.9, 0.98).

We set lr max as same as the initial learning rates of the Pre-LN Transformer in each corresponding experiment. Since Liu et al. (2019a) suggests that the learning rate warm-up stage can be removed using RAdam, we try this optimizer on the IWSLT14 De-En task. We use linear learning rate decay suggested by Liu et al. (2019a) and keep all other hyper-parameters to be the same as in other experiments.

Unsupervised Pre-training (BERT) We follow (Devlin et al., 2018) to use English Wikipedia corpus and BookCorpus for pre-training. As the dataset BookCorpus (Zhu et al., 2015) is no longer freely distributed. We follow the suggestions from (Devlin et al., 2018) to crawl and collect Book-Corpus4 on our own. The concatenation of two datasets contains roughly 3.4B words in total, which is comparable with the data corpus used in (Devlin et al., 2018) . We randomly split documents into one training set and one validation set. The training-validation ratio for pre-training is 199:1.

We use base model configuration in our experiments. Similar to the translation task, we train the Pre-LN BERT without the warm-up stage and compare it with the Post-LN BERT. We follow the same hyper-parameter configuration in Devlin et al. (2018) to train the Post-LN BERT using 10k warm-up steps with lr max = 1e -4 . For the Pre-LN BERT, we use linear learning rate decay starting from 3e -4 without the warm-up stage. We have tried to use a larger learning rate (such as 3e -4 ) for the Post-LN BERT but found the optimization diverged. All experiments are conducted on 32 NVIDIA Tesla P40 GPUs. 

4.2. EXPERIMENT RESULTS:



Machine:



5. CONCLUSION AND FUTURE WORK:

In this paper, we study why the learning rate warm-up stage is important in training the Transformer and show that the location of layer normalization matters. We show that in the original Transformer, which locates the layer normalization outside the residual blocks, the expected gradients of the parameters near the output layer are large at the beginning of the optimization. This leads to an unstable training when using a large learning rate. We further show that the Transformer which locates the layer normalization inside the residual blocks, can be trained without the warm-up stage and converges much faster. In the future, we will investigate other strategies of positioning the layer normalization, as well as the advantage of layer normalization to the Transformer from a theoretical perspective.

fine-tuning on each task for 6 times using different random seeds and compute the 95% confidence interval of validation accuracy.

B PROOF OF LEMMA 1:

Proof. Denote X = (X 1 , X 2 , ..., X d ) in which X i are i.i.d. Gaussian random variables with distribution N (0, σ 2 ). Denote ρ X (x) as the probability density function of

X 1 . Then E( ReLU(X) 2 2 ) = d i=1 E[ReLU(X i ) 2 ] = d i=1 E[ReLU(X i ) 2 |X i ≥ 0]P(X i ≥ 0) = d 2 E[ReLU(X 1 ) 2 |X 1 ≥ 0] = d 2 E[X 2 1 |X 1 ≥ 0] = d 2 +∞ -∞ x 2 ρ X|X>0 (x)dx = d 2 +∞ 0 x 2 2ρ X (x)dx = 1 2 σ 2 d.

C PROOF OF LEMMA 2:

Proof. At initialization, the layer normalization is computed as LN(v) = v-µ σ . It is easy to see that layer normalization at initialization projects any vector v onto the d -1-sphere of radius

√ d since LN(v) 2 2 = v-µ σ 2 = d k=1 (v k -µ) 2 σ 2 = d.

We first estimate the expected l 2 norm of each intermediate output

x post,1 l,i , • • • , x post,5 l,i

for l > 0. Using Xavier initialization, the elements in W V,l are i.i.d. Gaussian random variables sampled from N (0, 1/d). Since x post l,i 2 2 = d by the definition of Layer Normalization when l > 0, we have

EQUATION

2 ) + E( x post,1 l,i 2 2 ) + 2 n E( n j=1

x post l,j W V,l x post l,i

EQUATION

EQUATION

EQUATION

2 ) + E(

EQUATION

2 ) ≤ 2d (7)

and

EQUATION

EQUATION

EQUATION

EQUATION

EQUATION

2 ) ≥ E( x post l,i 2 2 ) = d.

Similarly, we have x post,3 l,i 2 2 = d by the definition of Layer Normalization. Again, for the ReLU activation function, the elements in W 1,l and W 2,l are i.i.d. Gaussian random variables sampled from N (0, 1/d). According to Lemma 1, we have

E( x post,4 l,i 2 2 ) = E( ReLU(x post,3 l,i W 1,l )W 2,l 2 2 ) (8) = E(E(E( ReLU(x post,3 l,i W 1,l )W 2,l 2 2 |x post,3 l,i , W 1,l )|x post,3 l,i )) (9) = E(E( ReLU(x post,3 l,i W 1,l ) 2 2 |x post,3 l,i )) = E( 1 2 x post,3 l,i 2 2 ) = d 2 (10)

Based on this, we can estimate the scale of E( x post,5 l,i

2 ) as follows.

EQUATION

EQUATION

EQUATION

2 ) + 2E(x post,3 l,i

x post,4 l,i

EQUATION

EQUATION

EQUATION

EQUATION

)

Using similar technique we can bound E( x pre l,i

2 ) for the Pre-LN Transformer. Since

EQUATION

2 ) + 2E(x pre,2 l,i

x pre l,i

) (14) = E( x pre l,i 2 2 ) + E( x pre,2 l,i 2 2 ) + 2 n E( n j=1

x pre,1 l,j W V,l x pre l,i

EQUATION

EQUATION

EQUATION

2 ) + E(

EQUATION

It is easy to see that we have E(

EQUATION

2 ) E(

EQUATION

EQUATION

2 ) + d. And similar to (11)-( 13),

EQUATION

EQUATION

Combining both, we have E(

EQUATION

EQUATION

EQUATION

2 ) ≤ (1 + 3l 2 )d by induction.

D PROOF OF LEMMA 3:

The proof of Lemma 3 is based on Lemma 4:

Lemma 4. Let α ∈ R d be a vector such that α 2 = 1, then the eigenvalue of I -α α is either 1 or 0.

Proof. Let {e 1 , ..., e d } be unit vectors such that e 1 = α and e i ⊥e j for all (i, j). Then we have e 1 (I -α α) = e 1 -e 1 α α = e 1 -α = 0 and e i (I -α α) = e i -e i α α = e i for i = 1. So e i are all the eigenvectors of I -α α, and their corresponding eigenvalues are (0, 1, 1, ..., 1). Hence we complete our proof.

Proof of Lemma 3. Denote y = x(I -1 d 1 1), where 1 = (1, 1, ..., 1) ∈ R d , then the layer normalization can be rewritten as

EQUATION

)

We explicitly calculate the Jacobian of layer normalization as

EQUATION

where δ ij = 1 when i = j and δ ij = 0 when i = j. In the matrix form,

EQUATION

and

EQUATION

Since the eigenvalue of the matrix (I -y y y 2

2:

) and (I -1 d 1 1) are either 1 or 0 (by Lemma 4), we have (I -y y

EQUATION

E PROOF OF THEOREM 1

The proof of Theorem 1 is based on Lemma 5: Lemma 5. Let Y be a random variable that is never larger than B. Then for all a < B,

EQUATION

Proof. Let X = B -Y , then X ≥ 0 and Markov's inequality tells us that

EQUATION

Hence

EQUATION

Proof of Theorem 1. We prove Theorem 1 by estimating each element of the gradient matrix. Namely, we will analyze ∂ L ∂W 2,L pq for p, q ∈ {1, ..., d}. The loss of the post-LN Transformer can be written as

EQUATION

Through back propagation, for each i ∈ {1, 2, ..., n} the gradient of L(x L+1,i ) with respect to the last layer's parameter W 2,L in the post-LN setting can be written as:

EQUATION

Here [ReLU(x post,3

EQUATION

which implies

EQUATION

Since we assume that all the derivatives are bounded, we have

∂L(x post L+1,i ) ∂x post L+1,i 2 2 = O(1). So | ∂L(x post L+1,i ) ∂W 2,L pq | 2 = O( J LN (x post,5 L,i ) 2 2 |[ReLU(x post,3 L,i W 1,L )] p | 2 ) (38) Since x post,3 L,i 2 2 = d, [x post,3 L,i

W 1,L ] p has distribution N (0, 1), using Chernoff bound we have

Pr[|[x post,3 L,i W 1,L ] p | ≥ a 0 ] ≤ exp(- a 2 0 2 ). So Pr[ReLU([x post,3 L,i W 1,L ] p ) 2 ≥ 2 ln 100d] ≤ 0.01 d .

Thus with probability at least 0.99, for all p = 1, 2, ..., d we have ReLU([x post,3

L,i W 1,L ] p ) 2 ≤ 2 ln 100d.

Since with probability 1 -δ( ),

EQUATION

for an arbitrary constant α 0 > 0, which equals

EQUATION

So according to union bound, with probability at least 0.99-δ( )-1+ -α0 we have

EQUATION

) ≤ O(

EQUATION

and

∂ L ∂W 2,L F = d p,q=1 | ∂ L ∂W 2,L pq | 2 = O( d 2 ln d α 0 )

.

The loss of the pre-LN Transformer can be written as

EQUATION

Using the same technique, in the pre-LN setting the gradient of L(x pre F inal,i ) with respect to the last layer's parameter W 2,L can be written as

EQUATION

So the absolute value of each component of the gradient is bounded by

| ∂L(x pre F inal,i ) ∂W 2,L pq | ≤ ∂L(x pre F inal,i ) ∂x pre F inal,i 2 J LN (x pre L+1,i ) 2 (0, 0, ..., [ReLU(x pre,4 L,i W 1,L )] p , ..., 0) 2 (45) = ∂L(x pre F inal,i ) ∂x pre F inal,i 2 J LN (x pre L+1,i ) 2 |[ReLU(x pre,4 L,i W 1,L )] p | (46) Since x pre,4

L,i 2 2 = d and [x pre,4 L,i W 1,L ] p obeys distribution N (0, 1), using Chernoff bound we have

Pr[|[x pre,4 L,i W 1,L ] p | ≥ a 0 ] ≤ exp(- a 2 0 2 ). So Pr[ReLU([x pre,4 L,i W 1,L ] p ) 2 ≥ 2 ln 100d] ≤ 0.01 d .

So with probability at least 0.99, for all p = 1, 2, ..., d we have ReLU([x pre,4 L,i W 1,L ] p ) 2 ≤ 2 ln 100d.

Since with probability 1 -δ( ),

EQUATION

According to union bound, with probability 0.99 -δ( ) -1+ -α0 we have

EQUATION

) ≤ O(

d ln d α0E x pre L+1,i 2 2 ) = O( ln d α0L ). So we have | ∂ L ∂W 2,L pq | 2 = | 1 n n i=1 ∂L(x pre F inal,i ) ∂W 2,L pq | 2 = O( ln d α 0 L ) (49) Thus ∂ L ∂W 2,L F = d p,q=1 | ∂ L ∂W 2,L pq | 2 ≤ O( d 2 ln d α0L ).

Take α 0 = 1 10 , we have that with probability at least 0.99 -δ( ) -0.9+ , for the Post-LN Transformer we have F EXTENSION TO OTHER LAYERS For simplicity, we denote x l = Concat(x l,1 , ..., x l,n ) ∈ R nd and x k l = Concat(x k l,1 , ..., x k l,n ) ∈ R nd for k = {1, 2, 3, 4, 5}. Then in the Post-LN Transformer, the gradient of the parameters in the l-th layer (take W 2,l as an example) can be written as 

EQUATION

where

J j i = diag σ x post,3 j,i w 1,j 1 , . . . , σ x post,3 j,i w 1,j d ∈ R d×d . ∂x post,3 j ∂x post,2 j =    J LN (x post,2 j,1

) . . . ), ..., J LN (x post,5 j,n

EQUATION

)), we have E 

EQUATION

If l is sufficiently large, the norm of J LN (x pre j,i ) and J LN (x pre,3 j,i ) are very small (of order O( 1 √ j )) as j is between l + 1 and L, which means the eigenvalues of matrix 

G EXAMPLES OF ( , δ)-BOUNDED RANDOM VARIABLES:

In this section we give an example of ( , δ)-bounded random variable. This example comes from Example 2.5 in (Wainwright, 2019) and we give a short description below.

If Z = (Z 1 , ..., Z n ) is a Gaussian vector with distribution N (0, I n ), then Y = Z 2 2 = n k=1 Z 2 k has distribution χ 2 n . And EY = n k=1 EZ 2 k = n A random variable X with mean µ = E[X] is called sub-exponential if there are non-negative parameters (ν, α) such that E[exp(λ(X -µ))] ≤ exp( ν 2 λ 2

2 ) for all |λ| < 1 α . The next proposition comes from Proposition 2.2 in (Wainwright, 2019) . Proposition 1 (Sub-exponential tail bound). Suppose that X is sub-exponential with parameters (ν, α). Then

EQUATION

)

and from Example 2.5 in (Wainwright, 2019), the χ 2 variable Y is sub-exponential with parameters (ν, α) = (2 √ n, 4). So we can derive the one-sided bound

EQUATION

So Y is ( , δ)-bounded with ∈ (0, 1) and δ = exp(-n 2 /8). We conduct experiments to study the gradient norm in the last FFN sub-layer in different sizes of the Transformer architecture at initialization to verify Theorem 1. We train 6-6/8-8/10-10/12-12/14-14 Post-LN/Pre-LN Transformer models, and record the gradient norm of the final FFN layer in different Transformer models. The results are plotted in Figure 6 . The x-axis is the size of the model, and the y-axis is the value of the gradient norm of W 2 in the final FFN sub-layer. It can be seen from the figure when the number of layers grows, the gradient norm remains in the Post-LN Transformer (around 1.6) and decreases in the Pre-LN Transformer. This observation is consistent well with our theory.

H EMPIRICAL VERIFICATION OF THE THEORETICAL FINDINGS:

Empirical verification of extended theory We record the gradient for each parameter for different mini-batches. For elements in a parameter matrix, we calculate their expected gradients and use the Frobenius norm of those values as the scale of the expected gradient of the matrix. Given the analysis above, we think our derived theory at the initialization stage is useful and consistent with the empirical studies above.

I LARGE GRADIENTS IN POST-LN TRANSFORMER HURTS THE OPTIMIZATION:

Theoretically, we find that the gradients of the parameters near the output layers are very large for the Post-LN Transformer and suggest using large learning rates to those parameters makes the training unstable. To verify whether using small-step updates mitigates the issue, we conduct a set of experiments that follows the setting in Section 3.3. There are two simple ways of "small-step updates", gradient norm clipping, and using small learning rates.

Experiments on using small learning rates We find using a very small but fixed learning rate can optimize the Post-LN Transformer (without the learning rate warm-up step) to a certain extent. We use a fixed learning rate of 1e -4 at the beginning of the optimization, which is much smaller than the lr max = 1e -3 in the paper. Please note that as the learning rates during training are small, the training converges slowly, and this setting is not very practical in real large-scale tasks. We plot the validation curve together with other baseline approaches in Figure 7 . We can see from the figure, the validation loss (pink curve) is around 4.3 in 27 epochs. This loss is much lower than that of the Post-LN Transformer trained using a large learning rate (blue curve). But it is still worse than the SOTA performance (green curve).

Experiments on using gradient norm clipping We find that using a small value of clip-norm can also optimize the Post-LN Transformer (without the learning rate warm-up stage). In the experiments, we find the gradient norm at initialization is about 5.00 and thus we clip the norm of the gradient update to 0.5. We plot the validation curve in Figure 7 (see gray curve). It can be seen from the figure that the performance is similar to the "small learning rate" experiment.

Discussions:

We think the experiments above help to verify our conclusion. Using a small learning rate/clip norm mitigates the instability of updating models with large gradient values. Although the theory we prove is based on some simplifications of the problem, the theoretical insights help us understand the optimization of different networks in practice.

:

An extremely small learning rate is practically used for the first update, e.g., 1e -7 .

One may notice that the orange curve is better than the blue curve on the validation loss but is worse on the BLEU score. This is because the BLEU score is defined on the translation results which is generated by step-wise decoding. There may be a big gap between the BLEU score and the validation loss when the model is not well-trained.

We mainly study the optimization of the Transformer layers and "last layer" here refers to the top FFN layer before the softmax operator. We did not focus on the parameters in the softmax layer as it is tied with the word embedding matrix, which is actually the input to the Transformer.

https://www.smashwords.com

https://www.smashwords.com

https://spacy.io

Figures and Tables:

Figure 1: (a) Post-LN Transformer layer; (b) Pre-LN Transformer layer.

k and σ 2 = 1 d d k=1 (v k -µ) 2 . Scale γ and bias vector β are parameters. Different orders of the sub-layers, residual connection and layer normalization in a Transformer layer lead to variants of Transformer architectures. One of the original and most popularly used

Figure 2: Performances of the models optimized by Adam and SGD on the IWSLT14 De-En task.

Based on Lemma 1, we have the following lemma to estimate the scale of the hidden states in different layers for the Post-LN Transformer and the Pre-LN Transformer.Lemma 2. At initialization, for the Post-LN Transformer, E( x post,all l > 0 and i. For the Pre-LN Transformer, (1+ l 2 )d ≤ E( x pre l,i22 ) ≤ (1 + 3l 2 )d for all l > 0 and i. Expectations are taken over the input and the randomness of initialization.Lemma 2 studies the expected norm of the hidden states in both Post-LN/Pre-LN Transformer. It is obviously that in the Post-LN Transformer, the norm of x post l,i is √ d and thus we study the norm of x post,5 l,i

Figure 3: Norm of expected gradients for Pre-LN/Post-LN Transformer

Figure 4: Performances of the models on the IWSLT14 De-En task and WMT14 En-De task

Translation We record the model checkpoints for every epoch during training and calculate the validation loss and BLEU score. The performance of the models at different checkpoints are plotted in Figure 4(a) -4(d).

Figure 5: Performances of the models on unsupervised pre-training (BERT) and downstream tasks Unsupervised Pre-training (BERT) We record validation loss of the model checkpoints and plot them in Figure5(a). Similar to the machine translation tasks, the learning rate warm-up stage can be removed for the Pre-LN model. The Pre-LN model can be trained faster. For example, the Post-LN model achieves 1.69 validation loss at 500k updates while the Pre-LN model achieves similar validation loss at 700k updates, which suggests there is a 40% speed-up rate. Note that T warmup (10k) is far less than the acceleration (200k) which suggests the Pre-LN Transformer is easier to optimize using larger learning rates. We also evaluate different model checkpoints on the downstream task MRPC and RTE (more details can be found in the appendix). The experiments results are plotted in Figure5(b) and 5(c). We can see that the Pre-LN model also converges faster on the downstream tasks.As a summary, the experiments on both machine translation and unsupervised pre-training tasks show that training the Pre-LN Transformer does not rely on the learning rate warm-up stage and can be trained much faster than the Post-LN Transformer.

∂ L ∂W 2,L F ≤ O(d √ ln d) and for the Pre-LN Transformer we have ∂ L ∂W 2,L F ≤ O(d ln d L )

to Lemma 2. Therefore, when we estimate the norm of ∂ L ∂W 2,l for post-LN transformer, there exists a term O( 2 3 (L-l)/2 ), which exponentially decreases as l goes smaller. Similarly, in the pre-LN Transformer, the gradient can be written asThe Jacobian matrices of the Pre-LN Transformer layers are:

1, and the norm of ∂ L ∂W 2,l for pre-LN transformer is independent of l when l is large.

Figure 6: Norm of expected gradients of W 2 in the last FFN sub-layer in different size of the Transformer architecture

Figure 3(a) and 3(b) shows those statistics for FFN sub-layers. The x-axis indexes different Transformer layers.It can be seen from the figure, the scale of the expected gradients grows along with the layer index for the Post-LN Transformer. On the contrary, the scale almost keeps the same for different layers in the Pre-LN Transformer. These observations are consistent with our theoretical findings.

max = 1e 3 ZRZDUPXS lr max = 5e 4 ZRZDUPXS lr max = 1e 4 ZRZDUPXS lr max = 1e 3 T warmup = 500 lr max = 5e 4 T warmup = 500 lr max = 1e 3 ZRZDUPXSJUDGLHQWQRUPFOLSSLQJ lr max = 1e 3 T warmup = 4000 lr max = 5e 4 T warmup = 4000

Figure 7: Performances of the models on the IWSLT14 De-En task.



Table content: <table><tr><td/><td>Post-LN Transformer v.s. Pre-LN Transformer</td></tr><tr><td>Post-LN Transformer</td><td>Pre-LN Transformer</td></tr></table>

Peer review (iclr2020):
Summary: The paper investigates the myth about layer normalization and learning rate warmup for the Transformer architecture. It shows, both theoretically and empirically, that putting the layer normalization in the residual blocks rather than between the residual blocks, could make a big difference to the scale of the gradient at the initialization stage. Pros: + A well-written paper with a good organization; notations are clear. + The proof I checked seem correct (but I didn*t check all of them). + Good experimental design that compares the pre-LN and post-LN Transformers in different settings/tasks. + Investigating the purpose and the theory behind using LN and learning rate warmup is a very interesting topic to me, as these modules are rarely used when training other kinds of deep nets (e.g., ConvNets, etc.). I will detail the cons below, along with my other questions/concerns/issues. ---------------------------------------- Questions/concerns: 1. One of the two major concerns I have is the novelty of this paper in terms of its methodology and empirical value to the community. The Pre-LN setting of Transformers has already been widely used. For instance, Baevski et al. [1] and Child et al. [2] are both well-known works that have applied the pre-LN setting and achieved SOTA results on various very challenging benchmarks. These papers also reported that using layer normalization before self-attention brought *more effective training*, which is one of the major empirical remarks that this paper made as well. 2. The second major concern I have is the connection the authors established between its theoretical findings and the empirical findings. While the post-LN Transformers may have larger gradient at higher-levels (as in, close-to-the-output levels), actually some (famous) prior works on Transformers have applied gradient clipping to their architecture, such as BERT [3] (https://github.com/google-research/bert/blob/master/optimization.py#L74), sparse transformer [2] and Transformer-XL [4]. But even when the gradient clipping is applied, learning rate warm-up still seems very helpful (and sometimes necessary), as was used in all of these works. Therefore, I think to further verify the theoretical hypotheses of the paper, the authors should at least also study whether (and to what degree) the very simple *gradient clipping* (or other gradient normalization techniques) solves the problem (which is a common solution to exploding gradients). 3. While Figure 3 is interesting to see, I don*t think it verifies Theorem 1 exactly. What it verifies is the *extending to other layers/parameters* paragraph (i.e., the gradient scale decreases with layers). Did you try training post-LN Transformers and pre-LN Transformer with different # of layers from scratch (i.e., different L)? According to Thm. 1, I think we should expect to see a plot where post-LN gradient expectation remains at the same level for all L, and pre-LN gradually decreases with L. 4. The proofs are based on the core assumption that we are *at initialization* (e.g., you assumed W_V entries are sampled from N(0, 1/d), that W_Q=W_K=0, and that the input data are normally distributed x ~ N(0, sigma^2 I_d)). How will the conclusion/derivation to change when these conditions are no longer met (e.g., after a few steps of warmup)? What do you expect to be the relationship between *the number of warmup steps* and solving the *gradient scale problem*, which you proved on these assumptions? ====================================== Some minor issues that didn*t impact the score: 1. You used delta = e^{-d epsilon^2 / 8} when discussing the tail bound, and later 3e^{-4}, 1e^{-4} for the learning rate. Just for notational consistency, maybe use 10^{-4} for the learning rate instead. 2. Referred to the wrong equation numbers (mentioned in another comment from me). 3. Appendix C: radius d -> radius sqrt{d} 4. Appendix C: *Similarly, we have | x_{l,i}^{post, 3} |_2^2 = d* should have an expected value. ====================================== Despite the potential lack of novelty on method, I do think investigating these myths and instabilities of training Transformers is a very interesting direction to pursue. I think this paper can be improved with more experimental settings to verify the claim it proposes (i.e., on the gradient, which brings a lot of different things to analyze/study/fix here). I put the paper slightly above the acceptance borderline. [1] https://arxiv.org/abs/1809.10853 [2] https://arxiv.org/abs/1904.10509 [3] https://arxiv.org/abs/1810.04805 [4] https://arxiv.org/abs/1901.02860