Neural Adaptive Sequential Monte Carlo

Introduction

Sequential Monte Carlo (SMC) is a class of algorithms that draw samples from a target distribution of interest by sampling from a series of simpler intermediate distributions. More specifically, the sequence constructs a proposal for importance sampling (IS) . SMC is particularly well-suited for performing inference in non-linear dynamical models with hidden variables, since filtering naturally decomposes into a sequence, and in many such cases it is the state-of-the-art inference method . Generally speaking, inference methods can be used as modules in parameter learning systems. SMC has been used in such a way for both approximate maximum-likelihood parameter learning and in Bayesian approaches such as the recently developed Particle MCMC methods .

Critically, in common with any importance sampling method, the performance of SMC is strongly dependent on the choice of the proposal distribution. If the proposal is not well-matched to the target distribution, then the method can produce samples that have low effective sample size and this leads to Monte Carlo estimates that have pathologically high variance . The SMC community has developed approaches to mitigate these limitations such as resampling to improve particle diversity when the effective sample size is low and applying MCMC transition kernels to improve particle diversity . A complementary line of research leverages distributional approximate inference methods, such as the extended Kalman Filter and Unscented Kalman Filter, to construct better proposals, leading to the Extended Kalman Particle Filter (EKPF) and Unscented Particle Filter (UPF) . In general, however, the construction of good proposal distributions is still an open question that severely limits the applicability of SMC methods.

This paper proposes a new gradient-based black-box adaptive SMC method that automatically tunes flexible proposal distributions. The quality of a proposal distribution can be assessed using the (intractable) Kullback-Leibler (KL) divergence between the target distribution and the parametrized proposal distribution. We approximate the derivatives of this objective using samples derived from SMC. The framework is very general and tractably handles complex parametric proposal distributions. For example, here we use neural networks to carry out the parameterization thereby leveraging the large literature and efficient computational tools developed by this community. We demonstrate that the method can efficiently learn good proposal distributions that significantly outperform existing adaptive proposal methods including the EKPF and UPF on standard benchmark models used in the particle filter community. We show that improved performance of the SMC algorithm translates into improved mixing of the Particle Marginal Metropolis-Hasting (PMMH) . Finally, we show that the method allows higher-dimensional and more complicated models to be accurately handled using SMC, such as those parametrized using neural networks (NN), that are challenging for traditional particle filtering methods .

The focus of this work is on improving SMC, but many of the ideas are inspired by the burgeoning literature on approximate inference for unsupervised neural network models. These connections are explored in section 6.

Sequential Monte Carlo

We begin by briefly reviewing two fundamental SMC algorithms, sequential importance sampling (SIS) and sequential importance resampling (SIR). Consider a probabilistic model comprising (possibly multi-dimensional) hidden and observed states $\bm{z}_{1:T}$ and $\bm{x}_{1:T}$ respectively, whose joint distribution factorizes as $p(\bm{z}_{1:T},\bm{x}_{1:T})=p(\bm{z}_{1})p(\bm{x}_{1}|\bm{z}_{1})\prod_{t=2}^{T}p(\bm{z}_{t}|\bm{z}_{1:t-1})p(\bm{x}_{t}|\bm{z}_{1:t},\bm{x}_{1:t-1})$. This general form subsumes common state-space models, such as Hidden Markov Models (HMMs), as well as non-Markovian models for the hidden state, such as Gaussian processes.

The choice of the proposal distribution in SMC is critical. Even when employing the resampling step, a poor proposal distribution will produce trajectories that, when traced backwards, quickly collapse onto a single ancestor. Clearly this represents a poor approximation to the true posterior $p(\bm{z}_{1:T}|\bm{x}_{1:T})$. These effects can be mitigated by increasing the number of particles and/or applying more complex additional MCMC moves , but these strategies increase the computational cost.

