Sequential Neural Models with Stochastic Layers

Introduction

Recurrent neural networks (RNNs) are able to represent long-term dependencies in sequential data, by adapting and propagating a deterministic hidden (or latent) state . There is recent evidence that when complex sequences such as speech and music are modeled, the performances of RNNs can be dramatically improved when uncertainty is included in their hidden states . In this paper we add a new direction to the explorer’s map of treating the hidden RNN states as uncertain paths, by including the world of state space models (SSMs) as an RNN layer. By cleanly delineating a SSM layer, certain independence properties of variables arise, which are beneficial for making efficient posterior inferences. The result is a generative model for sequential data, with a matching inference network that has its roots in variational auto-encoders (VAEs).

SSMs can be viewed as a probabilistic extension of RNNs, where the hidden states are assumed to be random variables. Although SSMs have an illustrious history , their stochasticity has limited their widespread use in the deep learning community, as inference can only be exact for two relatively simple classes of SSMs, namely hidden Markov models and linear Gaussian models, neither of which are well-suited to modeling long-term dependencies and complex probability distributions over high-dimensional sequences. On the other hand, modern RNNs rely on gated nonlinearities such as long short-term memory (LSTM) cells or gated recurrent units (GRUs) , that let the deterministic hidden state of the RNN act as an internal memory for the model. This internal memory seems fundamental to capturing complex relationships in the data through a statistical model.

This paper introduces the stochastic recurrent neural network (SRNN) in Section 3. SRNNs combine the gated activation mechanism of RNNs with the stochastic states of SSMs, and are formed by stacking a RNN and a nonlinear SSM. The state transitions of the SSM are nonlinear and are parameterized by a neural network that also depends on the corresponding RNN hidden state. The SSM can therefore utilize long-term information captured by the RNN.

We use recent advances in variational inference to efficiently approximate the intractable posterior distribution over the latent states with an inference network . The form of our variational approximation is inspired by the independence properties of the true posterior distribution over the latent states of the model, and allows us to improve inference by conveniently using the information coming from the whole sequence at each time step. The posterior distribution over the latent states of the SRNN is highly non-stationary while we are learning the parameters of the model. To further improve the variational approximation, we show that we can construct the inference network so that it only needs to learn how to compute the mean of the variational approximation at each time step given the mean of the predictive prior distribution.

In Section 4 we test the performances of SRNN on speech and polyphonic music modeling tasks. SRNN improves the state of the art results on the Blizzard and TIMIT speech data sets by a large margin, and performs comparably to competing models on polyphonic music modeling. Finally, other models that extend RNNs by adding stochastic units will be reviewed and compared to SRNN in Section 5.

Recurrent Neural Networks and State Space Models

Recurrent neural networks and state space models are widely used to model temporal sequences of vectors $\mathbf{x}_{1:T}=(\mathbf{x}_{1},\mathbf{x}_{2},\dots,\mathbf{x}_{T})$ that possibly depend on inputs $\mathbf{u}_{1:T}=(\mathbf{u}_{1},\mathbf{u}_{2},\dots,\mathbf{u}_{T})$. Both models rest on the assumption that the sequence $\mathbf{x}_{1:t}$ of observations up to time $t$ can be summarized by a latent state $\mathbf{d}_{t}$ or $\mathbf{z}_{t}$, which is deterministically determined ($\mathbf{d}_{t}$ in a RNN) or treated as a random variable which is averaged away ($\mathbf{z}_{t}$ in a SSM). The difference in treatment of the latent state has traditionally led to vastly different models: RNNs recursively compute $\mathbf{d}_{t}=f(\mathbf{d}_{t-1},\mathbf{u}_{t})$ using a parameterized nonlinear function $f$, like a LSTM cell or a GRU. The RNN observation probabilities $p(\mathbf{x}_{t}|\mathbf{d}_{t})$ are equally modeled with nonlinear functions. SSMs, like linear Gaussian or hidden Markov models, explicitly model uncertainty in the latent process through $\mathbf{z}_{1:T}$. Parameter inference in a SSM requires $\mathbf{z}_{1:T}$ to be averaged out, and hence $p(\mathbf{z}_{t}|\mathbf{z}_{t-1},\mathbf{u}_{t})$ and $p(\mathbf{x}_{t}|\mathbf{z}_{t})$ are often restricted to the exponential family of distributions to make many existing approximate inference algorithms applicable. On the other hand, averaging a function over the deterministic path $\mathbf{d}_{1:T}$ in a RNN is a trivial operation. The striking similarity in factorization between these models is illustrated in Figures 1(a) and 1(b).

