Root Mean Square Layer Normalization
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Layer normalization (LayerNorm) has been successfully applied to various deepneural networks to help stabilize training and boost model convergence becauseof its capability in handling re-centering and re-scaling of both inputs and weightmatrix.
However, the computational overhead introduced by LayerNorm makesthese improvements expensive and signiﬁcantly slows the underlying network, e.g. RNN in particular.
In this paper, we hypothesize that re-centering invariance in LayerNorm is dispensable and propose root mean square layer normalization, or RMSNorm .
RMSNorm regularizes the summed inputs to a neuron in one layeraccording to root mean square (RMS), giving the model re-scaling invariance property and implicit learning rate adaptation ability. RMSNorm is computationallysimpler and thus more efﬁcient than LayerNorm.
We also present partial RMSNorm, or p RMSNorm where the RMS is estimated from p % of the summed inputswithout breaking the above properties.
Extensive experiments on several tasksusing diverse network architectures show that RMSNorm achieves comparableperformance against LayerNorm but reduces the running time by 7% ∼ 64% ondifferent models. We will release our source code soon.
Introduction
How to train deep neural networks efﬁciently is a long-standing challenge.
To accelerate modelconvergence
Ba et al.
[3] propose the layer normalization (LayerNorm) which stabilizes the trainingof deep neural networks by regularizing neuron dynamics within one layer via mean and variancestatistics.
Due to its simplicity and requiring no dependencies among training cases, LayerNormhas been widely applied to different neural architectures which enables remarkable success onvarious tasks ranging from computer vision [17, 24], speech recognition [34] to natural languageprocessing [29].
In some cases, LayerNorm was found to be essential for successfully training amodel [5].
Besides, the decouple from batch-based samples endows LayerNorm with the superiorityover batch normalization (BatchNorm) [11] in handling variable-length sequences using RNNs.
Unfortunately, the incorporation of LayerNorm raises computational overhead.
Although this isnegligible to small and shallow neural models with few normalization layers, this problem becomesclearly severe when underlying networks grow larger and deeper.
As a result, the efﬁciency gainfrom faster and more stable training (in terms of number of training steps) is counter-balanced by anincreased computational cost per training step, which diminishes the net efﬁciency, as show in Figure1.
One major feature of LayerNorm that is widely regarded as contributions to the stabilization is itsre-centering invariance property: the summed inputs after LayerNorm remain intact when the inputsor weight matrix is shifted by some amount of noise. We argue that this mean normalization does notreduce the variance of hidden states or model gradients, and hypothesize that it has little impact onthe success of LayerNorm.
In this paper, we propose root mean square layer normalization (RMSNorm), which regularizes thesummed inputs to a neuron in one layer with the root mean square (RMS) statistic alone. RMSNorm
[Figure] reduces the amount of computation and increases efﬁciency over LayerNorm.
Despite the simplerformulation, the RMS normalizer helps reduce covariate shift of layer activations, ensuring invarianceto the re-scaling of both weights and datasets.
We also show the possibility of estimating RMS on asubset of the summed inputs, maintaining this invariance property. Assuming that the summed inputshave an independent identically distributed structure, we propose partial RMSNorm, where only theﬁrst p % summed inputs are utilized for RMS estimation.
We thoroughly examine our model on various tasks, including machine translation, image-captionretrieval and question answering.
Experimental results show that across different models, p RMSNormyields comparable performance against LayerNorm but shows superiority in terms of running speedwith a speed-up of 7% ∼ 64%. When estimating the RMS with partial (6.25%) summed inputs,p RMSNorm achieves competitive performance compared to RMSNorm.
Related Work
Deep neural networks suffer from the internal covariate shift issue [25], where a layer’s inputdistribution changes as previous layers are updated, which signiﬁcantly slows the training.
Onepromising direction to solve this problem is normalization.
Ioffe and Szegedy [11] introduce batchnormalization (BatchNorm) to stabilize activations based on mean and variance statistics estimatedfrom each training mini-batch.
Unfortunately, the reliance across training cases deprives BatchNormof the capability in handling variable-length sequences, though several researchers develop differentstrategies to enable it in RNNs [14, 7].
Instead, Salimans and Kingma [20] propose weight normalization (WeightNorm) to reparameterize weight matrix so as to decouple the length of weight vectorsfrom their directions.
Ba et al.
[3] propose layer normalization which differs from BatchNorm in thatstatistics are directly estimated from the same layer without accessing other training cases.
Due to itssimplicity and effectiveness, LayerNorm has been successfully applied to various deep neural models,and achieves state-of-the-art performance on different tasks [17, 34, 29, 5].
These studies pioneer the research direction that integrates normalization as a part of the modelarchitecture.
This paradigm ensures encouraging performance by shorting model convergence butat the cost of consuming more time for each running step.
To improve efﬁciency, Arpit et al.
[2]employ a data-independent method to approximately estimate mean and variance statistics, thusavoiding calculating batch statistics.
Ioffe [10] propose batch renormalization so as to reduce thedependence of mini-batches in BatchNorm.
Ulyanov et al.
[28] replace batch normalization withinstance normalization for image generation.
Hoffer et al.
[9] and Wu et al.
[31] observe that l 1 -normcan act as an alternative of variance in BatchNorm with the beneﬁt of fewer nonlinear operations andhigher computational efﬁciency.
Nevertheless, all these work still follow the original normalizationstructure and utilize mean statistic estimated from the whole summed inputs to handle re-centeringinvariance. Different from these related work, the proposed RMSNorm modiﬁes the normalization structure byremoving the re-centering operation and regularizing the summed inputs with RMS alone.
Our model only maintains the re-scaling invariance property which we ﬁnd can be inherited when the RMS isestimated from only subset of the summed inputs, partially inspired by the group normalization [32]. As a side effect, our model reduces the computational overhead and increases efﬁciency. Recently, Zhang et al.
[33] show that with careful initialization, residual networks can be trained as stable asthose with normalization. However, the approach mainly aims at improving residual networks andcan not be freely switched without modifying all initialization layers. Besides, it is not trivial tobe adapted to other general neural networks, such as RNNs where model depth expands along thevariable sequence length. By contrast, our model is simple, effective and can be used as a drop-inreplacement of LayerNorm.
Background
We brieﬂy review LayerNorm in this section based on a standard feed-forward neural network.
Givenan input vector x ∈ R m , a feed-forward network projects it into an output vector y ∈ R n through alinear transformation followed by a non-linear activation as follows:
[Equation]
where W is weight matrix, b is bias term which is usually initialized by 0, and f ( · ) is an element-wisenon-linear function.
a ∈ R n denotes the weight-summed inputs to neurons, which is also the targetof normalization.
This vanilla network suffers from internal covariate shift issue [11], where a layer’s input distributionchanges as previous layers are updated.
This could negatively affect the stability of parameters’gradients, delaying model convergence.
To reduce this shift, LayerNorm normalizes the summedinputs so as to ﬁx their mean and variance as follows:
[Equation]
where ¯ a i is the i -th value of vector ¯a ∈ R n , which acts as the normalized alternative of a i for layeractivation.
g ∈ R n is the gain parameter used to re-scale the standardized summed inputs, and is setto 1 at the beginning.
µ and σ 2 are the mean and variance statistic respectively estimated from rawsummed inputs a :
[Equation]
In this way, LayerNorm forces the norm of neurons to be decoupled from both the inputs and weightmatrix.
RMSNorm
A well-known explanation of the success of LayerNorm is its re-centering and re-scaling invarianceproperty. The former enables the model to be insensitive to shift noises on both inputs and weights,and the latter keeps the output representations intact when both inputs and weights are randomlyscaled. In this paper, we hypothesize that the re-scaling invariance is the reason for success of LayerNorm, rather than re-centering invariance.
We propose RMSNorm which only focuses on re-scaling invariance and regularizes the summedinputs simply according to the root mean square (RMS) statistic:
[Equation]
where RMS ( a
[Equation]
= (cid:118)(cid:117)(cid:117)(cid:116)
1 n n (cid:88) i = a 2 i .
Intuitively, RMSNorm simpliﬁes LayerNorm by totally removing the mean statistic in Eq.
(3) atthe cost of sacriﬁcing the invariance that mean normalization affords.
When the mean of summedinputs is zero, RMSNorm is exactly equal to LayerNorm.
Although RMSNorm does not re-centerthe summed inputs as in LayerNorm, we demonstrate through experiments that this property is notfundamental to the success of LayerNorm, and that RMSNorm is similarly effective. [Table]
RMS measures the quadratic mean of inputs, which in RMSNorm forces the summed inputs intoa √ n -scaled unit sphere.
By doing so, the output distribution remains regardless of the scaling ofinput and weight distributions, tackling the issue of internal covariate shift.
Although Euclidean normwhich only differs from RMS by a factor of √ n has been successfully explored [20], we empiricallyﬁnd that it does not work for layer normalization.
We hypothesize that scaling the sphere with thesize of the input vector is important because it makes the normalization more robust across vectors ofdifferent size.
As far as we know, the idea of employing RMS for neural network normalization hasnot been investigated before.
4.1 Invariance Analysis
Invariance measures whether model output after normalization changes highly in accordance withits input and weight matrix.
Ba et al.
[3] show that different normalization methods reveal differentinvariance properties, which contributes considerably to the model’s robustness.
In this section, wetheoretically examine the invariance properties of RMSNorm.
We consider the following general form of RMSNorm:
[Equation]
where (cid:12) denotes element-wise multiplication.
Our main results are summarized in Table 1.
RMSNorm is invariant to both weight matrix and input re-scaling, because of the following linearityproperty of RMS:
[Equation]
where α is a scale value.
Suppose the weight matrix is scaled by a factor of δ , i.e. W (cid:48)
=
δ W , thenthis change does not affect the ﬁnal layer output:
[Equation]
By contrast, if the scaling is only performed on individual weight vectors, this property does not holdanymore as different scaling factors break the linearity property of RMS.
Similarly, if we enforcea scale on the input with a factor of δ , i.e. x (cid:48) =
δ x , the output of RMSNorm remains through ananalysis analogous to that in Eq.
We can easily extend the equality to batch-based inputs as well asthe whole dataset.
Therefore, RMSNorm is invariant to the scaling of its inputs.
The main difference to LayerNorm is that RMSNorm is not re-centered and thus does not showsimilar linearity property for variable shifting.
It is not invariant to all re-centering operations.
4.2 Gradient Analysis
The above analysis only considers the effect of scaling inputs and the weight matrix on the layeroutput.
In a general setting, however, a RMSNorm-enhanced neural network is trained via standardstochastic gradient descent approach, where the robustness of model gradient is very crucial toparameters’ update and model convergence (see also Santurkar et al.
[21] who argue that the successof normalization methods does not come from the added stability to layer inputs, but due to increasedsmoothness of the optimization landscape).
In this section, we investigate the properties of modelgradients in RMSNorm.
Given a loss function L , we perform back-propagation through Eq.
(4) to obtain the gradient withrespect to parameters g , b as follows:
[Equation]
where v is short for the whole expression inside f ( · ) in Eq.
(4), and ∂ L / ∂ v is the gradient backpropagated from L to v .
Both gradients ∂
L / ∂ b and ∂ L / ∂ g are invariant to the scaling of inputs x andthe weight matrix W
(in the case of ∂ L / ∂ g because of the linearity property in Eq.
Besides, thegradient of g is proportional to the normalized summed inputs, rather than raw inputs.
This powersthe stability of the magnitude of g .
Unlike these vector parameters, the gradient of the weight matrix W is more complicated due to thequadratic computation in RMS.
Formally,
[Equation]
i , where
[Equation]
=
[Equation]
diag ( · ) denotes the diagonal matrix of input, ⊗ denotes the Kronecker product, and “ I ” indicatesidentity matrix.
For clarity, we explicitly use “ × ” to represent matrix multiplication.
The matrix term R associates the gradient of W with both inputs x and weight matrix W .
With a thorough analysis,we can demonstrate that this term is negatively correlated with both input and weight matrix scaling.
After assigning a scale of δ to either input x ( x (cid:48)
=
δ x ) or weight matrix ( W (cid:48) =
δ W ), we have
[Equation]
If we put the scaled term R (cid:48) back into Eq.
(9), we can easily prove that the gradient ∂ L / ∂ W isinvariant to input scaling, but keeps the negative correlation with weight matrix scaling.
Reducingthe sensitivity of gradient ∂ L / ∂ W to the scaling of inputs ensures its smoothness and improves thestability of learning.
On the other hand, the negative correlation acts as an implicit learning rateadaptor and dynamically controls the norm of gradients which avoids large-norm weight matrix andimproves model convergence.  p RMSNorm The re-scaling invariance property of RMSNorm ascribes to the linearity property of RMS. Considering that neurons in one layer often have independent identically distributed structure, we argue thatthe RMS can be estimated on a subset of these neurons rather than all of them.
We propose partial
RMSNorm ( p RMSNorm). Given the unnormalized input a , p RMSNorm infers the RMS statisticfrom ﬁrst- p % elements of a : RMS ( a ) = (cid:113) 1 m (cid:80) mi =1 a 2 i , where m = (cid:100) n · p (cid:101) denotes the number ofelements used for RMS estimation. Obviously, the linearity property still holds for the estimated
RMS, which indicates p RMSNorm shares the same invariance properties as RMSNorm as shown in Table 1.
RMS is a biased estimation of the RMS which is often inaccurate.
Though theoretically p RMSNormapproximates to RMSNorm, we observe gradient instability where the gradient tends to explode withsmall m . In practice, however, models with p RMSNorm can succeed in satisfactory convergencewith a partial ratio of 6.25%.
To test the efﬁciency of layer normalization across different implementations, we perform experimentswith Tensorﬂow [1], PyTorch [18] and Theano [27].
We add RMSNorm to different models, comparing against an unnormalized baseline and LayerNorm.
Unless otherwise noted, all speed-relatedstatistics are measured on one TITAN X (Pascal).
Reported time is averaged over 3 runs.
6.1 Machine Translation
Machine translation aims at transforming a sentence from one (source) language to another (target)language.
We focus on neural machine translation based on an attention-enhanced encoder-decoderframework.
We train two different models, a GRU-based RNNSearch [4] and a self-attentionbased neural Transformer [29] on WMT14 English-German translation task.
More details about theexperimental settings are {+listed in Appendix A.
+}
We ﬁrst experiment with RNNSearch.
In this experiment, we set the embedding size and hidden sizeto be 512 and 1024 respectively.
Normalization is added to the recurrent connections and feedforward
[Figure] [Figure] [Table] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table] [Figure] [Table] [Figure]
[Figure]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table] [Table] layers. Apart from RNNSearch without any normalization (Baseline) and with LayerNorm, we alsocompare against the same model equipped with L2-Norm (i.e. replacing RMS with L2-Norm), whichhas been observed to improve lexical selection [16].
Figure 2 illustrates the evolution of BLEU score on our development set after every 30k trainingsteps, and Table 2 summarizes the test results.
In short, both LayerNorm and RMSNorm outperformthe Baseline by accelerating model convergence: they reduce the number of training steps until convergence by about 50%, and improve test accuracy, with RMSNorm being comparable to LayerNorm. This supports our hypothesis that re-scaling invariance is the core property of LayerNorm, and that RMSNorm is an effective substitute.
Our results with L2-Norm show that it fails to improve themodel.
1 Results in Table 2 highlight the challenge that RNN with LayerNorm in Tensorﬂow suffersfrom serious computational inefﬁciency, where LayerNorm is slower than the Baseline by about 67%. In this respect, RMSNorm performs signiﬁcantly better, improving upon LayerNorm by ∼ 25%. Table 3 further lists translation results of different models implemented in Theano and Pytorch.
Overall, RMSNorm yields comparable translation quality compared with LayerNorm but incurs lesscomputational overhead, outperforming LayerNorm with speedups ranging from 11% ∼ 34%. Inaddition, we observe that though in theory the amount of computation in p RMSNorm is less thanthat in RMSNorm, p RMSNorm ( p = 6 . 25% ) sometimes tends to be slower. We ascribe this to thenon-optimal implementation of tensor slicing operation in these computational frameworks, whichcan be improved with speciﬁc low-level coding. In p RMSNorm, the partial ratio p directly controls the accuracy of estimated RMS, thereby affectingthe stability of model training.
Figure 3 shows the effect of p on model performance. Surprisingly, weﬁnd that the scale of p has little inﬂuence on the ﬁnal translation quality in RNNSearch: using a smallratio does not signiﬁcantly degenerate BLEU score. We set p to 6.25% for all following experiments. We also experiment with Transformer, which is based on self-attention, avoiding recurrent connections and allowing a higher degree of parallelization.
Still, layer normalization is an important part ofthe architecture.
We use an in-house Tensorﬂow implementation of the Transformer, and employ thebase setting as in [29] with all models trained for 300K steps. We treat Transformer with no normal- [Table]
[Table] [Figure] [Table] ization as our Baseline, and compare RMSNorm-enhanced Transformer with LayerNorm-equipped
Transformer.
Table 4 shows the results, from which we observe the importance of normalizationfor Transformer, without which training fails.
RMSNorm achieves BLEU scores comparable to LayerNorm, and yields a speedup of 7% ∼ 9%.
Compared with RNNSearch, the relative cost ofnormalization is lower because there are signiﬁcantly fewer sequential normalization operations in Transformer.
Effect of Normalization on Mean and Standard Deviation Table 5 shows the distribution of meanand standard deviation of hidden representations across token positions for an RNNSearch model.
Mean and standard deviation are unstable in the baseline, as observed by Ba et al. Due to theirnormalization properties, both RMSNorm and LayerNorm stabilize standard deviation. Although themean in RMSNorm is not normalized, in practice it is more stable than the mean of the baseline. Thissupports our hypothesis that RMSNorm stabilizes recurrent activations without the need to explicitlynormalize the mean.
6.2 CNN/Daily Mail Reading Comprehension
This reading comprehension task is a cloze-style question answering task, where models are requiredto answer a question regarding to a passage, and the answer is an anonymized entity from thepassage [8].
We train a bidirectional attentive reader model proposed by Hermann et al.
[8] on the CNN corpus.
More details about the experimental settings are given in Appendix A.2.
We compare
RMSNorm with both LayerNorm and BatchNorm.
Figure 4 and Table 6 show the results.
After normalizing RNN by BatchNorm with separate statisticsfor each time step in a sequence, both BatchNorm-LSTM and BatchNorm-Everywhere help speed upthe convergence of training process.
By contrast, LayerNorm and RMSNorm not only converge fasterthan BatchNorm, but also reach lower validation error rate, though p RMSNorm performs slightlyworse than RMSNorm.
Although in Figure 4 the performance of RMSNorm and LayerNorm iscomparable, RMSNorm is around 15% faster than LayerNorm as shown in Table {+6. 
+}
6.3 Image-Caption Retrieval
Image-caption retrieval is a cross-modal task aiming at learning a joint embedding space of imagesand sentences, which consists of two sub-tasks: image retrieval and caption retrieval .
The formerranks a set of images according to a query caption, and the latter ranks a set of captions basedon a query image.
We train an order-embedding model (OE) proposed by Vendrov et al. [30] on [Figure] [Figure]
Sym
[Figure] [Figure] [Figure] Existing
OE [30] 46.
[Figure]
[Figure]
[Figure]
[Figure]
Work
OE [3] 46.
[Figure]
[Figure]
[Figure]
[Figure]
9 OE + LayerNorm [3] 48.
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
OE + Baseline
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
7.7 This
OE + LayerNorm 47.
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
7.5 Work
OE + RMSNorm 48.
[Figure]
[Figure] [Figure] [Figure] [Figure]
[Figure] [Figure] 7.5 OE + p RMSNorm
[Table] the Microsoft COCO dataset [15] using their public source code in Theano. Model details aboutexperimental settings are provides in Appendix A.3.
We compare RMSNorm with two models: onewithout any normalization (Baseline) and one with LayerNorm.
[Table]
Figure 5 shows the R@K curve on validation set after every300 training steps, and Table 7 lists the ﬁnal test results.
Acrossall these metrics, RMSNorm and LayerNorm consistently outperform the Baseline in terms of model convergence as shownin Figure 5.
We observe that on the validation set, RMSNormslightly exceeds LayerNorm with respect to recall value.
Forthe ﬁnal test results as shown in Table 7, both RMSNorm and
LayerNorm improve the model performance, reaching higherrecall values (except LayerNorm on R@5) and lower mean rank, though RMSNorm reveals bettergeneralization than LayerNorm.
Besides, results in Table 8 show that RMSNorm accelerates trainingspeed by 40% ∼ 64% compared with LayerNorm, highlighting better efﬁciency of p RMSNorm.
6.4 Conclusion and Future Work
This paper presents RMSNorm, a novel normalization approach that normalizes the summed inputsaccording to the RMS.
RMSNorm preserves the re-scaling invariance property of LayerNorm buteschews the re-centering invariance property which contributes less to the model training.
Comparedwith LayerNorm, models with RMSNorm suffers from less computational overhead.
RMSNorm canbe easily applied to different model architectures as a drop-in replacement of LayerNorm.
Experimentson several NLP tasks show that RMSNorm is comparable to LayerNorm in quality, but acceleratesthe running speed.
Actual speed improvements depend on the framework, hardware, neural networkarchitecture and relative computational cost of other components, and we empirically observedspeedups of 7% ∼ 64% across different models and implementations.
Our efﬁciency improvementcome from simplifying the computation, and we thus expect them to be orthogonal to other meansof increasing training speed, such as low-precision arithmetic and GPU kernel fusion. We alsoexperimented with p RMSNorm which estimates the RMS on a subset of the summed inputs. Whiletheoretically faster, we did not observe empirical speed improvements for p RMSNorm. We leave itto future work to investigate if the performance can be improved via code optimization.
In the future, we would like to take more analysis about the success behind RMSNorm.
Inspired byrecent success of l 1 -norm for BatchNorm, we are also interested in exploring different norms for
RMSNorm, and in simplifying other normalization techniques such as BatchNorm.
A AppendixA.1 Machine Translation
We experiment on the WMT14 English-German translation task, where the training corpus consistsof 4.5M aligned sentence pairs.
We use newstest2013 as the development set for model selection,newstest2014 and newstest2017 as the test set.
We evaluate translation quality with case-sensitivedetokenized BLEU score reported by sacrebleu [19] 3 .
Byte pair encoding algorithm is applied toreduce out-of-vocabulary tokens with 32k merge operations [22].
All models are trained based onmaximum log-likelihood relaxed by label smoothing with a factor of 0.1.
A.2 CNN/Daily Mail Reading Comprehension
The model is trained on the CNN corpus based on the public source code in Theano [7].
We adoptthe top4 setting, where each passage in the pre-processed dataset contains at most 4 sentences.
Forfair comparison with LayerNorm, we only employ RMSNorm within LSTM.
We set hidden size of
LSTM to be 240.
Models are optimized via Adam optimizer [12] with a batch size of 64 and learningrate of 8 e − 5 .
A.3 Image-Caption Retrieval
In OE model, sentences are encoded through a GRU-based RNN [6] and images are represented bythe output of a pretrained VGGNet [26].
OE treats the caption-image pairs as a two-level partialorder, and trains the joint model using the pairwise ranking loss [13].
We adopt the 10crop feature from VGGNet as image representation, and set word embedding size and
GRU hidden size to be 300 and 1024 respectively.
All models are trained with Adam optimizer, with abatch size of 128 and learning rate of 1 e − 3 .
We employ Recall@K (R@K) values for evaluation, andreport averaged results on ﬁve separate test sets (each consisting of 1000 images and 5000 captions)as our ﬁnal test results.
[Table]