Layer Normalization

Abstract

Training state-of-the-art, deep neural networks is computationally expensive. One way to reduce the training time is to normalize the activities of the neurons. A recently introduced technique called batch normalization uses the distribution of the summed input to a neuron over a mini-batch of training cases to compute a mean and variance which are then used to normalize the summed input to that neuron on each training case. This significantly reduces the training time in feed-forward neural networks. However, the effect of batch normalization is dependent on the mini-batch size and it is not obvious how to apply it to recurrent neural networks. In this paper, we transpose batch normalization into layer normalization by computing the mean and variance used for normalization from all of the summed inputs to the neurons in a layer on a single training case. Like batch normalization, we also give each neuron its own adaptive bias and gain which are applied after the normalization but before the non-linearity. Unlike batch normalization, layer normalization performs exactly the same computation at training and test times. It is also straightforward to apply to recurrent neural networks by computing the normalization statistics separately at each time step. Layer normalization is very effective at stabilizing the hidden state dynamics in recurrent networks. Empirically, we show that layer normalization can substantially reduce the training time compared with previously published techniques.

1 Introduction

Deep neural networks trained with some version of Stochastic Gradient Descent have been shown to substantially outperform previous approaches on various supervised learning tasks in computer vision [Krizhevsky et al., 2012] and speech processing [Hinton et al., 2012]. But state-of-the-art deep neural networks often require many days of training. It is possible to speed-up the learning by computing gradients for different subsets of the training cases on different machines or splitting the neural network itself over many machines [Dean et al., 2012], but this can require a lot of communication and complex software. It also tends to lead to rapidly diminishing returns as the degree of parallelization increases. An orthogonal approach is to modify the computations performed in the forward pass of the neural net to make learning easier. Recently, batch normalization [Ioffe and Szegedy, 2015] has been proposed to reduce training time by including additional normalization stages in deep neural networks. The normalization standardizes each summed input using its mean and its standard deviation across the training data. Feedforward neural networks trained using batch normalization converge faster even with simple SGD. In addition to training time improvement, the stochasticity from the batch statistics serves as a regularizer during training.

Despite its simplicity, batch normalization requires running averages of the summed input statistics. In feed-forward networks with fixed depth, it is straightforward to store the statistics separately for each hidden layer. However, the summed inputs to the recurrent neurons in a recurrent neural network (RNN) often vary with the length of the sequence so applying batch normalization to RNNs appears to require different statistics for different time-steps. Furthermore, batch normalization cannot be applied to online learning tasks or to extremely large distributed models where the minibatches have to be small.

This paper introduces layer normalization, a simple normalization method to improve the training speed for various neural network models. Unlike batch normalization, the proposed method directly estimates the normalization statistics from the summed inputs to the neurons within a hidden layer so the normalization does not introduce any new dependencies between training cases. We show that layer normalization works well for RNNs and improves both the training time and the generalization performance of several existing RNN models.

2 Background

A feed-forward neural network is a non-linear mapping from a input pattern $𝕩$ to an output vector $y$. Consider the $l^{th}$ hidden layer in a deep feed-forward, neural network, and let $a^l$ be the vector representation of the summed inputs to the neurons in that layer. The summed inputs are computed through a linear projection with the weight matrix $W^l$ and the bottom-up inputs $h^l$ given as follows:

$$a_i^l=w_i^l{}_{}{}^{\top }{h}^l{h}_i^{l+1}=f(a_i^l+b_i^l)$$ (1)

where $f(\cdot )$ is an element-wise non-linear function and $w_i^l$ is the incoming weights to the $i^{th}$ hidden units and $b_i^l$ is the scalar bias parameter. The parameters in the neural network are learnt using gradient-based optimization algorithms with the gradients being computed by back-propagation.

One of the challenges of deep learning is that the gradients with respect to the weights in one layer are highly dependent on the outputs of the neurons in the previous layer especially if these outputs change in a highly correlated way. Batch normalization [Ioffe and Szegedy, 2015] was proposed to reduce such undesirable “covariate shift”. The method normalizes the summed inputs to each hidden unit over the training cases. Specifically, for the $i^{th}$ summed input in the $l^{th}$ layer, the batch normalization method rescales the summed inputs according to their variances under the distribution of the data

$${\overline{a}}_i^l=\frac{g_i^l}{{\sigma }_i^l}\left(a_i^l-{\mu }_i^l\right){\mu }_i^l=\underset{𝕩\sim P(𝕩)}{𝔼}\left[a_i^l\right]{\sigma }_i^l=\sqrt{\underset{𝕩\sim P(𝕩)}{𝔼}\left[{\left(a_i^l-{\mu }_i^l\right)}^2\right]}$$ (2)

where ${\overline{a}}_i^l$ is normalized summed inputs to the $i^{th}$ hidden unit in the $l^{th}$ layer and $g_i$ is a gain parameter scaling the normalized activation before the non-linear activation function. Note the expectation is under the whole training data distribution. It is typically impractical to compute the expectations in Eq. (2) exactly, since it would require forward passes through the whole training dataset with the current set of weights. Instead, $\mu$ and $\sigma$ are estimated using the empirical samples from the current mini-batch. This puts constraints on the size of a mini-batch and it is hard to apply to recurrent neural networks.

