SI χ O: Smoothing Inference with Twisted Objectives
Abstract
Sequential Monte Carlo (SMC) is an inference algorithm for state space models that approximates the posterior by sampling from a sequence of target distributions.
The target distributions are often chosen to be the filtering distributions, but these ignore information from future observations, leading to practical and theoretical limitations in inference and model learning.
We introduce SIXO , a method that instead learns target distributions that approximate the smoothing distributions, incorporating information from all observations.
The key idea is to use density ratio estimation to fit functions that warp the filtering distributions into the smoothing distributions.
We then use SMC with these learned targets to define a variational objective for model and proposal learning.
SIXO yields provably tighter log marginal lower bounds and offers more accurate posterior inferences and parameter estimates in a variety of domains.
Introduction
In this work we consider model learning and approximate posterior inference in probabilistic state space models.
Sequential Monte Carlo (SMC) is a general-purpose method for these problems [1–3] that produces an unbiased estimate of the marginal likelihood as well as latent state trajectories (i e. particles ) that can be used to approximate posterior expectations.
SMC can facilitate model learning via expectation-maximization or direct maximization of the marginal likelihood estimate [4, 5].
It can also be cast in a variational framework [6, 7] as a rich family of approximate posterior distributions that can be fit using stochastic gradient ascent and modern automatic differentiation methods [8–12].
The quality of SMC’s marginal likelihood and posterior estimates is driven by two design decisions: the choice of proposal distributions and target distributions .
The proposal distributions specify how particles propagate from one time step to the next while the target distributions specify how those particles are weighted and which ones survive to future time steps.
The most common SMC variant, filtering SMC, sets the targets to the filtering distributions , the conditional distributions over latent states x 1: t = ( x 1 , . . .
, x t ) given observations y 1: t = ( y 1 , . . .
, y t ) .
The central issue is that the filtering distributions do not incorporate information from future observations y t
+1: T .
Figure 1 illustrates why setting the target distributions to the filtering distributions can be problematic.
In this example, the latent states follow a simple Gaussian random walk, but only the last step is observed.
Thus, the filtering distributions reduce to the prior, a series of mean zero Gaussians, shown in Figure 1a.
If the observation is far from the prior, the filtering distribution suddenly jumps at time
T = 10 .
This is a recipe for disaster in SMC: the particles at time T − 1 will be distributed according to a mean zero Gaussian and very few will survive to the next time step, causing the variance of the
SMC estimator to explode.
Even if the proposals incorporate smoothing information, using filtering targets can cause particle degeneracy by resampling away high-quality particles, as seen in Figure 1c.
−    x t p θ ( x t | y 1 : t ) (a) Filtering distributions −    x t p θ ( x t | y 1 : T ) (b) Smoothing distributions y T
Timestep t      x t (c) FIVO particle lineages
Timestep t      x t (d) SIXO particle lineages y T t=
Suppose instead that the target distributions were the smoothing distributions —the conditional distributions over latents x 1: t given all observations y 1: T .
Figure 1b shows the smoothing distributions for the simple Gaussian random walk.
Unlike the filtering distributions, the smoothing distributions shift steadily toward the observation over time.
These slow, smooth changes are ideal for SMC:
Figure 1d shows many particles surviving from one step to the next, resulting in a low-variance SMC estimator.
In practice we do not have access to the smoothing distributions—if we did, there would be no need for SMC!
Here, we introduce a new method called SIXO:
Smoothing Inference with Twisted
Objectives.
SIXO provides a unified approach for learning model and proposal parameters, as well as a set of twisting functions that warp the filtering distributions into targets that better approximate the smoothing distributions [13].
Like its predecessor FIVO [8–10], SIXO uses a variational approach, deriving a lower bound to the marginal likelihood.
Unlike its predecessor, we prove that the SIXO bound can become tight, even with finitely many particles.
The key challenge with SIXO is learning the twisting functions.
We find that a simple density ratio estimation approach works best, and we propose an algorithm that interleaves twist updates with updates to the model and proposal.
Thus, SIXO offers a means of jointly learning model parameters,
SMC proposals, and targets for accurate posterior inference.
Finally, we give empirical evidence to support our theoretical claims.
Across a range of experiments with a Gaussian diffusion, a stochastic volatility model of currency exchange rates, and a HodgkinHuxley model of membrane potential in a neuron, SIXO consistently outperforms FIVO and related methods.
We dissect these results to illustrate how learning better targets enables more effective posterior inference and model learning.
Background
Consider modeling sequential data y 1: T ∈ Y T using latent variables x 1: T ∈ X T with Markovian structure, and let the joint distribution factorize as
[Equation]
with global parameters θ ∈ Θ .
We further assume that the conditional distributions p θ ( x t | x t − 1 ) and p θ ( y t | x t ) may depend nonlinearly on x t − 1 and x t respectively.
The marginal likelihood and posterior for this model class are not readily available from the joint distribution due to the intractable integral over the latents x 1: T , i.e.
[Equation]
cannot easily be computed due to the form of the conditional distributions.
2.1 Sequential Monte Carlo
Sequential Monte Carlo is an algorithm for inference in state-space models that approximates the posterior p θ ( x 1: T | y 1: T ) with a set of K weighted particles x 1: K 1: T .
These particles are constructed by approximately sampling from a sequence of target distributions { π t ( x 1: t ) }
Tt =1 , with the intuition that sampling from a series of distributions that gradually approach the posterior is easier than attempting to sample from it directly.
The targets are often only available up to an unknown normalizing constant
Z t , so SMC uses the unnormalized targets { γ t ( x 1: t ) } Tt =1 , which correspond to the normalized targets via π t ( x 1: t )
=
γ t ( x 1: t ) /Z t .
SMC repeats three steps: First, a set of latents are sampled from a proposal distribution q θ ( x kt | x kt − 1 , y 1: T ) conditional on the current particles x 1: K 1: t − 1 .
Then, each particle is weighted using the unnormalized target γ t ( x 1: t ) to form an empirical approximation of the normalized target distribution.
Finally, new particle trajectories x 1: K 1: t are drawn from this approximation to the normalized target.
Ideally the target distributions smoothly approach the posterior so that sampling from the target at time t + 1 is easy given samples from the target at time t .
As long as mild technical conditions are met and γ T ( x 1: T ) ∝ p θ ( x 1: T , y 1: T ) , SMC returns a consistent and unbiased estimate of the marginal likelihood p θ ( y 1: T ) and a set of weighted particles approximating the posterior p θ ( x 1: T | y 1: T ) [1–3].
For more details see Appendix B.1, and for a thorough treatment of SMC see [1–3].
2.2 Filtering SMC and Model Learning
The most commonly-used SMC algorithm is filtering SMC (FSMC), which sets the normalized targets to the filtering distributions, i.e. π t ( x 1: t ) =
p θ ( x 1: t | y 1: t ) and γ t ( x 1: t )
∝ p θ ( x 1: t , y 1: t ) .
Let (cid:98) Z FSMC ( θ , y 1: T ) be the marginal likelihood estimator returned from running filtering SMC with proposal distributions { q θ ( x t | x 1: t − 1 , y 1: t ) }
Tt =1 which may share parameters with p θ .
Previous work used filtering SMC to fit model parameters by ascending a lower bound on the log marginal likelihood called a filtering variational objective (FIVO)
The FIVO bound is derived using Jensen’s inequality and the unbiasedness of (cid:98) Z FSMC ,
[Equation]
and is optimized using stochastic gradient ascent in θ [8–10, 14].
2.3 Smoothing SMC via Twisting Functions
The main disadvantage of filtering SMC is that the filtering distributions only condition on observations up to the current timestep t , ignoring future observations y t
+1: T .
This creates situations where future observations are highly unlikely given the current latent trajectories, which in turn causes particle death, high variance normalizing constant estimates, and poor inference and model learning [8, 13, 15].
Performing smoothing SMC would resolve this issue by choosing the smoothing distributions as targets, i.e. π t ( x 1: t ) =
p θ ( x 1: t | y 1: T ) and γ t ( x 1: t ) ∝ p θ ( x 1: t , y 1: T ) .
Unfortunately, p θ ( x 1: t , y 1: T ) is not readily available from the model and computing it is roughly as hard as the original inference problem.
However, p θ ( x 1: t , y 1: T ) factors into the product of the filtering distributions, p θ ( x 1: t , y 1: t ) , and the lookahead distributions, p θ ( y t +1: T | x t ) (Appendix B.2).
If the lookahead distributions can be well-approximated by a series of “twisting” functions [13], { r ( y t +1: T , x t ) }
Tt =1 , then running SMC with targets γ t ( x 1: t )
=
p θ ( x 1: t , y 1: t ) r ( y t +1: T , x t ) would approximate smoothing SMC.
In this sense, the lookahead distributions are optimal twisting functions [13, 16].
Different twisting functions yield different SMC methods such as auxiliary particle filters and twisted particle filters [13, 17].
However, as long as the final unnormalized target γ T ( x 1: T ) is proportional to p θ ( x 1: T , y 1: T ) and regularity conditions are met, SMC will produce an unbiased estimate of the marginal likelihood, regardless of the choice of twisting functions [1, 3, 18].
Instead, the quality of the twisting functions affects the variance of SMC’s marginal likelihood estimate.
SIXO: Smoothing Inference with Twisted Objectives
Our goal is to fit models by optimizing a lower bound on their log marginal likelihood constructed using smoothing SMC.
To construct the lower bound, fix r ψ ( x T ) = 1 and let (cid:98) Z SIXO ( θ , ψ , y 1: T ) be the marginal likelihood estimator returned from running SMC with unnormalized targets { p θ ( x 1: t , y 1: t ) r ψ ( y t +1: T , x t ) }
Tt =1 and proposal distributions { q θ ( x t | x 1: t − 1 , y 1: T ) } Tt =1 .
Because the T th unnormalized target is p θ ( x 1: T , y 1: T ) , (cid:98) Z SIXO will be an unbiased estimator of the marginal likelihood p
θ ( y 1: T ) [1, 3].
This implies via Jensen’s inequality that
[Equation]
i.e. L SIXO ( θ , ψ , y 1: T ) is a lower bound on the log marginal likelihood log p θ ( y 1: T ) [14].
3.1 The Functional Form of the Twists
The structure of the lookahead distributions p θ ( y t +1: T | x t ) suggests a functional form for r ψ that accepts a single latent x t and produces distributions over all future observations y t
+1: T .
Because the twists will be evaluated once per particle and timestep in an SMC sweep, this functional form would lead to an algorithm with O ( T 2 ) complexity.
To reduce the complexity, we consider two methods: fixed-lag twisting and backwards twisting.
Fixed-lag twisting approximates the full lookahead distribution p θ ( y t
+1: T | x t ) using a fixed window of L observations, i.e. it models p θ ( y t +1:
t + L | x t ) [17, 19, 20].
We define the fixed-lag twisting functions { r ψ ( y t +1:
t + L , x t ) } T − 1 t =1 as a sequence of functions which accept x t ∈ X and produce a distribution over y t +1: t + L ∈ Y L .
This reduces the computational complexity to O ( TL ) , at the expense of only looking at L observations.
In our experiments we use an L = 1 twist that scores the next observation by approximating the one-step lookahead
[Equation]
with Gauss-Hermite quadrature [21].
We refer to this as the “quadrature twist”.
This approach is similar to the APF [17, 22], but uses numerical quadrature in place of sample-based integration.
Backwards twisting is motivated by rewriting the lookahead distributions using Bayes’ rule,
[Equation]
where we drop terms independent of x t because the twisting functions will be used to score particles in SMC.
Thus, we need only approximate p θ ( x t | y t
+1: T ) /p
θ ( x t ) .
The numerator p θ ( x t | y t +1: T )
1: procedure SIXO-DRE( y 1: T , θ 0 , ψ 0 , S , N , K ) 2: for s = 1 , . . .
, S do 3: ψ s = TWIST-UPDATE( θ s − 1 , ψ s − 1 , N ) 4: θ s = MODEL-UPDATE( y 1: T , θ s − 1 , ψ s , N , K ) 5: return θ S , ψ S 6: procedure TWIST-UPDATE( θ , ψ 0 , N ) 7: for i = 1 , . . .
, N do 8: ˜ x 1: T ∼ p θ ( x 1: T ) 9: x 1: T , y 1: T ∼ p θ ( x 1: T , y 1: T ) 10:
L DRE ( ψ ) = 1 T − 1 (cid:80) T − 1 t =1 log σ (log r ψ ( y t +1: T , x t ))
+ log(1 − σ (log r ψ ( y t +1: T , ˜ x t ))) 11:
Compute ψ
i using the gradients of L DRE evaluated at ψ i −12: return ψ N 13: procedure MODEL-UPDATE( y 1: T , θ 0 , ψ , N, K ) 14: for i = 1 , . . .
, N do 15: (cid:98) Z SIXO ( θ ) =
SMC( { p θ ( x 1: t , y 1: t ) r ψ ( y t +1: T , x t ) } T t =1 , { q θ ( x t | x t − 1 , y 1: T ) } T t =1 , K) 16:
Compute θ i using the biased gradients of (cid:98) Z SIXO evaluated at θ i −17: return θ N 18: procedure SMC( { γ t ( x 1: t ) }
Tt =1 , { q θ ( x t | x t − 1 , y 1: T ) }
Tt =1 , K) 19:
See Algorithm 2 in Appendix B.1.
is the reverse of the lookahead distributions—it is a distribution over a single latent conditioned on future observations.
This makes it possible to parameterize the twists using a recurrent function approximator (e.g. a recurrent neural network or RNN) run backwards across the observations
y 1: T to produce twist values for each timestep as a function of x kt .
We define the backwards twists { r ψ ( y t +1: T , x t ) } T − 1 t =1 as a sequence of positive, integrable, realvalued functions Y T − t × X
→ R + with parameters ψ ∈ Ψ .
Parameterizing backward twists with a recurrent function approximator results in O ( T ) complexity (see Appendix C.3) and allows the twist to condition on all future observations, making backwards twisting preferable to fixed-lag twisting.
3.2 Learning Twists
Ascending the Unified Objective One way to fit the twists, proposal, and model is to ascend L SIXO in the parameters of p θ , q θ , and r ψ , similar to FIVO [8–10].
The gradients of this objective include score-function terms that arise from the discrete resampling steps in SMC.
We refer to ascending
L SIXO with these unbiased gradients as SIXO-u. Because the resampling gradient terms have high variance, SIXO-u is impractical for complex settings [8, 11].
Lower-variance methods for estimating these gradients were explored by Lawson et al.
[11] but found to be ineffective.
We therefore seek an alternative method for training the twists.
For a detailed discussion and derivation of the gradient, see Appendix C.1.
Density Ratio Estimation Note that the optimal backwards twist is proportional to the ratio of a “backwards message” p θ ( x t | y t
+1: T ) and the latent marginal p θ ( x t ) (Equation 5).
Thus, we can learn the backwards twist using density ratio estimation (DRE) [23, 24].
DRE via classification estimates the ratio of two densities a ( x ) /b
( x ) by training a classifier to distinguish between samples from a and b .
If such a classifier is trained using the logit link function, then its raw output will approximate log a ( x ) − log b ( x ) up to a constant [24].
Using this approach, we interpret log r ψ ( y t
+1: T , x t ) as the logit of a Bernoulli classifier, which is trained to distinguish between samples from p θ ( x t , y t
+1: T ) and p θ ( x t ) p θ ( y t
+1: T ) , which are available from the model.
When trained in this way, log r ψ ( y t +1: T , x t ) will approximate log p θ ( x t | y t
+1: T ) − log p θ ( x t ) up to a constant, which can be ignored.
For details see Appendix C.2 and Sugiyama et al.
We use the DRE-learned twisting functions in an alternating scheme that first holds p θ , q θ fixed and updates r ψ using density ratio estimation, and then holds r ψ fixed and updates p θ and q θ by ascending a biased gradient estimator (no resampling terms) of L SIXO ( θ , ψ ) in θ .
We call the full alternating procedure for learning θ and ψ SIXO-DRE, see Algorithm 1.
3.3 The SIXO Bound Can Become Tight
Maddison et al.
[8] show that the FIVO bound can only become tight in models with uncommon dependency structures.
We show that the SIXO bound can become tight for any model in the class defined in Section 2.
Proposition 1.
Sharpness of the SIXO bound.
Let p ( x 1: T , y 1: T ) be a latent variable model with
Markovian structure as defined in Section 2, let Q be the set of possible sequences of proposal distributions indexed by parameters θ ∈ Θ , and let R be the set of possible sequences of positive, integrable twist functions indexed by parameters ψ ∈ Ψ .
Assume that { p ( x t | x t − 1 , y 1: T ) } Tt =1 ∈ Q and { p ( y t
+1: T | x t ) } T − 1 t =1 ∈ R .
Finally, assume L SIXO ( θ , ψ , y 1: T ) has the unique optimizer
[Equation]
Then the following holds:
q θ ∗ ( x t | x 1: t − 1 , y 1: T ) =
p ( x t | x 1: t − 1 , y 1: T ) for t = 1 , . .
, T , 2.
r ψ ∗ ( y t
+1: T , x t ) ∝ p ( y t +1: T | x t ) up to a constant independent of x t for t = 1 , . . .
, T − 1 , 3.
L K SIXO ( θ ∗ , ψ ∗ , y 1: T ) = log p ( y 1: T ) for any number of particles K ≥ 1 .
Proof.
See Appendix C.5.
This is an important advantage of our work—the SIXO objective is the first to recover the true marginal likelihood with a finite number of particles while also being tailored to sequential tasks.
Related Work
Standard references for SMC include Doucet and Johansen [1], Naesseth et al.
[2], and Del Moral [3] which provides a theoretical treatment of a generalization of SMC called Feynman-Kac formulae.
Pitt and Shephard [17] introduced the auxiliary particle filter, an early smoothing SMC method which constructs an estimate of the one-step backwards message p θ ( y t +1 | x t ) using simulations from the model.
Smoothing SMC in general is discussed thoroughly in Briers et al.
[15] and Del Moral et al.
Later, Whiteley and Lee [13] introduced twisted particle filters, which perform SMC on a twisted model that approximates the smoothing distributions by multiplying the model’s filtering distributions with “twisting” functions.
Our work extends the theoretical framework in Whiteley and
Lee [13] by proposing practical and effective methods for learning parametric twisting functions.
To make smoothing SMC computationally tractable, fixed-lag techniques use information from only a fixed window of future observations, as introduced in Clapp and Godsill [19] and surveyed in Lin et al.
For example, Park and Ionides [26] use simulations from the model to estimate fixed-lag twisting functions and Doucet et al.
[27] sample blocks of latents conditional on their observations via various Monte Carlo methods.
These methods suffer from computational complexity that grows with the window size, and fail to take advantage of all future observations.
Other methods use twisting functions which depend on all observations.
Most similar to our approach are Guarniero et al.
[16], Heng et al.
[28], and Zimmermann et al.
[29] which learn parametric twists using a Bellman-like decomposition of the lookahead distributions p ( y t +1: T | x t ) in terms of the same distributions one step into the future.
Our method instead uses independent DRE objectives for the twist at each timestep.
Del Moral and Murray [30] use Gaussian processes to approximate the twists, Lindsten et al.
[18] use traditional graphical model techniques such as loopy belief propagation and expectation propagation, and Ruiz and Kappen [31] use optimal control techniques.
None of these approaches consider model learning and their twist parameterizations and learning techniques are highly specialized to specific problem settings.
[Figure]
Fitting model parameters via stochastic gradient ascent on an evidence lower bound (ELBO) was introduced in Ranganath et al.
[32], Hoffman et al.
[33], Kingma and Welling [34], and later generalized to the Monte Carlo objectives (MCO) framework by Mnih and Rezende [14].
Since then, works have considered optimizing lower bounds defined by the normalizing constant estimators from multiple importance sampling [35], nested importance sampling, [36, 29], rejection sampling and
Hamiltonian Monte Carlo [12], filtering SMC [8–10], and smoothing SMC [11, 37, 38].
The prior work on smoothing SMC used an objective defined by forward filtering backwards smoothing [15] which suffers from the same particle degeneracy issues as filtering SMC and cannot become tight.
Kim et al.
[39] optimize the importance weighted autoencoder (IWAE) bound [35] using a baseline derived from future likelihood estimates, but do not use SMC or resampling in their bound.
Experiments
We experimentally explore our claims that:
The SIXO bound can become tight while FIVO cannot.
DRE-learned twists enable better posterior inference than filtering SMC.
Model learning with SIXO provides better parameter estimates than FIVO.
5.1 Gaussian Drift Diffusion
We first consider a one-dimensional Gaussian drift diffusion process with joint distribution
[Equation]
The single free model parameter is the drift α ∈ R , the state is x t ∈ R , and the observation is y T ∈ R .
Figures 1a and 1b show that for α = 0 the filtering and smoothing distributions in this model quickly diverge, which can lead to poor inference for filtering methods.
We compare joint model, proposal and twist learning using two variants of SIXO to variational inference with the IWAE bound [35] and FIVO with unbiased gradients [8–10].
All methods use an independent proposal at each time step parameterized as q t ( x t | x t − 1 , y T ) =
N ( x t ; f t ( x t − 1 , y T ) , σ 2 qt ) where f t is an affine function, a family which contains the optimal proposal.
SIXO-u uses twists parameterized as r t ( y T , x t ) =
N ( y T ; g t ( x t ) , σ 2 rt ) where g t is an affine function, a family which contains the true lookahead distributions.
The SIXO-DRE twist functions, log r t ( y T , x t ) , are parameterized as quadratic functions of x t , where the parameters of the quadratic function are generated by a neural network with inputs ( y T , t ) .
The true log density ratio will be quadratic in x t , so if the neural network is sufficiently flexible, the true log density ratio function can be obtained.
Figure 2a shows the convergence of the variational bound for each method.
As expected IWAE recovers a tight variational bound, whereas FIVO does not.
While SIXO-u does recover a tight variational bound, the high variance of the unbiased gradient estimator makes it slower to converge,
[Table]
and impractical for non-toy problems.
Conversely, SIXO-DRE achieves a tight bound while using the lower variance biased gradients.
This motivates its use in more complex, non-linear settings where the unbiased FIVO gradients are not practical.
Figure 2b shows that SIXO-u recovers the correct twist parameters.
More figures illustrating the convergence of θ and ψ are included in Appendix D.1.
In Figures 1c and 1d we compare particle trajectories under FIVO and SIXO-u.
We see that FIVO consistently proposes particles with high likelihood under the posterior distributions (identical to the smoothing distributions in this case) which are discarded by the resampling steps in filtering SMC.
In contrast, SIXO both proposes particles with high posterior likelihood and retains them through the resampling steps by properly scoring particles under the twisted target distributions.
These results empirically verify the theoretical claims made in Section 3.3.
5.2 Stochastic Volatility Model
We now apply SIXO to a stochastic volatility model (SVM) of monthly foreign exchange rates for
N = 22 currencies, over the period 9/2007 to 8/2017 [40].
The SVM generative model is defined as
[Equation]
with transition noise ν t ∼ N ( 0 , Q ) , observation noise ε t ∼ N ( 0 , I N × N ) , states x 1: T ∈ R T × N , and observations y 1: T ∈ R T × N .
All multiplications are performed element-wise, denoted by ⊙ .
The proposal, q θ , is structured as q θ ( x 1: T ) ∝ (cid:81)
Tt =1 N ( x t ; µ t , Σ t ) p θ ( x t | x t − 1 ) with means µ t ∈ R N and diagonal covariance matrices Σ t ∈ diag
R N + ) .
We compare four approaches: IWAE, FIVO,
SIXO with quadrature twist (SIXO-q), and SIXO with density ratio twist (SIXO-DRE).
Note that the observations for this model are dense in time, so we would expect filtering-based approaches to perform well.
See Appendix D.2 for more details.
Train Performance We first compare our methods in terms of their 4-particle log marginal likelihood lower bounds as in Naesseth et al.
[9], e.g. for IWAE we report the IWAE bound and for
FIVO we report the FIVO bound.
Even though SIXO-q only scores a single future observation, it still obtains a 7-nat improvement over FIVO.
SIXO-DRE, meanwhile, conditions on all future observations and obtains a 10-nat improvement over FIVO.
All SMC-based methods, however, significantly underperform IWAE which seems to indicate that IWAE’s unbiased gradients and use of smoothing information allow it to learn a proposal which is efficient at low numbers of particles.
This performance does not carry over to model learning or other high-dimensional problems, however, and is at odds with previous work on this model and dataset [8–10].
We also estimate the learned model’s training set marginal likelihood by computing a bootstrap particle filter’s (BPF) log marginal lower bound with 2,048 particles, denoted L 2048BPF [41].
Interestingly, a one-way ANOVA [42] does not reject the null hypothesis that the training set L 2048BPF means are all equal ( p = 0 . 26 ), suggesting that the log marginal likelihoods on training data are indistinguishable and training performance has saturated.
It is clear, however, that among the SMC methods SIXO performs the best inference and makes the most efficient use of particles as the L 4SIXO bounds are significantly higher than L 4FIVO for similar models.
Test Performance We also compare methods on a held-out test set to evaluate model learning.
We construct this test set using the same data source as the training set, but use the period of time since
Naesseth et al.
[9] was published, an extra 55 months.
Again, we report BPF log marginal lower bounds with 2,048 particles and find that SIXO-DRE outperforms IWAE, SIXO-q and FIVO.
All differences are statistically significant.
[Figure]
BPF
SIXO-bs
Resampling event y t
Time
[Figure]
ms
[Figure] [Figure]
[Figure]
[Figure]
M e m b r a n e p
o t
e n t i a l, m
V
[Table]
5.3 Hodgkin-Huxley Model
We conclude by comparing FIVO and SIXO on the Hodgkin-Huxley (HH) model of neural action potentials [43, 44].
A single neuron is represented with a four-dimensional state-space: the instantaneous membrane potential and the relative conductivity of three ion gates.
A noise-corrupted and subsampled membrane potential can be obtained using electrodes [45] or voltage imaging [46].
The state of the gates, however, is not observable, and must be inferred from the noisy potential recordings.
The physiological parameters governing the time-evolution of the system are also of interest, such as the base conductance of each of the ion channels.
We implement the HH model as a four-dimensional nonlinear state space model with Gaussian transition noise [47].
The observation is a single Gaussian-distributed value with mean equal to the instantaneous potential.
Unless otherwise stated, we subsample observations by a factor of 50 to simulate an acquisition frequency of 1 kHz (interval of 1 ms).
For more details, see Appendix D.3.
In this model action potentials, or spikes, are rare events that happen quickly and invoke a rapid change in the state.
Therefore, filtering-based inference is particularly disadvantageous as noisy observations may trigger erroneous spikes or “miss” true spikes.
We compare four methods: SIXO with a DRE twist and a learned smoothing proposal (SIXO-sm),
SIXO with a bootstrap proposal (SIXO-bs), FIVO with a learned filtering proposal (SIXO-fi), and
FIVO with a bootstrap proposal (FIVO-bs).
Inference In Figure 3 we compare HH trajectories generated by a BPF and SIXO-bs.
The BPF yields spurious spikes and misses the initiation of other spikes, illustrating the issue with filtering
SMC for HH inference.
SIXO-bs allows fewer spurious spikes to develop and fewer particles miss the onset of spikes.
SIXO-bs also achieves a log marginal lower bound of − 47 .
81 nats, higher than the − 49 . 24 nats achieved by the BPF, showing that it performs more effective inference.
In Table 2 we
[Figure]
[Table]
study this effect more extensively, comparing performance across different observation frequencies.
We see, as predicted, that the performance of SIXO is more consistent across observation frequencies than FIVO, supporting the claim that twists assist in inference when observations are sparse.
For more results, see Table 4 in Appendix D.3.
Model Learning We conclude by comparing FIVO and SIXO for parameter recovery in Figure 4 and Table 3.
We see that FIVO-bs converges to a poor parameter estimate, and recovers a poor variational bound.
FIVO-fi does not converge (details and full figure in Appendix D.3).
The SIXO methods recover much better parameter estimates, and achieve the highest bound values, with
SIXO-sm outperforming all other methods in terms of final performance and training stability.
Conclusions, Limitations, and Future Work
In this work we proposed a method of learning twisting functions for smoothing SMC via density ratio estimation.
Our approach ascends a lower bound on the log marginal likelihood that can theoretically become tight, a first for SMC objectives.
We verified our theoretical claims by experimentally demonstrating improvements over existing techniques in inference and model learning.
There are, however, important limitations to our approach.
Training and evaluating the twist requires additional computational effort compared to FIVO, and requires further hyperparameter tuning.
Although we consider SIXO to be mainly used for offline settings, extending SIXO to online settings could yield practical benefit.
Finally, there is a known pathology in DRE methods where the ratio may be poorly estimated if the difference between the densities in the ratio is very large [48].
Thus, new methods for learning the twist are important topics for future work.
Acknowledgements and Disclosure of Funding We thank Matt MacKay for helpful discussions and edits, and the anonymous reviewers for their feedback during the review process.
This work was supported by grants from the Simons Collaboration on the Global Brain (SCGB 697092), the NIH
BRAIN Initiative (U19NS113201 and R01NS113119).
Some of the computation for this work was made possible by Microsoft Education Azure cloud credits.
Dieterich Lawson was supported in part by a Stanford Data Science Fellowship.
A Table of Notation
[Table]
( y
[Table]
| x
[Table] [Table]
Proposal distributions { q θ ( x t | x t − 1 , y 1: T ) } Tt =
T conditional distributions on X
Twist functions { r ψ ( y t +1: T , x t ) } T − 1 t =
T − 1 positive integrable functions in X ×
Y T − t − 1 × Ψ → R ≥
Filtering distributions { p θ ( x 1: t | y 1: t ) } Tt =
T conditional distributions on X t
Smoothing distribution { p θ ( x 1: t | y 1: T ) } Tt =
T conditional distributions on X t
Unnormalized target distributions { γ t ( x 1: t ) } Tt =
T positive integrable functions in X t → R ≥
Normalized target distributions { π t ( x 1: t ) } Tt =
T distributions on X t t th normalizing constant
Z t
Positive real, π t ( x 1: t ) =
γ t ( x 1: t ) /Z
t
Bootstrap particle filter
BPF
SMC using p ( x t | x t − 1 ) as the proposal and no twist.
Filtering sequential Monte Carlo
FSMC
SMC with filtering distributions as targets.
[Table]
SIXO objective optimized w.r.t.
θ
[Table]
B Background B.1 Sequential Monte Carlo
Sequential Monte Carlo (SMC) is a popular method for sampling from posterior distributions with sequential structure.
For a thorough introduction we refer the reader to Doucet and Johansen [1] and
Naesseth et al.
We reproduce the general SMC algorithm in Algorithm 2.
1: procedure SMC( { γ t ( x 1: t ) } Tt =1 , { q t ( x t | x 1: t − 1 ) } Tt =1 , K ) 2: w 1: K= 1 , (cid:98) Z 0 =3: for t = 1 , . . .
, T do 4: for k = 1 , . . .
, K do 5: x kt ∼ q t ( x t | x k 1: t − 1 ) 6: x k 1: t = ( x k 1: t − 1 , x kt ) 7: α kt = γ t ( x k 1: t ) γ t ( x k 1: t − 1 ) q t ( x k t | x k 1: t − 1 ) 8: w kt = w kt − 1 α kt 9: (cid:92) Z t /Z
t − 1 = (cid:80)
Kk =1 w kt (cid:80)
Kk =1 w kt −10: (cid:98) Z t = (cid:98)
Z t − 1 ( (cid:92) Z t /Z
t − 1 ) 11: if should resample then 12: for k = 1 , . .
, K do 13: a kt ∼ Categorical ( w 1: K t ) 14: ˜ x k 1: t = x a kt 1: t 15: x 1: K 1: t = ˜ x 1: K 1: t 16: w 1: K t =17: return (cid:98) Z T , x 1: K 1: T
B.2 Factoring the Smoothing Distributions
Here we show that the smoothing distributions can be factored into the filtering distributions and the lookahead distributions.
Let p θ ( x 1: T , y 1: T ) be a model as defined in (1).
Then
[Equation]
where y t +1: T is conditionally independent of ( y 1: t , x 1: t − 1 ) given x t because of the Markov structure of p θ .
C Methods C.1 The Gradients of SMC
The original FIVO papers [8–10] used biased gradients to optimize L FSMC which ignore scorefunction gradient terms arising from the discrete resampling operations.
We use the same biased gradient estimator in SIXO-DRE when optimizing L SIXO ( θ , ψ , y 1: T ) in terms of θ .
We also examine
SIXO-u, where we use the full unbiased gradient estimator (see (17)).
We derive both estimators here.
Assume resampling occurs at each timestep, let A t = a 1: K t and X t = x 1: K t be the ancestor indices and latents for all particles at time t , and let A
=
( A 1 , . . .
, A T − 1 ) and X =
( X 1 , . . .
, X T ) be the full sequence of ancestor indices and latents.
We can then write the probability distribution over
X , A that defines SMC as
[Equation]
where α it = α it / ( (cid:80)
Kk
=1 α kt ) is the normalized incremental weight.
When the proposals and targets of
SMC are parametric functions of θ , both q and α will depend on θ .
To emphasize its dependence on θ when run with parametric twists and proposals we will write p SMC ( X , A ) as p ( X , A ; θ ) .
Then the gradient of L SIXO is defined as
[Equation]
(11) where we have rewritten (cid:98) Z SIXO ( θ , ψ , y 1: T ) as (cid:98) Z
( X , A , θ ) to emphasize its dependence on the random variables X and A and suppress its dependence on ψ and y 1: T .
The first step is to reparameterize the expectation in terms of continuous noise instead of X [49, 50, 34, 51, 52].
Assume X is from a reparameterizable distribution and let ϕ X ( θ , ϵ ) be a function that deterministically combines continuous noise ϵ and the parameters θ to produce a sample X .
Then we have
[Equation]
We will further abuse notation by writing (cid:98) Z ( ϕ X (
[Equation]
ϵ
[Equation]
, A
[Equation]
) as (cid:98) Z ( ϵ
[Equation]
Rewriting the expectation (12) as an integral gives
[Equation]
Assuming that we can differentiate under the integral allows us to break the integrand apart using the product rule as
[Equation]
[Equation]
The left hand term in the integrand of (14) equals E ϵ, A [ ∇ θ log (cid:98) Z ( ϵ, A , θ )] , the expectation of a gradient that can be estimated using simple Monte Carlo.
The right hand term is a “score-function gradient” [53] which can be rewritten using the fact that ∇ θ (log f ( θ ))
=
∇ θ ( f ( θ )) /f
( θ ) as
[Equation] [Equation]
which in turn equals the expectation
[Equation]
Writing both terms together gives the full unbiased gradient that is amenable to estimation with simple Monte Carlo,
[Equation]
Similar to prior work [8–11] we find that the term on the right hand side of (17) has prohibitively high variance which inhibits learning.
Dropping it gives the biased SMC gradient estimator used in
SIXO-DRE,
[Equation]
which can be estimated using open-source autodiff software [54].
The derivation above is adaptable for any resampling schedule that does not depend on the parameters (and by extension, the weights), but many common resampling schemes such as effective sample size resampling do not meet this requirement.
If the resampling scheme depends on the parameters of the model, it introduces additional gradient terms which are not described here.
Thus for all methods in experiments where compare to SIXO-u
we use a fixed resampling schedule.
C.2 Density Ratio Estimation
Density ratio estimation (DRE) considers estimating ratios of densities, e.g. a ( x ) /b
( x ) with a ( x ) and b ( x ) defined on the same probability space and b ( x ) > 0 for all x .
Instead of estimating a ( x ) and b ( x ) individually and then forming the ratio, an alternative approach is to directly estimate the odds that a given sample of x was drawn from a .
Let p ( x, z ) =
p ( z ) p ( x | z ) be an expanded generative model for x defined as
[Equation]
[Equation]
with α ∈ (0 , 1) .
We can now write the density ratio in terms of conditionals in this generative model,
[Equation]
Thus, the density ratio can be rewritten as proportional to the odds that x was drawn from a ( x ) instead of b ( x ) .
Density ratio estimation via classification suggests training a binary classifier with supervised learning to predict z given x [24, 23].
Let σ ( x ) = 1 / (1 + e − x ) be the sigmoid function and let g ψ ( x ) be a classifier trained with Bernoulli loss to maximize the log probability of a dataset z 1: N , x 1: N sampled IID from p ( x, z ) .
Specifically, ψ is fit by minimizing L DRE ( ψ ) , defined as
[Equation]
If trained in this way, the raw output of g ψ ( x ) will approximate the log-odds that x came from a ( x ) instead of b ( x ) ,
i
[Equation]
e.
[Equation]
The log of the density ratio can then be expressed as
[Equation]
Assuming a fixed α parameter, the log-ratio of densities is then proportional to the logit produced by g ψ up to an additive constant.
Thus as long as we can sample training pairs ( x, z ) from the expanded generative model above, we can estimate ratios of densities by training a binary classifier.
C.3 Backwards Density Ratio Twist Architecture
We now provide further implementation details for the the backwards DRE twist we propose.
The twist is defined through two neural networks, as illustrated in Figure 5.
First, a recurrent neural network (RNN) is passed backwards over the data to create a sequence of encodings, denoted e 2: T .
These backward encodings are computed upfront, before the SMC sweep, with a cost complexity that
[Figure]
is linear in
C.4 Alternating Density Ratio Twist Training
To train the density ratio twist functions we define a supervised maximum likelihood update that is applied offline from any update to the model and proposal parameters.
This update is shown in
Algorithm 1, labeled as DRE.
To define this update we use the approach in Section C.2 and set a ( x ) =
p θ ( x t | y t
+1: T ) and b ( x ) =
p θ ( x t ) .
To generate positive and negative examples for DRE, we first sample a set of M latent state and observation trajectories from the generative model, x 1: M 1: T , y 1: M 1: T ∼ p θ ( x 1: T , y 1: T ) .
Because these are samples from joint distribution they are also samples from the conditionals { p θ ( x t | y t
+1: T ) } T − 1 t
=
We will refer to these samples as positive samples .
We can then draw a second set of samples, but discard the observed data, ˜ x 1: M 1: T ∼ p θ ( x 1: T ) .
We will refer to these as negative samples .
The sets of positive and negative examples form the data on which we will train the twist classifier.
Generating examples sequentially in this manner is cheap and parallelizable, allowing us to use relatively large values of M .
To train the DRE twist, we first pass the RNN backwards over sampled synthetic observed data y m 2: T to generate the sequence of encodings e m 2: T (noting that we flip the resulting encodings so they are “forward” in time).
To evaluate the probability of a positive classification we concatenate the encoding e mt +1 with x mt at each timestep, feed the concatenation into the head MLP, and take the output as a positive example Bernoulli logit.
To evaluate the probability of a negative classification we take the same sequence of encodings, e m 2: T , concatenate them at each time step with ˜ x m 1: T − 1 , feed the result at each timestep into the head MLP, and take the result as a negative example Bernoulli logit.
These outputs are used to compute the cross-entropy loss as written in Algorithm 1 at each timestep, which we average across time and across positive and negative examples to create the final loss.
This approach allows us to estimate an entire sequence of ratios, { p θ ( x t | y t +1: T )
/p
θ ( x t ) } T − 1 t =1 , for positive and negative examples using a single RNN backwards pass.
C.5 The SIXO Bound Can Become Tight
Proposition 1.
(Reproduced from Section 3.3) Sharpness of the SIXO bound.
Let p ( x 1: T , y 1: T ) be a latent variable model with Markovian structure as defined in Section 2, let Q be the set of possible sequences of proposal distributions indexed by parameters θ ∈ Θ , and let R be the set of possible sequences of positive, integrable twist functions indexed by parameters ψ ∈ Ψ .
Assume that { p ( x t | x t − 1 , y 1: T ) } Tt =1 ∈ Q and { p ( y t
+1: T | x t ) } T − 1 t =1 ∈ R .
Finally, assume L SIXO ( θ , ψ , y 1: T ) has the unique optimizer θ ∗ , ψ ∗ = arg max
[Equation]
Then the following holds:
q θ ∗ ( x t | x 1: t − 1 , y 1: T ) =
p ( x t | x 1: t − 1 , y 1: T ) for t = 1 , . .
, T , 2.
r ψ ∗ ( y t
+1: T , x t ) ∝ p ( y t +1: T | x t ) up to a constant independent of x t for t = 1 , . . .
, T − 1 , 3.
L K SIXO ( θ ∗ , ψ ∗ , y 1: T ) = log p ( y 1: T ) for any number of particles K ≥ 1 .
Proof.
First we reproduce the proof in the main text that L SIXO ( θ , ψ , y 1: T ) ≤ p ( y 1: T ) for any setting of θ , ψ .
As previously, fix r ψ ( x T ) = 1 and let (cid:98) Z SIXO ( θ , ψ , y 1: T ) be the normalizing constant estimator returned by SMC run with targets { p ( x 1: t , y 1: t ) r ψ ( y t +1: T , x t ) } Tt =1 and proposals { q θ ( x t | x 1: t − 1 , y 1: T )) }
Tt =1 .
Then,
[Equation]
where (32) holds by Jensen’s inequality and the concavity of log , and (33) holds by the unbiasedness of SMC’s marginal likelihood estimator.
Because L SIXO ( θ , ψ , y 1: T ) ≤ p ( y 1: T ) and we assume L SIXO ( θ , ψ , y 1: T ) has a unique opti- mizer, any setting of θ and ψ that causes the bound to hold with equality must be θ ∗ , ψ ∗ = arg max θ , ψ L SIXO ( θ , ψ , y 1: T ) .
Thus, we conclude the proof by showing that the twists { p ( y t
+1: T | x t ) } T − 1 t
=
and proposals { p ( x t | x 1: t − 1 , y 1: T ) } Tt =1 cause the bound to hold with equality.
We proceed by induction on t , the timestep in the SMC sweep.
We will show that for t = 1 , . . .
, T , (cid:98) Z t = p ( y 1: T ) .
In proving this we will also show that for each t , w kt either equals 1 or p ( y 1: T ) for k = 1 , . . .
, K , depending on whether resampling occurred.
For t = 1 note that
γ 1 ( x
[Equation]
γ 0 ≜ 1 , 3.
and q 1 ( x k 1 )
=
p ( x k 1 | y 1: T ) .
Taken together this implies that the incremental weight α k 1 is
[Equation]
[Equation]
which does not depend on k .
Because w k 0 ≜ 1 , we have that w k 1 = w k 0 α k 1 = [Equation] 1: T ) for all k .
This in turn implies (cid:92)
[Equation]
[Equation]
which when combined with the fact that (cid:98) Z 0 ≜ 1 yields
[Equation]
If resampling occurs at the end of step 1, all weights w 1: Kwill be set to 1 .
Thus we have shown that (cid:98) Z 1 = p ( y 1: T ) and w 1: K= p ( y 1: T ) or 1 .
Now assume that (cid:98) Z t − 1 = p ( y 1: T ) and w 1: K t − 1 equals 1 or p ( y 1: T ) .
Again, we derive the incremental weights α kt by noting that
[Equation]
which yields α kt as
[Equation]
for k = 1 , . . .
, K .
Now there are two cases depending on the value of the weights at the previous timestep.
If w 1: K t − 1 = 1 , then w kt = w kt − 1 α kt = 1 for all k , implying that (cid:92) Z t /Z
t − 1 = 1 .
Alternatively, if w 1: K t − 1 = p ( y 1: T ) then w kt = p ( y 1: T ) for all k which also implies that (cid:92) Z t /Z
t − 1 = 1 .
Given that (cid:92) Z t /Z
t − 1 = 1 in both cases, and that (cid:98) Z t − 1 = p ( y 1: T ) , we have that
[Equation]
Finally, if resampling occurs then the weights w 1: K t will be set to 1.
Thus we have shown that (cid:98) Z t = p ( y 1: T ) and w 1: K t = p ( y 1: T ) or 1 for each t = 1 , . . .
, T .
We note that we have used r ψ ∗ ( y t +1: T , x t )
= p ( y t +1: T | x t ) .
In the case where r ψ ∗ ( y t
+1: T , x t ) ∝ p ( y t +1: T | x t ) , we have that (cid:98) Z t ∝ p ( y 1: T ) for t < T .
All the proportionality constants for each r ψ ∗ ( y t +1: T , x t ) cancel out in (41), and hence we still obtain (cid:98) Z T = p
( y 1: T ) as required.
Note that we have incidentally shown that all weights are equal at each step of SMC for the optimal proposals and twisting functions.
This implies that the variance of the importance weights is minimized (i.e. is 0), and if effective sample size is used to trigger resampling, resampling will never occur.
D Experiments
Code for reproduction of all experiments is released here: https://github.com/lindermanlab/ sixo .
D.1 Gaussian Drift Diffusion D.1.1 Model Details
The one-dimensional Gaussian drift-diffusion process has joint distribution:
[Equation]
where the free parameters of the model are θ = { α } ∈ Θ
= R , the state is x t ∈ X =
R , and the observed data are y 1: T = y T ∈ R .
Training data are sampled from this joint distribution with α = 1 .
Note that the distributions we show in Figure 1 were generated with α = 0 for clarity.
D.1.2 Analytic Forms
The t th marginal of the filtering distribution for t <
T is
[Equation]
The t th marginal of the smoothing distributions can be derived as follows:
[Equation]
Noting that the product of two Gaussian densities is also Gaussian:
[Equation]
allows us to combine the two Gaussian distributions in (44):
[Equation]
Hence the smoothing distribution is a Gaussian distribution with analytically computable mean and variance terms.
The filtering and smoothing distributions are equal at t =
T .
It is interesting to note that for α = y T T +1 , which is the maximum likelihood drift parameter for a single datapoint y T , the sequence of filtering and smoothing distributions have the same means.
However, the variances are different for all t < T ; in particular, the smoothing distribution variance peaks in the middle of the timeseries, whereas the variance of the filtering distribution is monotonically increasing for t < T , and then drops at t = T .
According to Proposition 1, we expect the proposal recovered by SIXO, q θ t , to match the conditional of the smoothing distribution:
[Equation]
Note that the mean is an affine function with bias equal to zero.
Furthermore, we also expect the optimal twist distribution to be equal to the true lookahead distribution:
[Equation]
D.1.3 SIXO Variants for the Gaussian Drift Diffusion
For all experiments we use a Gaussian proposal at each timestep, parameterized as q θ t ( x t | x t − 1 , y T ) =
N ( x t ; f t ( x t − 1 , y T ) , σ 2 qt ) where f t is a general affine function of x t − 1 and y T .
There are therefore 4 T − 1 proposal parameters to learn ( T biases, T y T coefficients, T − 1 x t − 1 coefficients, and T variances σqt ).
We test four variants of SIXO:
SIXO-u learns θ and ψ by gradient ascent on the unified bound given in (3) using the unbiased gradients (reparameterization and resampling gradients).
We parameterize the twists as r t ( y T , x t ) =
N ( y T ; g t ( x t ) , σ 2 rt ) for t < T , where g t is a learnable affine function and σ 2 rt is also learned. [Figure] B o und ga p l og p ( y 1 : T ) −
L b
SIXO-DRE as defined in Algorithm 1 learns θ by ascending the bound (3) using the biased reparameterization gradients as in (18).
The twist parameters ψ are then fit using a density ratio update.
The twist is parameterized as an MLP that produces the coefficients of a quadratic function over x t as a function of y T and t .
This quadratic function is evaluated at x t to compute the log r ψ value.
SIXO-a (not included in Figure 2) uses the analytic form for the twist as a function of θ and y T (specified in (45)).
There are no free parameters to learn for this twist, and θ is learned by ascending the bound (3) using the biased reparameterization gradients in (18).
SIXO-b (not included in Figure 2) learns θ and ψ by gradient ascent on the unified bound given in (3) using biased reparameterization gradients (18).
We parameterize the twists as r t ( y T , x t ) =
N ( y T ; g t ( x t ) , σ 2 rt ) for t < T , where g t is a learnable affine function and σ 2 rt is also learned.
Note that the true distributions lie within the variational families (assuming a sufficiently expressive
MLP for SIXO-DRE, which is not unreasonable).
In all of these models we initialize the parameter α = 0 .
D.1.4 Results
We compare the four different variants of SIXO to IWAE, FIVO with biased gradients, and FIVO with unbiased gradients.
For clarity, we omitted some of these comparisons from Figure 2 in the main text, but include them here in Figure 6.
Individual seeds for each experiment were run using two CPU cores, 8Gb of memory, and had a runtime of no longer than five hours.
We sample 1, synthetic trajectories of length T = 10 from the model and perform minibatch stochastic gradient descent using the ADAM optimizer [55] with a learning rate of 1 × 10 − 3 , a minibatch size of 32, and 10 particles.
In the case of SIXO-DRE we alternate between 1,000 updates to the model and proposal and 1,000 twist updates, where the tilt minibatch size is 64.
We report and plot the means across ten random seeds.
The variances for FIVO, FIVO-b, SIXO-b are too small to be seen.
The model drift α is initialized to zero, affine function weights are initialized to zero, affine function biases are initialized to zero, and σ 2 qt and σ 2 rt are initialized to one.
In Figure 6 we show the convergence of the bound across all seven methods we considered.
FIVO with unbiased gradients (FIVO) performs comparably to FIVO with biased gradients (FIVO-b), however, converges slightly more slowly.
We find that SIXO-b performed worse than all the other methods, similar to the results in Lawson et al.
We therefore omitted it from the main paper for brevity.
However, this result motivated us to find alternative methods.
[Figure] All of SIXO-u, SIXO-DRE and SIXO-a converge to the correct solution and achieve a tight variational bound (parameter convergence is shown in Figure 7 for SIXO-u).
SIXO-u converges to a tight bound most slowly, but SIXO-DRE and SIXO-a converge quickly at a rate similar to IWAE.
This result shows that SIXO is able to recover the optimal model, true posterior [Figure] and true twists [Figure]
and that it can recover a tight variational bound.
It also shows that SIXO-DRE converges in a rate commensurate with the best-case convergence of SIXO-a. Future work will investigate whether the use of a twist makes the reparameterization gradient less biased.
Figures 1c and 1d show the lineages for two SMC sweeps: one using a model and proposal learned using FIVO, and one using a model, proposal and twist learned using SIXO-u. We resample at every timestep using systematic resampling.
In each color we show the unweighted filtering particles prior to resampling (i.e. the propagated particles) at time t , and then the smoothing particles after performing a backwards pass from that timestep.
As such, we are showing the family of smoothing and predictive distributions.
We see that FIVO proposes particles towards the observation, but that these particles are often preferentially resampled back towards the prior distribution.
This results in gross particle degeneracy.
In contrast, the lineages for SIXO are nearly perfect, with every particle being resampled exactly once.
D.2 Stochastic Volatility Model (SVM) D.2.1 Model Details
The SVM models an N -dimensional state-space is defined as follows:
[Equation]
where the states and observations are defined as:
[Equation]
the transition and observation noise terms are defined as:
[Equation]
and all multiplications are performed element-wise.
The model has free parameters defined as:
[Equation]
[Equation]
[Equation]
[Equation]
} , where µ ∈
[Equation]
[Equation]
such that there are 4 N model parameters.
The model learning objective is to recover the free parameters, θ , given observed data y .
The data we consider are the monthly returns from N =currencies over the period from 9/2007 to 8/2017, transformed into the log domain.
As a result,
X = Y
[Equation] × 22 .
D.2.2 Results
We use a fixed proposal per timestep parameterized as a Gaussian perturbation to p θ as in [9],
[Equation]
Thus the parameters specific to q are the means µ t ∈ R N and covariance matrices Σ t ∈ diag( R N + ) for t = 1 , . .
, T .
SIXO-q uses a parameterless quadrature twist, so we skip twist optimization and use biased reparameterization gradients of the SIXO bound to fit the model and proposal.
For the quadrature twist we use Gauss-Hermite quadrature with degree five.
For SIXO-DRE, we model the twist using the backwards RNN method introduced in Section 3.2.
The twist is parameterized with a one-layer LSTM with 128 hidden units and a one-layer MLP with 128 hidden units.
The twist is learned using the alternating DRE method described in Section 3.2.
We generate a batch of 32,000 synthetic trajectories from the current model, and perform minibatch stochastic gradient descent using the ADAM optimizer [55] with a learning rate of 3 × 10 − 3 and a minibatch size of 64.
We apply 1,000 twist updates (corresponding to two epochs) and then apply 1,000 updates to the model and proposal.
The model and proposal parameters are updated using the
ADAM optimizer [55] with a learning rate of 1 × 10 − 4 .
We use four particles per SMC sweep, and average across four datasets per model and proposal update.
Individual FIVO and SIXO-q seeds for each experiment were run using four CPU cores, 24Gb of memory, and had a wallclock time of no longer than twenty four hours.
SIXO-DRE had a longer runtime of five days, however, competitive results were achieved within two days.
Train Bound Performance We compare four methods: IWAE (learning θ ), FIVO (learning θ ),
SIXO-q (learning θ with a quadrature twist) and SIXO-DRE (learning θ and ψ using a DRE twist).
All methods (other than IWAE) use the biased resampling gradients of (18).
We show the median and quartiles across five random seeds.
Each µ n is initialized from N (0 , 0 . 3) (with n ∈ 1 , . .
ϕ n is learned in the unconstrained space R , and is transformed to [0 , 1] by passing the raw ϕ through a hyperbolic tangent function.
The unconstrained ϕ n ’s are initialized from N (arctanh(0 . 1) , 0 . 3) , and log β n is initialized from N (log 1 . 0 , 0 . 3) .
Finally, log Q n is initialized from N (log 1 . 0 , 0 . 3) .
The train bound value we report is taken as the average train bound after 75% of training.
We also report
L 2048BPF as the bound evaluated using a BPF with 2,048 particles.
This tests the performance of the learned model opposed to inference performance.
A one-way ANOVA [42] rejected the null hypothesis that the mean train bounds are equal ( p < 0 . 0001 ), and a post-hoc Dunnett T3 test [56] found that all methods are statistically significantly different ( p < 0 .
01 for all pairs).
For the L 2048BPF values, a one-way ANOVA failed to reject the null hypothesis that the train bounds are equal ( p = 0 . 26 ), so all entries are bolded.
This methodology was recommended in Sauder and DeMars [57].
Test Set Performance We also report the performance on a held-out dataset constructed from the new data since Naesseth et al.
[9] was published.
The test L 2048BPF we report is the bound evaluated using a BPF with 2,048 particles, averaged across all checkpoints after 75% of training.
A one-way
ANOVA rejected the null hypothesis that the mean test bounds are equal ( p < 0 . 0001 ), and a post-hoc
Dunnett T3 test found the mean SIXO-DRE bound to be the highest ( p < 0 .
01 for all pairs).
D.3 Hodgkin-Huxley Model
We provide a brief overview of the model here for completeness, but refer the reader to Chapter 5. of Dayan and Abbott [44] for more detailed information.
D .3.1 Model Details
The HH model is a physiologically grounded model of neural action potentials [43] defined through a set of four nonlinear differential equations.
Each neuron is defined by four state variables: the instantaneous membrane potential v ( t ) ∈ R , the potassium channel activation n ( t ) ∈ [0 , 1] , sodium channel activation m ( t ) ∈ [0 , 1] , and sodium channel inactivation h ( t ) ∈ [0 , 1] .
The channel states represent the aggregated probability that the given channel is active.
The state evolves according to:
[Equation]
[Equation]
The membrane capacitance C m is often defined to be 1 . 0 .
The first term, i ext ( t ) is the externally injected current.
The second term represents the net current through cell membrane due to the potential difference between the intracellular and extracellular mediums, often referred to as the leakage current .
This current is a function of the membrane capacitance, g L , and the potential of the extracellular medium, E L , where the potential difference across the membrane is then v ( t ) − E L .
The third term represents the net current through the membrane as a result of the potassium channels as a function of the channel state n ( t ) , the channel capacitance g K , and the potassium reversal potential
E K .
The final term represents the current through the membrane as a result of the sodium channel states, both m ( t ) and h ( t ) , the sodium channel conductance g Na , and the sodium reversal potential
E Na .
The channel states evolve according to:
[Equation]
[Equation]
where α n , α m , α h , β n , β m , β h are all fixed scalar functions of the membrane potential.
We discretize this continuous-time differential equation into a discrete-time latent variable model by integrating using Euler integration with an integration timestep of 0 .
02 ms (similarly to Huys and Paninski [47]).
This defines the deterministic time-evolution of the neural state.
We follow a similar approach as Huys and Paninski [47] and add Gaussian random noise to each of the four states at each timestep.
The membrane potential is unconstrained, and so we can add noise directly.
The gate states, however, are constrained to the range [0 , 1] .
Huys and Paninski [47] use truncated Gaussian noise to avoid pushing the state outside the constrained range.
We use a different approach and transform the constrained states into an unconstrained state by applying the inverse sigmoid function to the raw gate value.
This has the effect of modifying the variance of the perturbation in constrained space as a function of the state (heteroscedastic noise in constrained space).
However, this hetereoscedasticity carries with it a favorable intuition.
The magnitude of the noise term is reduced (after being pushed through a sigmoid) close the limits.
This means that the same noise kernel provides smaller perturbations close to the extremes, while still retaining a larger permissible perturbations in the mid-range.
The model is integrated with a timestep of 0.02ms.
We add zero-mean Gaussian noise to the potential with variance scaled by the integration timestep, σ 2 v = 9mVs − 1 × 0 . 02ms
= 0 .
18 µ V .
The unconstrained gate variables are perturbed by zero-mean Gaussian noise with variance also scaled by the integration timestep, σ 2 { n,m,h } = 0 .
1s − 1 × 0 .
02ms
= 0 . 000002 .
Observations are sampled from a Gaussian emission distribution centered on the current membrane potential with variance 25mV.
Observations are generated every 50 timesteps, corresponding to an acquisition rate of 1kHz.
We initialize the potential according to a Gaussian distribution with mean equal to -65mV and a standard deviation of 25mV. The unconstrained gate variables are initialized from a Gaussian distribution with mean defined by an estimate of the steady-state value x = sigmoid ( α x ( − 65) / α x ( − 65) − β x ( − 65) ) , where x represents the n , m or h states.
To iterate the model, we first constrain the state by passing the gate variables through a sigmoid.
The potential is already unconstrained and so requires no transform.
We then iterate the model given the constrained state.
The iterated state is then unconstrained by passing the gate states through the logit function (inverse of the sigmoid function).
The noise term is then added to the iterated, unconstrained state.
Observations are then generated by sampling from the emission distribution every 50 steps.
We generate traces with 2,048 timesteps, corresponding to approximately 40ms.
D.3.2 Results
For the experiments presented in Section 5.3 we use either a bootstrap, filtering or smoothing proposal.
The bootstrap proposal is defined as q θ ( x t | x t − 1 , y 1: T ) =
p θ ( x t | x t − 1 ) .
The filtering proposal runs an RNN forward over the observed data to create an encoding of all previous observations.
This encoding is then concatenated with a Transformer-style embedding of the elapsed time since the last observation [58], and is input into an MLP to produce the residual density with is multiplied with the prior density from which particles are then proposed (a type of resq proposal [59]).
The smoothing proposal runs two separate RNNs forwards and backwards over the data to encode all previous and future observations.
Both RNN encodings and the time encoding are concatenated and used as input to an MLP.
For SIXO, the twists are parameterized by an RNN running backwards over the observations. The hidden state of the RNN is concatenated with a time encoding and fed into an
MLP to produce the twist value.
All RNNs in these experiments have a hidden state of size 32, and all MLPs have two hidden layers of size 32. Individual seeds for each experiment were run using three CPU cores, 9Gb of memory, and had a runtime time of no longer than five days.
Inference In Figure 3 we compare the inference performance of a BPF (equivalent to FIVO-bs), to SIXO-bs.
The true generative model was used in both cases.
This highlights the advantage of using a learned twist under a fixed proposal.
In Table 4 we compare different inference methodologies in the true model across different observation intervals and number of particles.
FIVO-fi performs poorly, and drops off dramatically with the number of particles and sampling interval.
The performance of SIXO (both bootstrap and smoothing) outperforms FIVO-bs, and is more consistent across the number of particles.
Model Learning For the experiments presented in Figure 4 and Table 3 in Section 5.3 we learn the constant external current input, i ext in (52).
We generate 10,000 training sequences, and evaluate on 64 validation and test sequences, all generated with θ True = i ext = 13mV .
We use five seeds, deterministically initialized i ext to a value uniformly spaced between 1 . 3 µ A and 37 .
7 µ A , corresponding to a relative error of between − 0 .
9 and 1 . 9 .
We use K = 4 particles per SMC sweep during learning, when evaluating the model and proposal gradients (cf.
Line 11 of Algorithm 1), and average across four sequences per parameter update.
We use the same proposal and twist [Table] architectures described above.
We take 400 steps in the model per 100 steps in the twist.
We perform a hyperparameter search over the learning rate for the model, proposal [Table] and twist [Table]
and a learning rate decay schedule that halves the learning rate at specific steps.
The optimal parameters are shown in
Table 5.
We selected the best hyperparamters based on validation-set accuracy
[Table]
and then evaluated the
[Figure]
[Table]