Can we combine the best of both worlds, and make the stochastic state transitions of SSMs nonlinear whilst keeping the gated activation mechanism of RNNs? Below, we show that a more expressive model can be created by stacking a SSM on top of a RNN, and that by keeping them layered, the functional form of the true posterior distribution over $\mathbf{z}_{1:T}$ guides the design of a backward-recursive structured variational approximation.

Stochastic Recurrent Neural Networks

We define a SRNN as a generative model $p_{\theta}$ by temporally interlocking a SSM with a RNN, as illustrated in Figure 2(a). The joint probability of a single sequence and its latent states, assuming knowledge of the starting states $\mathbf{z}_{0}=\mathbf{0}$ and $\mathbf{d}_{0}=\mathbf{0}$, and inputs $\mathbf{u}_{1:T}$, factorizes as

The stochastic variables $\mathbf{z}_{1:T}$ of the nonlinear SSM cannot be analytically integrated out to obtain $\mathcal{L}(\theta)$ in (2). Instead of maximizing $\mathcal{L}$ with respect to $\theta$, we maximize a variational evidence lower bound (ELBO) $\mathcal{F}(\theta,\phi)=\sum_{i}\mathcal{F}_{i}(\theta,\phi)\leq\mathcal{L}(\theta)$ with respect to both $\theta$ and the variational parameters $\phi$ . The ELBO is a sum of lower bounds $\mathcal{F}_{i}(\theta,\phi)\leq\mathcal{L}_{i}(\theta)$, one for each sequence $i$,

where $A=\{\mathbf{u}_{1:T},\mathbf{z}_{0},\mathbf{d}_{0}\}$ is a notational shorthand. Each sequence’s approximation $q_{\phi}$ shares parameters $\phi$ with all others, to form the auto-encoding variational Bayes inference network or variational auto encoder (VAE) shown in Figure 2(b). Maximizing $\mathcal{F}(\theta,\phi)$ – which we call “training” the neural network architecture with parameters $\theta$ and $\phi$ – is done by stochastic gradient ascent, and in doing so, both the posterior and its approximation $q_{\phi}$ change simultaneously. All the intractable expectations in (4) would typically be approximated by sampling, using the reparameterization trick or control variates to obtain low-variance estimators of its gradients. We use the reparameterization trick in our implementation. Iteratively maximizing $\mathcal{F}$ over $\theta$ and $\phi$ separately would yield an expectation maximization-type algorithm, which has formed a backbone of statistical modeling for many decades . The tightness of the bound depends on how well we can approximate the $i=1,\ldots,N$ factors $p_{\theta}(\mathbf{z}_{1:T_{i}}^{i},\mathbf{d}_{1:T_{i}}^{i}|\mathbf{x}_{1:T_{i}}^{i},A^{i})$ that constitute the true posterior over all latent variables with their corresponding factors $q_{\phi}(\mathbf{z}_{1:T_{i}}^{i},\mathbf{d}_{1:T_{i}}^{i}|\mathbf{x}_{1:T_{i}}^{i},A^{i})$. In what follows, we show how $q_{\phi}$ could be judiciously structured to match the posterior factors.

This choice exactly approximates one delta-function by itself, and simplifies the ELBO by letting them cancel out. By further taking the outer average in (4), one obtains

2 Exploiting the temporal structure

The true posterior distribution of the stochastic states $\mathbf{z}_{1:T}$, given both the data and the deterministic states $\mathbf{d}_{1:T}$, factorizes as $p_{\theta}(\mathbf{z}_{1:T}|\mathbf{d}_{1:T},\mathbf{x}_{1:T},\mathbf{u}_{1:T},\mathbf{z}_{0})=\prod_{t}p_{\theta}(\mathbf{z}_{t}|\mathbf{z}_{t-1},\mathbf{d}_{t:T},\mathbf{x}_{t:T})$. This can be verified by considering the conditional independence properties of the graphical model in Figure 2(a) using d-separation . This shows that, knowing $\mathbf{z}_{t-1}$, the posterior distribution of $\mathbf{z}_{t}$ does not depend on the past outputs and deterministic states, but only on the present and future ones; this was also noted in . Instead of factorizing $q_{\phi}$ as a mean-field approximation across time steps, we keep the structured form of the posterior factors, including $\mathbf{z}_{t}$’s dependence on $\mathbf{z}_{t-1}$, in the variational approximation