3 Layer normalization

We now consider the layer normalization method which is designed to overcome the drawbacks of batch normalization.

Notice that changes in the output of one layer will tend to cause highly correlated changes in the summed inputs to the next layer, especially with ReLU units whose outputs can change by a lot. This suggests the “covariate shift” problem can be reduced by fixing the mean and the variance of the summed inputs within each layer. We, thus, compute the layer normalization statistics over all the hidden units in the same layer as follows:

$${\mu }^l=\frac{1}{H}\sum _{i=1}^H{a}_i^l{\sigma }^l=\sqrt{\frac{1}{H}\sum _{i=1}^H{\left(a_i^l-{\mu }^l\right)}^2}$$ (3)

where $H$ denotes the number of hidden units in a layer. The difference between Eq. (2) and Eq. (3) is that under layer normalization, all the hidden units in a layer share the same normalization terms $\mu$ and $\sigma$, but different training cases have different normalization terms. Unlike batch normalization, layer normaliztion does not impose any constraint on the size of a mini-batch and it can be used in the pure online regime with batch size 1.

3.1 Layer normalized recurrent neural networks

The recent sequence to sequence models [Sutskever et al., 2014] utilize compact recurrent neural networks to solve sequential prediction problems in natural language processing. It is common among the NLP tasks to have different sentence lengths for different training cases. This is easy to deal with in an RNN because the same weights are used at every time-step. But when we apply batch normalization to an RNN in the obvious way, we need to to compute and store separate statistics for each time step in a sequence. This is problematic if a test sequence is longer than any of the training sequences. Layer normalization does not have such problem because its normalization terms depend only on the summed inputs to a layer at the current time-step. It also has only one set of gain and bias parameters shared over all time-steps.

In a standard RNN, the summed inputs in the recurrent layer are computed from the current input ${𝕩}^t$ and previous vector of hidden states ${𝕙}^{t-1}$ which are computed as ${𝕒}^t=W_{hh}{h}^{t-1}+W_{xh}{𝕩}^t$. The layer normalized recurrent layer re-centers and re-scales its activations using the extra normalization terms similar to Eq. (3):

$${𝕙}^t=f\left[\frac{𝕘}{{\sigma }^t}\odot \left({𝕒}^t-{\mu }^t\right)+𝕓\right]{\mu }^t=\frac{1}{H}\sum _{i=1}^H{a}_i^t{\sigma }^t=\sqrt{\frac{1}{H}\sum _{i=1}^H{\left(a_i^t-{\mu }^t\right)}^2}$$ (4)

where $W_{hh}$ is the recurrent hidden to hidden weights and $W_{xh}$ are the bottom up input to hidden weights. $\odot$ is the element-wise multiplication between two vectors. $𝕓$ and $𝕘$ are defined as the bias and gain parameters of the same dimension as ${𝕙}^t$.

In a standard RNN, there is a tendency for the average magnitude of the summed inputs to the recurrent units to either grow or shrink at every time-step, leading to exploding or vanishing gradients. In a layer normalized RNN, the normalization terms make it invariant to re-scaling all of the summed inputs to a layer, which results in much more stable hidden-to-hidden dynamics.

4 Related work

Batch normalization has been previously extended to recurrent neural networks [Laurent et al., 2015, Amodei et al., 2015, Cooijmans et al., 2016]. The previous work [Cooijmans et al., 2016] suggests the best performance of recurrent batch normalization is obtained by keeping independent normalization statistics for each time-step. The authors show that initializing the gain parameter in the recurrent batch normalization layer to 0.1 makes significant difference in the final performance of the model. Our work is also related to weight normalization [Salimans and Kingma, 2016]. In weight normalization, instead of the variance, the L2 norm of the incoming weights is used to normalize the summed inputs to a neuron. Applying either weight normalization or batch normalization using expected statistics is equivalent to have a different parameterization of the original feed-forward neural network. Re-parameterization in the ReLU network was studied in the Path-normalized SGD [Neyshabur et al., 2015]. Our proposed layer normalization method, however, is not a re-parameterization of the original neural network. The layer normalized model, thus, has different invariance properties than the other methods, that we will study in the following section.

5 Analysis

In this section, we investigate the invariance properties of different normalization schemes.

5.1 Invariance under weights and data transformations

The proposed layer normalization is related to batch normalization and weight normalization. Although, their normalization scalars are computed differently, these methods can be summarized as normalizing the summed inputs $a_i$ to a neuron through the two scalars $\mu$ and $\sigma$. They also learn an adaptive bias $b$ and gain $g$ for each neuron after the normalization.

$h_i=f({\displaystyle \frac{g_i}{{\sigma }_i}}\left(a_i-{\mu }_i\right)+b_i)$ (5)

Note that for layer normalization and batch normalization, $\mu$ and $\sigma$ is computed according to Eq. 2 and 3. In weight normalization, $\mu$ is 0, and $\sigma ={\Vert w\Vert }_2$.

