T OWARDS B INARY -V ALUED G ATES FOR R OBUST LSTM T RAINING
Anonymous authors
Paper under double-blind review
A BSTRACT
Long Short-Term Memory (LSTM) is one of the most widely used recurrent structures in sequence modeling.
Its goal is to use gates to control the information ﬂow (e.g., whether to skip some information/transformation or not) in the recurrent computations, although its practical implementation based on soft gates only partially achieves this goal and is easy to overﬁt.
In this paper, we propose a new way for LSTM training, which pushes the values of the gates towards 0 or 1.
By doing so, we can (1) better control the information ﬂow: the gates are mostly open or closed, instead of in a middle state; and (2) avoid overﬁtting to certain extent: the gates operate at their ﬂat regions, which is shown to correspond to better generalization ability.
However, learning towards discrete values of the gates is generally difﬁcult.
To tackle this challenge, we leverage the recently developed GumbelSoftmax trick from the ﬁeld of variational methods, and make the model trainable with standard backpropagation.
Experimental results on language modelling and machine translation show that (1) the values of the gates generated by our method are more reasonable and intuitively interpretable, and (2) our proposed method generalizes better and achieves better accuracy on test sets in all tasks.
Moreover, the learnt models are not sensitive to low-precision approximation and low-rank approximation of the gate parameters due to the ﬂat loss surface.
I NTRODUCTION
Recurrent neural networks (RNN) (Hochreiter, 1998) are widely used in sequence modeling tasks, such as language modeling (Kim et al., 2016; Jozefowicz et al., 2016), speech recognition (Zhang et al., 2016), time series prediction (Xingjian et al., 2015), machine translation (Wu et al., 2016;
Britz et al., 2017), image captioning (Vinyals et al., 2015; Xu et al., 2015), and image generation (Villegas et al., 2017).
To address the long-term dependency and gradient vanishing problem of conventional RNN, long short-term memory (LSTM) (Gers et al., 1999; Hochreiter & Schmidhuber, 1997b) was proposed, which introduces gate functions to control the information in a recurrent unit: a forget gate function to determine how much previous information should be excluded for the current step, an input gate function to ﬁnd relevant signals to be absorbed into the hidden context, and an output gate function for prediction and decision making.
For ease of optimization, in practical implementation, one usually uses element-wise sigmoid function to mimic the gates, whose outputs are soft values between 0 and 1.
By using such gates, LSTM usually performs much better than conventional RNN.
However, the beneﬁts come with the cost of introducing many more parameters in the gates, which makes the training of a LSTM model inefﬁcient and easy to overﬁt (Krueger et al., 2016; Zaremba et al., 2014; Semeniuta et al., 2016).
In this paper, we explore a new way to train LSTM by pushing the values of the gates to the boundary of their ranges (0 , 1) 1 .
Pushing the values of the gates to 0/1 has certain advantages.
First, it well aligns with the original purpose of the development of gates: to get the information in or skip by “opening” or “closing” the gates during the recurrent computation.
Second, training LSTM towards binary-valued gates can make the learnt model generalize better.
According to Hochreiter &
Schmidhuber (1997a); Haussler et al. Keskar et al. ; Chaudhari et al. (2016), a model lying in a ﬂat region of the loss surface is likely to generalize well, since any small perturbation to the model makes little ﬂuctuation to the loss.
Training LSTM towards binary-valued gates means seeking a set of parameters to make the values of the gates approaching zero or one, namely residing in the ﬂat region of the sigmoid function.
Simple deductions show that this also corresponds to the ﬂat region of the overall loss surface.
Technically, pushing the outputs of the gates towards such discrete values is challenging.
A straightforward approach is to sharpen the sigmoid function by a smaller temperature.
However, this is equivalent to rescaling the input and cannot guarantee the value of the learnt gate to be close to 0 or 1.
To tackle this challenge, in this paper, we leverage the Gumbel-Softmax trick that Jang et al.
; Maddison et al.
(2016) recently develop for variantional methods.
The trick aims to generate approximated samples for categorical latent variables in a stochastic computational graph, e.g., variational autoencoder, brings convenience to using reparametrization tricks, and thus leads to efﬁcient learning.
Speciﬁcally, during training, we apply the Gumbel-Softmax trick to the gates to approximate the values sampled from the Bernoulli distribution given by the parameters, and train the LSTM model with standard backpropagation methods.
We call this method Gumbel-Gate LSTM ( G 2 -LSTM).
We conduct three experiments on two tasks (language modelling and machine translation) to verify our proposed method.
We have the following observations from experimental results:
• Our model generalizes well: In all tasks, we achieve superior performance to baseline algorithms on the test sets, and the gap between training and test is effectively reduced.
• Our model is not sensitive due to its ﬂat loss surface: We apply several model compression algorithms to the parameters in the gates, including low-precision approximation and lowrank approximation, and all results show that our learnt models are better.
• The gates in our learnt model are meaningful and intuitively interpretable after visualization.
Furthermore, our model can automatically learn the boundaries inside the sentences.
The organization of the paper is as follows.
We introduce related work in Section 2 and propose our learning algorithm in Section 3.
Experiments are reported in Section 4 and future work is discussed in the last section.
R ELATED W ORK 2. L OSS SURFACE AND GENERALIZATION
The concept of sharp and ﬂat minima has been ﬁrst discussed in Hochreiter & Schmidhuber (1997a); Haussler et al.
Intuitively, a ﬂat minimum x of a loss f ( · ) corresponds to the point for which the function f varies slowly in a relatively large neighborhood of x .
In contrast, a sharp minimum x is such that the function f increases rapidly in a small neighborhood of x .
The sensitivity of the loss function at sharp minima negatively impacts the generalization ability of a trained model on new data.
Recently, several papers discuss how to modify the training process and to learn a model in a ﬂat region so as to obtain better generalization ability.
Keskar et al.
(2016) show by using smallbatch training, the learnt model is more likely to converge to a ﬂat region rather than a sharp one.
Chaudhari et al.
(2016) propose a new objective function considering the local entropy and push the model to be optimized towards a wide valley.
2. D ROPOUT IN R ECURRENT N EURAL N ETWORK
Dropout is one of the most standard tricks used in deep learning to improve generalization ability.
For recurrent neural networks, Zaremba et al.
(2014) and Semeniuta et al.
(2016) apply dropout to feed-forward connections and recurrent units of RNNs.
In Zoneout (Krueger et al., 2016), the values of the hidden states and memory cells are randomly either maintained by their previous value or updated as usual, which introduces stochastic identity connections between subsequent time steps.
Different from dropout, which is to regularize the training of a deep neural network by randomly dropping nodes/edges to prevent co-adaptations, our method is to bias the optimization process and ensure to ﬁnd a model in a ﬂat region to avoid overﬁtting. Therefore, our method is complementary to dropout in RNN, and actually in our experiments our method is well combined with dropout.
Jang et al. ; Maddison et al.
(2016) develop a continuous relaxation of discrete random variables in stochastic computational graphs. The main idea of the method is that the multinomial distribution can be represented according to Gumbel-Max trick, thus can be approximated by GumbelSoftmax distribution. In detail, given a probability distribution over k categories with parameter π 1 , π 2 · · · π k , the Gumbel-Softmax trick approximately samples the categorical variable according to: [Equation]
where τ is the temperature and q i is independently sampled from Gumbel distribution: q i = − log( − log( U i )) , U i ∼ Uniform (0 , 1) . By using the Gumbel-Softmax trick, we can generate sample y = ( y 1 , ..., y k ) to approximate the categorical distribution. Furthermore, as the randomness q is independent of π (which is usually deﬁned by a set of parameters), we can use reparameterization trick to optimize the model parameters using standard backpropagation algorithms.
Gumbel-Softmax trick has been adopted in several applications such as variation autoencoder (Jang et al., 2016), generative adversarial net (Kusner & Hern ´ andez-Lobato, 2016), and language generation (Subramanian et al., 2017). To the best of our knowledge, this is the ﬁrst work to introduce the Gumbel-Softmax trick in LSTM for robust training purpose.  T HE P ROPOSED T RAINING A LGORITHM
In this section, we present a new and robust training algorithm for LSTM by learning towards binaryvalued gates.
3. B ACKGROUND
Recurrent neural networks process an input sequence { x 1 , x 2 , . . .
, x T } sequentially and construct a corresponding sequence of hidden states/representations { h 1 , h 2 , . . .
, h T } .
In single-layer recurrent neural networks, the hidden states { h 1 , h 2 , . . .
, h T } are used for prediction or decision makings.
In deep (stacked) recurrent neural networks, the hidden states in layer k are used as inputs to layer k + 1 .
In recurrent neural networks, each hidden state is trained (implicitly) to remember and emphasize task-relevant aspects of the preceding inputs, and to incorporate new inputs via a recurrent operator,
T , which converts the previous hidden state and presents input into a new hidden state, e.g.,
[Equation]
where W h , W x and b are parameters.
Long short-term memory RNN (LSTM) (Hochreiter & Schmidhuber, 1997b) is a carefully designed recurrent structure.
In addition to the hidden state h t used as a transient representation of state at timestep t , LSTM introduces a memory cell c t , intended for internal long-term storage.
c t and h t are computed via three gate functions.
The forget gate function f t directly connects c t to the memory cell c t − 1 of the previous timestep via an element-wise multiplication.
Large values of the forget gates cause the cell to remember most (if not all) of its previous values.
The other gates control the ﬂow of information in input ( i t ) and output ( o t ) of the cell.
Each gate function has a weight matrix and a bias vector; we use subscripts f , i and o to denote parameters for the forget gate function, the input gate function and the output gate function respectively, e.g., the parameters for the forget gate function is denoted by W xf , W hf , and b f .
With above notations, an LSTM is formally deﬁned as follows:
[Equation]
where σ ( · ) represents the sigmoid function and (cid:12) is the element-wise product.
3. T RAINING LSTM G ATES T OWARDS B INARY V ALUES
The LSTM unit requires much more parameters than the simple RNN unit, and makes it hard to generalize.
As we can see from Eqn (2) - (7), a large percentage of the parameters are used to compute the gate (sigmoid) functions.
If we can push the outputs of the gates to the saturation area of the sigmoid function (i.e., towards 0 or 1), the loss function with respect to the parameters in the gates will be ﬂat: if the parameters in the gates perturb, the change to the output of the gates is small due to the sigmoid operator (see Figure 1), and then the change to the loss is little, which means the ﬂat region of the loss. As discussed in Chaudhari et al.
(2016), minima in a ﬂat region is more likely to generalize better, and thus toward binary-valued gates will lead to better generalization.
However, the task of training towards binary-valued gates is quite challenging.
One straightforward idea is to sharpen sigmoid function by using a smaller temperature, i.e., 1 / (1 + e − x/τ ) .
However, this equals to rescale the parameters in the gates which cannot guarantee the outputs to be close to the boundary after training.
In this work, we leverage the recently developed Gumbel-Softmax trick. This trick is efﬁcient in approximating discrete distributions, and is one of the widely used methods to learn discrete random variables in stochastic computational graphs.
We ﬁrst provide a proposition about the approximation ability of this trick for Bernoulli distribution, which will be used in our proposed algorithm.
Proposition 1.
Assume σ ( · ) is the sigmoid function.
Given α ∈ R and temperature τ > 0 , we deﬁne random variable D α ∼ B ( σ ( α )) where B ( σ ( α )) is the Bernoulli distribution with parameter σ ( α ) , and deﬁne G ( α, τ ) =
σ ( α +log U − log(1 − U ) τ ) where U ∼ Uniform (0 , 1) .
Then the following inequalities hold for arbitrary (cid:15) ∈ (0 , 1 2 ) ,
[Equation]
Proof.
Since σ − 1 ( x ) = log( x 1 − x ) , we have P ( G ( α, τ ) ≥ 1 − (cid:15) )
= P ( α +log
U − log(1 − U ) τ ) ≥ log( 1 (cid:15) − 1))
=
P ( e α − τ log( 1 (cid:15) − 1) ≥ 1 − UU ) =
P ( U ≥ 1+ e α − τ log(1 (cid:15) − 1) )
=
σ ( α − τ log( 1 (cid:15) − 1)) .
Considering that sigmoid function is monotonically increasing and 14 -Lipschitz continuous, we have P ( D α = 1) − P ( G ( α, τ ) ≥ 1 − (cid:15) )
=
σ ( α ) − σ ( α − τ log( 1 (cid:15) − 1)) ≥ 0 and P ( D α = 1) − P ( G ( α, τ ) ≥ 1 − (cid:15) )
=
σ ( α ) − σ ( α − τ log( 1 (cid:15) − 1)) ≤ τ 4 log( 1 (cid:15) − 1) ≤ τ 4 log( 1 (cid:15) ) .
We omit the proof for (9) as it is almost identical to the proof of (8).
We can see from the above proposition, the distribution of G ( α, τ ) can be considered as an approximation of Bernoulli distribution B ( σ ( α )) .
The rate of convergence is characterized by (8) and (9).
When the temperature τ approaches positive zero, we directly obtain the following property which is also proved by Maddison et al.
[Equation]
We apply this method into the computation of the gate.
Imagine an one-dimensional gate σ ( α ( θ )) where α is a scalar parameterized by θ , and assume the model will produce a larger loss if the output of the gate is close to one, and produce a smaller loss if the gate value is close to zero. If we can repeatedly sample the output of the gate using G ( α ( θ ) , τ ) = σ ( α ( θ )+
log U − log(1 − U ) τ ) and estimate the loss, any gradient-based algorithm will push the parameter θ such that the output value of the gate is close to zero in order to minimize the expected loss. By this way, we can optimize towards the binary-valued gates.
As the gate function is usually a vector-valued function, we extend the notations into a general form: Given α ∈ R d and τ > 0 , we deﬁne G ( α, τ ) =
σ ( α +log U − log(1 − U ) τ ) , where U is a vector and each element u i in U is independently sampled from Uniform (0 , 1) , i = 1 , 2 , . . .
, d .
In particular, we only push the outputs of input gates and forget gates towards binary values as the output gates usually need ﬁne-granularity information for decision making which makes binary values less desirable (to justify this, we conducted similar experiments and observed the performance drops if we push the output gates to 0/1).
We call our proposed learning method Gumbel-Gate LSTM ( G 2 -LSTM), which works as follows during training:
[Equation]
[Equation] [Equation]
[Equation]
[Equation] [Equation] In the forward pass, we ﬁrst independently sample values for U in each time step, then update LSTM units using Eqn (11) - (16) and calculate the loss, e.g., negative log likelihood loss.
In the backward pass, as G is continuous and differentiable with respect to the parameters and the loss is continuous and differentiable with respect to G , we can use any standard gradient-based method to update the model parameters.
E XPERIMENTS 4. S ETTINGS
We tested the proposed training algorithm on two tasks – language modelling and machine translation.
4.1. L ANGUAGE M ODELLING
Language modelling is a very basic task for LSTM.
We used the Penn Treebank corpus which contains about 1 million words.
The task is to train an LSTM model to correctly predict the next [Table] [Table] Our
[Table] word conditioned on previous words. A model is evaluated by the prediction perplexity: smaller the perplexity, better the prediction. We followed the practice in Zaremba et al.
(2014) to set up the model architecture for LSTM: a stacked two-layer LSTM with 1500 units, with weights initialized uniformly in [ − 0 .
04 , +0 . 04] .
The mini-batch size is set to 20.
The model was trained for 14 epochs with learning rate 1.0 using standard stochastic gradient decent, after which the learning rate was reduced by a factor of 1.15 at the beginning of each epoch.
Gradient norms were clipped at 10.0 and the dropout rate is set to 0.65.
Our training code for G 2 -LSTM was also based on the code released by Zaremba et al.
We found the temperature τ used in G 2 -LSTM is not very sensitive.
We set the temperature to 0.9 and did not change during training.
4.1. M ACHINE T RANSLATION
We used two datasets for experiments on neural machine translation (NMT): (1) IWSLT German → English translation dataset (Cettolo et al., 2014), widely adopted in machine learning community (Bahdanau et al., 2016; Wiseman & Rush, 2016a; Ranzato et al., 2015).
The training/validation/test sets contains about 153 k/ 7 k/ 7 k sentence pairs respectively, with words preprocessed into sub-word units using byte pair encoding (BPE) (Sennrich et al., 2016).
We chose 25 k most frequent sub-word units as vocabulary for both German and English.
(2) English → German translation dataset in WMT’14, which is also commonly used as a benchmark task to evaluate different NMT models (Bahdanau et al., 2014; Wu et al., 2016;
Gehring et al., 2017).
The training set contains 4.5M English → German sentence pairs, Newstest 2014 is used as the test set, and the concatenation of Newstest 2012 and Newstest2013 is used as the validation set.
Similarly, BPE was used to form a vocabulary of most frequent 30 k sub-word units for both language.
In both datasets, we removed the sentences with more than 64 sub-word units in training.
For German → English dataset, we adopted a stacked two-layer encoder-decoder framework.
We set the size of word embedding and hidden state to 256.
As amount of data in the English → German dataset is much larger, we adopted a stacked three-layer encoder-decoder framework and set the size of word embedding and hidden state to 512 and 1024 respectively.
The ﬁrst layer of the encoder was bi-directional.
We also used dropout in training stacked LSTM as in Zaremba et al.
(2014), with dropout value determined via validation set performance.
For both experiments, we set the temperature τ for G 2 -LSTM to 0.9, which is the same as in the language model task.
The minibatch size was 32/64 for German → English/English → German respectively.
All models were trained with AdaDelta (Zeiler, 2012) on one M40 GPU.
Both gradient clipping norms were set to 2.0.
We [Figure] used tokenized case-sensitive BLEU (Papineni et al., 2002) 2 as evaluation measure.
The beam size is set to 5 during the inference step.
4. E XPERIMENTAL R ESULTS
The experimental results are shown in Table 1 and 2.
First, we compare our training algorithm with the baseline algorithm.
From the results, we can see that our learnt models are better than the baseline models.
In language modelling task, we outperform the baseline algorithm for 1.1 points in terms of perplexity.
For machine translation, we outperform the baseline for 0.95 and 0.54 points in terms of BLEU score for German → English and English → German dataset respectively.
Note that the only difference between G 2 -LSTM and baseline is the training algorithm, while they adopt the same model structure.
Thus, better results of
G 2 -LSTM demonstrate the effectiveness of our proposed training method. Second, training and validation loss curves of the baseline and G 2 -LSTM are shown in Fig. 2 for the two small tasks. Both curves show that the gap between training and validation is effectively reduced using our algorithm.
As shown in Fig.
2(b), the baseline LSTM achieves its lowest validation loss around the 18th epoch and begin to overﬁt after that, while the validation loss of G 2 -LSTM still drops even in the 30th epoch.
This clearly shows G 2 -LSTM generalizes better.
Third, we also list the performance of previous works in literature, which may adopt different model architectures or settings.
For language modelling, we obtain comparable results as Zoneout (Krueger et al., 2016).
For German → English translation, the two-layer stacked encoder-decoder model we learnt outperforms all previous works and achieves state-of-the-art performance.
For English → German translation, our result is worse than GNMT (Wu et al., 2016) as they used a stacked eight-layer LSTM encoder-decoder model while we only used a three-layer one.
[Table]
4. S ENSITIVENESS A NALYSIS
We conducted a set of experiments to test how sensitive our learnt models were if their gate parameters were compressed.
We considered two ways of parameter compression [Table] [Table] [Table]
Low-Precision Compression
We compressed parameters in the input and forget gates to lower precision.
Doing so the model can be compressed to a relatively small size.
In particular, we applied round and clip operations to the parameters of the input and forget gates [Table] [Equation]
[Equation]
We tested two settings of low-precision compression.
In the ﬁrst setting (named as Round ), we rounded the parameters using Eqn (17).
In this way, we reduced the support set of the parameters in the gate.
In the second setting (named as Round & Clip ), we further clipped the rounded value to a ﬁxed range using Eqn (18) and thus restricted the number of different values.
We set the round parameter r = 0 . 1 and the clip parameter c = 0 . 2 for the task of language modeling, and c = 1 .and r = 0 .
5 for neural machine translation.
As a result, parameters of input gates and forget gates in language modeling can only take values from ( 0 . 0 , ± 0 . 1 , ± 0 . 2 ), and ( 0 . 0 , ± 0 . 5 , ± 1 . 0 ) for machine translation.
Low-Rank Compression We compressed parameter matrices of the input/forget gates to lowerrank matrices through single value decomposition.
Doing so can reduce model size and lead to faster matrix multiplication.
Given that the hidden states of the task of language modeling were of much larger dimension than that of neural machine translation, we set rank = 64 / 128 for language modeling and rank = 16 / 32 for neural machine translation.
We summarize the results in Table 3-5.
More comprehensive results can be found in Appendix A.
From Table 3, we can see that for language modelling both the baseline and our learnt model are quite robust to low-precision compression, but our model is much more robust and signiﬁcantly outperforms the baseline with low-rank approximation.
Even setting rank = 64 (roughly 12x compression rate of the gates), we still get 80.4 perplexity, while the perplexity of the baseline model increases from 78.4 to 94.7, i.e., becoming 20% worse.
For machine translation, our proposed method is always better than the baseline model, no matter for low-precision or low-rank compression.
Even if setting rank = 16 (roughly 8x/32x compression rate of the gates for German → English and
English → German respectively), we still get roughly comparable translation accuracy to the baseline model with full parameters.
All results show that the models trained with our proposed method are more insensitive.
4. V ISUALIZATION OF THE G ATES
In addition to compare the ﬁnal accuracy in previous two subsections, we further look inside the learnt models and check the values of the gates.
To well verify the effectiveness of our proposed G 2 -LSTM, we did a set of experiments to show the values of gates we have learnt are near the boundary and are reasonable, based on the model learnt from German → English translation task.
We show the value distribution of the gates trained using classic LSTM and G 2 -LSTM.
To achieve this, we sampled 10000 sentence pairs from the training set of German → English and fed them into the learnt models.
We got the output value vectors of the input/forget gates in both the encoder and decoder.
We recorded the value of each element in the output vectors and plotted the value distribution in Figure 3 and Figure 4.
From the ﬁgure, we can see that although both LSTM and G 2 -LSTM work reasonably well in practice, the output values of the gates are very different.
In LSTM, the distributions of the gate [Figure] [Figure]
values are relatively uniform and have no clear concentration. In contrast, the values of the input gates of G 2 -LSTM are concentrated in the region close to 1, which suggests that our learnt model tries to keep most information from the input words; the values of the forget gates are concentrated in the boundary regions (i.e., either the region close to 0 or the region close to 1).
This observation shows that our training algorithm meets our expectation and successfully pushes the gates to 0/1.
Besides the overall distribution of gate values over a sampled set of training data, here we provide a case study for a sampled sentence.
As it is hard to go deep into individual dimensions of a hidden state, we just calculated the average value of the output vector of the input and forget gate functions for each word.
In particular, for each word, we focused on the average value of input/forget gate functions in the ﬁrst layer and check whether the average is reasonable.
We plot the heatmap of the English sentence part in Figure 5.
More visualizations can be found in Appendix B. First, we can see that our G 2 -LSTM does not drop information in the input gate function, since the average value is relatively large for all words.
In contrast, the average values of the input gates of LSTM are sometimes small (less than 0.5), even for the meaningful word like “data”.
As those words are not included into LSTM, they cannot be effectively encoded and decoded, thus lead to bad translation result.
Second, for G 2 -LSTM, most of the words with small values for forget gates are function words (e.g., conjunctions and punctuations) or the boundaries in clauses.
That is, our training algorithm indeed ensures the model to forget information on the boundaries inside the sentences, and reset the hidden states with new inputs.
C ONCLUSION AND F UTURE W ORK
In this paper, we have designed a new training algorithm for LSTM by leveraging the recently developed Gumbel-Softmax trick.
Our training algorithm can push the values of the input and forget gates to 0 or 1, leading to robust LSTM models.
Experiments on language modelling and machine translation have demonstrated the effectiveness of the proposed training algorithm.
We will explore following directions in the future.
First, we only tested with shallow LSTM models in this paper.
We will apply our algorithm to deeper models (e g., 8+ layers) in larger datasets.
Second, we considered the tasks of language modeling and machine translation.
We will study more applications such as question answering and text summarization.
A E XTRA E XPERIMENTS ON S ENSITIVENESS
We did an extra set of experiments on language modeling to show our models are more insensitive than the baseline model, no matter what the parameters ( c, r in low precision compression, rank in low rank compression) are.
The results are shown in Table 6 and Table 7.
[Table]
[Table]
B E XAMPLES