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 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 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.