Table 1 highlights the following invariance results for three normalization methods.

	Weight matrix	Weight matrix	Weight vector	Dataset	Dataset	Single training case
	re-scaling	re-centering	re-scaling	re-scaling	re-centering	re-scaling
Batch norm	Invariant	No	Invariant	Invariant	Invariant	No
Weight norm	Invariant	No	Invariant	No	No	No
Layer norm	Invariant	Invariant	No	Invariant	No	Invariant

Table 1: Invariance properties under the normalization methods.

Weight re-scaling and re-centering: First, observe that under batch normalization and weight normalization, any re-scaling to the incoming weights $w_i$ of a single neuron has no effect on the normalized summed inputs to a neuron. To be precise, under batch and weight normalization, if the weight vector is scaled by $\delta$, the two scalar $\mu$ and $\sigma$ will also be scaled by $\delta$. The normalized summed inputs stays the same before and after scaling. So the batch and weight normalization are invariant to the re-scaling of the weights. Layer normalization, on the other hand, is not invariant to the individual scaling of the single weight vectors. Instead, layer normalization is invariant to scaling of the entire weight matrix and invariant to a shift to all of the incoming weights in the weight matrix. Let there be two sets of model parameters $\theta$, ${\theta }^{\prime }$ whose weight matrices $W$ and $W^{\prime }$ differ by a scaling factor $\delta$ and all of the incoming weights in $W^{\prime }$ are also shifted by a constant vector $𝜸$, that is $W^{\prime }=\delta W+\mathrm{𝟏}{𝜸}^{\top }$. Under layer normalization, the two models effectively compute the same output:

${𝕙}^{\prime }=$ $f({\displaystyle \frac{𝕘}{{\sigma }^{\prime }}}\left(W^{\prime }𝕩-{\mu }^{\prime }\right)+𝕓)=f({\displaystyle \frac{𝕘}{{\sigma }^{\prime }}}\left((\delta W+\mathrm{𝟏}{𝜸}^{\top })𝕩-{\mu }^{\prime }\right)+𝕓)$

$=$ $f({\displaystyle \frac{𝕘}{\sigma }}\left(W𝕩-\mu \right)+𝕓)=𝕙.$ (6)

Notice that if normalization is only applied to the input before the weights, the model will not be invariant to re-scaling and re-centering of the weights.

Data re-scaling and re-centering: We can show that all the normalization methods are invariant to re-scaling the dataset by verifying that the summed inputs of neurons stays constant under the changes. Furthermore, layer normalization is invariant to re-scaling of individual training cases, because the normalization scalars $\mu$ and $\sigma$ in Eq. (3) only depend on the current input data. Let ${𝕩}^{\prime }$ be a new data point obtained by re-scaling $𝕩$ by $\delta$. Then we have,

$h_i^{\prime }=$ $f({\displaystyle \frac{g_i}{{\sigma }^{\prime }}}\left(w_i^{\top }{𝕩}^{\prime }-{\mu }^{\prime }\right)+b_i)=f({\displaystyle \frac{g_i}{\delta \sigma }}\left(\delta w_i^{\top }𝕩-\delta \mu \right)+b_i)=h_i.$ (7)

It is easy to see re-scaling individual data points does not change the model’s prediction under layer normalization. Similar to the re-centering of the weight matrix in layer normalization, we can also show that batch normalization is invariant to re-centering of the dataset.

5.2 Geometry of parameter space during learning

We have investigated the invariance of the model’s prediction under re-centering and re-scaling of the parameters. Learning, however, can behave very differently under different parameterizations, even though the models express the same underlying function. In this section, we analyze learning behavior through the geometry and the manifold of the parameter space. We show that the normalization scalar $\sigma$ can implicitly reduce learning rate and makes learning more stable.

5.2.1 Riemannian metric

The learnable parameters in a statistical model form a smooth manifold that consists of all possible input-output relations of the model. For models whose output is a probability distribution, a natural way to measure the separation of two points on this manifold is the Kullback-Leibler divergence between their model output distributions. Under the KL divergence metric, the parameter space is a Riemannian manifold.

The curvature of a Riemannian manifold is entirely captured by its Riemannian metric, whose quadratic form is denoted as $d{s}^2$. That is the infinitesimal distance in the tangent space at a point in the parameter space. Intuitively, it measures the changes in the model output from the parameter space along a tangent direction. The Riemannian metric under KL was previously studied [Amari, 1998] and was shown to be well approximated under second order Taylor expansion using the Fisher information matrix:

$d{s}^2$ $={\mathrm{D}}_{\mathrm{KL}}[P(y|𝕩;\theta )\parallel P(y|𝕩;\theta +\delta )]\approx {\displaystyle \frac{1}{2}}{\delta }^{\top }F(\theta )\delta ,$ (8)

$F(\theta )$ $=\underset{𝕩\sim P(𝕩),y\sim P(y|𝕩)}{𝔼}\left[{\displaystyle \frac{\partial \log P(y|𝕩;\theta )}{\partial \theta }}{{\displaystyle \frac{\partial \log P(y|𝕩;\theta )}{\partial \theta }}}^{\top }\right],$ (9)

