S TATE S PACE LSTM M ODELS WITH P ARTICLE MCMC I NFERENCE
Anonymous authors
Paper under double-blind review
A BSTRACT
Long Short-Term Memory (LSTM) is one of the most powerful sequence models.
Despite the strong performance, however, it lacks the nice interpretability as in state space models.
In this paper, we present a way to combine the best of both worlds by introducing State Space LSTM (SSL) models that generalizes the earlier work Zaheer et al.
(2017) of combining topic models with LSTM.
However, unlike Zaheer et al.
(2017), we do not make any factorization assumptions in our inference algorithm.
We present an efﬁcient sampler based on sequential Monte
Carlo (SMC) method that draws from the joint posterior directly.
Experimental results conﬁrms the superiority and stability of this SMC inference algorithm on a variety of domains.
I NTRODUCTION
State space models (SSMs), such as hidden Markov models (HMM) and linear dynamical systems (LDS), have been the workhorse of sequence modeling in the past decades From a graphical model perspective, efﬁcient message passing algorithms (Stratonovich, 1960; Kalman, 1960) are available in compact closed form thanks to their simple linear Markov structure.
However, simplicity comes at a cost: real world sequences can have long-range dependencies that cannot be captured by Markov models; and the linearity of transition and emission restricts the ﬂexibility of the model for complex sequences.
A popular alternative is the recurrent neural networks (RNN), for instance the Long Short-Term
Memory (LSTM) (Hochreiter & Schmidhuber, 1997) which has become a standard for sequence modeling nowadays.
Instead of associating the observations with stochastic latent variables, RNN directly deﬁnes the distribution of each observation conditioned on the past, parameterized by a neural network.
The recurrent parameterization not only allows RNN to provide a rich function class, but also permits scalable stochastic optimization such as the backpropagation through time (BPTT) algorithm.
However, ﬂexibility does not come for free as well: due to the complex form of the transition function, the hidden states of RNN are often hard to interpret.
Moreover, it can require large amount of parameters for seemingly simple sequence models (Zaheer et al., 2017).
In this paper, we propose a new class of models State Space LSTM (SSL) that combines the best of both worlds.
We show that SSLs can handle nonlinear, non-Markovian dynamics like RNNs, while retaining the probabilistic interpretations of SSMs.
The intuition, in short, is to separate the state space from the sample space.
In particular, instead of directly estimating the dynamics from the observed sequence, we focus on modeling the sequence of latent states, which may represent the true underlying dynamics that generated the noisy observations.
Unlike SSMs, where the same goal is pursued under linearity and Markov assumption, we alleviate the restriction by directly modeling the transition function between states parameterized by a neural network.
On the other hand, we bridge the state space and the sample space using classical probabilistic relation, which not only brings additional interpretability, but also enables the LSTM to work with more structured representation rather than the noisy observations.
Indeed, parameter estimation of such models can be nontrivial.
Since the LSTM is deﬁned over a sequence of latent variables rather than observations, it is not straightforward to apply the usual BPTT algorithm without making variational approximations.
In Zaheer et al.
(2017), which is an instance of SSL, an EM-type approach was employed: the algorithm alternates between imputing the latent states and optimizing the LSTM over the imputed sequences.
However, as we show below, the inference implicitly assumes the posterior is factorizable through time.
This is a restrictive assumption since the beneﬁt of rich state transition brought by the LSTM may be neutralized by breaking down the posterior over time.
We present a general parameter estimation scheme for the proposed class of models based on sequential Monte Carlo (SMC) (Doucet et al., 2001), in particular the Particle Gibbs (Andrieu et al., 2010).
Instead of sampling each time point individually, we directly sample from the joint posterior without making limiting factorization assumptions.
Through extensive experiments we verify that sampling from the full posterior leads to signiﬁcant improvement in the performance.
Related works Enhancing state space models using neural networks is not a new idea.
Traditional approaches can be traced back to nonlinear extensions of linear dynamical systems, such as extended or unscented Kalman ﬁlters (Julier & Uhlmann, 1997), where both state transition and emission are generalized to nonlinear functions.
The idea of parameterizing them with neural networks can be found in Ghahramani & Roweis (1999), as well as many recent works (Krishnan et al., 2015;
Archer et al., 2015; Johnson et al., 2016; Krishnan et al., 2017; Karl et al., 2017) thanks to the development of recognition networks (Kingma & Welling, 2014; Rezende et al., 2014).
Enriching the output distribution of RNN has also regain popularity recently.
Unlike conventionally used multinomial output or mixture density networks (Bishop, 1994), recent approaches seek for more ﬂexible family of distributions such as restricted Boltzmann machines (RBM) (Boulanger-Lewandowski et al., 2012) or variational auto-encoders (VAE) (Gregor et al., 2015; Chung et al., 2015).
On the ﬂip side, there have been studies in introducing stochasticity to recurrent neural networks.
For instance, Pachitariu & Sahani (2012) and Bayer & Osendorfer (2014) incorporated independent latent variables at each time step; while in Fraccaro et al.
the RNN is attached to both latent states and observations.
We note that in our approach the transition and emission are decoupled, not only for interpretability but also for efﬁcient inference without variational assumptions.
On a related note, sequential Monte Carlo methods have recently received attention in approximating the variational objective (Maddison et al., 2017; Le et al., 2017; Naesseth et al., 2017).
Despite the similarity, we emphasize that the context is different: we take a stochastic EM approach, where the full expectation in E-step is replaced by the samples from SMC.
In contrast, SMC in above works is aimed at providing a tighter lower bound for the variational objective.
B ACKGROUND
In this section, we provide a brief review of some key ingredients of this paper.
We ﬁrst describe the
SSMs and the RNNs for sequence modeling, and then outline the SMC methods for sampling from a series of distributions.
2. S TATE SPACE MODELS
Consider a sequence of observations x 1: T = ( x 1 , . .
, x T ) and a corresponding sequence of latent states z 1: T = ( z 1 , . . .
, z T ) .
The SSMs are a class of graphical models that deﬁnes probabilistic dependencies between latent states and the observations.
A classical example of SSM is the (Gaussian)
LDS, where real-valued states evolve linearly over time under the ﬁrst-order Markov assumption.
Let x t ∈ R d and z t ∈ R k , the LDS can be expressed by two equations:
[Equation]
where A ∈ R k × k , C ∈ R d × k , and Q and R are covariance matrices of corresponding sizes.
They are widely applied in modeling the dynamics of moving objects, with z t representing the true state of the system, such as location and velocity of the object, and x t being the noisy observation under zero-mean Gaussian noise.
We mention two important inference tasks (Koller & Friedman, 2009) associated with SSMs.
The ﬁrst tasks is ﬁltering : at any time t , compute p ( z t | x 1: t ) , i.e . the most up-to-date belief of the state z t conditioned on all past and current observations x 1: t .
The other task is smoothing , which computes p ( z t | x 1: T ) , i.e .
the update to the belief of a latent state by incorporating future observations.
One of the beauties of SSMs is that these inference tasks are available in closed form, thanks to the simple Markovian dynamics of the latent states.
For instance, the forward-backward algorithm (Stratonovich, 1960), the Kalman ﬁlter (Kalman, 1960), and RTS smoother (Rauch et al., 1965) are widely appreciated in the literature of HMM and LDS.
Having obtained the closed form ﬁltering and smoothing equations, one can make use of the EM algorithm to ﬁnd the maximum likelihood estimate (MLE) of the parameters given observations.
In the case of LDS, the E-step can be computed by RTS smoother and the M-step is simple subproblems such as least-squares regression.
We refer to Ghahramani & Hinton (1996) for a full exposition on learning the parameters of LDS using EM iterations.
2. R ECURRENT NEURAL NETWORKS
RNNs have received remarkable attention in recent years due to their strong benchmark performance as well as successful applications in real-world problems.
Unlike SSMs, RNNs aim to directly learn the complex generative distribution of p ( x t | x 1: t − 1 ) using a neural network, with the help of a deterministic internal state s t : [Equation]
[Equation]
) where RNN ( · , · ) is the transition function deﬁned by a neural network, and g ( · ) is an arbitrary differentiable function that maps the RNN state s t to the parameter of the distribution of x t .
The ﬂexibility of the transformation function allows the RNN to learn from complex nonlinear nonGaussian sequences.
Moreover, since the state s t is a deterministic function of the past observations x 1: t − 1 , RNNs can capture long-range dependencies, for instance matching brackets in programming languages (Karpathy et al., 2015).
The BPTT algorithm can be used to ﬁnd the MLE of the parameters of RNN ( · , · ) and g ( · ) .
However, although RNNs can, in principle, model long-range dependencies, directly applying BPTT can be difﬁcult in practice since the repeated application of a squashing nonlinear activation function, such as tanh or logistic sigmoid, results in an exponential decay in the error signal through time.
LSTMs (Hochreiter & Schmidhuber, 1997) are designed to cope with the such vanishing gradient problems, by introducing an extra memory cell that is constructed as a linear combination of the previous state and signal from the input.
In this work, we also use LSTMs as building blocks, as in
Zaheer et al.
2. S EQUENTIAL M ONTE C ARLO
Sequential Monte Carlo (SMC) (Doucet et al., 2001) is an algorithm that samples from a series of potentially unnormalized densities π 1 ( z 1 ) , . .
, π T ( z 1: T ) .
At each step t , SMC approximates the target density π t with P weighted particles using importance distribution f
( z t | z 1: t − 1 ) :
[Equation]
[Equation]
where α pt is the importance weight of the p -th particle and δ x is the Dirac point mass at x .
Repeating this approximation for every t leads to the SMC method, outlined in Algorithm 1.
The key to this method lies in the resampling, which is implemented by repeatedly drawing the ancestors of particles at each step.
Intuitively, it encourages the particles with a higher likelihood to [Figure] survive longer, since the weight reﬂects the likelihood of the particle path.
The ﬁnal Monte Carlo estimate (4) consists of only survived particle paths, and sampling from this point masses is equivalent to choosing a particle path according to the last weights α T .
We refer to Doucet et al.
Andrieu et al.
(2010) for detailed proof of the method.
S TATE S PACE LSTM MODELS
In this section, we present the class of State Space LSTM (SSL) models that combines interpretability of SSMs and ﬂexibility of LSTMs.
The key intuition, motivated by SSMs, is to learn dynamics in the state space, rather than in the sample space.
However, we do not assume transition in the state space is linear, Gaussian, or Markovian.
Existing approaches such as the extended Kalman ﬁlter (EKF) attempted to work with a general nonlinear transition function.
Unfortunately, additional ﬂexibility also introduced extra difﬁculty in the parameter estimation:
EKF relies heavily on linearizing the nonlinear functions.
We propose to use
LSTM to model the dynamics in the latent state space, as they can learn from complex sequences without making limiting assumptions.
The BPTT algorithm is also well established so that no additional approximation is needed in training the latent dynamics.
Generative process Let h ( · ) be the emission function that maps a latent state to a parameter of the sample distribution.
As illustrated in Figure 1 (a), the generative process of SSL for a single sequence is:
• For t = 1 , . . .
, T : 1.
Perform LSTM transition: s t =
LSTM ( s t − 1 , z t − 1 ) 2.
Draw latent state:
z t ∼ p ( z ; g ( s t )) 3.
Draw observation: x t ∼ p ( x ; h ( z t ))
The generative process speciﬁes the following joint likelihood, with a similar factorization as SSMs except for the Markov transition:
[Equation]
where p ω ( z t | z 1: t − 1 ) = p ( z t ; g ( s t )) , ω is the set of parameters of LSTM ( · , · ) and g ( · ) , and φ is the parameters of h ( · ) .
The structure of the likelihood function is better illustrated in Figure 1 (b), where each latent state z t is dependent to all previous states z 1: t − 1 after substituting s t recursively.
This allows the SSL to have non-Markovian state transition, with parsimonious parameterization thanks to the recurrent structure of LSTMs.
Parameter estimation We continue with a single sequence for the ease of notation.
A variational lower bound to the marginal data likelihood is given by
[Equation]
where q ( z 1: T ) is the variational distribution.
Following the (stochastic) EM approach, iteratively maximizing the lower bound w.r.t.
q and the model parameters ( ω, φ ) leads to the following updates:
• E-step: The optimal variational distribution is given by the posterior:
[Equation]
In the case of LDS or HMM, efﬁcient smoothing algorithms such as the RTS smoother or the forward-backward algorithm are available for computing the posterior expectations of sufﬁcient statistics.
However, without Markovian state transition, although the forward messages can still be computed, the backward recursion can no longer evaluated or efﬁciently approximated.
• S-step: Due to the difﬁculties in taking expectations, we take an alternative approach to collect posterior samples instead:
[Equation]
given only the ﬁltering equations.
We discuss the posterior sampling algorithm in detail in the next section.
• M-step:
Given the posterior samples z (cid:63) 1: T , which can be seen as Monte Carlo estimate of the expectations, the subproblem for ω and φ are
[Equation]
which is exactly the MLE of an LSTM, with z (cid:63) 1: T serving as the input sequence, and the MLE of the given emission model.
Having seen the generative model and the estimation algorithm, we can now discuss some instances of the proposed class of models.
In particular, we provide two examples of SSL, for continuous and discrete latent states respectively.
Example 1 (Gaussian SSL) Suppose z t and x t are real-valued vectors.
A typical choice of the transition and emission is the Gaussian distribution:
[Equation]
where g µ ( · ) and g σ ( · ) map to the mean and the covariance of the Gaussian respectively, and similarly h µ ( · ) and h σ ( · ) .
For closed form estimates for the emission parameters, one can further assume
[Equation]
where C is a matrix that maps from state space to sample space, and R is the covariance matrix with appropriate size.
The MLE of φ =
( C, b, R ) is then given by the least squares ﬁt.
Example 2 (Topical SSL, (Zaheer et al., 2017)) Consider x 1: T as the sequence of websites a user has visited.
One might be tempted to model the user behavior using an LSTM, however due to the enormous size of the Internet, it is almost impossible to even compute a softmax output to get a discrete distribution over the websites.
There are approximation methods for large vocabulary problems in RNN, such as the hierarchical softmax (Morin & Bengio, 2005).
However, another interesting approach is to operate on a sequence with a “compressed” vocabulary, while learning how to perform such compression at the same time.
Let z t be the indicator of a “topic”, which is a distribution over the vocabulary as in Blei et al.
Accordingly, deﬁne
[Equation]
[Equation]
where W is a matrix that maps LSTM states to latent states, b is a bias term, and φ z t is a point in the probability simplex.
If z t lies in a lower dimension than x t , the LSTM is effectively trained over a sequence z 1: T with a reduced vocabulary.
On the other hand, the probabilistic mapping between z t and x t is interpretable, as it learns to group similar x t ’s together.
The estimation of φ is typically performed under a Dirichlet prior, which then corresponds to the MAP estimate of the Dirichlet distribution (Zaheer et al., 2017).
I NFERENCE WITH P ARTICLE G IBBS
In this section, we discuss how to draw samples from the posterior (7), corresponding to the S-step of the stochastic EM algorithm:
[Equation]
Assuming the integration and normalization can be performed efﬁciently, the following quantities can be evaluated in the forward pass without Markov state transition:
[Equation]
[Equation]
The task is to draw from the joint posterior of z 1: T only given access to these forward messages .
One way to circumvent the tight dependencies in z 1: T is to make a factorization assumption, as in
Zaheer et al.
More concretely, the joint distribution is decomposed as (factorization assumption
[Equation]
where z prev1: t − 1 is the assignments from the previous inference step.
However, as we conﬁrm in the experiments, this assumption can be restrictive since the ﬂexibility of LSTM state transitions is offset by considering each time step independently.
In this work, we propose to use a method based on SMC, which is a principled way of sampling from a sequence of distributions.
More importantly, it does not require the model to be Markovian (Frigola et al., 2013; Lindsten et al., 2014).
As described earlier, the idea is to approximate the posterior (15) with point masses, i.e ., weighted particles.
Let f ( z t | z 1: t − 1 , x t ) be the importance density, and P be the number of particles.
We then can run Algorithm 1 with π t ( z 1: t ) =
p ( x 1: t , z 1: t ) being the unnormalized target distribution at time t , where the weight becomes
[Equation]
As for the choice of the proposal distribution f ( · ) , one could use the transition density p ω ( z t | z 1: t − 1 ) , in which case the algorithm is also referred to as the bootstrap particle ﬁlter.
An alternative is the predictive distribution, a locally optimal proposal in terms of variance (Andrieu et al., 2010):
[Equation]
which is precisely one of the available forward messages:
[Equation]
[Equation]
Notice the similarity between terms in (19) and (20).
Indeed, with the choice of predictive distribution as the proposal density, the importance weight simpliﬁes to
[Equation]
which is not a coincidence that the name collides with the message α t .
Interestingly, this quantity no longer depends on the current particle z pt .
Instead, it marginalizes over all possible particle assignments of the current time step.
This is beneﬁcial computationally since the intermediate terms from (20) can be reused in (22).
Also note that the optimal proposal relies on the fact that the normalization in (20) can be performed efﬁciently, otherwise the bootstrap proposal should be used.
After a full pass over the sequence, the algorithm produces Monte Carlo approximation of the posterior and the marginal likelihood:
[Equation]
The inference is completed by a ﬁnal draw from the approximate posterior,
[Equation]
which is essentially sampling a particle path indexed by the last particle.
Speciﬁcally, the last particle z pT is chosen according to the ﬁnal weights α T , and then earlier particles can be obtained by tracing backwards to the beginning of the sequence according to the ancestry indicators a pt at each position.
Since SMC produces a Monte Carlo estimate, as the number of particles P → ∞ the approximate posterior (23) is guaranteed to converge to the true posterior for a ﬁxed sequence.
However, as the length of the sequence T increases, the number of particles needed to provide a good approximation grows exponentially.
This is the well-known depletion problem of SMC (Andrieu et al., 2010).
One elegant way to avoid simulating enormous number of particles is to marry the idea of MCMC with SMC (Andrieu et al., 2010).
The idea of such Particle MCMC (PMCMC) methods is to treat the particle estimate
ˆ p ( · ) as a proposal, and design a Markov kernel that leaves the target distribution invariant.
Since the invariance is ensured by the MCMC, it does not demand SMC to provide an accurate approximation to the true distribution, but only to give samples that are approximately distributed according to the target.
As a result, for any ﬁxed P > 0 the PMCMC methods ensure the target distribution is invariant.
We choose the Gibbs kernel that requires minimal modiﬁcation from the basic SMC.
The resulting algorithm is Particle Gibbs (PG), which is a conditional SMC update in a sense that a reference path z ref1: T with its ancestral lineage is ﬁxed throughout the particle propagation of SMC.
It can be shown that this simple modiﬁcation to SMC produces a transition kernel that is not only invariant, but also ergodic under mild assumptions.
In practice, we use the assignments from previous step as the reference path.
The ﬁnal algorithm is summarized in Algorithm 2.
Combined with the stochastic
EM outer iteration, the ﬁnal algorithm is an instance of the particle SAEM (Lindsten, 2013;
Sch ¨ on et al., 2015), under non-Markovian state transition.
We conclude this section by deriving forward messages for the previous examples.
Example 1 (Gaussian SSL, continued)
The integration and normalization preserves normality thanks to the Gaussian identify.
The messages are given by
[Equation]
[Equation]
where
V = (cid:0) [ g σ ( s t )] − 1 + C T R − 1 C (cid:1) − 1 .
Example 2 (Topical SSL, continued) Let θ t = softmax ( Ws t + b ) .
Since the distributions are discrete, we have
[Equation]
where ◦ denotes element-wise product.
Note that the integration for α t corresponds to a summation in the state space.
It is then normalized across P particles to form a weight distribution.
For γ t the normalization is performed in the state space as well, hence the computation of the messages are manageable.
[Figure]
E XPERIMENTS
We now present empirical studies for our proposed model and inference (denoted as SMC) in order to establish that (1) SSL is ﬂexible in capturing underlying nonlinear dynamics, (2) our inference is accurate yet easily applicable to complicated models, and (3) it opens new avenues for interpretable yet nonlinear and non-Markovian sequence models, previously unthinkable.
To illustrate these claims, we evaluate on (1) synthetic sequence tracking of varying difﬁculties, (2) language modeling, and (3) user modeling utilizing complicated models for capturing the intricate dynamics.
For SMC inference, we gradually increase the number of particles P from 1 to K during training.
Software & hardware
All the algorithms are implemented on TensorFlow (Abadi et al., 2016).
We run our experiments on a commodity machine with Intel R (cid:13)
Xeon R (cid:13) CPU E5-2630 v4 CPU, 256GB
RAM, and 4 NVidia R (cid:13) Titan X (Pascal) GPU.
5. S YNTHETIC S EQUENCE T RACKING
To test the ﬂexibility of SSL, we begin with inference using synthetic data.
We consider four different dynamics in 2D space: (i) a straight line, (ii) a sine wave, (iii) a circle, and (iv) a swiss role.
Note that we do not add additional states such as velocity, keeping the dynamics nonlinear except for the ﬁrst case.
Data points are generated by adding zero mean Gaussian noise to the true underlying dynamics.
The true dynamics and the noisy observations are plotted in the top row of Figure 2.
The ﬁrst 60% of the sequence is used for training and the rest is left for testing.
The middle and bottom row of Figure 2 show the result of SSL and vanilla LSTM trained for same number of iterations until both are sufﬁciently converged.
The red points refer to the prediction of z t after observing x 1: t , and the green points are blind predictions without observing any data.
We can observe that while both methods are capturing the dynamics well in general, the predictions of
LSTM tend to be more sensitive to initial predictions.
In contrast, even when the initial predictions are not incorrect, SSL can recover in the end by remaining on the latent dynamic.
5. L ANGUAGE M ODELING
For Topical SSL, we compare our SMC inference method with the factored old algorithm (Zaheer et al., 2017) on the publicly available Wikipedia dataset, where documents with less than 500 words are excluded and the most frequent 200k word types are retained. We train on the ﬁrst 60% of the
[Figure]
documents and test on the rest, using the same settings in Zaheer et al.
Figure 3 shows the test perplexity (lower is better) and number of nonzeros in the learned word topic count matrix (lower is better).
In all cases, the SMC inference method consistently outperforms the old factored method. For comparison, we also run LSTM with the same number of parameters, which gives the lowest test perplexity of 1942.26. However, we note that LSTM needs to perform expensive linear transformation for both embedding and softmax at every step, which depends linearly on the vocabulary size V . In contrast, SSL only depends linearly on number of topics K (cid:28) V .
[Figure]
Ablation study We also want to explore the beneﬁt of the newer inference as dataset size increases.
We observe that in case of natural languages which are highly structured the gap between factored approximation and accurate SMC keeps reducing as dataset size increases.
But as we will see in case of user modeling when the dataset is less structured, the factored assumption leads to poorer performance. Also when the data size is ﬁxed and the number of topics are varying, the SMC algorithm gives better perplexity compared to the old algorithm. Therefore we the SMC inference is consistently better in various settings.
Visualizing particle paths In Figure 5, we show the particle paths on a snippet of an article about a music album 1 .
As we can see from the top row, which plots the particle paths at the initial iteration, the model proposed a number of candidate topic sequences since it is uncertain about the latent semantics yet.
However, after 100 epochs, as we can see from the bottom row, the model is much more conﬁdent about the underlying topical transition.
Moreover, by inspecting the learned parameters φ of the probabilistic emission, we can see that the topics are highly concentrated on topics related to music and time.
This conﬁrms our claim about ﬂexible sequence modeling while retaining interpretability.
[Figure]
5. U SER M ODELING
We use an anonymized sample of user search click history to measure the accuracy of different models on predicting users future clicks.
An accurate model would enable better user experience by presenting the user with relevant content.
The dataset is anonymized by removing all items appearing less than a given threshold, this results in a dataset with 100K vocabulary and we vary the number of users from 500K to 1M.
Algorithm # Users 500k 1M
Factored
SMC
We ﬁx the number of topics at 500 for all user experiments.
We used the same setup to the one used in the experiments over the Wikipedia dataset for parameters.
The dataset is less structured than the language modeling task since users click patterns are less predictable than the sequence of words which follow deﬁnite syntactic rules.
As shown in table 1, the beneﬁt of new inference method is highlighted as it yields much lower perplexity than the factored model.
D ISCUSSIONS & C ONCLUSION
In this paper we revisited the problem of posterior inference in Latent LSTM models as introduced in Zaheer et al.
We generalized their model to accommodate a wide variety of state space models and most importantly we provided a more principled Sequential Monte-Carlo (SMC) algorithm for posterior inference.
Although the newly proposed inference method can be slower, we showed over a variety of dataset that the new SMC based algorithm is far superior and more stable.
While computation of the new SMC algorithm scales linearly with the number of particles, this can be naively parallelized.
In the future we plan to extend our work to incorporate a wider class of dynamically changing structured objects such as time-evolving graphs.
R