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 flow (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 overfit. 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 flow: the gates are mostly open or closed, instead of in a middle state; and (2) avoid overfitting to certain extent: the gates operate at their flat regions, which is shown to correspond to better generalization ability. However, learning towards discrete values of the gates is generally difficult. To tackle this challenge, we leverage the recently developed Gumbel-Softmax trick from the field of variational methods, and make the model trainable with standard backpropagation. Experimental results on language modeling 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 flat loss surface. Recurrent neural networks (RNN) BID10 are widely used in sequence modeling tasks, such as language modeling BID19 BID17 , speech recognition BID44 , time series prediction BID40 , machine translation BID2 , image captioning BID35 BID41 , and image generation BID34 .To address the long-term dependency and gradient vanishing problem of conventional RNN, long short-term memory (LSTM) BID7 BID12 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 find 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 benefits come with the cost of introducing many more parameters in the gates, which makes the training of a LSTM model inefficient and easy to overfit BID20 BID42 BID30 .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 1 The output of a gate function is usually a vector. For simplicity, in the paper, we say "pushing the output of the gate function to 0/1" when meaning "pushing each dimension of the output vector of the gate function to either 0 or 1". We also say that each dimension of the output vector of the gate function is a gate, and say a gate is open/closed if its value is close to 1/0. towards binary-valued gates can make the learnt model generalize better. According to BID11 BID9 BID18 BID4 , a model lying in a flat region of the loss surface is likely to generalize well, since any small perturbation to the model makes little fluctuation 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 flat region of the sigmoid function. Simple deductions show that this also corresponds to the flat 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 values of the learnt gates to be close to 0 or 1. To tackle this challenge, in this paper, we leverage the Gumbel-Softmax trick that BID15 and BID23 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 efficient learning. Specifically, 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 modeling 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 flat 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.2 RELATED WORK The concept of sharp and flat minima has been first discussed in BID11 BID9 . Intuitively, a flat 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 flat region so as to obtain better generalization ability. BID18 show by using smallbatch training, the learnt model is more likely to converge to a flat region rather than a sharp one. BID4 propose a new objective function considering the local entropy and push the model to be optimized towards a wide valley. Dropout is one of the most standard tricks used in deep learning to improve generalization ability. For recurrent neural networks, BID42 and BID30 apply dropout to feed-forward connections and recurrent units of RNNs. In Zoneout BID20 , 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 find a model in a flat region to avoid overfitting. Therefore, our method is complementary to dropout in RNNs, and actually in our experiments our method is well combined with dropout. BID15 and BID23 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 Gumbel-Softmax 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: where τ is the temperature and q i is independently sampled from Gumbel distribution: DISPLAYFORM0 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 defined 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 BID15 , generative adversarial net BID21 , and language generation BID33 . To the best of our knowledge, this is the first work to introduce the Gumbel-Softmax trick in LSTM for robust training purpose. In this section, we present a new and robust training algorithm for LSTM by learning towards binaryvalued gates. 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 making. 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., DISPLAYFORM0 where W h , W x and b are parameters. Long short-term memory RNN (LSTM) BID12 ) 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 flow 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 are denoted by W xf , W hf , and b f . With the above notations, an LSTM is formally defined as follows: DISPLAYFORM1 where σ(·) represents the sigmoid function and is the element-wise product. 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) - FORMULA3 , 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 flat: if the parameters in the gates perturb, the change to the output of the gates is small due to the sigmoid operator (see FIG0 , and then the change to the loss is little, which means the flat region of the loss. As discussed in BID4 , minima in a flat 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., f W,b (x) = σ((W x + b)/τ ), where τ < 1 is the temperature. However, it is computationally equivalent to f W ,b (x) = σ(W x + b ) by setting W = W/τ and b = b/τ . Then using a small temperature is equivalent to rescale the initial parameters as well as the gradients to a larger range. Usually, using an initial point in a large range with a large learning rate will harm the optimization process, and apparently 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 efficient in approximating discrete distributions, and is one of the widely used methods to learn discrete random variables in stochastic computational graphs. We first 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 define random variable D α ∼ B(σ(α)) where B(σ(α)) is the Bernoulli distribution with parameter σ(α), and define G(α, τ ) = σ( DISPLAYFORM0 where U ∼ Uniform(0, 1). Then the following inequalities hold for arbitrary ∈ (0, 1 2 ), DISPLAYFORM1 DISPLAYFORM2 Proof. Since σ −1 (x) = log( DISPLAYFORM3 Considering that sigmoid function is monotonically increasing and DISPLAYFORM4 . 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 BID23 , DISPLAYFORM5 We apply this method into the computation of the gates. 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(α(θ), τ ) = σ( DISPLAYFORM6 ) 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: DISPLAYFORM7 , 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 fine-granularity information for decision making which makes binary values less desirable (to justify this, we conducted similar experiments and observed a performance drop when pushing the output gates to 0/1 together with the input gates and the forget gates).We call our proposed learning method Gumbel-Gate LSTM (G 2 -LSTM), which works as follows during training: DISPLAYFORM8 In the forward pass, we first independently sample values for U in each time step, then update LSTM units using Eqn FORMULA0 - FORMULA0 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. 4.1 SETTINGSWe tested the proposed training algorithm on two tasks -language modeling and machine translation. Language modeling 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 word Table 2 : Performance comparison on machine translation (BLEU) English→German task BLEU German→English task BLEU Existing end-to-end system RNNSearch-LV BID16 19.40 BSO BID38 26.36 MRT BID32 20.45 NMPT (Huang et al.) 28.96 Global-att BID22 20.90 NMPT+LM (Huang et al.) 29.16 GNMT 24.61 ActorCritic BID1 conditioned on previous words. A model is evaluated by the prediction perplexity: smaller the perplexity, better the prediction. We followed the practice in BID26 to set up the model architecture for LSTM: a stacked three-layer LSTM with drop-connect BID36 on recurrent weights and a variant of averaged stochastic gradient descent (ASGD) (Polyak & Juditsky, 1992) for optimization. Our training code for G 2 -LSTM was also based on the code released by BID26 2 . We found the temperature τ used in G 2 -LSTM is not very sensitive. We set the temperature to 0.9 and followed all configurations in BID26 . We added neural cache model BID8 on the top of our trained language model to further improve the perplexity. We used two datasets for experiments on neural machine translation (NMT): (1) IWSLT2014 German→English translation dataset BID3 , widely adopted in machine learning community BID1 BID37 BID29 . The train- ing/validation/test sets contains about 153k/7k/7k sentence pairs respectively, with words preprocessed into sub-word units using byte pair encoding (BPE) BID31 . We chose 25k 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 BID0 BID6 . 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 30k 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 first layer of the encoder was bi-directional. We also used dropout in training stacked LSTM as in BID42 , with dropout value determined via validation set performance. For both experiments, we set the temperature τ for G 2 -LSTM to 0.9, which was 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 used tokenized case-sensitive BLEU BID27 3 as evaluation measure. The beam size is set to 5 during the inference step. The experimental results are shown in TAB0 and 2.First, we compare our training method with two algorithms. For the first algorithm (we call it Baseline), we remove the Gumble-Softmax trick and train the model using standard optimization methods. For the second algorithm (we call it Sharpened Sigmoid), we use a sharpened sigmoid function as described in Section 3.2 by setting τ = 0.2 and check whether such trick can bring better generalization. From the results, we can see that our learnt models are better than all baseline models. In language modeling task, we outperform the baseline algorithms for 0.7/1.1 points (1.2/1.4 points without continuous cache pointer) in terms of test perplexity. For machine translation, we outperform the baselines for 0.95/2.22 and 0.54/0.79 points in terms of BLEU score for German→English and English→German dataset respectively. Note that the only difference between G 2 -LSTM and the baselines 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 FIG1 for the two small tasks. Both curves show that the gap between training and validation is effectively reduced using our algorithm. As shown in FIG1 , the baseline LSTM achieves its lowest validation loss around the 18th epoch and begins to overfit after that, while the validation loss of G 2 -LSTM still drops even in the 30th epoch. This clearly shows that 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 modeling, we obtain the best performance as far as we know. 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 as they used a stacked eight-layer LSTM encoder-decoder model while we only used a three-layer one. 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. 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. DISPLAYFORM0 We tested two settings of low-precision compression. In the first setting (named as Round), we rounded the parameters using Eqn (17). In this way, we reduced the support set of the parameters in the gates. In the second setting (named as Round & Clip), we further clipped the rounded value to a fixed range using Eqn FORMULA0 and thus restricted the number of different values. As the two tasks are far different, we set the round parameter r = 0.2 and the clip parameter c = 0.4 for the task of language modeling, and set c = 1.0 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.2, ±0.4), and (0.0, ±0.5, ±1.0) for machine translation. More comprehensive results on different choices of hyperparameters can be found in Appendix A. 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 TAB2 , we can see that for language modeling both the baseline and our learnt model are quite robust to low-precision compression, but our model is much more robust and significantly outperforms the baseline with low-rank approximation. Even setting rank = 64 (roughly 12x compression rate of the gates), we still get 56.0 perplexity, while the perplexity of the baseline model increases from 52.8 to 65.5, i.e., becoming 24% 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 less sensitive to parameter compression. In addition to compare the final 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 distributions in Figure 3 and Figure 4 .From the figures, 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 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 first layer and check whether the average is reasonable. We plot the heatmap of the English sentence part in FIG3 . 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 values are 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. 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 modeling and machine translation have demonstrated the effectiveness of the proposed training algorithm. We will explore following directions in the future. First, we have only tested with shallow LSTM models in this paper. We will apply our algorithm to deeper models (e.g., 8+ layers) and test on larger datasets. Second, we have considered the tasks of language modeling and machine translation. We will study more applications such as question answering and text summarization. Third, we are cleaning and refactoring the code and will release the training code to public soon. We did an extra set of experiments on language modeling to show our model is less sensitive than the baseline model, no matter what the hyperparameters (c, r in low-precision compression, rank in low-rank compression) are. The results are shown in TAB7 .
We propose a new algorithm for LSTM training by learning towards binary-valued gates which we shown has many nice properties. Propose a new "gate" function for LSTM to enable the values of the gates towards 0 or 1.  The paper aims to push LSTM gates to be binary by employing the recent Gumbel-Softmax trick to obtain end-to-end trainable categorical distribution.