where, $\delta$ is a small change to the parameters. The Riemannian metric above presents a geometric view of parameter spaces. The following analysis of the Riemannian metric provides some insight into how normalization methods could help in training neural networks.

5.2.2 The geometry of normalized generalized linear models

We focus our geometric analysis on the generalized linear model. The results from the following analysis can be easily applied to understand deep neural networks with block-diagonal approximation to the Fisher information matrix, where each block corresponds to the parameters for a single neuron.

A generalized linear model (GLM) can be regarded as parameterizing an output distribution from the exponential family using a weight vector $w$ and bias scalar $b$. To be consistent with the previous sections, the log likelihood of the GLM can be written using the summed inputs $a$ as the following:

$\log P(y|𝕩;w,b)$ $={\displaystyle \frac{(a+b)y-\eta (a+b)}{\varphi }}+c(y,\varphi ),$ (10)

$𝔼[y|𝕩]=f(a+b)$ $=f(w^{\top }𝕩+b),\text{Var}[y|𝕩]=\varphi f^{\prime }(a+b),$ (11)

where, $f(\cdot )$ is the transfer function that is the analog of the non-linearity in neural networks, $f^{\prime }(\cdot )$ is the derivative of the transfer function, $\eta (\cdot )$ is a real valued function and $c(\cdot )$ is the log partition function. $\varphi$ is a constant that scales the output variance. Assume a $H$-dimensional output vector $𝕪=[y_1,y_2,\mathrm{\cdots },y_H]$ is modeled using $H$ independent GLMs and $\log P(𝕪|𝕩;W,𝕓)={\sum }_{i=1}^H\log P(y_i|𝕩;w_i,b_i)$. Let $W$ be the weight matrix whose rows are the weight vectors of the individual GLMs, $𝕓$ denote the bias vector of length $H$ and $\mathrm{vec}(\cdot )$ denote the Kronecker vector operator. The Fisher information matrix for the multi-dimensional GLM with respect to its parameters $\theta ={[w_1^{\top },b_1,\mathrm{\cdots },w_H^{\top },b_H]}^{\top }=\mathrm{vec}({[W,𝕓]}^{\top })$ is simply the expected Kronecker product of the data features and the output covariance matrix:

$F(\theta )$ (12)

We obtain normalized GLMs by applying the normalization methods to the summed inputs $a$ in the original model through $\mu$ and $\sigma$. Without loss of generality, we denote $\overline{F}$ as the Fisher information matrix under the normalized multi-dimensional GLM with the additional gain parameters $\theta =\mathrm{vec}({[W,𝕓,𝕘]}^{\top })$:

(13)

${\chi }_i$ $=𝕩-{\displaystyle \frac{\partial {\mu }_i}{\partial w_i}}-{\displaystyle \frac{a_i-{\mu }_i}{{\sigma }_i}}{\displaystyle \frac{\partial {\sigma }_i}{\partial w_i}}.$ (14)

Implicit learning rate reduction through the growth of the weight vector: Notice that, comparing to standard GLM, the block ${\overline{F}}_{ij}$ along the weight vector $w_i$ direction is scaled by the gain parameters and the normalization scalar ${\sigma }_i$. If the norm of the weight vector $w_i$ grows twice as large, even though the model’s output remains the same, the Fisher information matrix will be different. The curvature along the $w_i$ direction will change by a factor of $\frac{1}{2}$ because the ${\sigma }_i$ will also be twice as large. As a result, for the same parameter update in the normalized model, the norm of the weight vector effectively controls the learning rate for the weight vector. During learning, it is harder to change the orientation of the weight vector with large norm. The normalization methods, therefore, have an implicit “early stopping” effect on the weight vectors and help to stabilize learning towards convergence.

Learning the magnitude of incoming weights: In normalized models, the magnitude of the incoming weights is explicitly parameterized by the gain parameters. We compare how the model output changes between updating the gain parameters in the normalized GLM and updating the magnitude of the equivalent weights under original parameterization during learning. The direction along the gain parameters in $\overline{F}$ captures the geometry for the magnitude of the incoming weights. We show that Riemannian metric along the magnitude of the incoming weights for the standard GLM is scaled by the norm of its input, whereas learning the gain parameters for the batch normalized and layer normalized models depends only on the magnitude of the prediction error. Learning the magnitude of incoming weights in the normalized model is therefore, more robust to the scaling of the input and its parameters than in the standard model. See Appendix for detailed derivations.

6 Experimental results

We perform experiments with layer normalization on 6 tasks, with a focus on recurrent neural networks: image-sentence ranking, question-answering, contextual language modelling, generative modelling, handwriting sequence generation and MNIST classification. Unless otherwise noted, the default initialization of layer normalization is to set the adaptive gains to $1$ and the biases to $0$ in the experiments.

6.1 Order embeddings of images and language

 

(a) Recall@1

(b) Recall@5

(c) Recall@10

Figure 1: Recall@K curves using order-embeddings with and without layer normalization.

MSCOCO
	Caption Retrieval				Image Retrieval
