Title: RECURRENT NORMALIZATION PROPAGATION

Abstract: We propose an LSTM parametrization that preserves the means and variances of the hidden states and memory cells across time. While having training benefits similar to Recurrent Batch Normalization and Layer Normalization, it does not need to estimate statistics at each time step, therefore, requiring fewer computations overall. We also investigate the parametrization impact on the gradient flows and present a way of initializing the weights accordingly. We evaluate our proposal on language modelling and image generative modelling tasks. We empirically show that it performs similarly or better than other recurrent normalization approaches, while being faster to execute.

1. INTRODUCTION:

Recurrent neural network have shown remarkably good performances for sequential modelling tasks including machine translation (Bahdanau et al., 2015) , visual captioning (Xu et al., 2015; Yao et al., 2015) or question answering (Hermann et al., 2015) . However, such models remain notoriously hard to train with gradient backpropagation. As the number of time steps in the input sequence increases, the contractive or expanding effects associated with the state-to-state transformation at each time step can shrink or grow exponentially, leading respectively to vanishing or exploding gradients (Hochreiter, 1991; Bengio et al., 1994; Pascanu et al., 2012) . In particular, with gradient vanishing, states at a given time are not influenced by changes happening much earlier in the sequence, preventing the model from learning long-term dependencies.

While the long-term dependencies problem is unsolvable in absolute (Hochreiter, 1991; Bengio et al., 1994) , different RNN parameterizations, such as LSTM or GRU (Hochreiter & Schmidhuber, 1997; Cho et al., 2014) can help mitigate it. Furthermore, the LSTM parametrization has been recently extended to include layer-wise normalization (Cooijmans et al., 2016; Ba et al., 2016) , building upon Batch Normalization (BN) (Ioffe & Szegedy, 2015) . By normalizing the hidden state distributions to a fix scale and shift through the different time steps, normalized LSTMs have been shown to ease training, resulting in a parametrization that converges faster than a standard LSTM.

However, normalized LSTM introduces extra-computations as it involves standardizing the hidden states, enforcing their means and variances at each time step. By contrast, we propose an LSTM reparametrization that allows by construction to cheaply preserve the normalization of the hidden states through time. Our approach can be seen as the recurrent counterpart to the recent normalization propagation applied in feed-forward network (Arpit et al., 2016) . It results in faster training convergence similar to Layer Normalization (LN) and Recurrent Batch Normalization while requiring fewer operations per time step and generalizing naturally to variable length sequences.

In addition, we investigate the impact of our parametrization, and more generally of normalized LSTM, on the vanishing and exploding gradient problems. We observe that layer-wise normalization provides a direct way to orient LSTM behaviour toward either gradient explosion or vanishing, and therefore biases the LSTM either towards reliably storing bits of information throughout time or allowing it to be more sensitive to new input changes.

We empirically validate our proposal on character-level language modelling on the Penn Treebank corpus (Marcus et al., 1993) and on image generative modelling, applying our normalisation to the DRAW architecture (Gregor et al., 2015) .

The paper is structured as follows: section 2 provides a brief overview of the Batch-Normalized LSTM, in section 3 we derive our Normalized LSTM, section 4 investigates the impact of such normalization on the gradient flow, section 5 presents some experimental results, and we conclude in section 5.

2. PRE-REQUISITES:

2.1 BN-LSTM Batch-Normalized Long Short-Term Memory (BN-LSTM) (Cooijmans et al., 2016) is a reparametrization of LSTM that takes advantage of Batch Normalization (BN) to address the Covariate Shift (Shimodaira, 2000) occurring between time steps. Changes in the LSTM output at one time-step are likely to cause correlated changes in the summed inputs of the sequence next timesteps. This Temporal Covariate Shift can slow down the training process as the parameters of the model must not only be updated to minimize the cost of the task at hand but also adapt to the changing distribution of the inputs. In other words, the latter time steps in an LSTM need to account for the shifting distribution of the previous hidden states. BN-LSTM proposes to reduce this temporal covariate shift by fixing the mean and the variance at each time step, relying on the BN transform

EQUATION

where E[x], Var[x] are the activation mean and variance estimated from the mini-batch samples.