The conclusion is that the proposal should be chosen with care. The optimal choice for an unconstrained proposal that has access to all of the observed data at all times is the intractable posterior distribution $q_{\phi}(\bm{z}_{1:T}|\bm{x}_{1:T})=p_{\theta}(\bm{z}_{1:T}|\bm{x}_{1:T})$. Given the restrictions imposed by the factorization, this becomes $q(\bm{z}_{t}|\bm{z}_{1:t-1},\bm{x}_{1:t})=p(\bm{z}_{t}|\bm{z}_{1:t-1},\bm{x}_{1:t})$, which is still typically intractable. The bootstrap filter instead uses the prior $q(\bm{z}_{t}|\bm{z}_{1:t-1},\bm{x}_{1:t})=p(\bm{z}_{t}|\bm{z}_{1:t-1},\bm{x}_{1:t-1})$ which is often tractable, but fails to incorporate information from the current observation $\bm{x}_{t}$. A halfway-house employs distributional approximate inference techniques to approximate $p(\bm{z}_{t}|\bm{z}_{1:t-1},\bm{x}_{1:t})$. Examples include the EKPF and UPF . However, these methods suffer from three main problems. First, the extended and unscented Kalman Filter from which these methods are derived are known to be inaccurate and poorly behaved for many problems outside of the SMC setting . Second, these approximations must be applied on a sample by sample basis, leading to significant additional computational overhead. Third, neither approximation is tuned using an SMC-relevant criterion. In the next section we introduce a new method for adapting the proposal that addresses these limitations.

Adapting Proposals by Descending the Inclusive KL Divergence

The gradient of the negative KL divergence with respect to the parameters of the proposal distribution takes a simple form,

The expectation over the posterior can be approximated using samples from SMC. One option would use the weighted sample trajectories at the final time-step of SMC, but although asymptotically unbiased such an estimator would have high variance due to the collapse of the trajectories. An alternative, that reduces variance at the cost of introducing some bias, uses the intermediate ancestral trees i.e. a filtering approximation (see the supplementary material for details),

The simplicity of the proposed approach brings with it several advantages and opportunities.

Online and batch variants. Since the derivatives distribute over time, it is trivial to apply this update in an online way e.g. updating the proposal distribution every time-step. Alternatively, when learning parameters in a batch setting, it might be more appropriate to update the proposal parameters after making a full forward pass of SMC. Conveniently, when performing approximate maximum-likelihood learning the gradient update for the model parameters $\theta$ can be efficiently approximated using the same sample particles from SMC (see supplementary material and Algorithm 1). A similar derivation for maximum likelihood learning is also discussed in .

Efficiency of the adaptive proposal. In contrast to the EPF and UPF, the new method employs an analytic function for propagation and does not require costly particle-specific distributional approximation as an inner-loop. Similarly, although the method bears similarity to the assumed-density filter (ADF) which minimizes a (local) inclusive KL, the new method has the advantage of minimizing a global cost and does not require particle-specific moment matching.

Training complex proposal models. The adaptation method described above can be applied to any parametric proposal distribution. Special cases have been previously treated by . We propose a related, but arguably more straightforward and general approach to proposal adaptation. In the next section, we describe a rich family of proposal distributions, that go beyond previous work, based upon neural networks. This approach enables adaptive SMC methods to make use of the rich literature and optimization tools available from supervised learning.

Flexibility of training. One option is to train the proposal distribution using samples from SMC derived from the observed data. However, this is not the only approach. For example, the proposal could be trained using data sampled from the generative model instead, which might mitigate overfitting effects for small datasets. Similarly, the trained proposal does not need to be the one used to generate the samples in the first place. The bootstrap filter or more complex variants can be used.

Flexible and Trainable Proposal Distributions Using Neural Networks

The proposed adaption method can be applied to any parametric proposal distribution. Here we briefly describe how to utilize this flexibility to employ powerful neural network-based parameterizations that have recently shown excellent performance in supervised sequence learning tasks . Generally speaking, applications of these techniques to unsupervised sequence modeling settings is an active research area that is still in its infancy and this work opens a new avenue in this wider research effort.

In a nutshell, the goal is to parameterize $q_{\phi}(\bm{z}_{t}|\bm{z}_{1:t-1},\bm{x}_{1:t})$ – the proposal’s stochastic mapping from all previous hidden states $\bm{z}_{1:t-1}$ and all observations (up to and including the current observation) $\bm{x}_{1:t}$, to the current hidden state, $\bm{z}_{t}$ – in a flexible, computationally efficient and trainable way. Here we use a class of functions called Long Short-Term Memory (LSTM) that define a deterministic mapping from an input sequence to an output sequence using parameter-efficient recurrent dynamics, and alleviate the common vanishing gradient problem in recurrent neural networks . The distributions $q_{\phi}(\bm{z}_{t}|\bm{h}_{t})$ can be a mixture of Gaussians (a mixture density network (MDN) ) in which the mixing proportions, means and covariances are parameterised through another neural network (see the supplementary for details on LSTM, MDN, and neural network architectures).