Model	R@1	R@5	R@10	Mean r	R@1	R@5	R@10	Mean r
Sym [Vendrov et al., 2016]	45.4		88.7	5.8	36.3		85.8	9.0
OE [Vendrov et al., 2016]	46.7		88.9	5.7	37.9		85.9	8.1
OE (ours)	46.6	79.3	89.1	5.2	37.8	73.6	85.7	7.9
OE + LN	48.5	80.6	89.8	5.1	38.9	74.3	86.3	7.6

Table 2: Average results across 5 test splits for caption and image retrieval. R@K is Recall@K (high is good). Mean r is the mean rank (low is good). Sym corresponds to the symmetric baseline while OE indicates order-embeddings.

In this experiment, we apply layer normalization to the recently proposed order-embeddings model of Vendrov et al. [2016] for learning a joint embedding space of images and sentences. We follow the same experimental protocol as Vendrov et al. [2016] and modify their publicly available code to incorporate layer normalization which utilizes Theano [Team et al., 2016]. Images and sentences from the Microsoft COCO dataset [Lin et al., 2014] are embedded into a common vector space, where a GRU [Cho et al., 2014] is used to encode sentences and the outputs of a pre-trained VGG ConvNet [Simonyan and Zisserman, 2015] (10-crop) are used to encode images. The order-embedding model represents images and sentences as a 2-level partial ordering and replaces the cosine similarity scoring function used in Kiros et al. [2014] with an asymmetric one.

We trained two models: the baseline order-embedding model as well as the same model with layer normalization applied to the GRU. After every 300 iterations, we compute Recall@K (R@K) values on a held out validation set and save the model whenever R@K improves. The best performing models are then evaluated on 5 separate test sets, each containing 1000 images and 5000 captions, for which the mean results are reported. Both models use Adam [Kingma and Ba, 2014] with the same initial hyperparameters and both models are trained using the same architectural choices as used in Vendrov et al. [2016]. We refer the reader to the appendix for a description of how layer normalization is applied to GRU.

Figure 1 illustrates the validation curves of the models, with and without layer normalization. We plot R@1, R@5 and R@10 for the image retrieval task. We observe that layer normalization offers a per-iteration speedup across all metrics and converges to its best validation model in 60% of the time it takes the baseline model to do so. In Table 2, the test set results are reported from which we observe that layer normalization also results in improved generalization over the original model. The results we report are state-of-the-art for RNN embedding models, with only the structure-preserving model of Wang et al. [2016] reporting better results on this task. However, they evaluate under different conditions (1 test set instead of the mean over 5) and are thus not directly comparable.

6.2 Teaching machines to read and comprehend

Figure 2: Validation curves for the attentive reader model. BN results are taken from [Cooijmans et al., 2016].

In order to compare layer normalization to the recently proposed recurrent batch normalization [Cooijmans et al., 2016], we train an unidirectional attentive reader model on the CNN corpus both introduced by Hermann et al. [2015]. This is a question-answering task where a query description about a passage must be answered by filling in a blank. The data is anonymized such that entities are given randomized tokens to prevent degenerate solutions, which are consistently permuted during training and evaluation. We follow the same experimental protocol as Cooijmans et al. [2016] and modify their public code to incorporate layer normalization which uses Theano [Team et al., 2016]. We obtained the pre-processed dataset used by Cooijmans et al. [2016] which differs from the original experiments of Hermann et al. [2015] in that each passage is limited to 4 sentences. In Cooijmans et al. [2016], two variants of recurrent batch normalization are used: one where BN is only applied to the LSTM while the other applies BN everywhere throughout the model. In our experiment, we only apply layer normalization within the LSTM.

The results of this experiment are shown in Figure 2. We observe that layer normalization not only trains faster but converges to a better validation result over both the baseline and BN variants. In Cooijmans et al. [2016], it is argued that the scale parameter in BN must be carefully chosen and is set to 0.1 in their experiments. We experimented with layer normalization for both 1.0 and 0.1 scale initialization and found that the former model performed significantly better. This demonstrates that layer normalization is not sensitive to the initial scale in the same way that recurrent BN is.

6.3 Skip-thought vectors

 
 

(a) SICK($r$)

(b) SICK(MSE)

(c) MR

(d) CR

(e) SUBJ

(f) MPQA

Figure 3: Performance of skip-thought vectors with and without layer normalization on downstream tasks as a function of training iterations. The original lines are the reported results in [Kiros et al., 2015]. Plots with error use 10-fold cross validation. Best seen in color.

Skip-thoughts [Kiros et al., 2015] is a generalization of the skip-gram model [Mikolov et al., 2013] for learning unsupervised distributed sentence representations. Given contiguous text, a sentence is encoded with a encoder RNN and decoder RNNs are used to predict the surrounding sentences. Kiros et al. [2015] showed that this model could produce generic sentence representations that perform well on several tasks without being fine-tuned. However, training this model is time-consuming, requiring several days of training in order to produce meaningful results.

Method	SICK($r$)	SICK($\rho$)	SICK(MSE)	MR	CR	SUBJ	MPQA
Original [Kiros et al., 2015]	0.848	0.778	0.287	75.5	79.3	92.1	86.9
Ours	0.842	0.767	0.298	77.3	81.8	92.6	87.9
Ours + LN	0.854	0.785	0.277	79.5	82.6	93.4	89.0
Ours + LN $†$	0.858	0.788	0.270	79.4	83.1	93.7	89.3

