Title: VARIATIONAL BI-LSTMS

Abstract: Recurrent neural networks like long short-term memory (LSTM) are important architectures for sequential prediction tasks. LSTMs (and RNNs in general) model sequences along the forward time direction. Bidirectional LSTMs (Bi-LSTMs), which model sequences along both forward and backward directions, generally perform better at such tasks because they capture a richer representation of the data. In the training of Bi-LSTMs, the forward and backward paths are learned independently. We propose a variant of the Bi-LSTM architecture, which we call Variational Bi-LSTM, that creates a dependence between the two paths (during training, but which may be omitted during inference). Our model acts as a regularizer and encourages the two networks to inform each other in making their respective predictions using distinct information. We perform ablation studies to better understand the different components of our model and evaluate the method on various benchmarks, showing state-of-the-art performance.

1. INTRODUCTION:

Recurrent neural networks (RNNs) have become the standard models for sequential prediction tasks, having achieved state of the art performance in a number of applications that includes sequence prediction, language translation, machine comprehension, and speech synthesis (Arik et al., 2017; Wang et al., 2017; Mehri et al., 2016; Sotelo et al., 2017) . RNNs model temporal data by encoding a given arbitrary-length input sequentially, at each time step combining some transformation of the current input with the encoding from the previous time step. This encoding, referred to as the RNN hidden state, summarizes all previous input tokens.

Viewed as "unrolled" feedforward networks, RNNs can become arbitrarily deep depending on the input sequence length, and use a repeating module to combine the input with the previous state at each time step. Consequently, they suffer from the vanishing/exploding gradient problem (Pascanu et al., 2012) . This problem has been addressed through architectural variants like the long short-term memory (LSTM) (Hochreiter & Schmidhuber, 1997) and the gated recurrent unit (GRU) (Chung et al., 2014) . These architectures add a linear path along the temporal sequence which allows gradients to flow more smoothly back through time.

Various regularization techniques have also been explored to improve RNN performance and generalization. Dropout (Srivastava et al., 2014) regularizes a network by randomly dropping hidden units during training. However, it has been observed that using dropout directly on RNNs is not as effective as in the case of feed-forward networks. To combat this, Zaremba et al. (2014) propose to instead apply dropout on the activations that are not involved in the recurrent connections (Eg. in a multi-layer RNN); Gal & Ghahramani (2016) propose to apply the same dropout mask through an input sequence during training. In a similar spirit to dropout, Zoneout (Krueger et al., 2016) proposes to choose randomly whether to use the previous RNN hidden state.

The aforementioned architectures model sequences along the forward direction of the input sequence. Bidirectional-LSTM, on the other hand, is a variant of LSTM that simultaneously models each sequence in both the forward and backward directions. This enables a richer representation of data, since each token's encoding contains context information from the past and the future. It has been shown empirically that bidirectional architectures generally outperform unidirectional ones on many sequence-prediction tasks. However, the forward and backward paths in Bi-LSTMs are trained separately and the benefit usually comes from the combined hidden representation from both paths. In this paper, our main idea is to create a dependence between the two paths that acts as a regulariza-tion during training, but doesn't hinder inference even in the absence of the other path (the backward path in practical scenarios). We note that recently proposed methods like TwinNet Serdyuk et al. (2017) and Z-forcing Sordoni et al. (2017) are similar in spirit to this idea. In our approach, we use a variational auto-encoder (VAE; Kingma & Welling (2014) ) that takes as input the hidden states from the two paths of the Bi-LSTM and maps them to a shared hidden representation of the VAE at each time step. The samples from the VAE hidden state are then used both for reconstructing the LSTM hidden states and feeding forward to the next hidden state. In this way, we create a channel between the two paths that acts as a regularization for learning better representations. We refer to the resulting model as a Variational Bi-LSTM.

Below, we describe Variational Bi-LSTMs in detail and then demonstrate empirically their ability to model complex sequential distributions. In experiments, we obtain state-of-the-art or competitive performance on the tasks of Penn Treebank, IMDB, TIMIT, Blizzard, and Sequential MNIST.

2. VARIATIONAL BI-LSTM:

Bi-LSTM is a powerful architecture for sequential tasks because it models temporal data both in the forward and backward directions. For this it uses two LSTMs that are generally learned independently of each other; the richer representation results from combining the hidden states of these LSTMs, where combination is often by concatenation. The idea behind variational Bi-LSTMs is to create a channel of information exchange between the two LSTMs that helps the model to learn better representations. We create this dependence using a variational auto-encoder (VAE). This enables us to take advantage of the fact that VAE allows for sampling from a prior during inference. For sequence prediction tasks like language generation, while one can use Bi-LSTMs during training, there is no straightforward way to employ the full bidirectional model during inference -this would involve, eg, generating a sentence starting at both its beginning and end. In such cases, the VAE allows us to sample from the prior at inference time to make up for the absence of the backward LSTM. Now we describe our variational Bi-LSTM model formally. Let X = {x (i) } N i=1 be a dataset consisting of N i.i.d. sequential data samples of continuous or discrete variables. For notational convenience, we will henceforth drop the superscript i indexing samples. For each sample sequence x = (x 1 , . . . , x T ), the hidden state of the forward LSTM is given by:

h t = - → f (x t , h t-1 , z t , bt ).

The hidden state of the backward LSTM is given by,

b t = ← - f (x t , b t+1 ).

In both cases, the function f represents the standard LSTM updates, modified to account for the additional arguments.

In the forward LSTM model, we introduce three latent random variables, z t , bt , and ht-1 , where z t depends on h t-1 and b t during training, and bt and ht-1 depend only on z t (see figure 1 -left, for a graphical representation). Note that so far, bt and ht-1 are simply latent vectors drawn from conditional distributions p ψ and p ξ , respectively, to be defined below. However, as explained in Section 2.1 (see also dashed lines in figure 1 -left), we will encourage these to lie near the manifolds of backward and forward LSTM states, respectively.

By design, the joint conditional distribution over latent variables z t and bt with parameters θ and ψ factorizes as p θ (z t |x 1:t-1 , z 1:t-1 )p ψ ( bt |z t ). This factorization enables us to formulate several helpful auxiliary cots, as defined in the next subsection. Further, p η (x t+1 |x 1:t , z 1:t , bt ) defines the generating model, which induces the distribution over the next observation given the previous states and the current input. During training, each step t is composed of an encoder which receives both the past and future summary via h t-1 and b t respectively, and a decoder that generates ht-1 and bt which are forced to be close enough to h t-1 and b t using two auxiliary reconstruction costs (dashed lines). This dependence between backward and forward LSTM through the latent random variable encourages the forward LSTM to learn a richer representation. During inference, the backward LSTM is removed. In this case, z t is sampled from the prior as in a typical VAE, which in our case, is defined as a function of h t-1 .

Then the marginal likelihood of each individual sequential data sample x can be written as

EQUATION

where q φ (z t |x) is the conditional inference model and Γ = {φ, θ, ψ, η} is the set of all parameters of the model. Here, we assume that all conditional distributions belong to parametrized families of distributions which can be evaluated and sampled from efficiently.

Note that the joint distribution in equation ( 1) is intractable. Kingma & Welling (2014) demonstrated how to maximize a variational lower bound, L Γ , of the data log likelihood instead, which is given by

EQUATION

where D KL is the Kullback-Leibler (KL) divergence between the approximate posterior and the conditional prior (see the appendix). This is the approach we take.

2.1. TRAINING AND INFERENCE:

In the proposed variational Bi-LSTM, the latent variable z t is inferred as

EQUATION