Given an input sequence X = (x 1 , x 2 , . . . , x T ), the BN-LSTM defines a sequence of hidden states h t and memory cell states c t according to

    ĩt ft õt gt     = BN(W x x t ; γ x , β x ) + BN(W h h t-1 ; γ h , β h ) + b

(2)

EQUATION

where using the different feature channels rather than mini-batch samples in order to naturally generalize to variable length sequences.

W h ∈ R d h ×4d h , W x ∈ R dx×4d h , b ∈ R 4d h and the initial states h 0 ∈ R d h , c 0 ∈ R d h are

2.2. NORMALIZATION PROPAGATION:

While increasing the training convergence speed relatively to a standard LSTM (Cooijmans et al., 2016) , BN-LSTM needs to perform more computations per sample as it requires to compute 3x the BN transform at each time step.

On the other hand, Normalization Propagation (Norm Prop) (Arpit et al., 2016) aims at preserve the normalization of the input throughout the network. Unlike BN, the normalization doesn't rely on the statistics of the mini-batch. Instead, it is the structure of the network itself that maintains the normalization. We therefore propose an LSTM reparametrization that preserves the normalization through the different time steps in order to avoid those extra computation.

3. NORMALIZED LSTM:

While Norm Prop properties are appealing for recurrent models, its application to LSTM is not straightforward due to the memory cell structure. In this section we show how to derive a LSTM reparametrization that preserves normalization of the state h t through time.

3.1. CONSTRUCTION OF THE NORMALIZED LSTM:

Following (Arpit et al., 2016; Salimans & Kingma, 2016) , we will attempt to ensure, through an analytical reparametrization, that several intermediate quantities in the computation remain approximately standardized. We first compensate for the distribution changes induced by the weight matrices in the gates and cell candidate

EQUATION

)

where ||W •,i || 2 is the vector of L2-norm of each line of the matrix and γ x and γ h are the trainable rescaling factors that restore the representation power lost in the rescaling of the weight matrices.

To preserve the constant error carousel mechanism of the LSTM, we use the usual cell update,

EQUATION

6) Let us now construct an approximate analytical estimate of Var(c t ). The evolution of c t through time can bee seen as a geometric series, with σ( ft ) as constant ratio. Since σ(•) is upper-bounded by (and in practice smaller than) 1, c t will converge in expectation to a fixed value. This is the reason why in BN-LSTM the mini-batch statistics converge to a fixed value after a few time steps (Cooijmans et al., 2016) . Moreover, if we consider that ĩt , ft , gt and c t-1 are (as a rough approximation) independent1 , we can use the variance product rule of two independent random variables X and Y

Var[XY ] = Var[X] Var[Y ] + Var[X]E[Y ] 2 + Var[Y ]E[X] 2 (7) to compute Var[c t ].

Considering that E[tanh(g t )] ≈ 0 and assuming that the cell has converged i.e.

Var[c t ] = Var[c t-1 ], we have Var[c t ] = Var[tanh(g t )] Var[σ( ĩt )] + E[σ( ĩt )] 2 1 -Var[σ( ft )] -E[σ( ft )] 2 (8)

We can therefore analytically or numerically compute the mean and variance of each of those elements, assuming that both input x t and hidden state h t-1 are independent drawn from N (0, 1)

E[i t ] = E[σ(γ x z x + γ h z h )] (9) Var[i t ] = Var[σ(γ x z x + γ h z h )] (10) E[g t ] = E[tanh(γ x z x + γ h z h )] (11) Var[g t ] = Var[tanh(γ x z x + γ h z h )]

(12) where z x , z h ∼ N (0, 1). The statistics of the gates o t and f t can be computed in a similar way. We can then compute the value to which Var[c t ] converges. Using this variance estimate, we compensate c t in order to compute the next hidden state h t

EQUATION

Since we assumed that Var[h t-1 ] = 1, to ensure that we need to correct for the variance induced by the product of the tanh with the output gate. Using again the variance product rule (equation 7) we obtain

Var[h t ] = Var tanh γ c c t Var[c t ] (Var[σ(õ t )] + E[σ(õ t )] 2 ) (14)

We can estimate this variance through similar computation than equation 12. Scaling h t with 1/ Var[h t ] ensure that its variance is 1 and so the propagation is maintained throughout the recurrence.