Experiments

The goal of the experiments is three fold. First, to evaluate the performance of the adaptive method for inference on standard benchmarks used by the SMC community with known ground truth. Second, to evaluate the performance when SMC is used as an inner loop of a learning algorithm. Again we use an example with known ground truth. Third, to apply SMC learning to complex models that would normally be challenging for SMC comparing to the state-of-the-art in approximate inference.

In order to evaluate the effectiveness of our adaptive SMC method, we tested our method on a standard nonlinear state-space model often used to benchmark SMC algorithms . The model is given by Eq. 3, where $\theta=(\sigma_{v},\sigma_{w})$. The posterior distribution $p_{\theta}(\bm{z}_{1:T}|\bm{x}_{1:T})$ is highly multi-modal due to uncertainty about the signs of the latent states.

The experiments investigated how the new proposal adaptation method performed in comparison to standard methods including the bootstrap filter, EKPF, and UKPF. In particular, we were interested in the following questions: Do rich multi-modal proposals improve inference? For this we compared a Gaussian proposal with a diagonal Gaussian to a mixture density network with three components (-MD-). Does a recurrent parameterization of the proposal help? For this we compared a non-recurrent neural network with 100 hidden units (-NN-) to a recurrent neural network with 50 LSTM units (-RNN-). Can injecting information about the prior dynamics into the proposal improve performance (similar in spirit to for variational methods)? To assess this, we parameterized proposals for $v_{t}$ (process noise) instead of $z_{t}$ (-f-), and let the proposal have access to the prior dynamics $f(z_{t-1},t)$ .

For all experiments, the parameters in the non-linear state-space model were fixed to $(\sigma_{v},\sigma_{w})=(\sqrt{10},1)$. Adaptation of the proposal was performed on 1000 samples from the generative process at each iteration. Results are summarized in Fig. 1 and Table 1 (see supplementary material for additional results). Average run times for the algorithms over a sequence of length 1000 were: 0.782s bootstrap, 12.1s EKPF, 41.4s UPF, 1.70s NN-NASMC, and 2.67s RNN-NASMC, where EKPF and UPF implementations are provided by . Although these numbers should only be taken as a guide as the implementations had differing levels of acceleration.

The new adaptive proposal methods significantly outperform the bootstrap, EKPF, and UPF methods, in terms of ESS, RMSE and the variance in the LML estimates. The multi-modal proposal outperforms a simple Gaussian proposal (compare RNN-MD-f to RNN-f) indicating multi-modal proposals can improve performance. Moreover, the RNN outperforms the non-recurrent NN (compare RNN to NN). Although the proposal models can effectively learn the transition function, injecting information about the prior dynamics into the proposal does help (compare RNN-f to RNN). Interestingly, there is no clear cut winner between the EKPF and UPF, although the UPF does return LML estimates that have lower variance . All methods converged to similar LMLs that were close to the values computed using large numbers of particles indicating the implementations are correct.

2 Inference in the Cart and Pole System

As a second and more physically meaningful system we considered a cart-pole system that consists of an inverted pendulum that rests on a movable base . The system was driven by a white noise input. An ODE solver was used to simulate the system from its equations of motion. We considered the problem of inferring the true position of the cart and orientation of the pendulum (along with their derivatives and the input noise) from noisy measurements of the location of the tip of the pole. The results are presented in Fig. 2. The system is significantly more intricate than the model in Sec. 5.1, and does not directly admit the usage of EKPF or UPF. Our RNN-MD proposal model successfully learns good proposals without any direct access to the prior dynamics.

3 Bayesian learning in a Nonlinear SSM

SMC is often employed as an inner loop of a more complex algorithm. One prominent example is Particle Markov Chain Monte Carlo , a class of methods that sample from the joint posterior over model parameters $\theta$ and latent state trajectories, $p(\theta,\bm{z}_{1:T}|\bm{x}_{1:T})$. Here we consider the Particle Marginal Metropolis-Hasting sampler (PMMH). In this context SMC is used to construct a proposal distribution for a Metropolis-Hasting (MH) accept/reject step. The proposal is formed by sampling a proposed set of parameters e.g. by perturbing the current parameters using a Gaussian random walk, then SMC is used to sample a proposed set of latent state variables, resulting in a joint proposal $q(\theta^{*},\bm{z}_{1:T}^{*}|\theta,\bm{z}_{1:T})=q(\theta^{*}|\theta)p_{\theta^{*}}(\bm{z}_{1:T}^{*}|\bm{x}_{1:T})$. The MH step uses the SMC marginal likelihood estimates to determine acceptance. Full details are given in the supplementary material.