in which [µ q,t , σ 2 q,t ] = f φ (h t-1 , b t ) where f φ is a multi-layered feed-forward network with Gaussian outputs. We assume that the prior over z t is a diagonal multivariate Gaussian distribution given by p θ (z

EQUATION

for a fully connected network f θ . This is important because, during generation (see Figure 1 -right, for a graphical representation), we will not have access to the backward LSTM. In this case, as in a VAE, we will sample from the prior for z t . Since we define the prior to be a function of h t-1 , the forward LSTM is encouraged during training to learn the dependency due to the backward hidden state b t .

The latent variable bt is meant to model information coming from the future of the sequence. Its conditional distribution is given by

EQUATION

where [µ b,t , σ 2 b,t ] = f ψ (z t ) for a fully connected neural network f ψ (See Figure 1 (a)). To encourage the encoding of future information in bt , we maximize the probability of the true backward hidden state, b t , under the distribution p ψ , as an auxiliary cost during training. In this way we treat bt as a predictor of b t , similarly to what was done by Sordoni et al. (2017) .

To capture information from the past in the latents, we similarly use ht-1 as a predictor of h t-1 . This is accomplished by maximizing the probability of the latter under the conditional distribution of the former, log p ξ ( ht-1 |z t ), as another auxiliary cost, where

EQUATION

Here, [µ h,t , σ 2 h,t ] is the output of a fully-connected neural network f ξ taking z t as input. The auxiliary costs arising from distributions p ξ and p ψ teach the variational Bi-LSTM to encode past and future information into the latent space of z.

We assume that parameters of the generating distribution p η (x t+1 |x 1:t , z 1:t , bt ) are computed via MLP, taking the form of either a Gaussian distribution output in the continuous case or categorical proportions output in the discrete (ie, one-hot) prediction case (See Figure 1(b) ).

All the parameters in Γ and ξ are updated based on backpropagation through time (Rumelhart et al., 1988) using the reparameterization trick (Kingma & Welling, 2014) , where the gradients are computed by differentiating of the following function:

L(x; Γ, ξ) = T t=0 E z 1:T ∼q φ (z|x) E bt∼pψ(b|z1:t) log p(x t+1 |x 1:t , z 1:t , bt )+ α log p ψ (b t |z t ) + β log p ξ (h t-1 |z t ) - D KL (q φ (z t |x) p θ (z t |x 1:t-1 , z 1:t-1 )).

Here, α and β are non-negative real numbers. We improve training convergence with a trick for the variational Bi-LSTM, which we refer to as skip gradient, meant to ease learning of the latent variables. It is well known that autoregressive decoder models tend to ignore their stochastic variables (Bowman et al., 2015) . Skip gradient is a technique to encourage that relevant summaries of the past and the future are encoded in the latent space. The idea is to skip the gradient of the stochastic operations with respect to the recurrent units through time. To achieve this, at each time step, a mask drawn from a Bernoulli distribution governs whether to skip the gradient or to backpropagate it for a given data point.

3. EXPERIMENTAL RESULTS:

In this section demonstrate the effectiveness of our proposed model on several tasks. We present experimental results obtained when training the Variational Bi-LSTM on various sequential datasets: Penn Treebank (PTB), IMDB, TIMIT, Blizzard, and Sequential MNIST. Our main goal is to ensure that the model proposed in Section 2 can benefit from a generated relevant summary of the future that yields competitive results. In all experiments, we train all the models using ADAM optimizer (Kingma & Ba, 2014) and we set all MLPs in Section 2 to have one hidden layer with leaky-ReLU hidden activation. All the models are implemented using Theano (Theano Development Team, 2016) and the code is available at https://anonymous.url.

Blizzard: Blizzard is a speech model dataset with 300 hours of English, spoken by a single female speaker. We report the average log-likelihood for half-second sequences (Fraccaro et al., 2016) . In our experimental setting, we use 1024 hidden units for MLPs, 1024 LSTM units and 512 latents. Our model is trained using learning rate of 0.001 and minibatches of size 32 and we set α = β = 1. A fully factorized multivariate Gaussian distribution is used as the output distribution. The final lower bound estimation on TIMIT can be found in Table 1 . ≥ 11991 ≥ 60550 Z-Forcing (Sordoni et al., 2017) ≥ 14315 ≥ 68852 Variational Bi-LSTM ≥ 17319 ≥ 73315 TIMIT: Another speech modeling dataset is TIMIT with 6300 English sentences read by 630 speakers. Like the work done in (Fraccaro et al., 2016) , our model is trained on raw sequences of 200 dimensional frames. In our experiments, we use 1024 hidden units, 2048 LSTM units and 128 latent variables, and batch size of 128. We train the model using learning rate of 0.0001, α = 1 and β = 0.

The average log-likelihood for the sequences on test can be found in Table 1 .

Sequential MNIST: We use the MNIST dataset which is binarized according to (Murray & Salakhutdinov, 2009) and we downloaded in binrized from (Larochelle, 2011) . Our best model consists of 1024 hidden units, 1024 LSTM units and 256 latent variables. We train the model using a learning rate of 0.0001 and a batch size of 32. To reach the negative log-likelihood reported in Table 2 , we set α = 0.001 and β = 0. ≈ 84.55 NADE (Uria et al., 2016) 88.33 EoNADE-5 2hl (Raiko et al., 2014) 84.68 DLGM 8 (Salimans et al., 2014) ≈ 85.51 DARN 1hl (Gregor et al., 2015) ≈ 84.13 BiHM (Bornschein et al., 2015) ≈ 84.23 DRAW (Gregor et al., 2015) ≤ 80.97 PixelVAE (Gulrajani et al., 2016) ≈ 79.02 Prof. Forcing (Goyal et al., 2016) 79.58 PixelRNN (1-layer) (Oord et al., 2016) 80.75 PixelRNN (7-layer) (Oord et al., 2016) 79.20 Z-Forcing (Sordoni et al., 2017) ≤ 80.09 Variational Bi-LSTM ≤ 79.78 IMDB: It is a dataset consists of 350000 movie reviews (Diao et al., 2014) in which each sentence has less than 16 words and the vocabulary size is fixed to 16000 words. In this experiment, we use 500 hidden units, 500 LSTM units and latent variables of size 64. The model is trained with a batch size of 32 and a learning rate of 0.001 and we set α = β = 1. The word perplexity on valid and test dataset is shown in Table 3 .

PTB: Penn Treebank (Marcus et al. (1993) ) is a language model dataset consists of 1 million words. We train our model with 1024 LSTM units, 1024 hidden units, and the latent variables of size 128. We train the model using a standard Gaussian prior, a learning rate of 0.001 and batch size of 50 and 

4. ABLATION STUDIES:

The goal of this section is to study the importance of the various components in our model to avoid any triviality. The experiments are as follows:

1. Reconstruction loss on h t vs activity regularization on h t The authors of Merity et al. (2017) study the importance of activity regularization (AR) on the hidden states on LSTMs given as,

EQUATION

However, since our model's reconstruction term on h t can be decomposed as,

EQUATION

we perform experiments to confirm that the gains in our approach is not due to the 2 regularization alone since our regularization encapsulates an 2 term along with the dot product term.

We use activity regularization using hyperparameter α ∈ {0.001, 1, 4, 8, 16} in place of reconstruction term in our model and study the test perplexity. The results are shown in table 5 . We find that in all the cases performance using activity regularization is worse compared with our best model shown in table 3 . Interestingly we found that the model with fixed prior performed similarly in the case of PTB, but hurt performance in the other cases, which can be explained given their large KL divergence values in the original experiments.

5. RELATED WORK:

Variational auto-encoders (Kingma & Welling, 2014) can be easily combined with many deep learning models. They have been applied in the feed-forward setting but they have also found usage in RNNs to better capture variation in sequential data (Sordoni et al., 2017; Fraccaro et al., 2016; Chung et al., 2015; Bayer & Osendorfer, 2014) . VAEs consists of several muti-layer neural networks as probabilistic encoders and decoders and training is based on the gradient on log-likelihood lower bound (as the likelihood is in general intractable) of the model parameters Γ along with a reparametrization trick. The derived variational lower-bound L Γ for an observed random variable x is:

EQUATION

where D KL denotes the Kullback-Leibler divergence and p θ is the prior over a latent variable z. The KL divergence tem can be expressed as the difference between the entropy of q φ and the prior and fortunately, it can be computed and differentiated without estimation for some distribution families like Gaussians. Although maximizing the log-likelihood corresponds to minimizing the KL divergence, we have to ensure that the resulting q φ remains far enough from an undesired equilibrium state where q φ is almost everywhere equal to the prior over latent variables. Combining recurrent neural networks with variational auto encoders can lead to powerful generative models that are capable of capturing the variations in data, however, they suffer badly from this optimization issue as discussed by Bowman et al. (2015) .

Recently, VAEs have also been applied to Bi-LSTMs by Sordoni et al. (2017) through a technique called Z-forcing. It is a powerful generative autoregressive model which is trained using the following variational evidence lower-bound L(x; θ, φ, ξ) = E q φ (z |x) log p θ (x t+1 |x 1:t , z 1:t ) -D KL (q φ (z t |x) (p θ (z t |x 1:t-1 , z 1:t-1 ))

plus an auxiliary cost as a regularizer which is defined as log p ξ (b t |z t ). They show that the auxiliary cost helps in improving the final performance; however during inference the backward reconstructions have not been used in their approach. In our ablation study section below, we show experimentally that this connection is important towards improving the performance of Bi-LSTMs as is the case in our model.

6. CONCLUSION:

Variational Bi-LSTMs are powerful autoregressive generative models that are capable of learning better representations by creating a channel to exchange the information of the past and the future. Moreover, the conditional distribution over the backward LSTM variables is learned that can lead to better learning results in practice. Furthermore, Variational Bi-LSTM model acts as a regularizer and makes both networks be informative enough to perform well on different benchmark problems taken from the literature.

:



Figures and Tables:

Figure1: Graphical description of our proposed variational Bi-LSTM model during train phase (left) and inference phase (right). During training, each step t is composed of an encoder which receives both the past and future summary via h t-1 and b t respectively, and a decoder that generates ht-1 and bt which are forced to be close enough to h t-1 and b t using two auxiliary reconstruction costs (dashed lines). This dependence between backward and forward LSTM through the latent random variable encourages the forward LSTM to learn a richer representation. During inference, the backward LSTM is removed. In this case, z t is sampled from the prior as in a typical VAE, which in our case, is defined as a function of h t-1 .

Figure 2: Evolution of the average of log-likelihood during training of Variational Bi-LSTMs with and without using skip gradient and auxiliary costs on PTB and Blizzard.

Figure3: Evolution of the bits per character on PTB validation with sampling latent variables z from N (0, I) during training or using a fixed vector which we set to be the mean of latent variables. Interestingly, not sampling from prior during inference does not hurt the final performance on PTB.

The average of log-likelihood per sequence on Blizzard and TIMIT testset

Table content: <table><tr><td>Model</td><td colspan="2">Blizzard TIMIT</td></tr><tr><td>RNN-Gauss</td><td>3539</td><td>-1900</td></tr><tr><td>RNN-GMM</td><td>7413</td><td>26643</td></tr><tr><td>VRNN-I-Gauss</td><td colspan="2">≥ 8933 ≥ 28340</td></tr><tr><td>VRNN-Gauss</td><td colspan="2">≥ 9223 ≥ 28805</td></tr><tr><td>VRNN-GMM</td><td colspan="2">≥ 9392 ≥ 28982</td></tr><tr><td>SRNN (smooth+res q )</td><td/><td/></tr></table>

The average of negative log-likelihood on sequential MNIST

Table content: <table><tr><td>Models</td><td>Seq-MNIST</td></tr><tr><td>DBN 2hl (Germain et al., 2015)</td><td/></tr></table>

Word perplexity on IMDB on valid and test sets

Table content: <table><tr><td>Model</td><td>Valid Test</td></tr><tr><td>Gated Word-Char</td><td>70.60 70.87</td></tr><tr><td colspan="2">Z-Forcing (Sordoni et al., 2017) 56.48 65.68</td></tr><tr><td>Variational Bi-LSTM</td><td>51.43 51.60</td></tr></table>

Bits Per Character (BPC) on PTB valid and test sets

Table content: <table><tr><td>Model</td><td colspan="2">Valid Test</td></tr><tr><td>Unregularized LSTM</td><td colspan="2">1.47 1.36</td></tr><tr><td>Weight noise</td><td colspan="2">1.51 1.34</td></tr><tr><td>Norm stabilizer</td><td colspan="2">1.46 1.35</td></tr><tr><td>Stochastic depth</td><td colspan="2">1.43 1.34</td></tr><tr><td>Recurrent dropout</td><td colspan="2">1.40 1.29</td></tr><tr><td>Zoneout (Krueger et al. (2016))</td><td colspan="2">1.36 1.25</td></tr><tr><td>RBN (Cooijmans et al. (2016))</td><td>-</td><td>1.32</td></tr><tr><td>H-LSTM + LN (Ha et al. (2016))</td><td colspan="2">1.28 1.25</td></tr><tr><td>3-HM-LSTM + LN (Chung et al., 2016)</td><td>-</td><td>1.24</td></tr><tr><td>2-H-LSTM + LN (Ha et al. (2016))</td><td colspan="2">1.25 1.22</td></tr><tr><td>Z-Forcing</td><td colspan="2">1.29 1.26</td></tr><tr><td>Variational Bi-LSTM</td><td colspan="2">1.26 1.23</td></tr></table>

Perplexity on IMDB using different coefficient γ for activity regularization

Table content: <table><tr><td>γ</td><td>0.001</td><td>1.</td><td>4.</td><td>8.</td><td>16.</td></tr><tr><td colspan="6">Test perplexity 56.07 60.74 69.97 77.24 86.72</td></tr></table>

KL divergence of the Variational Bi-LSTMIn our variational Bi-LSTM model, we propose to have the encoder prior over z t as a function of the previous forward LSTM hidden state h t-1 . This is done to omit the need of the backward LSTM during inference because it is unavailable in practical scenarios since predictions are made in the forward direction. However, to study whether the model learns to use this encoder or not, we record the KL divergence value of the best validation model for the various datasets. The results are reported in table 6. We can see that the KL divergence values are large in the case of IMDB, TIMIT and Blizzard datasets, but small in the case of Seq-MNIST and PTB. To further explore, we ran experiments on these datasets with fixed standard Gaussian prior like in the case of traditional VAE.

Table content: <table><tr><td colspan="5">Dataset PTB Seq-MNIST IMDB TIMIT Blizzard</td></tr><tr><td>KL</td><td>0.001</td><td>0.02</td><td>0.18</td><td>3204.71 3799.79</td></tr><tr><td colspan="5">2. Use of parametric encoder prior vs. fixed Gaussian prior</td></tr></table>

Peer review (iclr2018):
 This paper builds a sequential deep generative model with (1) an inference network parameterized by an RNN running from the future to the past and (2) an explicit representation of the hidden state of the backward RNN in the generative model. The model is validated on held-out likelihood via the ELBO on text, handwriting, speech and images. It presents good emprical results and works at par with or better than many other baselines considered. The main source of novelty here the choice made in the transition function of z_t to also incorporate an explicit variable to models the hidden state of the backward RNN at inference time and use that random variable in the generative process. This is a choice of structural prior for the transition function of the generative model that I think lends it more expressivity realizing the empirical gains obtained. I found the presentation of both the model and learning objective to be confusing and had a hard time following it. The source of my confusion is is that .tilde{b} (the following argument applies equivalently to .tilde{h}) is argued to be a latent variable. Yet it is not inferred (via a variational distribution) during training. Please correct me if I*m wrong but I believe that an easier to understand way to explain the model is as follows: both .tilde{b} and .tilde{h} should be presented as *observed* random variables during *training* and latent at inference time. Training then comprises maximizing the marginal likelihood of the data *and* maximizing the conditional likelihood of the two observed variables(via p_psi and p_eta; conditioned on z_t). Under this view, setting beta to 0 simply corresponds to not observing .tilde{h_t}. alpha can be annealed but should never be set to anything less than 1 without breaking the semantics of the learned generative model. Consider Figure 1(b). It seems that the core difference between this work and [Chung et. al] is that this work parameterizes q(Z_t) using x_t....x_T (via a backward RNN). This choice of inference network can be motivated from the point of view of building a better approximation to the structure of the posterior distribution of Z_t under the generative model. Both [Fracarro et. al] and [Krishnan et. al] (https://arxiv.org/pdf/1609.09869.pdf) use RNNs from x_T to x_1 to train sequential state space models. [Gao et. al] (https://arxiv.org/pdf/1605.08454.pdf) derive an inference network with a block-diagonal structure motivated by correlations in the posterior distribution. Incorporating a discussion around this idea would provide useful context for where this work stands amongst the many sequential deep generative models in the literature. Questions for the authors: * How important is modeling .tilde{h_t} in TIMIT, Blizzard and IMDB? * Did you try annealing the KL divergence in the PTB experiment. Based on the KL divergence you report it seems the latent variable is not necessary. Overall, I find the model to be interesting and it performs well empirically. However, the text of the paper lacks a bit of context and clarity that makes understanding it challenging to understand in its current form.