3.2. PROPOSED REPARAMETRIZATION:

Using equations 5, 6 and 13, we propose the following reparametrization of the LSTM, simply called the Normalized LSTM 

EQUATION

where Var[c t ] and Var[h t ] are computed using equations 8 and 14, respectively. Those two variances are estimated at the initialization of the network (eq. 10 to eq. 12), and are then kept fixed during the training as in Norp Prop. γ x , γ h and γ c are parameters learned via gradient descent.

Note that the reparametrization of equation 15 is identical to Weight Normalization (Weight Norm) (Salimans & Kingma, 2016) . The main difference comes from equation 17, where we compensate for the variance of c t , the tanh and σ(õ t ), which ensures a normalized propagation. Overall, this reparametrization is equivalent in spirit to the BN-LSTM, but it benefits from the same advantages that Norm Prop has over BN: There is no dependence on the mini-batch size and the computation is the same for training and inference. Also, the rescaling of the matrices W x and W h can be done before the recurrence, leading to computation time closer to a vanilla LSTM.

3.3. WEIGHTS INITIALIZATION:

With such reparametrization of the weight matrices, one can think that the scale of the initialization of the weights doesn't matter in the learning process anymore. It is actually true for the forward and backward computation of the layer

EQUATION

)

∂y i ∂x = aW i ||aW i || 2 = W i ||W i || 2 (19)

and since the variance of both forward and backward passes is fixed, using an initialization scheme such as Glorot (Glorot & Bengio, 2010) doesn't make sense with Norm Prop. However, the update of the parameters is affected by their scale:

EQUATION

The scale of the parameters affect the learning rate of the layer: the bigger the weights, the smaller the update. This induces a regularization effect in Norm Prop that is also present in BN (Ioffe & Szegedy, 2015) . However, this could possibly be an issue for such parametrization: different initializations lead to different learning rates, and it is true even with adaptive step rules, such as Adam (Kingma & Ba, 2014) . Moreover, the parameters that are not normalized (such as γ and b) aren't affected by this effect, and so they are not regularized. This is the reason why forcing the weight matrices to have a unit L2 norm of the lines, as proposed in Arpit et al. (2016) , helps the training procedure.

To still benefit from the reduction of the learning rate, which is know to ease the optimization (Vogl et al., 1988) , we propose to simply force the unit L2 norm of the lines of the matrices and combine it with a global learning rate decay schedule.

4. GRADIENT PROPAGATION IN NORMALIZED LSTM:

In this section we study the gradient flow in the Normalized LSTM. Since this reparametrization is similar to the BN-LSTM, the analysis we do here can be transposed to the BN-LSTM case.

4.1. THE EXPLODING AND VANISHING GRADIENTS PROBLEM:

Given an input sequence X = (x 1 , x 2 , . . . , x T ), we consider a recurrent network, parametrized by θ, that defines a sequence of hidden states h t = f θ (h t-1 , x t ) and cost function L which evaluates the model performance on a given task. Such network is usually trained using backpropagation through time, where the backpropagation is applied on the time-unrolled model. The chain rule can be applied in order to compute the derivative of the loss L with respect to parameters θ.

EQUATION

The factors ∂h k ∂ht = k≤l≤t ∂h l ∂h l-1 transports the error "in time" from step t back to step k and are also the cause of vanishing or exploding gradient in RNN (Pascanu et al., 2012) . Indeed, if the Jacobian ∂h l ∂h l-1 has singular value different from 1, the factor ∂h k ∂ht , which is a product of t -k Jacobian matrices will either explode or vanish.

4.2. GRADIENT OF THE NORMALIZED LSTM:

To study the gradient propagation of the Normalized LSTM, we first need to derive it. Using equation 15-17, we can write the gradient of h t with respect to h t-1

EQUATION

)

EQUATION

As we can see in equation 23 with the normalization, the gradient depends not only on the derivative of the cell candidate, the gates and the output tanh, but also on on the variance of h t and c t .

If we assume that h t-1 and x t are independent, we can compute the variance of c t . Neglecting the weight matrices and the effect of the gates, we can write from equations 8 and 14

EQUATION