In this experiment, we evaluate our method in a PMMH sampler on the same model from Section 5.1 following .Only the prior proposal is compared, since Sec. 5.1 shows the advantage of our method over EKPF/UPF. A random walk proposal is used to sample $\theta=(\sigma_{v},\sigma_{w})$, $q(\theta^{*}|\theta)=\mathcal{N}(\theta^{*}|\theta,\text{diag}([0.15,0.08]))$. The prior over $\theta$ is set as $\mathcal{IG}(0.01,0.01)$. $\theta$ is initialized as $(10,10)$, and the PMMH is run for 500 iterations.

Two of the adaptive models considered section 5.1 are used for comparison (RNN-MD and RNN-MD-f) , where “-pre-” models are pre-trained for 500 iterations using samples from the initial $\theta=(10,10)$. The results are shown in Fig. 3 and were typical for a range of parameter settings. Given a sufficient number of particles ($N=100$), there is almost no difference between the prior proposal and our method. However, when the number of particles gets smaller ($N=10$), NASMC enables significantly faster burn-in to the posterior, particularly on the measurement noise $\sigma_{w}$ and, for similar reasons, NASMC mixes more quickly. The limitation with the NASMC-PMMH is that the model needs to continuously adapt as the global parameter is sampled, but note this is still not as costly as adapting on a particle-by-particle basis as is the case for the EKPF and UPF.

4 Polyphonic Music Generation

Finally, the new method is used to train a latent variable recurrent neural network (LV-RNN) for modelling four polymorphic music datasets of varying complexity . These datasets are often used to benchmark RNN models because of their high dimensionality and the complex temporal dependencies involved at different time scales . Each dataset contains at least 7 hours of polyphonic music with an average polyphony (number of simultaneous notes) of 3.9 out of 88. LV-RNN contains a recurrent neural network with LSTM layers that is driven by i.i.d. stochastic latent variables ($\bm{z}_{t}$) at each time-point and stochastic outputs ($\bm{x}_{t}$) that are fed back into the dynamics (full details in the supplementary material). Both the LSTM layers in the generative and proposal models are set as 1000 units and Adam is used as the optimizer. The bootstrap filter is compared to the new adaptive method (NASMC). 10 particles are used in the training. The hyperparameters are tuned using the validation set . A diagonal Gaussian output is used in the proposal model, with an additional hidden layer of size 200. The log likelihood on the test set, a standard metric for comparison in generative models , is approximated using SMC with 500 particles. The results are reported in Table 2.Results for RNN-NADE are separately provided for reference, since this is a different model class. The adaptive method significantly outperforms the bootstrap filter on three of the four datasets. On the piano dataset the bootstrap method performs marginally better. In general, the NLLs for the new methods are comparable to the state-of-the-art although detailed comparison is difficult as the methods with stochastic latent states require approximate marginalization using importance sampling or SMC.

Comparison of Variational Inference to the NASMC approach

Conclusion

This paper developed a powerful method for adapting proposal distributions within general SMC algorithms. The method parameterises a proposal distribution using a recurrent neural network to model long-range contextual information, allows flexible distributional forms including mixture density networks, and enables efficient training by stochastic gradient descent. The method was found to outperform existing adaptive proposal mechanisms including the EKPF and UPF on a standard SMC benchmark, it improves burn in and mixing of the PMMH sampler, and allows effective training of latent variable recurrent neural networks using SMC. We hope that the connection between SMC and neural network technologies will inspire further research into adaptive SMC methods. In particular, application of the methods developed in this paper to adaptive particle smoothing, high-dimensional latent models and adaptive PMCMC for probabilistic programming are particular exciting avenues.

SG is generously supported by Cambridge-Tübingen Fellowship, the ALTA Institute, and Jesus College, Cambridge. RET thanks the EPSRC (grants EP/G050821/1 and EP/L000776/1). We thank Theano developers for their toolkit, the authors of for releasing the source code, and Roger Frigola, Sumeet Singh, Fredrik Lindsten, and Thomas Schön for helpful suggestions on experiments.

References