In (7) all time steps are taken into account when constructing the variational approximation at time $t$; this can therefore be seen as a smoothing problem. In our experiments we also consider filtering, where only the information up to time $t$ is used to define $q_{\phi}(\mathbf{z}_{t}|\mathbf{z}_{t-1},\mathbf{d}_{t},\mathbf{x}_{t})$. As the parameters $\phi$ are shared across time steps, we can easily handle sequences of variable length in both cases.

As both the generative model and inference network factorize over time steps in (1) and (7), the ELBO in (6) separates as a sum over the time steps

where $q_{\phi}^{*}(\mathbf{z}_{t-1})$ denotes the marginal distribution of $\mathbf{z}_{t-1}$ in the variational approximation to the posterior $q_{\phi}(\mathbf{z}_{1:t-1}|\widetilde{\mathbf{d}}_{1:T},\mathbf{x}_{1:T},\mathbf{z}_{0})$, given by

We can interpret (9) as having a VAE at each time step $t$, with the VAE being conditioned on the past through the stochastic variable $\mathbf{z}_{t-1}$. To compute (8), the dependence on $\mathbf{z}_{t-1}$ needs to be integrated out, using our posterior knowledge at time $t-1$ which is given by $q_{\phi}^{*}(\mathbf{z}_{t-1})$. We approximate the outer expectation in (8) using a Monte Carlo estimate, as samples from $q_{\phi}^{*}(\mathbf{z}_{t-1})$ can be efficiently obtained by ancestral sampling. The sequential formulation of the inference model in (7) allows such samples to be drawn and reused, as given a sample $\mathbf{z}_{t-2}^{(s)}$ from $q_{\phi}^{*}(\mathbf{z}_{t-2})$, a sample $\mathbf{z}_{t-1}^{(s)}$ from $q_{\phi}(\mathbf{z}_{t-1}|\mathbf{z}_{t-2}^{(s)},\widetilde{\mathbf{d}}_{t-1:T},\mathbf{x}_{t-1:T})$ will be distributed according to $q_{\phi}^{*}(\mathbf{z}_{t-1})$.

3 Parameterization of the inference network

where we used the parameterization of the prior distribution in (3). We estimate the integral required to compute \widehat{\text{\boldmath\mu}}_{t}^{(p)} by reusing the samples that were needed for the Monte Carlo estimate of the ELBO in (8). This predictive prior mean can then be used in the parameterization of the mean of the variational approximation $q_{\phi}(\mathbf{z}_{t}|\mathbf{z}_{t-1},\mathbf{a}_{t})$,

Results

In this section the SRNN is evaluated on the modeling of speech and polyphonic music data, as they have shown to be difficult to model without a good representation of the uncertainty in the latent states . We test SRNN on the Blizzard and TIMIT raw audio data sets (Table 1) used in . The preprocessing of the data sets and the testing performance measures are identical to those reported in . Blizzard is a dataset of 300 hours of English, spoken by a single female speaker. TIMIT is a dataset of 6300 English sentences read by 630 speakers. As done in , for Blizzard we report the average log-likelihood for half-second sequences and for TIMIT we report the average log likelihood per sequence for the test set sequences. Note that the sequences in the TIMIT test set are on average 3.1s long, and therefore 6 times longer than those in Blizzard. For the raw audio datasets we use a fully factorized Gaussian output distribution. Additionally, we test SRNN for modeling sequences of polyphonic music (Table 2), using the four data sets of MIDI songs introduced in . Each data set contains more than 7 hours of polyphonic music of varying complexity: folk tunes (Nottingham data set), the four-part chorales by J. S. Bach (JSB chorales), orchestral music (MuseData) and classical piano music (Piano-midi.de). For polyphonic music we use a Bernoulli output distribution to model the binary sequences of piano notes. In our experiments we set $\mathbf{u}_{t}=\mathbf{x}_{t-1}$, but $\mathbf{u}_{t}$ could also be used to represent additional input information to the model.

