On Layer Normalization in the Transformer Architecture
Authors: Ruibin Xiong
Yunchang Yang
Di He
Kai Zheng
Shuxin Zheng
Huishuai Zhang
Yanyan Lan
Liwei Wang
Tie-Yan Liu
Venue:  (2020)
Keywords: Transformer
BERT
Layer Normalization
Natural Language Processing

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.

Under review as a conference paper at ICLR 2020
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 Transformer model, 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 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 learn- ing 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 param- eters near the output layer are large. Then using large learning rate those gradients makes the training unstable. The warm-up up stage is practically helpful to avoid this problem. Such an analysis motivates us to investigate slightly mod- ified 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 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.
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 usu- ally referred to as the Transformer with Post-Layer Normalization (Post-LN). This architecture has achieved state-of- oftthe--ttttttttt- performance in many tasks including language modeling (Dai et al. 2019] Al-Rfou et al.] 2018) and machine translation (Vaswani et al. 2017 Dehghani 2018 Edunov ct al. 2018). U Unsupervised pre-trained models based on the Post-LN Transformer architecture also
show impressive performance in many downstream tasks (Radford al. 2019 Devlin et al. 2018
Yang et al. 2019).
Although achieves great success, people usually need deal with the optimization of the Post-
LN Transformer more carefully than convolutional networks (He et al. [0066] Tan & Le 2019) or other sequence-to-sequenec models (Sutskever et al., 2014] Gehring et al.) 2017). In particular,
to train the model from scratch, any gradient- optimizatin approach requires learning rate warm-up stage: The optimization starts from using an extremely small learning rate, e.g.. 1e and then gradually increases it to a pre-defined maximum value in 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). well 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-paramen parameters is costly in training large-scale models, e.g., BERT (Devlin et al. 2018) or XLNet (Yang et al.) 2019).
empirical study, show that such
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

---

Under review as a conference paper at ICLR 2020
the warm-up stage happens in the first several iterations, we investigate the optimization behavior 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 improved model and unstable. Using warm-up stage and training the model from small learning rates practically avoid this problem.
make the optimization process
As the location of the layer normalization plays crucial role in controlling the gradient scales, we investigate whether there are some other ways of positioning the layer normalization that lead to better-normalizer gradients. In par- ticular, we study another variant, the Trans- former 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 final- layer normalization before prediction (Please see Figure for the differences between the two variants of the Transformer architectures). In this paper, we show that the gradients are well- behaved without any exploding or vanishing at initialization for the Pre-LN Transformer both theoretically and empirically.
X1+1
X
Layer Norm
acdiition
I-N
Layer Norm
Layer
addition
Multi-Hend
Muht-HHad Attercion
Norm
Given the gradients are well- boll-behaved in the Pre- LN Transformer,
natural to consider re- moving the learning rate warm-up stage during training. We conduct extensive experiments, including IWSLT14 German-Englist transla- tion, WMT14 English-Germen translation, and BERT pre- pre-training tasks. We show that, in all tasks, the learning rate warm-up stage can be safely removed and thus, the number of hyper- parameter reduced. Furthermore, we observe that the loss decays faster for the Pre-LN Trans- former 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.
X0
(a)
(b)
Figure 1: (a) Post-LN Transformer layer; (b) Pre- LN Transformer layer.
Our contributions summarized 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 essential in training the Post-LN Transformer. Sec- ond, 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 largely reduced.
2 RELATED WORK
Duchi al. 2011 Tieleman & Hinton 2012) are popularly used in optimizing deep neural networks. For convolutional neural net- works and recurrent neural networks, 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 al., (019). The learning rate warm-up stage has been shown be essen- tial in dealing with some specific problems, e.g., the large-batch training. Goyal et al. (0117); He ctt 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 large learning rate in the beginning usually leads to poor performance. Training with a learning rate warm-up stage can eliminate the performance gap.
Gradient descent-basse based methods (Kingma & Ba 2014 Zeiler 2012:
However, when optimizing the Post-LN Transformer models, as far as we know, in almost all pre-
vious works Vaswani al. 2017 Devlin et al. 2018 Dai al. 2019) Radford et al. 2019 Lu
2

---