In both cases, the variance depends explicitly on the value of the different γ: The bigger the γ, the higher the variance. Neglecting again the weight matrices, we can now write the equations of the cell candidates g t and the gates i t , o t and f t with respect to h t-1

EQUATION

)

The gradients of o t and f t can be computed similarly. The effect of the γ here is double: They appear both in the activation function, where they control the saturation regime, and γ h also appears as a multiplicative term in the gradient. They should therefore be small enough to prevent the activation from saturating too much, but at the same time γ h can't be too small, because it can also make the gradients vanish. Putting it all together, we have

∂h t ∂h t-1 = ∂o t ∂õ t γ h a t + o t ∂a t ∂ã t γ c Var[c t ] γ h ∂i t ∂ ĩt g t + i t ∂g t ∂ gt + ∂f t ∂ ft c t-1 (28)

In this equations we can see that the different γ directly scale the gradient, and they also control the saturation of the activation functions. Bad initialization of γ could thus lead to saturation or explosion regimes. Figure 1 shows the norm of the gradient with respect to γ x and γ h in a simulated LSTM. As we can see, one important parameter is the ratio between γ h and γ x : They control most of the propagation of the gradients. If γ x > γ h , the network will focus more on the input and so the gradients will tend to vanish more. On the other hand, if γ h > γ x , the network will tend have less vanishing gradients, but will focus less on its inputs. 

5. EXPERIMENTS:



5.1. CHARACTER-LEVEL LANGUAGE MODELLING:

The first task we explore is character-level language modelling on the Penn Treebank corpus (Marcus et al., 1993) . The goal is to predict the next character of the sequence given the previous ones. We use the same splits as Mikolov et al. (2012) and the same training procedure as Cooijmans et al. (2016) , i.e. we train on sequences of length 100, with random starting point. The model is a 1000 units LSTM followed by a Softmax classifier. We use orthogonal initialization for the weight matrices. Because Norm Prop requires normalized inputs, we multiply the one-hot inputs vector with an untrained but fixed orthogonal matrix. This tricks does not only help the optimization of Norm Prop, but also all other variants.

To compare the convergence properties of Norm Prop against LN and BN, we first ran experiments using Adam (Kingma & Ba, 2014) with learning rate 2e-3, exponential decay of 1e-3 and gradient clipping at 1.0. As explained in section 3.3, we rescale the matrices such that they have a unit norm on the lines. For Norm Prop, we use γ x = γ h = 2 and γ c = 1, for LN all the γ = 1.0 and for BN all the γ = 0.1. The results are presented in Table 1 and in Figure 2 . To show the potential of Norm Prop against other state-of-the-art system, we followed Ha et al. (2016) and apply dropout on both the input and output layer (p = 0.1) and recurrent dropout inside the LSTM (p = 0.1). We also used the Batch Data Normalization scheme presented by Arpit et al. (2016) , so we standardize each input example using the mini-batch statistics and use population statistics at inference time. Finally, we also reduce the learning rate decay to 1e-4, to compensate for the fact that a network with dropout needs more time to train. The results are presented in Table 2 .

Model:

As we can see in Figure 2 and in Table 1 , Norm Prop compares really well against the other reparametrization. Also Norm Prop is roughly 30 % computationally faster2 than BN and LN. LN shows better optimization performances, but also overfits more. We also see that both optimization and generalization are better than the ones from Weight Norm, which shows the importance of compensating for the variance of c t and h t . Moreover, although Norm Prop doesn't combine well with 

Model Test:

Recurrent Dropout LSTM (Semeniuta et al., 2016) 1.301 Zoneout LSTM (Krueger et al., 2016) 1.27 Layer Norm LSTM (Ha et al., 2016) 1.267 HyperLSTM (Ha et al., 2016) 1.265 Norm Prop LSTM (ours)

1.262 Layer Norm HyperLSTM (Ha et al., 2016) 1.250 dropout in feed-forward networks (Arpit et al., 2016) , it works will with recurrent dropout, as we can see in Table 2 . We believe it is because recurrent dropout is less affecting its output distribution than dropout in feed forward networks, because we copy the variable at the previous time step instead of setting it to 0. With such regularization, Norm Prop compares well with other state-of-the-art approaches.

5.2. DRAW:

The second task we explore is a generative modelling task on binarized MNIST (Larochelle & Murray, 2011) using the Deep Recurrent Attentive Writer (DRAW) (Gregor et al., 2015) architecture. DRAW is a variational auto-encoder, where both encoder and decoder are LSTMs, and has two attention mechanisms to select where to read and where to write.

We use Jörg Bornschein's implementation3 , with the same hyper-parameters as Gregor et al. (2015) , ie the read and write size are 2x2 and 5x5 respectively, the number of glimpses is 64, the LSTMs have 256 units and the dimension of z is 100. We use Adam with learning rate of 1e-2, exponential decay of 1e-3 and mini-batch size of 128. We use orthogonal initialization and force the norm of the lines of the matrices to be 1. For Norm Prop, we use γ x = γ h = γ c = 0.5. The test variational bound for the first 100 epochs is presented in Figure 3 .

As we can see in Figure 3 , both Weight Norm and Norm Prop outperform the baseline network by a significant margin. Also, as expected, Norm Prop performs better than Weight Norm, showing one again the importance of the compensation of the variance of c t and h t . 

Model DRAW:

Baseline (ours) 84.30 Layer Norm (Ba et al., 2016) 82.09 Weight Norm (ours) 81.98 Norm Prop (ours) 81.17 

6. CONCLUSION:

Based on the BN-LSTM, we have shown how to build a Normalized LSTM that is able to preserve the variance of its output at each time step, by compensating for the variance of the cell and the hidden state. Such LSTM can be seen as the Norm Prop version of the BN-LSTM, and thus benefits from the same advantages that Norm Prop has over BN, while being way faster to compute. Also, we propose a scheme to initialize the weight matrices that takes into account the reparametrization. Moreover, we have derived the gradients of this LSTM and pointed out the importance of the initialization of the rescaling parameters. We have validated the performances of the Normalized LSTM on two different tasks, showing similar performances than BN-LSTM and LN-LSTM, while being significantly faster in computation time. Also, unlike the feed-forward case, this architecture works well with recurrent dropout, leading to close to state-of-the-art performances on the character-level language modelling task.

Future work includes trying this architecture on more challenging tasks and also studying the impact of not keeping the variance estimates of the cell and the hidden states fixed during the learning process.

:

This assumption is strong, but we don't have any easy way to model the covariance between those terms without estimating it from the data.

The GPU used is a NVIDIA GTX 750.

https://github.com/jbornschein/draw

Figures and Tables:

Figure 1: Norm of the gradients for one time step in an LSTM with respect to γ x and γ h (simulation). Left: γ c = 0.1. Right: γ c = 1.0.

Figure 2: Perplexity (bits-per-character) on sequences of length 100 from the Penn Treebank corpus. The dashed lines are the training curves, and the solid ones are the validation curves.

Figure 3: Test negative log-likelihood on binarized MNIST.

Perplexity (bits-per-character) on sequences of length 100 from the Penn Treebank validation set, and training time (seconds) per epoch.

Table content: <table><tr><td/><td colspan="2">Validation Time</td></tr><tr><td>Baseline</td><td>1.455</td><td>386</td></tr><tr><td>Weight Norm</td><td>1.438</td><td>402</td></tr><tr><td>Batch Norm</td><td>1.433</td><td>545</td></tr><tr><td>Layer Norm</td><td>1.439</td><td>530</td></tr><tr><td>Norm Prop</td><td>1.422</td><td>413</td></tr></table>

Perplexity (bits-per-character) of the full Penn Treebank test sequence.

Table content: <table/>

Table 3 shows the test variational bound after 200 epochs of training. Norm Prop also compares favorably against LN.

Table content: <table/>

Test variational log likelihood (nats) after 200 epochs of training.

Table content: <table/>

Peer review (iclr2017):
I think this build upon previous works, in the attempt of doing something similar to batch norm specific for RNNs. To me the experiments are not yet very convincing, I think is not clear this works better than e.g. Layer Norm or not significantly so. I*m not convinced on how significant the speed up is either, I can appreciate is faster, but it doesn*t feel like order of magnitude faster. The theoretical analysis also doesn*t provide any new insights. All in all I think is good incremental work, but maybe is not yet significant enough for ICLR.