Table 1 compares the average log-likelihood per test sequence of SRNN to the results from . For RNNs and VRNNs the authors of test two different output distributions, namely a Gaussian distribution (Gauss) and a Gaussian Mixture Model (GMM). VRNN-I differs from the VRNN in that the prior over the latent variables is independent across time steps, and it is therefore similar to STORN . For SRNN we compare the smoothing and filtering performance (denoted as smooth and filt in Table 1), both with the residual term from (12) and without it (10) (denoted as $\text{Res}_{q}$ if present). We prefer to only report the more conservative evidence lower bound for SRNN, as the approximation of the log-likelihood using standard importance sampling is known to be difficult to compute accurately in the sequential setting . We see from Table 1 that SRNN outperforms all the competing methods for speech modeling. As the test sequences in TIMIT are on average more than 6 times longer than the ones for Blizzard, the results obtained with SRNN for TIMIT are in line with those obtained for Blizzard. The VRNN, which performs well when the voice of the single speaker from Blizzard is modeled, seems to encounter difficulties when modeling the 630 speakers in the TIMIT data set. As expected, for SRNN the variational approximation that is obtained when future information is also used (smoothing) is better than the one obtained by filtering. Learning the residual between the prior mean and the mean of the variational approximation, given in (12), further improves the performance in 3 out of 4 cases.

Polyphonic music.

Table 2 compares the average log-likelihood on the test sets obtained with SRNN and the models introduced in . As done for the speech data, we prefer to report the more conservative estimate of the ELBO in Table 2, rather than approximating the log-likelihood with importance sampling as some of the other methods do. We see that SRNN performs comparably to other state of the art methods in all four data sets. We report the results using smoothing and learning the residual between the mean of the predictive prior and the mean of the variational approximation, but the performances using filtering and directly learning the mean of the variational approximation are now similar. We believe that this is due to the small amount of data and the fact that modeling MIDI music is much simpler than modeling raw speech signals.

Related work

A number of works have extended RNNs with stochastic units to model motion capture, speech and music data . The performances of these models are highly dependent on how the dependence among stochastic units is modeled over time, on the type of interaction between stochastic units and deterministic ones, and on the procedure that is used to evaluate the typically intractable log likelihood. Figure 4 highlights how SRNN differs from some of these works.

In STORN (Figure 4(a)) and DRAW the stochastic units at each time step have an isotropic Gaussian prior and are independent between time steps. The stochastic units are used as an input to the deterministic units in a RNN. As in our work, the reparameterization trick is used to optimize an ELBO.

Conclusion

This work has shown how to extend the modeling capabilities of recurrent neural networks by combining them with nonlinear state space models. Inspired by the independence properties of the intractable true posterior distribution over the latent states, we designed an inference network in a principled way. The variational approximation for the stochastic layer was improved by using the information coming from the whole sequence and by using the $\text{Res}_{q}$ parameterization to help the inference network to track the non-stationary posterior. SRNN achieves state of the art performances on the Blizzard and TIMIT speech data set, and performs comparably to competing methods for polyphonic music modeling.

Acknowledgements

We thank Casper Kaae Sønderby and Lars Maaløe for many fruitful discussions, and NVIDIA Corporation for the donation of TITAN X and Tesla K40 GPUs. Marco Fraccaro is supported by Microsoft Research through its PhD Scholarship Programme.

References

Appendix A Experimental setup

The sampling rate is 16KHz and the raw audio signal is normalized using the global mean and standard deviation of the traning set. We split the raw audio signals in chunks of 2 seconds. The waveforms are then divided into non-overlapping vectors of size 200. The RNN thus runs for 160 steps2s$\cdot$16Khz / 200 = 160. The model is trained to predict the next vector ($\mathbf{x}_{t}$) given the current one ($\mathbf{u}_{t}$). During training we use backpropagation through time (BPTT) for 0.5 seconds, i.e we have 4 updates for each 2 seconds of audio. For the first 0.5 second we initialize hidden units with zeros and for the subsequent 3 chunks we use the previous hidden states as initialization.

For Blizzard we split the data using 90% for training, 5% for validation and 5% for testing. For testing we report the average log-likelihood per 0.5s sequences. For TIMIT we use the predefined test set for testing and split the rest of the data into 95% for training and 5% for validation. The training and testing setup are identical to the ones for Blizzard. For TIMIT the test sequences have variable length and are on average $3.1s$, i.e. more than 6 times longer than Blizzard.

A.2 Polyphonic music

We use the same model architecture as in the speech modeling experiments, except for the output Bernoulli variables used to model the active notes. We reduced the number of parameters in the model to 300 deterministic hidden units for the GRU networks, and 100 stochastic units whose distributions are parameterized with neural networks with 1 layer of 500 units.