Table 3: Skip-thoughts results. The first two evaluation columns indicate Pearson and Spearman correlation, the third is mean squared error and the remaining indicate classification accuracy. Higher is better for all evaluations except MSE. Our models were trained for 1M iterations with the exception of ($†$) which was trained for 1 month (approximately 1.7M iterations)

In this experiment we determine to what effect layer normalization can speed up training. Using the publicly available code of Kiros et al. [2015] , we train two models on the BookCorpus dataset [Zhu et al., 2015]: one with and one without layer normalization. These experiments are performed with Theano [Team et al., 2016]. We adhere to the experimental setup used in Kiros et al. [2015], training a 2400-dimensional sentence encoder with the same hyperparameters. Given the size of the states used, it is conceivable layer normalization would produce slower per-iteration updates than without. However, we found that provided CNMeM is used, there was no significant difference between the two models. We checkpoint both models after every 50,000 iterations and evaluate their performance on five tasks: semantic-relatedness (SICK) [Marelli et al., 2014], movie review sentiment (MR) [Pang and Lee, 2005], customer product reviews (CR) [Hu and Liu, 2004], subjectivity/objectivity classification (SUBJ) [Pang and Lee, 2004] and opinion polarity (MPQA) [Wiebe et al., 2005]. We plot the performance of both models for each checkpoint on all tasks to determine whether the performance rate can be improved with LN.

The experimental results are illustrated in Figure  3. We observe that applying layer normalization results both in speedup over the baseline as well as better final results after 1M iterations are performed as shown in Table  3. We also let the model with layer normalization train for a total of a month, resulting in further performance gains across all but one task. We note that the performance differences between the original reported results and ours are likely due to the fact that the publicly available code does not condition at each timestep of the decoder, where the original model does.

6.4 Modeling binarized MNIST using DRAW

Figure 4: DRAW model test negative log likelihood with and without layer normalization.

We also experimented with the generative modeling on the MNIST dataset. Deep Recurrent Attention Writer (DRAW) [Gregor et al., 2015] has previously achieved the state-of-the-art performance on modeling the distribution of MNIST digits. The model uses a differential attention mechanism and a recurrent neural network to sequentially generate pieces of an image. We evaluate the effect of layer normalization on a DRAW model using 64 glimpses and 256 LSTM hidden units. The model is trained with the default setting of Adam [Kingma and Ba, 2014] optimizer and the minibatch size of 128. Previous publications on binarized MNIST have used various training protocols to generate their datasets. In this experiment, we used the fixed binarization from Larochelle and Murray [2011]. The dataset has been split into 50,000 training, 10,000 validation and 10,000 test images.

Figure 4 shows the test variational bound for the first 100 epoch. It highlights the speedup benefit of applying layer normalization that the layer normalized DRAW converges almost twice as fast than the baseline model. After 200 epoches, the baseline model converges to a variational log likelihood of 82.36 nats on the test data and the layer normalization model obtains 82.09 nats.

6.5 Handwriting sequence generation

Figure 5: Handwriting sequence generation model negative log likelihood with and without layer normalization. The models are trained with mini-batch size of 8 and sequence length of 500,

The previous experiments mostly examine RNNs on NLP tasks whose lengths are in the range of 10 to 40. To show the effectiveness of layer normalization on longer sequences, we performed handwriting generation tasks using the IAM Online Handwriting Database [Liwicki and Bunke, 2005]. IAM-OnDB consists of handwritten lines collected from 221 different writers. When given the input character string, the goal is to predict a sequence of x and y pen co-ordinates of the corresponding handwriting line on the whiteboard. There are, in total, 12179 handwriting line sequences. The input string is typically more than 25 characters and the average handwriting line has a length around 700.

We used the same model architecture as in Section (5.2) of Graves [2013]. The model architecture consists of three hidden layers of 400 LSTM cells, which produce 20 bivariate Gaussian mixture components at the output layer, and a size 3 input layer. The character sequence was encoded with one-hot vectors, and hence the window vectors were size 57. A mixture of 10 Gaussian functions was used for the window parameters, requiring a size 30 parameter vector. The total number of weights was increased to approximately 3.7M. The model is trained using mini-batches of size 8 and the Adam [Kingma and Ba, 2014] optimizer.

The combination of small mini-batch size and very long sequences makes it important to have very stable hidden dynamics. Figure 5 shows that layer normalization converges to a comparable log likelihood as the baseline model but is much faster.

6.6 Permutation invariant MNIST

In addition to RNNs, we investigated layer normalization in feed-forward networks. We show how layer normalization compares with batch normalization on the well-studied permutation invariant MNIST classification problem. From the previous analysis, layer normalization is invariant to input re-scaling which is desirable for the internal hidden layers. But this is unnecessary for the logit outputs where the prediction confidence is determined by the scale of the logits. We only apply layer normalization to the fully-connected hidden layers that excludes the last softmax layer.

Figure 6: Permutation invariant MNIST 784-1000-1000-10 model negative log likelihood and test error with layer normalization and batch normalization. (Left) The models are trained with batch-size of 128. (Right) The models are trained with batch-size of 4.