Under review as a conference paper at ICLR 2020
et al. 2019), the learning rate warm-up stage is essential and critical for training. Popel & Bojar 2018) investigated the influence 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 concurrent and independent work (Liu et al. 2019ap, 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 rectify the variance of adaptive learning rate by 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 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) self- attention sub-layer and the position-wel -iiee feed-forward network sub-layer. Residual connection (He) 2016)) and layer normalization (Lei Ba et al. (016) 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 key- value pairs (Vaswani et al. 2017). The self-attention sub-layer uses scaled dot-product attention, which is defined as: Attention Q,K, K, V) = softmax (QT VV where d is the dimensionality of the hidden representations, and Q (Query). K (Key), V (Value) are specified the hidden representa- tions 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
Multi-head (Q, K, V) Concat(headı,,-' .headd )Wo
head, Attention(QW KWW VWY),
where w Rdxdx W Rdxdx WV Rdxdd, and W Hdy are project parameter matrices, H is the number of heads. dx and dy are the dimensionalities of Key and Value. Without any confusion, given a sequence of vectors (I1, .... En), we use MultiHeadAAt(rr, [I1, T2. 2x,)) as the multi-head self-attention mechanism on position i which considers the attention from n; to the entire sequence.
Position-wise FFN sub-layer In addition to the self-attention sub-layer, each Transformer layer contains fully connected network, which is applied to each position separately and identically. This sub-layer is two-layer of vectors hi,..., An, the computation of a position-wise FFN sub-layer on any h; is defined as:
network with ReLU activation function. Given a sequence
FFN(h,) ReLU 6')w² 6²,
where W W 6t and 62 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 the layer normalization is computed as LayerNorm(r)
- + B, in which μt, are the
mean and standard deviation of the elements in v, i.e., / = Ei= 2kk and 0² E Ei -2)= Scale y and bias vector / are parameters.
Different orders of the sub-layers, residual connection and layer normalization in Transformer layer lead to variants of Transformer architectures. One of the original and most popularly used
3

---

Under review as a conference paper at ICLR 2020
Table 1: Post-LN Transformer v.s. Pre-IN Transformer
Post-LN Transformer
Pre-LN Transformer
MultHeadAAt(222 post 2 (g(e))
T pre.l LayerNorm (x)))
p Ti,i + post,1
2 MultiHeadAtt(ar mre.1 [x]1 ,' ,.
post.3 LayerNorm (x.
a Ls
pre,3
mi,6 + 2i, x.pre
xpre
post,1 (a
poet, W 6..))2 +
L LayerNorm (2f pre.3)
pre,
p ReLU(xZ WL. + (l.) 1211
post
LayerNorm(zZ,, (
a 4,41, = prees + 1,
pre,3
144,,
1,5
Final LayerNorm: Finat, LayerNorm 22 L+1.))
architecture for the Transformer and BERT (Vaswani al. 2017: Devlin et al. 2018) follows "self- attention (FFN) sub-layer + residual connection + layer normalization", which we call the Trans- former with Post-Layer normalization (Post-LN Transformer), as illustrated in Figure
Post-LN Transformer Denote Fi, as the input of the 1-th Transformer layer at position i, where
Fy, is a real-valued vector of dimension d, i 1. n, = 1, ..., L. n is the length of the sequence and L is the number of layers. For completeness, we define 20, as the input embedding at position i which usually a combination of word embedding and positional embedding. The computations inside the I-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 Table (left), where W1.1 W2,, bad and 122. are parameters of the FFN sub-layer in the I-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). learning rate warm-up stage for the Post-LN Transformer is critical. Specifically, denote the learning rate of the (-th iteration as Ir((). Denote the maximum learning rate during training mmm- Given predefined time frame T Twarmups the learning rate scheduler for the first T warmuu
iterations is defined as (Vaswani et al. 2018)
t
lr(t) Twarmup
Immass warmup-
(4)
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 zerd and then linearly increases to Irmax in T wanmuu 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 to taaimttt-Enggllee 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 wanmmpp- 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) set hyper-parameter
be (0.9, 0.98) in Adam. We also test different Ir Imaz for both optmmizers For Adam, we set Ir Immaz 5e- or le 3 and for SGD, we set Iimas 5e-3 or - When the warm-up stage is used, we set T Twarmup 4000 as suggested by the original paper (Vaswani et al. 2017). To study the second aspect, we set T Twarmup to be 1/500/4000 ("1" refers to the no warm-up setting) and use Irmax = 5e or 1e with Adam. For all experiments, same inverse square root learning rate scheduler is used after the warm-up stage. We use both validation loss and BLEU (Papineni et al. (002) as the evaluation measure of the model performance. All other details can found in the Appendix.
An extremely small learning practically used for the first update, e.g., 1e
4

---

Under review as a conference paper at ICLR 2020
10
30
Validation Loss
20
6
10
3
5
11
13
15
9
11
13
15
Epochs
Epochs
Kinaz 1e wio warm-up
-
Mrnax 1e
Toarmuo 5000
Vmax le Tuarmup 4000
Miax 5e wio warm-up
Trax Se
mammuu 500
Imax Se 4, Tearmup 4000
(a) Loss/BLEU on the IWSLT14 De-En En task (Adam)
35
10
30
25
Validation
20
15
11
13
16
13
15
Epochs
Epochs
Mruus le wio warm-up
Mmax 1e_3, Twanmpp = 4000
-* Irmax Se-3, wlo warm-up
Vruss Se Twarmpp 4000
(b) Loss/BLEU the IWSLT14 De-En task (SGD)
Figure 2: Performances of the models optimized by Adam and SGD the IWSLT14 De-En task.
Result We record the model checkpoints for every epoch during training and calculate the val- idation loss and BLEU score. The performance of the models trained with Adam and SGD plotted Figure 2(a) and Figure [bb) The x-axis is the epoch number and the y-axis is the BLEU score/validation loss. "w/o warm-up" indicates "without the warm-up stage" while "w/ warm-up" indicates "with the warm-up stage".
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 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 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 warrmups which means sppeemerrrmeter in training the Post-LN Transformer. For example, when setting T Tarmmpp 500, the learned models with Adam achieve only 31.16 16 and 2.77 in term of BLEU score for Ir Imaax
T warmup is an important
5e and le-3 respectively P
Discussion First, we can see that the learning rate warm-up stage significantly helps the opti- mization of the Post-LN Transformer and also significantly affects the final performance. Such warm-up stage brings additional efforts on hyper-parameter tuning which is computationally expen- sive 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,
'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 big gap between the BLFUU score and the validation loss when the model not well-trained.
5

---

Under review as a conference paper at ICLR 2020
which may slow down the optimization process. Liu 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 this, they design 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 particular optimizer.
variant of Adam optimizer, RAdam. However, according
3.3 UNDERSTANDING THE TRANSFORMER AT INITIALIZATION
We can see that the Post-LN Transformer cannot be trained with 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. E()) as the loss function of the whole sequence, || ll2 and || as the l2 norm (spectral norm) and the Frobenius norm, LN(x)) as the standard layer normalization with scale 7 1 and bias 3 O, and ( @LN(2) de as the Jacobian matrix of LN(x). Let (.) 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 matrix of size nin X Rout. the Xavier initialization sets the value of each element by indepen- 2 dently sampling from Gaussian distribution N(0, 7tn +o..) ). 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 WK.,, WV. W a to be d d. Second, we initialize the parameter matrices the self-attention sub-layer W and W.,. to be zero matrices. In this setting, the attention is uniform distribution at initialization and MultiHeadAAt( (T [ar,,, #7,2 ... .xi..)) can be simplified as Σー Ti, WVV. 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 reason- able 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. Pre-LN Transformer 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 [ and Figure
2018 Liu et al., 20195). Wang et al. (001)) suggested that when stacking more layers, the
For both architectures, each L,, 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 pre- diction, the loss function is defined as L( = log softmax, (Wemb) post ())) for the Post- LN Transformer and ((FFFatt,)
=2+1;))
log( peffttta,,, (Wemb__pre (inaa,,)) for the Pre-LN Transformer, where softmax,, is the probability of ground truth token y; outpuited by the softmax distribution and Wemb 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 2 is called (e, bb-uunde if with probability at least 1 6. Z-EZ S e, where e 0 and 0 < 8 < 1.

---

Under review as a conference paper at ICLR 2020
Intuitively, if the random variable Z is (eoundbod then with a high probability its realization will not get too far away from its expectation. For example, if Y a d-dimensional standard Gaus- sian random vector, then Z ||Y 2 is bbbbmdddd with 8 exp(-de² /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 (Gradients of the last layer in the Transformer). Assume that >>= 112 and 1 1l are (e. 8) -bounded for all i, where e and = ô(e) are small numbers. Then with probability at least 0.99 å 0.94c. for the Post-LN Transformer with L layers, the gradient of the parameters the last layer satisfies
aC law.z || S O(dvInd)
and for the Pre-LN Transformer with L layers,
Oc aaw.. 11. SO
d Ind
From Theorem 1, we can see that for the Post-LN Transformer, the scale of the gradients to the last FFN layer of order Ody'In d) which independent of L.. For the Pre-LN Transformer, the seale of the gradients is much smaller. We first study the forward propagation of the Post-LNN Transformer and the Pre-LN Transformer. Lemma will be served as basic tool to prove the main theorem and other lemmas.
Lemma 1. X E Rdd is a Gaussian vector, X ~ (T)) then (([[ReLLU(X))33) = 102d.
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.
post.5 (2)= 3d for all 1 > 0 and i. Lemma 2. At initialization, for the Post-LN Transformer, B(IaaT, (12)
For the Pre-LN Transformer, (1 + 1 )dd s [R[[ 122) S (1+ 2 for for all l 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-LLN Transformer. It post is v and thus we study the norm
is obviously that in the Post-LN Transformer, the norm of Lt,,
of post, 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 3. For x: E Rdd, we have [J (x) Ol vd in which JLN( aLN(z) |[]
dw
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. 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 L, and thus the gradients of all parameters will be normalized by LL.
the scale of the inputs to the layer normalization is independent
Extending to other layers/parameters We have provided 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 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.
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.

---

Under review as a conference paper at ICLR 2020
Empirical study We also conduct experiments to study the gradients at initialization for the PostLN/PrreLLN Transformer in real scenarios. The model and training configuration exactly fol- lows Section 3.2. The experiments initialized model, we record the hidden states in the Jost-LNNPrre-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-batcher For elements in a parameter matrix, we calculate their expected gradients and the Frobenius norm of those values the scale of the expected gradient of the matrix. Figure 3(a) and ((bb shows those statistics for FFN sub-layers. The x-axis indexes different Transformer layers. 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.
repeated ten times using different random seeds. Given an
the scale almost keeps the same for different layers the Pre-LN Transformer. These observations
are consistent with our theoretical findings. More analysis can be found in Appendix H.
1.2
Pro-LN jn.l)
Pro (HI)
Gradient Expectation
Poxi-UN (ril)
Gradient Expeciation
PostUUN (ini)
1.0
Fost (ottor up)
1.25
wart
0.8
1.00
0.6
0.75
0.4
0.50
0.2
0.0
3
5
2
3
Layer
Layer
W in the FFN sub-layers
(b) W in the FFN sub-layers
Figure 3: Norm of expected gradients for Pre-LN/Post-LN Transformer
We further study the gradient statistics for the Post-LN Transformer after the warm-up stage with Adam. It can seen from the figure that the scale 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 careful learning rate scheduling in the beginning. Since the gradients are large for some layers, using 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 al. (2017) use the learning rate warm-up stage. We conduct experiments for the Pre-LN Transformer to test whether the learning rate warm- up stage can be removed and how to set learning rate schedulers.
4.1 EXPERIMENT SETTINGS
We conduct our experiments two widely used tasks: the IWSLT14 German-to-English (De-En) task and the WMT14 English-to-German (En-De) task. For IWSLT14 De-En task, 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.
Machine Translation
For training the Pre-LN Transformer, 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, set to be 7e 7/1./1 respectively. Both learning rates are decayed at the 6-th epoch followed by the inverse square root learning rate scheduler.
two experiments in which the initial learning rates
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, set the number of the warm-up stage to be 4000
8

---

Under review as a conference paper at ICLR 2020
Validation Loss
Pre-NN (RAdim warmsup)
Fost
Pre-LN JAdem
BLEU
20
Paxi-IN (RAdim wo wertup)
10
Pro LN (RAdam w/o wam--us)
Post UN (Adam w/ wam-upl
Pre LN (Adam wo wamm-pi
0
Epochs
Epochs
(a) Validation Loss (IWSLT)
(b) BLEU (IWSLT)
Pos--E Moree
25
Validation Loss
(Frus 76"
20
(Fres
BLEU
15
PN Wrows 71
10
Post-UN Wrawe 1.50-))
PredNN (Poow 7o
PretN (Poww 1.50 3
5
11
15
17
19
3
7
11 Epochs
13
15
17 19
Epochs
(c) Validation Loss (WMT)
(d) BLEU (WMT)
Figure 4: Performances of the models on the IWSLT14 De-En task and WMT14 En-De task
following Vaswani et al. (001)) 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 Irmaz 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 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 Book Corpus (Zhu et al.) 2015) is no longer freely distributed. We follow the suggestions from (Devlin et al., 2018) to crawl and collect Book- Corpus on our own. The concatenation of two datasets contains roughly 3.4B words in total, which comparable with the data corpus used in (Devlin et al. 2018). We randomly split documents into one training set and one validation set. The coaning-validatioon ratio for pre- pre-training is 199:1.
We use base model configuration 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
configuration in Devlin et al. (2018) to train the Post-LN BERT using 10k warm-up steps with Irmaz le- For the Pre-LN BERT,
linear learning rate decay starting from 34 without the warm-up stage. We have tried to use larger learning rate (such as 3e 3e-4) for the Post-LN BERT but found the optimization diverged. All experiments conducted 32 NVIDIA Tesla P40 GPUs.
4.2 EXPERIMENT RESULTS
Machine Translation We record the model checkpoints for every epoch during training and cal- culate the validation loss and BLEU score. The performance of the models at different checkpoints are plotted in Figure 4(a) 4())
whtpp://www.smmsswrrttti

---

Under review as a conference paper at ICLR 2020
First, can see from the figure, the learning rate warm-up stage is not critical anymore for train- ing the Pre-LN Transformer and the performance of the learned model is competitive. For example, on the IWSLT14 De-En task, the BLEU score and validation loss of the Pre-LN Transformer can achieve around 34 and 4, which are comparable with the performance of the Post-LN Transformer.
Second, the Pre-LN Transformer converges faster than the Post-LN Transformer using the same Ima.. On the IWSLT14 De-En task, the 9-th checkpoint of the Pre-LN Transformer achieves nearly the same performance (validation Loss/BLEU score) 15-th checkpoint of the Post-LN Transformer. Similar observations can be found in the WMT14 En-De task. The first model checkpoint of the Pre- LN Transformer can achieve BLEU score near 20, As a comparison, the BLEU score of the first checkpoint of the Post-LN Transformer is less than 10.
Third, compared with RAdam, we find that the change of the position of layer normalization "dom- inates" the change of the optimizer. According to our experiments on the IWSLT14 14 De-En task, we can see that although RAdam trains the Post-LN Transformer well without the warm-up stage, has little difference with Adam when training the Pre-LN Transformer.
0.90
0.74
Pasl-N
2.0
0.88
0.72
Validation Loss
1.9
0.85
0.70
0.84
5 0.68
0.82
1.7
0.80
0.68
0.78
0.64
1.6
0.76
PioLN
0.62
Pr
100
300
500
730
900
200
400
600
800
200
400
600
800
Pre-training Staps (Thousands)
Pre-training Steps (Thousands)
Pre-training Steps (Thousands)
(a) Validation Loss on BERT
(b) Accuracy on MRPC
(c) Accuracy on RTE
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 Figure 5(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 500k updates while the Pre-LN model achieves similar validation loss at 700k updates, which suggests there is 40% speed-up rate. Note that T Twarmup (10k) is far less than the acceleration (200k) which suggests the Pre-I.N Transformer 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 Figure 5(b) and 5(c) We can see that the Pre-LN model also converges faster on the downstream tasks.
As 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.
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 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, will investigate other strategies of positioning the layer normalization, as well as the advantage of layer normalization to the Transformer from a theoretical perspective.
10

---

Under review as a conference paper at ICLR 2020
REFERENCES
Rami Al-Rfou, Dokook Choe, Noah Constant, Mandy Guo, and Llion Jones. Character-level lan- guage modeling with deeper self-attention. arXiv preprint 04444, 2018.
Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. arXiv preprint 1809. 10853, 2018.
Dzmitry Bahdanau, KyungHyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. 2017.
Luisa Bentivogli, Ido Dagan, Hoa Trang Dang, Danilo Giampiccolo, and Bernardo. Magnini. The fifth PASCAL recognizing textual entailment challenge. 2009.
Zihang Dai, Zhilin Yang, Yiming Yang, William W Cohen, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. arXiv preprint arXiv:1901. 0208600 2019.
Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In International Conference on Machine Learning, pp. 933-941, 2017.
Mostafa Dehghami, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. arXiv preprint arXiv: urXiv: 0007/000790099 2018.
Jacob Devlin, Ming-Wos Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805. 2018.
William B Dolan and Chris. Brockett. Automatically constructing corpus of sentential paraphrases. In Proceedings of the International Workshop on Paraphrasing. 2005.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 1200u):2121-21553 2011.
Sergey Edunov, Myle Ott, Michael Auli, and David Grangier. Understanding back-translation at scale. arXiv preprint /808.
2018.
Jonas Gehring, Michael Auli, David Grangier, Denis Yarats, and Yann N Dauphin. Convolutional sequence to sequence learning. In International Conference on Machine Learning, pp. 1243- 1252, 2017.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 249-256, 2010.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, An- drew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in 1 hour. arXiv preprint arXiv: 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross Girshick. Mask r-cnn. In Proceedings of the IEEE international conference on computer vision, pp. 2961-2969, 2017.
Tong He, Zhi Zhang, Hang Zhang, Zhongyue Zhang, Junyuan Xie, and Mu Li. Bag of tricks for image classification with convolutional neural networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 558 558-567, 2019.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: /4/2.6980, 2014.
Guillaume Klein, Yoon Kim, Yuntian Deng, Vincent Nguyen, Jean Senellart, and Alexander Rush. Opennmt: Neural machine translation toolkit. In Proceedings of the 13th Conference of the Association for Machine Translation in the Americas (Volume 1: Research Papers), volume 1, pp. 177-184, 2018
11

---

Under review as a conference paper at ICLR 2020
Philipp Koehn, Hieu Hoang, Alexandra Birch, Chris Callison Burchh, Marcello Federico, Nicola Bertoldi, Brooke Cowan, Wade Shen, Christine Moran, Richard Zens, et al. Moses: Open source toolkit for statistical machine translation. In Proceedings of the 45th annual meeting of the as- sociation for computational linguistics companion volume proceedings of the demo and poster sessions, pp. 177-180, 2007.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.064500 2016.
Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. arXiv preprint arXiv: /90i::/90..000 1908.03265, 2019a.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 1907. I 1692,, 2019b.
Yiping Lu, Zhuohan Li, Di He, Zhiqing Sun, Bin Dong, Tao Qin, Liwei Wang, and Tie- Yan Liu. Understanding and improving transformer from a multi-particle dynamic system point of view. arXiv preprint orXiv:1906 02062, 2019.
Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. Bleu: a method for automatic evaluation of machine translation. In Proceedings of the 40th annual meeting association for computational linguistics, pp. 311-318. Association for Computational Linguistics, 2002.
Martin Popel and Ondicj Bojar. Training tips for the transformer model. The Prague Bulletin of Mathematical Linguistics, 1 110(1):33770, 2018.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
Sebastian Ruder. An overview of gradient descent optimization algorithms. 16099004074477 2016.
arXiv preprint
Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. arXiv preprint arXiv: (.t.:1500.07999 07909, 2015.
Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. In ACL,, 2016.
Ilya Sutskever, Oriol Vinyals, and Quoc V Le. Sequence to sequence learning with neural networks. In Advances neural information processing systems, pp. 3104-3112, 2014.
Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. Rethink- ing the inception architecture for computer vision. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 2818-2826, 2016.
Mingxing Tan and Quoe Le. Efficientnet: Rethinking model scaling for convolutional neural networks. arXiv preprint 19Xiv: 17946, 2019.
Tijmen Tieleman and Geoffrey Hinton. Lecture 6.5-rmsprop, coursera: Neural networks for machine learning. University of Toronto, Technical Report, 2012.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Ashish Vaswani, Samy Bengio, Eugene Brevdo, Francois Chollet, Aidan N. Gomez, Stephan Gouws, Llion Jones, Fukasz. Kaiser, Nal Kalchbrenner, Niki Parmar, Ryan Sepassi, Noam Shazeer, and Jakob Uszkoreit. Tensor2tensor for neural machine translation. CoRR, abs/1803.07416, 2018 URL http::/ttt/aiiia org/ aaaaabbs/1000 011//1113.00166666611111
Martin Wainwright. High- dimensiimli statistics: A non-asymptotic viewpoint, volume 48. Cam- bridge University Press, 2019.
12

---

Under review as a conference paper at ICLR 2020
Qiang Wang, Bei Li, Tong Xiao, Jingbo Zhu, Changliang Li, Derek F Wong, and Lidia S Chao. Learning deep transformer models for machine translation. arXiv preprint arXi:1190..11877, 2019.
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Ruslan Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding. arXiv preprint arXiv:1906.08237, 2019.
Yang You, Zhao Zhang, Cho-Jui Hsieh, James Demmel, and Kurt Keutzer. Imagenet training in minutes. In Proceedings of the 47th International Conference on Parallel Processing, pp. 1. ACM, 2018.
Matthew D Zeiler. Adadelta: an adaptive learning rate method. arXiv preprint arXiv: 2012.
Yukun Zhu, Ryan Kiros, Richard Zemel, Ruslan Salakhutdinov. Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. arXiv preprint arXiv:/506.000 067.44 2015.
13

---

Under review as a conference paper at ICLR 2020
Appendix
A EXPERIMENTAL SETTINGS
A.1 MACHINE TRANSLATION
Experiment on Section The training/validationneest sets of the IWSLT14 German-to-Ennlish (De-En) task contain about 153K/7K/7K sentence pairs, respectively. We use vocabulary of 10K tokens based on a joint source and target byte pair encoding (BPE) (Sennrich et al.) 2013). All of experiments use Transformer architecture with 6-layer encoder and 6-layer decoder. The size of embedding is set to 512, the size of hidden nodes in attention sub-layer and position-wise feed-forward network sub-layer smoothed cross entropy is used as the objective function by setting E = 0.1 (Szegedy et al. (2016), and we apply dropout with ratio 0.1. The batch size set be 4096 tokens. When decode translation results from the model during inference, we set beam size as 5 and the length penalty as 1.2.
4. Label set 512 and 1024, and the number of heads is
Experiment Section 4 The configuration of IWLST11 De-En task is the same in Section 3. For the WMT14 En-De task, we replicate the setup of (Vaswani et al. 2017). which consists of about 4.5M training parallel sentence pairs, and uses 37K vocabulary based on a joint source and target BPE. Newstest2013 is used as the validation set, and Newstest2014 is used as the test set. One of the basic configurations of the Transformer architecture is the base setting, which consists of 6-1ayer encoder and 6-layer decoder. The size of the hidden nodes and embeddings are set to 512. The number of heads Label smoothed cross entropy is used as the objective function by setting 0.1. The batch size is set to be 8192 tokens per GPU on 16 NVIDIA Tesla P40 GPUs.
A.2 UNSUPERVISED PRETRAINING
We follow Devlin et al. (018) to use English Wikipedia corpus and BookCorpus for the pre- pre-training. As the dataset BookCorpus (Zhu al. 2015) is no longer freely distributed. We follow the sugges- tions from Devlin et al. (2018) to crawl and collect BookCorpus on our own. The concatenation of two datasets includes roughly 3.4B words in total, which is comparable with the data corpus used in Devlin et al. (2018). We first segment documents into sentences with Spacy Then, we normal- ize, lower-case, and tokenize texts using Moses (Koehn et al. (200)) and apply BPE(SSennrih et al.) 2016). We randomly split documents into one training set and one validation set. The training- validation ratio for pre-training is 199:1.
The base model in Devlin et al. (2018) consists of 12 Transformer layers. The size of hidden nodes and embeddings are set to 768, and the number of heads set to 12.
A.3 GLUE DATASET
MRPC The Microsoft Research Paraphrase Corpus (Dolan Brockett 2005) is corpus of sen- tence pairs automatically extracted from online news sources, with human annotations for whether the sentences in the pair are semantically equivalent, and the task is to predict the equivalence. The performance is evaluated by the accuracy.
The Recognizing Textual Entailment (RTE) datasets come from a series of annual textual entailment challenges (Bentivogli at. 2009). The task to predict whether sentences in a sentence pair are entailment. The performance is evaluated by the accuracy.
RTE
Fine-tuning on GLUE tasks We use the validation set for evaluation. To fine-tune the models, following [Devlin al. (2018); Liu et al. (2019b), we search the optimization hyper-parameterss in a search space including different batch sizes (16/32), learning rates (le-5 le 1e-4) and number of epochs (3-8). We find that the validation accuracy are sensitive to random seeds, so we repeat
"btts:/wwww.mmasswrrk.cck
ihups://ppcy.io
14

---

Under review as a conference paper at ICLR 2020
fine-tuning on each task for 6 times using different random seeds and compute the 95% confidence interval validation accuracy.
B PROOF OF LEMMA
Proof. Denote X = (X1, X2.... Xa) in which X, are i.i.d. Gaussian random variables with distri-
bution .V (0) Denote px (a) as the probability density function of X1. Then |[| (|ReLU(X)]]))
E (X,,) = -i (x, 2 ofp(X, 2 0) = / (eLU(())) 2 0] =
x 2 0] = 4 t px)x>u(z)dz $ Jo (x) f
C PROOF OF LEMMA
Proof. initialization, the layer normalization is computed LN(u) Vーn It is easy see that layer normalization at initialization projects any vector e onto the d - 1-sphere of radius vd since (1h d.
[IN((>) = 11 Σi=
5²
We first estimate the expected 12 norm of each intermediate output "1,4 Using Xavier initialization. the elements in VV1. are i.i.d. Gaussian random variables sampled from (0, 1/d). Since
p for >> 0. I,
||0|3 d by the definition of Layer Normalization when O, we have
E [[| [| 212) = B(I1 past 12
+ 2B
(
post,1 post T T1,3
(5)
past 12
past,11
112)
2 E
post WV post
(6)
post B(I (f 12:)
B(
post (12) + E(1
post
2d
(7)
and [[ E(Iaa pest
post.2 (3) = E(1 113) + E(11=],
+ (13) = E(I post (13)
E(I
n Ei=
post
133)
2
d.
Similarly, we have ||=r, d by the definition of Layer Normalization. Again, for the ReLU activation function, the elements in W and W221 are i..dd Gaussian random variables sampled we have
from (0, 1/d). According to Lemma] 1
E [(I 112)
B( | (o W 2
(8)
((E((B[| ReeC[] (1 post, & 1 11g
post,3 w (appoar.)))
(9)
d
= E(E([|ReLU(z Post.3 w 1.
(ost. EE post,3 (12)
(10)
2
Based on this, we can estimate the scale of [[[ post. (12) follows.
B(Iaa post.5 (2]]
E(I1 1,5
3
(22)
E [[ post,4
(12) + 2
post,3 poet,1
(11)
E(Ila p
(2))
E
post, 1 (13) +
2 E
ReLU 1 poot,
post,3
(12)
d
BlIlagpot.. 112) - 01 4
(13) = d
+
d
(13)
2
2
15

---

Under review as a conference paper at ICLR 2020
Using similar technique can bound E(I|=];, (l3) for the Pre-LN Transformer. Since
E [1 [ 123) = |
[|
(
+
2 21,
pre.2 P
(14)
B(I pre.2
+
E
pre.l W
V. xppp
(15)
=( +
pre.2
B(
1,3
pre
E(I
pre,
(16)
It is easy to that have =[ - 12:) < B( (11)-(13),
v B(I 12: + d. And similar to
B(I 11)
+ E =[| 11 + 2E (at ( p p
(17)
B(1|=f,, 12) + IE +|,+ 11
(18)
B(1|=f,, pre.3 12 + d
(19)
2
Combining both, we have |E[|= 112) + 1 B(1=44i.11 S B(| 112) + 2.. Then we have
(1+ + 4) s | [ (12) s (1 ) yd by induction.
D PROOF OF LEMMA 3
The proof of Lemma 3 is based on Lemma 4:
Lemma 4. Let E Rd be vector such that ||a:||2 then the eigenvalue of T a. o is either or O.
Proof. Let {eI,..., Ea} be unit vectors such that ei and e; for all (i,j). Then we have ei - a aTa)===== =1 e a e1 0 and ea(T - a a) = e; e;a" a 6; for + 1. So e are all the eigenvectors complete our proof.
a a, and their corresponding eigenvalues are (0, 1, 1, ..., 1). Hence
Proof of Lemma 3. Denote y = r(1 (1T1), 1), where 1 = (1, (I.,,.....) ization can be rewritten
1) E Rd, then the layer normal-
LN(x); = Vi y
(20)
We explicitly calculate the Jacobian of layer normalization as
8ig V Ek v2 Vi / E M0
OLN(x);
0
y1
(21)
Oy;
oy:
1/ Dim V
Σ k=
Vh
=v 8zg 11y11
Nilli
Vd
DeMl 1
(5
(22)
1l
where 8G = 1 when i = and dzz
O when i + j. In the matrix form,
oLN(x)
Vd
y
(T
(23)
Oy
1
16

---

Under review as a conference paper at ICLR 2020
and
JLN(N) (a)
(24)
dr
(L oy
(25)
dy
dr
y )1 1).
Vd 11v122
(1
(26)
Since the eigenvalue of the matrix
(I
y and (I - 1 1) are either or 0 (by Lemma 4), we 1 !
have [] x²x Tvll l| = (1) and [I //2 1)ll2 = (1). So the spectral norm of J JN
[]]x())]22 0r vd
vd |
0
(27)
|
E PROOF OF THEOREM 1
The proof Theorem based on Lemma 5:
Lemma 5. Let Y be a random variable that never larger than B. Then for all a < B.
B Y] B a
Pr[Y S al
(28)
Proof. Let = B - Y, then X and Markov's inequality tells us that
E
Pr[X > B a] 5
(29)
B
Hence
Pr|Y I a] < E Y]
(30)
B - a
Proof of Theorem I. We prove Theorem by estimating each element of the gradient matrix. ac for p.q E (1....dd. The loss the post-LN Transformer can Namely, we will analyze TW2.L be written
L bm
post
post
(31)
n
Through back propagation, for each E (1,2,..., n) the gradient of ((xL+1,) with respect to the last layer's parameter 2 in the post-LN setting can be written as:
aL(xpoot Ow2,1
post aC(xport.) papa
Dxpt Ox post,4 L.t
(32)
Dusts
post.
aw2.
ol(xtort.))
post.4
JLN (2
(33)
Oxyoat
aw2.t
J (
(oat..))0,0....... 0, ( 3ost. 111
Ipp 0)
(34)
Oxpost
17

---

Under review as a conference paper at ICLR 2020
Here Re means the p-th element of (a So the absolute value 0C(x of
can be bounded by
OW
(L(x post Ow2,"
aL(x Oxt L+1,
(22.. 1221(, o,... [22, w1.....L........1.1..111 (l..,).
0)
(35)
af(22+1, port
112 LN (
(36)
post drlti L+1
which implies
OC(aL aw pq
HC (+11,)
11 IN
post.
[eeL (1 Wp.h))pp
post.3
(37)
(C(x
Since we assume that all the derivatives are bounded, we have
||
((1. So
8n post
aC(I (+1s) OO[IJEN ( (e wh }np)
(38)
aW?1"
2.1
Since ||=2,,, 113 = d, [=2.t
has distribution N(0, 1), using Chernoff bound we have W1 1
Pr[l WLLL),, n1 2 aol S exp(( _
a
2
So
0.01
Pr post, W 2 2 In 100d] S
d
Thus with probability at least 0.99, for all p 1,2, .... d we have to 2 In 100d.
Since with probability 1 8(e),
1 El
I e, we have |
post,5
12 1 (1 +
post.
B post,5 Izg. Using Lemma 5, we have
(1+))Blarg post,
B|
E
post.5
Pr Pr|l Post 1|2 <
post,5
L,i
(39)
(1+ ao) E0) post,
1+c ag
for arbitrary constant oxp 0, which equals
Pr Prlll=too 2 Bl
post,5
(40)
1 a0
So according to union bound, with probability at least 0.99-5(c) 5(c) 1+c-00 we have
aw
o
[| LN poot, 112
(Re..
post,3 111.1 p))
< O
2d 100d
O dnnd
= 0
Ind
popt,
"13
post.
00
So we have
ac
1
lnd
2
0
(41)
aW pq
OW-2.
n1
aW2,1 pq
n
2w2.1.
pq
00
and
ac
ac
d'
12
011
0W2,L
aW2.L 2.
ag
p.9=1
18

---

Under review as a conference paper at ICLR 2020
The loss of the pre-LN Transformer be written
I
pre Pinal,n
-
1
Σ C(
((Faat)))))
Final,e)
(42)
n1. im1
Using the same technique, in the pre-LN setting the gradient of [(x layer's parameter W2L can be written as
Final,;) with respect to the last
(f(Ftiiat,) Ou De Or dx Final,
(43)
da Ory OWit
aflxFFmatt
JLN (1, o..,, R2LL ( lp..... 0)T
(44)
Ox
So the absolute value of each component of the gradient is bounded by
aL(aFFnntd
al(rFimatt)) pre Final,i
Final,)
O,... ... (ReLU(Z 0)ll2
(45)
@L(rFinatd Final,)
||
(46)
Dre Fénal,i
Since ||= d and
(.r In obeys distribution N(00 1). using Chernoff bound we have
( 2 anl exp((- ad
So
0.01
Pr|ReLU((u (rr (. 1.1. 2=>= > 1,)),,)))
2ln 100d] S
d
So with probability at least 0.99, for all p= 1, 2, ..., we have ( In 100d.
Since with probability 8(e), 11 11 -1 Ella
I 6, have | 1 (1 +
e pre a.
Using Lemma 5, we have
PrIlI=ZA,, & Bl
(1 + 1 pre
1.+1
E 2
pre
(47)
(1 + 0rp |E L+,,i
+ c a0
which equals
Pr[l Pril pre
B0BBi=2 pre 2+1,1
> 1 1 a0
(48)
al(x 12
According to union bound, with probability 0.99 - -5(c) 1te-00 we have
OW3
O (A "LLL) 0
2d 100d
0
dlnd
0
Ind
an 2
aoL
So we have
Oc aW2.
a )
1 O Ind
(49)
aW2.
anL
Thus T
E d
B
12
< O(
Ind Ln
p.
aW
Take Og 10 we have that with probability least 0.99 8(c) 0.ste for the Post-LN Trans-
former we have a lle I O(dvInd) and for the Pre-LN Transformer we have || 12, ll <
až
Old
Ind I
19

---

Under review as a conference paper at ICLR 2020
F EXTENSION TO OTHER LAYERS
For simplicity, denote xі Co,, Fi,n) Rnd and x for k {1, 3, 5). Then in the Post-LN Transformer, the gradient of the parameters in the 1-th layer (take W as an example) can be written as
(2 z(,n) E Rnd
σC OW2,
ac
L duft
Ouffi OW2,
0xyt1 post Dujo
0ust1 Ox On On D On On dxy
II
where
dxp j=1+1
The Jacobian matrices of the Post-LN Transformer layers are:
J LL (1
Ox Ourpoat.
(50)
JLN (2gfart
W2.
J
WIa
Oxposst Or.
=
+
3
J
W1
(51)
)) Rddd
where J diag or
p J.1
(w))))
popet (wd" 1.j
J JIN (2000.2)
Onpast.3
(52)
J.N J
0ap p
1
tWVS
TAM! +
(53)
Dupost
LWVS
Using Hölder's inequality, we have
@xppr post
Ox
Dappost:
Ox post,3
Oxpostt
Ell grs
ll2 SE
||l||
Ilall
|12||
(54)
fast..
post
dz post,
On
0uj+1
On
Ox poot,3
On post,2
E
||
113 Ony 3ost,
E
(55)
Or
poot
Oz poot,
Da post
Since
0r,+1 Dugorts
diag (J LN post,5 ),... J LN (1. post.5));
we have
Bu Ür poxt,
3,1
when || poets
concentrates around its expectation E E|aa p
[|3 which equals
J.1
d according to Lemma 2. Therefore, when we estimate the norm of øW?7 for post-LN transformer,
BC
there exists term ( (L-11/22 ) which exponentially decreases as I goes smaller. Similarly, in the pre-LN Transformer, the gradient can be written as
ac ac Ox Final OW2J Peraa Oupre
Oxere Pre
pre Oxpre
Oxpre Ox
Oxpre Oxgre..3 Oxpr On
II
where
OWVD
j=1+1
The Jacobian matrices of the Pre-LN Transformer layers
20

---

for Under review as a conference paper at ICLR 2020
W22
+
I
W2.
W13
J LLN JLN (2ftn
wis
J JLN (aggr
Wws
Awvs
JIN(
Dufren Oulfre
WVV
wwvs
JLN (afmn)
(56)
If l sufficiently large, the norm of J LN (1gg)) and J JLN (aff are very small (of order O
(3))
as j is between / + and L, which means the eigenvalues of matrix du Duf
Duprais au are close
and
to 1. Then we can see that E]
Bagrers |l and [E]
axpra
ac l2 are nearly 1, and the norm aw2r Bejre
pre-LN transformer is independent of / when is large.
G EXAMPLES OF ( BOOUNDDD RANDOM VARIABLES
In this section we give an example (c, bbuunddd random variable. This example comes from Example 2.5 in (Wainwright) 2019) and give short description below.
Z = (Z1,... Zn) is a Gaussian vector with distribution N (0, In), then Y = || Σソ= 2 has distribution x2. And EY = Σi=1 EZI = n2
A random variable X with mean μ B[X] is called sub-exponential if there are non-negative parameters (v,a) a)) such that Ejexp(A(X - - exp( x²x² for all >| < The next proposition
comes from Proposition 2.2 in (Wainwright 2019).
Proposition 1 (Sub-exponential tail bound). Suppose that X subbexxonential with parameters (u, a). Then
P[X- - / 2 exp( 2 2
f S < 2 and fort > a
(57)
exp(
and from Example 2.5 in WWainwrigt 2019), the x2 variable Y is sub-exponential with parameters (u, a) (2 4). So we can derive the one-sided bound
P|Y 2 nd] I exp(-nc²/8) for all (0,1) So Y is (e, -bbounde with E (0, 1) and exp(-ne2/8)
(58)
H EMPIRICAL VERIFICATION OF THE THEORETICAL FINDINGS
As our theory is derived based on several simplifications of the problem, we conduct experiments to study whether our theoretical insights on the gradients are consistent with what we observe in real scenarios. We empirically study the gradients initialization for both Post-LN/Pre-LNN Transformer on the IWSLT14 De-En task. The general model and training configuration exactly follow Section 3.2. The experiments are repeated ten times using different random seeds.
Empirical verification of concentration property Given initialized model, we record the hid- den states in the Post-LN/Pre- Transformer and find that the norm of the hidden states satisfies the concentration property (0.1,0.125)-bounne,
21

---

Under review as a conference paper at ICLR 2020
0.30
Fro UN
Fom LN
Gradient Expeciation
Gradient Expectation
0.25
1.5
0.20
0.16
1.0
0.5
0.06
0.00
0.0
1D
#Layer
#Laaee
(a) Pre-LN Transformer
(b) Post-LN Transformer
Figure 6: Norm of expected gradients of W in the last FFN sub-layer in different size of the Transformer architecture
Empirical verification of Theorem 1 Post-LN Transformer, the scale of the gradient norm in the last FFN sub-layer remains the same. On the contrary, the scale of the gradient norm in the last FFN sub-layer of the Pre-LN Transformer decreases as the size (total depth) of the model grows.
Theorem the theory suggests that for any sizes of the
conduct experiments to study the gradient norm the last FFN sub-layer in different sizes of the Transformer architecture at initialization to verify Theorem 1. We train 6-6/8- 101--8/1010-1-- 10/00-0/12-1111 12/14- 14 Post-LN/Pre-LM. Transformer models, and record the gradient norm of the final FFN layer in different Transformer models. The results model, and the y-axis is the value of the gradient norm of W 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 theory.
plotted in Figure 6. The x-axis is the size of the
Empirical verification of extended theory We record the gradient for each parameter for differ- ent mini-batches. For elements a parameter matrix, we calculate their expected gradients and use the Frobenius norm of those values the scale of the expected gradient of the matrix. Figure (3a) (bb shows those statistics for FFN sub-layers. The x-axis indexes different Transformer layers. 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.
Given the analysis above, we think our derived theory at the initialization stage is useful and consis- tent with the empirical studies above.
I LARGE GRADIENTS IN POST-LN TRANSFORMER HURTS THE OPTIMIZATION
Theoretically, 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, conduct 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 fixed learning rate of le at the beginning of the optimization, which is much smaller than the Irmez = 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 Figure 7. We can see from the figure, the validation loss (pink curve) around 4.3 in 27 epochs. This loss is much lower than that of the
22

---

Under review as a conference paper at ICLR 2020
10
30
Validation Loss
BLEU
6
10
3
11
15 Epochs
13
19
21
23
27
5
11
13
15
21
23
25
Epochs
IFmax -, warm-up
Immax 1e
Tnarmup
Imax
Twammup
Irma 5e w/o warm-up
Irmex 5e
Tmarmup 500
Iraax
5e
Twarmup 4000
Irmax 10
warm-up
Irmax 1e w/o warm-up, gradient norm clipping
Figure 7: Performances of the models on the IWSLT14 De-En task.
Post-LN Transformer trained using 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 ex- periments, 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 verify our conclusion. Using small learning rate/clip norm mitigates the instability of updating models with large gradient values. Although the theory prove is based on some simplifications of the problem, the theoretical insights help us understand the optimization of different networks in practice.
23


Reviews:
[{"review_id": "SJxoWm9qdr", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nPros:\n  + A well-written paper with a good organization; notations are clear.\n  + The proof I checked seem correct (but I didn't check all of them).\n  + Good experimental design that compares the pre-LN and post-LN Transformers in different settings/tasks.\n  + 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.).\n\nI will detail the cons below, along with my other questions/concerns/issues.\n\n----------------------------------------\n\nQuestions/concerns:\n\n1. 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.\n\n2. 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).\n\n3. 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.\n\n4. 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?\n\n======================================\n\nSome minor issues that didn't impact the score:\n\n1. 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.\n\n2. Referred to the wrong equation numbers (mentioned in another comment from me). \n\n3. Appendix C: radius d -> radius \\sqrt{d}\n\n4. Appendix C: \"Similarly, we have \\| x_{l,i}^{post, 3} \\|_2^2 = d\" should have an expected value.\n\n======================================\n\nDespite 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.\n\n\n\n[1] https://arxiv.org/abs/1809.10853\n[2] https://arxiv.org/abs/1904.10509\n[3] https://arxiv.org/abs/1810.04805\n[4] https://arxiv.org/abs/1901.02860", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly."}, "extra_text": {}}, {"review_id": "BJlyv2xIqH", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "1. Specific problem tackled by the paper:\n\nMoving the LayerNorm layer to be inside the residual connection in a stack of transformers can remove the need for learning rate warm up. This paper provides a theoretical motivation for doing this.\n\n2. Motivation of the paper:\n\nThe authors motivate the problem clearly, performing experiments to demonstrate the problem. Their experiments provide good context for the problem they are solving and acts as a solid reference point for their (and other peoples' future) work. To be more convincing  the authors should have performed multiple runs and shown the standard deviations across runs.\n\n3. Claims of the paper:\n\nThe authors claim that layer norm should be placed inside the residual connection (Pre-LN) rather than outside (Post-LN) it. The authors show, theoretically, that in a Post-LN transformer the magnitude of the gradients in a transformer scale with the number of layers and that the magnitude of the hidden states scale linearly with the layers that that they are output by. While in a Pre-LN transformer the magnitudes of the gradients and states are independent of the number of layers.\n\ni.e. During training a Post-LN transformer is likely to have weaker gradients in the lower (closer to input layers) than at the output layers.\n\nThis theory is backed by an empirical study. It is good that these experiments are repeated ten times however the authors should show the standard deviations too.\n\nThe authors show machine translation results, demonstrating that using Pre-LN rather than Post-LN leads to faster convergence, however the models converge to the same result. However, there is benefit in not having to design a training schedule. Wang et al. have also shown the benefits of Pre-LN rather than Post-LN transformers for machine translation.\n\n4. Decision (accept or reject) with one or two key reasons for this choice and reasons for the decision.\n\nWeak accept.\n\nPros:\n(1) The authors provide theory that supports the use of Pre-LN rather than Post-LN transformers. Using Pre-LN rather than Post-LN transformers may save a lot of time by avoiding hyper-parameter tuning, without loss in performance and this looks very easy to implement.\n(2) The paper is well organised and the problem is very well motivated.\n(3) The experiments are sufficient. They could be improved with repeats and by showing the standard deviations (as mentioned above).\n\nCons:\n(1) This work is incremental. Wang et al. have already shown that Pre-LN transformers are better that Post-LN transformers when the network has many transformer layers, they explain theoretically why this is the case. \n(2) While the organisation of the paper is good, the paper is not well written. The gramma is poor.\n(3) The paper is very long for an incremental improvement.\n\n5. Additional feedback with the aim to improve the paper. \n\n[a] Ideally, it would be good to see repeats for Figure 2 and the standard deviations for Figure 3.\n[b] Without reading the appendix it is not clear where the assumption that W^Q and W^K are zero is used. Making some connection with how this assumption relates to the lemmas would be useful. Additionally, you should explain what this means qualitatively, because this makes the assumption more acceptable. I am assuming that it means that the attention is uniform.\n[c] In Lemma two you are comparing the magnitudes for the input in the Pre-LN and the output in the Post-LN transformer according to how x_{l,i}^post and x_{l, i}^pre are defined in Table 1.\n[d] In Figure 3(b) the gradients are clearly decreasing with the number of layers, are there any comments on this? In the limit this could cause vanishing gradients?\n[e] On the surface Figure 2 and 4 appear to contradict. Is the difference a result of using RAdam? If so, this should be made very clear. If not, why are the results contradictory?\n\nMore minor comments:\n[1] Many gramma errors.\n[2] Figure 1 is referenced before explaining what an FFN is. Figure 1 could also be enhanced by labelling Post-LN as previous work and Pre-LN as current work.\n[3] MultiHeadAtt is not well defined. Multi-Head( ), Attention( ) and Head( ) each take three arguments, while MutliHeadAtt( ) takes two. It would be worth connecting these.\n[4]\"sub-layer --> ...\" here sub-layer is not defined. Should this say self-attention sub-layer?\n[5] Where does equation (4) come from? There should be a citation and/or explanation.\n[6] The BLEU score is not defined (this could be with a footnote).\n[7] Top of page 7: \"As most of the parameters are initialized by Gaussian distributions\" --> Using the word \"most\" is very vague. The authors should be specific about which parameters they are referring to.\n[8] There is a good balance of equations in the main text with most of the proofs in the appendix.\n[9] Inconsistent use of LN and LayerNorm, both are used.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "BJgoO9NCtr", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\n=========Update========\nI appreciate the authors' response and additional appendix sections connecting theory and practice, verifying the claims about the final layer gradients scaling with L. I have also read the other reviews. I'm still not completely convinced that the multi-layer analysis well explains the reality, but because of the other contributions I am updating my score, which would be a 5/10 on the old ICLR scale.\n=======================\n\nThis paper investigates the placement of Layer Normalization in the transformer architecture. The authors show that the Pre-LN placement leads to better behaved gradients as the network gets larger. This in turn allows them to remove the warmup stage of the learning rate schedule, leading to a faster and simpler training procedure. They examine a simplification of the attention layer analytically, and show a scaling with the number of layers that occurs with the Pre-LN placement but not the Post-LN placement. Finally, they demonstrate the effectiveness of the transformer changes both for machine translation and in BERT pretraining.\n\nEven though the novelty here is limited, Pre-LN placement has been used in prior work, the potential for accelerating future research is large. In general I think the impact of this kind of research, exploring improvements to commonly used methods that add no additional complexity or even simplify, is underappreciated in the reviewing process. Still, I have some concerns about the relation between the analytic investigation of the gradient norms and the empirical results that are presented, and I am concerned that the analytical results are used to imply something stronger than they actually show.\n\nAt some level, it seems like the theoretical results have come along for the ride but do not clearly demonstrate that there is a problem with Post-LN and that this problem is fixed by switching to Pre-LN, or at least the relationship is not clearly explained. In the empirical study and central to the paper’s narrative is the fact that Post-LN normalization leads to a vanishing gradient problem where gradient in later layers is substantially larger in magnitude than in early layers at initialization, whereas the Pre-LN placement does not suffer from this problem. The point mentioned in the main text (theorem 1) is that the final FFN layer has gradient magnitude that scales at most like 1/sqrt(L) in the Pre-LN case and independent of L in the Post-LN case. This alone does not imply any of the observed empirical behavior because if all of the layers were scaled by this factor 1/sqrt(L) then the same problems would persist for the Pre-LN network. More relevant to the findings is how the scaling of the gradient changes throughout the layers which is examined in appendix section F, where the gradient norm of the Post-LN network can be upper bounded by an expression in which one of the terms scales approximately as (2/3)^(L-l)/2, whereas in the Pre-LN network the scaling is explained to be independent of layer index l. The connection between the expression that scales in the upper bound and the actual gradient norm is tenuous and there multiple places where the argument could break down (upper bound not being tight, scaling of other terms canceling out, validity of the simplifying assumptions being used). It would be useful to verify this sqrt(2/3) scaling on the data from the empirical study that is shown, is the decay shown in figure 3 geometric with this factor. Also the fact that the expectation near the bottom of page 20 is approximately 2/3 needs to be explained as it’s not obvious where that comes from.\n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper investigates layer normalization and learning rate warmup in transformers, demonstrating that placing layer norm inside the residual connection (pre-LN) leads to better behaved gradients than post-LN placement. Doing so allows the learning rate warm-up stage to be removed, leading to faster training.

Reviewers were mildly positive about the submission, commenting on the interesting insight provided about transformers, as well as the clear, focused motivation and contribution.

However they also stated that it seem rather incremental of a contribution, as pre-LN placement has been introduced before, and found it confusingly written at times.

R2 clearly read it very closely, and had many detailed comments and discussions with authors and other reviewers. They had concerns about the relationship of this work with gradient clipping. The authors deserve credit for quickly investigating this in further experiments. Interestingly, the found that even with gradient clipping, post-LN models still needed the learning rate warm-up stage, although this issue went away with smaller clipped values or much lower learning rates. Overall, R2 appears to find the paper’s motivation very compelling, but the insights incomplete and not fully satisfactory, while all reviewers find the novelty rather limited.

I think a future submission that forges closer connections between the empirical findings and the theoretical interpretations would be of a great interest to the community, but in its current form is probably unsuitable for publication at ICLR 2020.

Decision: Reject