All the models were trained using 55000 training data points and the Adam [Kingma and Ba, 2014] optimizer. For the smaller batch-size, the variance term for batch normalization is computed using the unbiased estimator. The experimental results from Figure 6 highlight that layer normalization is robust to the batch-sizes and exhibits a faster training convergence comparing to batch normalization that is applied to all layers.

6.7 Convolutional Networks

We have also experimented with convolutional neural networks. In our preliminary experiments, we observed that layer normalization offers a speedup over the baseline model without normalization, but batch normalization outperforms the other methods. With fully connected layers, all the hidden units in a layer tend to make similar contributions to the final prediction and re-centering and re-scaling the summed inputs to a layer works well. However, the assumption of similar contributions is no longer true for convolutional neural networks. The large number of the hidden units whose receptive fields lie near the boundary of the image are rarely turned on and thus have very different statistics from the rest of the hidden units within the same layer. We think further research is needed to make layer normalization work well in ConvNets.

7 Conclusion

In this paper, we introduced layer normalization to speed-up the training of neural networks. We provided a theoretical analysis that compared the invariance properties of layer normalization with batch normalization and weight normalization. We showed that layer normalization is invariant to per training-case feature shifting and scaling.

Empirically, we showed that recurrent neural networks benefit the most from the proposed method especially for long sequences and small mini-batches.

Supplementary Material

Application of layer normalization to each experiment

This section describes how layer normalization is applied to each of the papers’ experiments. For notation convenience, we define layer normalization as a function mapping $LN:{}_{}{}^D\to {}^D$ with two set of adaptive parameters, gains $𝜶$ and biases $𝜷$:

$LN(𝐳;𝜶,𝜷)={\displaystyle \frac{(𝐳-\mu )}{\sigma }}\odot 𝜶+𝜷,$ (15)

$\mu ={\displaystyle \frac{1}{D}}{\displaystyle \sum _{i=1}^D}{z}_i,\sigma =\sqrt{{\displaystyle \frac{1}{D}}{\displaystyle \sum _{i=1}^D}{(z_i-\mu )}^2},$ (16)

where, $z_i$ is the $i^{th}$ element of the vector $𝐳$.

Teaching machines to read and comprehend and handwriting sequence generation

The basic LSTM equations used for these experiment are given by:

$\left(\begin{array}{c}{𝐟}_t\\ {𝐢}_t\\ {𝐨}_t\\ {𝐠}_t\end{array}\right)$ $=$ ${𝐖}_h{𝐡}_{t-1}+{𝐖}_x{𝐱}_t+b$ (17)

${𝐜}_t$ $=$ $\sigma ({𝐟}_t)\odot {𝐜}_{t-1}+\sigma ({𝐢}_t)\odot \text{tanh}({𝐠}_t)$ (18)

${𝐡}_t$ $=$ $\sigma ({𝐨}_t)\odot \text{tanh}({𝐜}_t)$ (19)

The version that incorporates layer normalization is modified as follows:

$\left(\begin{array}{c}{𝐟}_t\\ {𝐢}_t\\ {𝐨}_t\\ {𝐠}_t\end{array}\right)$ $=$ $LN({𝐖}_h{𝐡}_{t-1};{𝜶}_1,{𝜷}_1)+LN({𝐖}_x{𝐱}_t;{𝜶}_2,{𝜷}_2)+b$ (20)

${𝐜}_t$ $=$ $\sigma ({𝐟}_t)\odot {𝐜}_{t-1}+\sigma ({𝐢}_t)\odot \text{tanh}({𝐠}_t)$ (21)

${𝐡}_t$ $=$ $\sigma ({𝐨}_t)\odot \text{tanh}(LN({𝐜}_t;{𝜶}_3,{𝜷}_3))$ (22)

where ${𝜶}_i,{𝜷}_i$ are the additive and multiplicative parameters, respectively. Each ${𝜶}_i$ is initialized to a vector of zeros and each ${𝜷}_i$ is initialized to a vector of ones.

Order embeddings and skip-thoughts

These experiments utilize a variant of gated recurrent unit which is defined as follows:

$\left(\begin{array}{c}{𝐳}_t\\ {𝐫}_t\end{array}\right)$ $=$ ${𝐖}_h{𝐡}_{t-1}+{𝐖}_x{𝐱}_t$ (23)

${\widehat{𝐡}}_t$ $=$ $\text{tanh}({\mathrm{𝐖𝐱}}_t+\sigma ({𝐫}_t)\odot ({\mathrm{𝐔𝐡}}_{t-1}))$ (24)

${𝐡}_t$ $=$ $(1-\sigma ({𝐳}_t)){𝐡}_{t-1}+\sigma ({𝐳}_t){\widehat{𝐡}}_t$ (25)

Layer normalization is applied as follows:

$\left(\begin{array}{c}{𝐳}_t\\ {𝐫}_t\end{array}\right)$ $=$ $LN({𝐖}_h{𝐡}_{t-1};{𝜶}_1,{𝜷}_1)+LN({𝐖}_x{𝐱}_t;{𝜶}_2,{𝜷}_2)$ (26)

${\widehat{𝐡}}_t$ $=$ $\text{tanh}(LN({\mathrm{𝐖𝐱}}_t;{𝜶}_3,{𝜷}_3)+\sigma ({𝐫}_t)\odot LN({\mathrm{𝐔𝐡}}_{t-1};{𝜶}_4,{𝜷}_4))$ (27)

${𝐡}_t$ $=$ $(1-\sigma ({𝐳}_t)){𝐡}_{t-1}+\sigma ({𝐳}_t){\widehat{𝐡}}_t$ (28)

just as before, ${𝜶}_i$ is initialized to a vector of zeros and each ${𝜷}_i$ is initialized to a vector of ones.

Modeling binarized MNIST using DRAW

The layer norm is only applied to the output of the LSTM hidden states in this experiment:

The version that incorporates layer normalization is modified as follows:

$\left(\begin{array}{c}{𝐟}_t\\ {𝐢}_t\\ {𝐨}_t\\ {𝐠}_t\end{array}\right)$ $=$ ${𝐖}_h{𝐡}_{t-1}+{𝐖}_x{𝐱}_t+b$ (29)

${𝐜}_t$ $=$ $\sigma ({𝐟}_t)\odot {𝐜}_{t-1}+\sigma ({𝐢}_t)\odot \text{tanh}({𝐠}_t)$ (30)

${𝐡}_t$ $=$ $\sigma ({𝐨}_t)\odot \text{tanh}(LN({𝐜}_t;𝜶,𝜷))$ (31)

where $𝜶,𝜷$ are the additive and multiplicative parameters, respectively. $𝜶$ is initialized to a vector of zeros and $𝜷$ is initialized to a vector of ones.

Learning the magnitude of incoming weights

We now compare how gradient descent updates changing magnitude of the equivalent weights between the normalized GLM and original parameterization. The magnitude of the weights are explicitly parameterized using the gain parameter in the normalized model. Assume there is a gradient update that changes norm of the weight vectors by ${\delta }_g$. We can project the gradient updates to the weight vector for the normal GLM. The KL metric, ie how much the gradient update changes the model prediction, for the normalized model depends only on the magnitude of the prediction error. Specifically,

under batch normalization:

$d{s}^2={\displaystyle \frac{1}{2}}\mathrm{vec}{({[0,0,{\delta }_g]}^{\top })}^{\top }\overline{F}(\mathrm{vec}({[W,𝕓,𝕘]}^{\top })\mathrm{vec}({[0,0,{\delta }_g]}^{\top })={\displaystyle \frac{1}{2}}{\delta }_g^{\top }\underset{𝕩\sim P(𝕩)}{𝔼}\left[{\displaystyle \frac{\text{Cov}[𝕪|𝕩]}{{\varphi }^2}}\right]{\delta }_g.$ (32)

Under layer normalization:

$d{s}^2=$ ${\displaystyle \frac{1}{2}}\mathrm{vec}{({[0,0,{\delta }_g]}^{\top })}^{\top }\overline{F}(\mathrm{vec}({[W,𝕓,𝕘]}^{\top })\mathrm{vec}({[0,0,{\delta }_g]}^{\top })$

$=$ (33)

Under weight normalization:

$d{s}^2=$ ${\displaystyle \frac{1}{2}}\mathrm{vec}{({[0,0,{\delta }_g]}^{\top })}^{\top }\overline{F}(\mathrm{vec}({[W,𝕓,𝕘]}^{\top })\mathrm{vec}({[0,0,{\delta }_g]}^{\top })$

$=$ (34)

Whereas, the KL metric in the standard GLM is related to its activities $a_i=w_i^{\top }𝕩$, that is depended on both its current weights and input data. We project the gradient updates to the gain parameter ${\delta }_{g}i$ of the $i^{th}$ neuron to its weight vector as ${\delta }_{gi}\frac{w_i}{{\Vert w_i\Vert }_2}$ in the standard GLM model:

${\displaystyle \frac{1}{2}}\mathrm{vec}{({[{\delta }_{gi}{\displaystyle \frac{w_i}{{\Vert w_i\Vert }_2}},0,{\delta }_{gj}{\displaystyle \frac{w_j}{{\Vert w_i\Vert }_2}},0]}^{\top })}^{\top }F({[w_i^{\top },b_i,w_j^{\top },b_j]}^{\top })\mathrm{vec}({[{\delta }_{gi}{\displaystyle \frac{w_i}{{\Vert w_i\Vert }_2}},0,{\delta }_{gj}{\displaystyle \frac{w_j}{{\Vert w_j\Vert }_2}},0]}^{\top })$

$=$ ${\displaystyle \frac{{\delta }_{gi}{\delta }_{gj}}{2{\varphi }^2}}\underset{𝕩\sim P(𝕩)}{𝔼}\left[\text{Cov}(y_i,y_j|𝕩){\displaystyle \frac{a_i{a}_j}{{\Vert w_i\Vert }_2{\Vert w_j\Vert }_2}}\right]$ (35)

The batch normalized and layer normalized models are therefore more robust to the scaling of the input and its parameters than the standard model.