Title: NAS-X: Neural Adaptive Smoothing via Twisting
Venue: NeurIPS 2023
Accepted: yes
Keywords: sequence models

probabilistic inference

reweighted wake-sleep

sequential monte carlo

smoothing

mechanistic models
Abstract: Sequential latent variable models (SLVMs) are essential tools in statistics and machine learning, with applications ranging from healthcare to neuroscience. As their flexibility increases, analytic inference and model learning can become challenging, necessitating approximate methods. Here we introduce neural adaptive smoothing via twisting (NAS-X), a method that extends reweighted wake-sleep (RWS) to the sequential setting by using smoothing sequential Monte Carlo (SMC) to estimate intractable posterior expectations. Combining RWS and smoothing SMC allows NAS-X to provide low-bias and low-variance gradient estimates, and fit both discrete and continuous latent variable models. We illustrate the theoretical advantages of NAS-X over previous methods and explore these advantages empirically in a variety of tasks, including a challenging application to mechanistic models of neuronal dynamics. These experiments show that NAS-X substantially outperforms previous VI- and RWS-based methods in inference and model learning, achieving lower parameter error and tighter likelihood bounds.

NASχ : Neural Adaptive Smoothing via Twisting

Abstract

Sequential latent variable models (SLVMs) are essential tools in statistics and machine learning, with applications ranging from healthcare to neuroscience. As their flexibility increases, analytic inference and model learning can become challenging, necessitating approximate methods. Here we introduce neural adaptive smoothing via twisting (NAS-X), a method that extends reweighted wake-sleep (RWS) to the sequential setting by using smoothing sequential Monte Carlo (SMC) to estimate intractable posterior expectations. Combining RWS and smoothing SMC allows NAS-X to provide low-bias and low-variance gradient estimates, and fit both discrete and continuous latent variable models. We illustrate the theoretical advantages of NAS-X over previous methods and explore these advantages empirically in a variety of tasks, including a challenging application to mechanistic models of neuronal dynamics. These experiments show that NAS-X substantially outperforms previous VI- and RWS-based methods in inference and model learning, achieving lower parameter error and tighter likelihood bounds.

1 Introduction

Sequential latent variable models (SLVMs) are a foundational model class in statistics and machine learning, propelled by the success of hidden Markov models [1] and linear dynamical systems [2]. To model more complex data, SLVMs have incorporated nonlinear conditional dependencies, resulting in models such as sequential variational autoencoders [3-7], financial volatility models [8], and biophysical models of neural activity [9]. While these nonlinear dependencies make SLVMs more flexible, they also frustrate inference and model learning, motivating the search for approximate methods.

One popular method for inference in nonlinear SLVMs is sequential Monte Carlo (SMC), which provides a weighted particle approximation to the true posterior and an unbiased estimator of the marginal likelihood. For most SLVMs, SMC is a significant improvement over standard importance sampling, providing an estimator of the marginal likelihood with variance that grows linearly in the length of the sequence rather than exponentially [10]. SMC's performance, however, depends on having a suitable proposal. The optimal proposal is often intractable, so in practice, proposal parameters are learned from data. Broadly, there are two approaches to proposal learning: variational inference [11, 12] and reweighted wake-sleep [13, 14].

Variational inference (VI) methods for SLVMs optimize the model and proposal parameters by ascending a lower bound on the log marginal likelihood that can be estimated with SMC, an approach called variational SMC [15-17]. Recent advances in variational SMC have extended it to work with smoothing SMC [18, 19], a crucial development for models where future observations are strongly related to previous model states. Without smoothing SMC, particle degeneracy can cause high variance of the lower bound and its gradients, resulting in unstable learning [18, 10].

Reweighted wake-sleep (RWS) methods instead use SMC's posterior approximation to directly estimate gradients of the log marginal likelihood [14, 20]. This approach can be interpreted as descending the inclusive KL divergence from the true posterior to the proposal. Notably, RWS methods work with discrete latent variables, providing a compelling advantage over VI methods that typically resort to high-variance score function estimates for discrete latent variables.

In this work, we combine recent advances in smoothing variational SMC with the benefits of reweighted wake-sleep. To this end, we introduce neural adaptive smoothing via twisting (NASX), a method for inference and model learning in nonlinear SLVMs that uses smoothing SMC to approximate intractable posterior expectations. The result is a versatile, low-bias and low-variance estimator of the gradients of the log marginal likelihood suitable for fitting proposal and model parameters in both continuous and discrete SLVMs.

After introducing NAS-X, we present two theoretical results that highlight the advantages of NAS-X over other RWS-based methods. We also demonstrate NAS-X's performance empirically in model learning and inference in linear Gaussian state-space models, discrete latent variable models, and high-dimensional ODE-based mechanistic models of neural dynamics. In all experiments, we find that NAS-X substantially outperforms several VI and RWS alternatives including ELBO, IWAE, FIVO, SIXO. Furthermore, we empirically show that our method enjoys lower bias and lower variance gradients, requires minimal additional computational overhead, and is robust and easy to train.

2 Background

This work considers model learning and inference in nonlinear sequential latent variable models with Markovian structure; i.e. models that factor as

with latent variables x 1: T ∈ X T , observations y 1: T ∈ Y T , and global parameters θ ∈ Θ . By nonlinear, we mean latent variable models where the parameters of the conditional distributions p θ ( x t | x t -1 ) and p θ ( y t | x t ) depend nonlinearly on x t -1 and x t , respectively.

Estimating the marginal likelihood p θ ( y 1: T ) and posterior p θ ( x 1: T | y 1: T ) for this model class is difficult because it requires computing an intractable integral over the latents,

We begin by introducing two algorithms, reweighted wake-sleep [14, 13] and smoothing sequential Monte Carlo [19], that are crucial for understanding our approach.

2.1 Reweighted Wake-Sleep

Reweighted wake-sleep (RWS, [13, 14]) is a method for maximum marginal likelihood in LVMs that estimates the gradients of the marginal likelihood using self-normalized importance sampling. This is motivated by Fisher's identity, which allows us to write the gradients of the marginal likelihood as a posterior expectation,

as proved in Appendix 8.1. The term inside the expectation is computable with modern automatic differentiation tools [21, 22], but the posterior p θ ( x 1: T | y 1: T ) is unavailable. Thus, SNIS is used to form a biased but consistent Monte Carlo estimate of Eq. (2) [23]. Specifically, SNIS draws N IID

samples from a proposal distribution, q ϕ ( x 1: T | y 1: T ) and weights them to form the estimator

The variance of this estimator is reduced as q ϕ approaches the posterior [23], so RWS also updates q ϕ by minimizing the inclusive Kullback-Leibler (KL) divergence from the posterior to the proposal. Crucially, the gradient for this step can also be written as the posterior expectation

as derived in Appendix 8.2. This allows RWS to estimate Eq. (4) using SNIS with the same set of samples and weights as Eq. (3),

Importantly, any method that provides estimates of expectations w.r.t. the posterior can be used for gradient estimation within the RWS framework, as we will see in the next section.

2.2 Estimating Posterior Expectations with Smoothing Sequential Monte Carlo

As we saw in Eqs. (2) and (4), key quantities in RWS can be expressed as expectations under the posterior. Standard RWS uses SNIS to approximate these expectations, but in sequence models the variance of the SNIS estimator can scale exponentially in the sequence length. In this section, we review sequential Monte Carlo (SMC) [10, 24], an inference algorithm that can produce estimators of posterior expectations with linear or even sub-linear variance scaling.

SMC approximates the posterior p θ ( x 1: T | y 1: T ) with a set of N weighted particles x 1: N 1: T constructed by sampling from a sequence of target distributions { π t ( x 1: t ) } T t =1 . Since these intermediate targets are often only known up to some unknown normalizing constant Z t , SMC uses the unnormalized targets { γ t ( x 1: t ) } T t =1 , where π t ( x 1: t ) = γ t ( x 1: t ) /Z t . Provided mild technical conditions are met and γ T ( x 1: T ) ∝ p θ ( x 1: T , y 1: T ) , SMC returns weighted particles that approximate the posterior p θ ( x 1: T | y 1: T ) [10, 24]. These weighted particles can be used to compute biased but consistent estimates of expectations under the posterior, similar to SNIS.

SMC repeats the following steps for each time t :

Sample latents x 1: N 1: t from a proposal distribution q ϕ ( x 1: t | y 1: T ) .

Weight each particle using the unnormalized target γ t to form an empirical approximation ˆ π t to the normalized target distribution π t .

Draw new particles x 1: N 1: t from the approximation ˆ π t (the resampling step).

By resampling away latent trajectories with low weight and focusing on promising particles, SMC can produce lower variance estimates than SNIS. For a thorough review of SMC, see Doucet and Johansen [10], Naesseth et al. [24], and Del Moral [25].

Filtering vs. Smoothing The most common choice of unnormalized targets γ t are the filtering distributions p θ ( x 1: t , y 1: t ) , resulting in the algorithm known as filtering SMC or a particle filter. Filtering SMC has been used to estimate posterior expectations within the RWS framework in neural adaptive sequential Monte Carlo (NASMC) [20], but a major disadvantage of filtering SMC is that it ignores future observations y t +1: T . Ignoring future observations can lead to particle degeneracy and high-variance estimates, which in turn causes poor model learning and inference [17, 18, 26, 27].

We could avoid these issues by using the smoothing distributions as unnormalized targets, choosing γ t ( x 1: t ) = p θ ( x 1: t , y 1: T ) , but unfortunately the smoothing distributions are not readily available from the model. We can approximate them, however, by observing that p θ ( x 1: t , y 1: T ) is proportional to the filtering distributions p θ ( x 1: t , y 1: t ) times the lookahead distributions p θ ( y t +1: T | x t ) . If the lookahead distributions are well-approximated by a sequence of twists { r ψ ( y t +1: T , x t ) } T t =1 , then running SMC with targets γ t ( x 1: t ) = p θ ( x 1: t , y 1: t ) r ψ ( y t +1: T , x t ) approximates smoothing SMC [26].

Twist Learning We have reduced the challenge of obtaining the smoothing distributions to learning twists that approximate the lookahead distributions. Previous twist-learning approaches include maximum likelihood training on samples from the model [18, 28] and Bellman-type losses motivated by writing the twist at time t recursively in terms of the twist at time t +1 [18, 29]. For NAS-X we use density ratio estimation (DRE) via classification to learn the twists, as introduced in Lawson et al. [19]. This method is motivated by observing that the lookahead distribution is proportional to a ratio of densities up to a constant independent of x t ,

Results from the DRE via classification literature [30] provide a way to approximate this density ratio: train a classifier to distinguish between samples from the numerator p θ ( x t | y t +1: T ) and denominator p θ ( x t ) . Then, the pre-sigmoid output of the classifier will approximate the log of the ratio in Eq. (6). For an intuitive argument for this fact see Appendix 8.3, and for a full proof see Sugiyama et al. [30].

In practice, it is not possible to sample directly from p θ ( x t | y t +1: T ) . Instead, Lawson et al. [19] sample full trajectories from the model's joint distribution, i.e. draw x 1: T , y 1: T ∼ p θ ( x 1: T , y 1: T ) , and discard unneeded timesteps, leaving only x t and y t +1: T which are distributed marginally as p θ ( x t , y t +1: T ) . Training the DRE classifier on data sampled in this manner will approximate the ratio p θ ( x t , y t +1: T ) /p θ ( x t ) p θ ( y t +1: T ) , which is equivalent to Eq. (6), see Appendix 8.3.

3 NAS-X: Neural Adaptive Smoothing via Twisting

The goal of NAS-X is to combine recent advances in smoothing SMC with the advantages of reweighted wake-sleep. Because SMC is a self-normalized importance sampling algorithm, it can be used to estimate posterior expectations and therefore the model and proposal gradients within a reweighted wake-sleep framework. In particular, NAS-X repeats the following steps:

Draw a set of N trajectories x (1: N ) 1: T and weights w (1: N ) 1: T from a smoothing SMC run with model p θ , proposal q ϕ , and twist r ψ .

Use those trajectories and weights to form estimates of gradients for the model p θ and proposal q ϕ , as in reweighted wake-sleep. Specifically, NAS-X computes the gradients of the inclusive KL divergence for learning the proposal q ϕ as

and computes the gradients of the model p θ as

Update the twists r ψ using density ratio estimation via classification.

A full description is available in Algorithms 1 and 2.

A key design decision in NAS-X is the specific form of the gradient estimators. Smoothing SMC provides two ways to estimate expectations of test functions with respect to the posterior: both the timestept and timestepT approximations of the target distribution could be used, in the latter case by discarding timesteps after t . Specifically,

where δ ( a ; b ) is a Dirac delta of a located at b and ( x 1: T ) 1: t denotes selecting the first t timesteps of a timestepT particle; due to SMC's ancestral resampling step these are not in general equivalent. For NAS-X we choose the timet approximation of the posterior to lessen particle degeneracy, as in NASMC[20]. Note, however, that in the case of NASMC this amounts to approximating the posterior with the filtering distributions, which ignores information from future observations. In the case of NAS-X, the intermediate distributions directly approximate the posterior distributions because of the twists, a key advantage that we explore theoretically in Section 3.1 and empirically in Section 5.

Algorithm 1: NAS-X

3.1 Theoretical Analysis of NAS-X

In this section, we state two theoretical results that illustrate NAS-X's advantages over NASMC, with proofs given in Appendix 7.

Proposition 1. Consistency of NAS-X's gradient estimates. Suppose the twists are optimal so that 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 . Let ˆ ∇ θ log p θ ( y 1: T ) be NAS-X's weighted particle approximation to the true gradient of the log marginal likelihood ∇ θ log p θ ( y 1: T ) . Then ˆ ∇ θ log p θ ( y 1: T ) a.s. → ∇ θ log p θ ( y 1: T ) as N →∞ .

Proposition 2. Unbiasedness of NAS-X's gradient estimates. Assume that proposal distribution q ϕ ( x t | x 1: t -1 , y 1: T ) is optimal so that q ϕ ( x t | x 1: t -1 , y 1: T ) = p ( x t | x 1: t -1 , y 1: T ) for t = 1 , . . . , T , and the twists r ψ ( y t +1: T , x t ) are optimal so that 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 . Let ˆ ∇ θ log p θ ( y 1: T ) be NAS-X's weighted particle approximation to the true gradient of the log marginal likelihood, ∇ θ log p θ ( y 1: T ) . Then, for any number of particles, E [ ˆ ∇ θ log p θ ( y 1: T )] = log p θ ( y 1: T ) .

4 Related Work

VI Methods There is a large literature on model learning via stochastic gradient ascent on an evidence lower bound (ELBO) [31, 32, 4, 33]. Subsequent works have considered ELBOs defined by the normalizing constant estimates from multiple importance sampling [34], nested importance sampling, [35, 36], rejection sampling, and Hamiltonian Monte Carlo [37]. Most relevant to our work is the literature that uses SMC's estimates of the normalizing constant as a surrogate objective. There are a number of VI methods based on filtering [17, 16, 15] and smoothing SMC [18, 38, 39, 19, 27], but filtering SMC approaches can suffer from particle degeneracy and high variance [18, 19].

Reweighted Wake-Sleep Methods The wake-sleep algorithm was introduced in Hinton et al. [13] as a way to train deep directed graphical models. Bornschein and Bengio [14] interpreted the wake-sleep algorithm as self-normalized importance sampling and proposed reweighted wakesleep, which uses SNIS to approximate gradients of the inclusive KL divergence and log marginal likelihood. Neural adaptive sequential Monte Carlo (NASMC) extends RWS by using filtering SMC to approximate posterior expectations instead of SNIS [20]. To combat particle degeneracy, NASMC approximates the posterior with the filtering distributions, which introduces bias.

NAS-X vs. SIXO Both NAS-X and SIXO [19] leverage smoothing SMC with DRE-learned twists, but NAS-X uses smoothing SMC to estimate gradients in an RWS-like framework while SIXO uses it within a VI-like framework. Thus, NAS-X follows biased but consistent estimates of the log marginal likelihood while SIXO follows unbiased estimates of a lower bound on the log marginal likelihood. It is not clear a-priori which approach would perform better, but we provide empirical evidence in Section 5 that shows NAS-X is more stable than SIXO and learns better models and proposals. In

0

100

200

300

400

500

600

700

Proposal training steps (1000s)

10

-

1

10

0

10

1

10

2

log

p

(

y

1:

T

)

-L

128

NAS-X

NASMC

ELBO

IWAE

RWS

SIXO

FIVO

0

100

200

300

400

500

600

700

Proposal training steps (1000s)

10

-

2

10

-

1

10

0

Proposal relative error

0

10

20

30

40

50

Model timesteps

0

.

30

0

.

35

0

.

40

0

.

45

0

.

50

0

.

55

0

.

60

0

.

65

σ

2

value

addition to these empirical advantages, NAS-X can fit discrete latent variable models while SIXO would require high-variance score function estimates.

5 Experiments

We empirically validate the following advantages of NAS-X:

By using the approximate smoothing distributions as targets for proposal learning, NAS-X can learn proposals that match the true posterior marginals, while NASMC and other baseline methods cannot. We illustrate this in Section 5.1, in a setting where the true posterior is tractable. We illustrate the practical benefits on inference in nonlinear mechanistic models in Section 5.3.

By optimizing the proposal within the RWS framework ( e.g., descending the inclusive KL), NAS-X can perform learning and inference in discrete latent variable models, which SIXO cannot. We explore this in Section 5.2.

We explore the practical benefits of this combination in a challenging setting in Section 5.3, where we show NAS-X can fit ODE-based mechanistic models of neural dynamics with 38 model parameters and an 18-dimensional latent state.

In addition to these experiments, we analyze the computational complexity and wall-clock speed of each method and the bias and variance of the gradient estimates in Sections 15 and 16 of the Appendix.

5.1 Linear Gaussian State Space Model

We first consider a one-dimensional linear Gaussian state space model with joint distribution

In Figure 1, we compare NAS-X against several baselines (NASMC, FIVO, SIXO, RWS, IWAE, and ELBO) by evaluating log marginal likelihood estimates (left panel) and recovery of the true posterior (middle and right panels). For all methods we learn a mean-field Gaussian proposal factored over time, q ( x 1: T ) = ∏ T t =1 q t ( x t ) = ∏ T t =1 N ( x t ; µ t , σ 2 t ) , with parameters µ 1: T and σ 2 1: T corresponding to the means and variances at each time-step. For twist-based methods, we parameterize the twist as a quadratic function in x t whose coefficients are functions of the observations and time step. We chose this form to match the functional form of the analytic log density ratio. For details, see Section 9 in the Appendix. NAS-X outperforms all baseline methods, achieving a tighter lower bound on the log-marginal likelihood and lower parameter error.

In the right panel of Figure 1, we compare the learned proposal variances against the true posterior variance, which can be computed in closed form. See Section 9 for comparison of proposal means; we do not report this comparison in the main text since all methods recover the posterior mean. This comparison gives insight into NAS-X's better performance. NASMC's learned proposal overestimates the posterior variance and fails to capture the true posterior distribution, because it employs a filtering

approximation to the gradients of the proposal distribution. On the other hand, by using the twisted targets, which approximate the smoothing distributions, to estimate proposal gradients, NAS-X recovers the true posterior.

5.2 Switching Linear Dynamical Systems

To explore NAS-X's ability to handle discrete latent variables, we consider a switching linear dynamical system (SLDS) model [40, 41]. Specifically, we adapt the recurrent SLDS example from Linderman et al. [41] in which the latent dynamics trace ovals in a manner that resembles cars racing on a NASCAR track. There are two coupled sets of latent variables: a discrete state z t , with K = 4 possible values, and a two-dimensional continuous state x t that follows linear dynamics that depend on z t . The observations are a noisy projection of x t into a ten-dimensional observation space. There are 1000 observations in total. For the proposal, we factor q over both time and the continuous and discrete states. The continuous distributions are parameterized by Gaussians, and categorical distributions are used for the discrete latent variables. For additional details on the proposal and twist, see Section 10 in the Appendix and for details on the generative model see Linderman et al. [41].

(a) Comparing learned dynamics of NAS-X, NASMC, and Laplace EM on rSLDS.

-

2

0

2

x

1

-

1

.

5

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

1

.

5

x

2

Ground truth

-

2

0

2

x

1

-

1

.

5

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

1

.

5

x

2

NAS-X

-

2

0

2

x

1

-

1

.

5

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

1

.

5

x

2

NASMC

-

2

0

2

x

1

-

1

.

5

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

1

.

5

x

2

Laplace EM

(b) Train

L

1024

BPF

for rSLDS.

Method

σ

2

O

= 0

.

001

σ

2

O

= 0

.

01

σ

2

O

= 0

.

1

NAS-X

19

.

837

±

0

.

0234

8

.

63

±

0

.

0015

-

2

.

79

±

0

.

0009

NASMC

19

.

834

±

0

.

0018

8

.

53

±

0

.

001

-

2

.

874

±

0

.

0007

Laplace EM

19

.

154

±

0

.

057

8

.

54

±

0

.

0039

-

2

.

765

±

0

.

0012

RWS

17

.

148

0

.

087

6

.

314

0

.

023

5

.

78

0

.

0026

±

±

-

±

We present qualitative results from model learning and inference in the top panel of Figure 2. We compare the learned dynamics for NAS-X, NASMC, and a Laplace-EM algorithm designed specifically for recurrent state space models [42]. In each panel, we plot the vector field of the learned dynamics and the posterior means, with different colors corresponding to the four discrete states. NAS-X recovers the true dynamics accurately. In the Table in Figure 2, we quantitatively compare the model learning performances across these three approaches by running a bootstrap proposal with the learned models and the true dynamics and observation variances. We normalize the bounds by the sequence length. NAS-X outperforms or performs on par with both NASMC and Laplace EM across the different observation noises σ 2 O . See Section 10, for additional results on inference.

5.3 Biophysical Models of Neuronal Dynamics

For our final set of experiments we consider inference and parameter learning in Hodgkin-Huxley (HH) models [9, 43] - mechanistic models of voltage dynamics in neurons. These models use systems of coupled nonlinear differential equations to describe the evolution of the voltage difference across a neuronal membrane as it changes in response to external stimuli such as injected current. Un-

derstanding how voltage propagates throughout a cell is central to understanding electrical signaling and computation in the brain.

Voltage dynamics are governed by the flow of charged ions across the cell membrane, which is in turn mediated by the opening and closing of ion channels and pumps. HH models capture the states of these ion channels as well as the concentrations of ions and the overall voltage, resulting in a complex dynamical system with many free parameters and a high dimensional latent state space. Model learning and inference in this setting can be extremely challenging due to the dimensionality, noisy data, and expensive and brittle simulators that fail for many parameter settings.

Model Description We give a brief introduction to the models used in this section and defer a full description to the appendix. We are concerned with modeling the potential difference across a neuronal cell membrane, v , which changes in response to currents flowing through a set of ion channels, c ∈ C . Each ion channel c has an activation which represents a percentage of the maximum current that can flow through the channel and is computed as a nonlinear function g c of the channel state λ c , with g c ( λ c ) ∈ [0 , 1] . This activation specifies the time-varying conductance of the channel as a fraction of the maximum conductance of the channel, g c . Altogether, the dynamics for the voltage v can be written as

where c m is the specific membrane capacitance, I ext is the external current applied to the cell, S is the cell membrane surface area, c ion is the ion transported by channel c , and E c ion is that ion's reversal potential. In addition to the voltage dynamics, the ion channel states { λ c } c ∈C evolve as

where A ( v ) and b ( v ) are nonlinear functions of the membrane potential that produce matrices and vectors, respectively. Together, equations (11) and (12) define a conditionally linear system of first-order ordinary differential equations (ODEs), meaning that the voltage dynamics are linear if the channel states are known and vice-versa.

Following Lawson et al. [19], we augment the deterministic dynamics with zero-mean, additive Gaussian noise to the voltage v and unconstrained gate states logit( λ c ) at each integration time-step. The observations are produced by adding Gaussian noise to the instantaneous membrane potential.

Proposal and Twist Parameterization For all models in this section we amortize proposal and twist learning across datapoints. SIXO and NAS-X proposals used bidirectional recurrent neural networks (RNNs) [44, 45] with a hidden size of 64 units to process the raw observations and external current stimuli, and then fed the processed observations, previous latent state, and a transformer positional encoding [46] into a 64-unit single-layer MLP that produced the parameters of an isotropic Gaussian distribution over the current latent state. Twists were similarly parameterized with an RNN run in reverse across observations, combined with an MLP that accepts the RNN outputs and latent state and produces the twist values.

Integration via Strang Splitting The HH ODEs are stiff , meaning they are challenging to integrate at large step sizes because their state can change rapidly. While small step sizes can ensure numerical stability, they also make methods prohibitively slow. For example, many voltage dynamics of interest unfold over hundreds of milliseconds, which could take upwards of 40,000 integration steps at the standard 0.005 milliseconds per step. Because running our models for even 10,000 steps would be too costly, we developed new numerical integration techniques based on an explicit Strang splitting approach that allowed us to stably integrate at 0.1 milliseconds per step, a 20-time speedup [47]. For details, see Section 13 in the Appendix.

5.3.1 Hodgkin-Huxley Inference Results

First, we evaluated NAS-X, NASMC, and SIXO in their ability to infer underlying voltages and channel states from noisy voltage observations. For this task we sampled 10,000 noisy voltage traces from a probabilistic Hodgkin-Huxley model of the squid giant axon [9], and used each method to

4

8

16

32

64

128

256

-169

-181

-193

-206

Log-likelihood lower bound

IWAE

FIVO

NASMC

SIXO

NAS-X

-

50

0

50

Voltage (mV)

SIXO

True Voltage

Inferred Voltage

Resampling Event

0

5

10

15

20

25

30

35

40

45

50

Time (ms)

-

50

0

50

Voltage (mV)

NAS-X

(a)

Proposals were trained with 4 particles and evaluated across a range of particle numbers. RWS performed too poorly to be included.

(b) Inferred voltage traces for SIXO and NAS-X.

(top) SIXO generates a high number of resampling events leading to particle degeneracy and a single mistimed spike. (bottom) NASX perfectly infers the latent voltage with no mistimed spikes, and resamples infrequently. See Fig. 9 in Appendix for NASMC's traces.

train proposals (and twists for NAS-X and SIXO) to compute the marginal likelihood assigned to the data under the true model. As in [19], we sampled trajectories of length 50 milliseconds, with a single noisy voltage observation every millisecond. The stability of the Strang splitting based ODE integrator allowed us to integrate at dt = 0 . 1 ms, meaning there were 10 latent states per observation.

In Figure (3a) we plot the performance of proposals and twists trained with 4 particles and evaluated across a range of particle numbers. All methods perform roughly the same when evaluated with 256 particles, but with lower numbers of evaluation particles the smoothing methods emerge as more particle-efficient than the filtering methods. To achieve NAS-X's inference performance with 4 particles, NASMC would need 256 particles, a 64x increase, and NAS-X is also on average 2x more particle-efficient than SIXO.

In Figure 3b we further investigate these results by examining the inferred voltage traces of NAS-X and SIXO. SIXO accurately infers the timing of most spikes but resamples at a high rate, which can lead to particle degeneracy and poor bound performance. NAS-X correctly infers the voltage across the whole trace with no spurious or mistimed spikes and almost no resampling events, indicating it has learned a high-quality proposal that does not generate poor particles that must be resampled away. These qualitative results support the quantitative results in Figure 3a: SIXO's high resampling rate and NASMC's filtering approach lead to lower bound values.

These results highlight a benefit of RWS-based methods over VI methods: when the model is correctly specified, it can be beneficial to have a more deterministic proposal. Empirically, we find that maximizing the variational lower bound encourages the proposal to have high entropy, which in this case resulted in SIXO's poorer performance relative to NAS-X. In the next section, we explore the implications of this on model learning.

5.3.2 Hodgkin-Huxley Model Learning Results

In this section, we assess NAS-X and SIXO's ability to fit model parameters in a more complex, biophysically realistic model of a pyramidal neuron from the mouse visual cortex. This model was taken from the Allen Institute Brain Atlas [48] and includes 9 different voltage-gated ion channels as well as a calcium pump/buffer subsystem and a calcium-gated potassium ion channel. In total, the model had 38 free parameters and an 18-dimensional latent state space, in contrast to the 1 free parameter and 4-dimensional state space of the model considered by Lawson et al. [19]. For full details of the models, see Appendix Section 12.

We fit these models to voltage traces gathered from a real mouse neuron by the Allen Institute, but downsampled and noised the data to simulate a more common voltage imaging setting. We ran a hyperparameter sweep over learning rates and initial values of the voltage and observation noise variances (270 hyperparameter settings in all), and selected the best performing model via early stopping on the train log marginal likelihood lower bound. Each hyperparameter setting was run for

0

20

40

60

80

100

120

140

160

180

200

Time (ms)

True Voltage

SIXO Samples

NAS-X Samples

Method

L

32

BPF

Cross-Corr.

%Runs Failed

NAS-X

-

686

.

4

±

6

.

8

0

.

76

±

0

.

15

2

.

74

±

0

.

1

6258

±

11

18

.

9

SIXO

660

.

6

4

.

4

1

.

88

0

.

41

1

.

8

0

.

2

6055

22

25

.

2

-

±

±

±

±

5 seeds, and each seed was run for 2 days on a single CPU core with 7 Gb of memory. Because of the inherent instability of these models, many seeds failed, and we discarded hyperparameter settings with more than 2 failed runs.

In Figure 4 (bottom), we compare NAS-X and SIXO-trained models with respect to test set loglikelihood lower bounds as well as biophysically relevant metrics. To compute the biophsyical metrics, we sampled 32 voltage traces for each input stimulus trace in the test set, and averaged the feature errors over the samples and test set. NAS-X better captures the number of spikes, an important feature of the traces, and attains a higher cross correlation. Both methods capture the resting voltage well, although SIXO attains a slightly lower error and outperforms NAS-X in terms of log-likelihood lower bound.

Training instability is a significant practical challenge when fitting mechanistic models. Therefore, we also include the percentage of runs that failed for each method. SIXO's more entropic proposals more frequently generate biophysically implausible latent states, causing the ODE integrator to return NaNs. In contrast, fewer of NAS-X's runs suffer from numerical instability issues, a great advantage when working with mechanistic models.

6 Conclusion

In this work we presented NAS-X, a new method for model learning and inference in sequential latent variable models, that combines reweighted wake-sleep framework and approximate smoothing SMC. Our approach involves learning twist functions to use in smoothing SMC, and then running smoothing SMC to approximate gradients of the log marginal likelihood with respect to the model parameters and gradients of the inclusive KL divergence with respect to the proposal parameters. We validated our approach in experiments including model learning and inference for discrete latent variable models and mechanistic models of neural dynamics, demonstrating that NAS-X offers compelling advantages in many settings.

References

Leonard E Baum and Ted Petrie. Statistical inference for probabilistic functions of finite state markov chains. The annals of mathematical statistics , 37(6):1554-1563, 1966.

Rudolph Emil Kalman. A new approach to linear filtering and prediction problems. ASME Journal of Basic Engineering , 82(1):35-45, March 1960.

Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International Conference on Machine Learning , pages 1278-1286. PMLR, 2014.

Diederik Kingma and Max Welling. Auto-encoding variational Bayes. In 2nd International Conference on Learning Representations , 2014.

Rahul G Krishnan, Uri Shalit, and David Sontag. Deep Kalman filters. arXiv preprint arXiv:1511.05121 , 2015.

Junyoung Chung, Kyle Kastner, Laurent Dinh, Kratarth Goel, Aaron C Courville, and Yoshua Bengio. A recurrent latent variable model for sequential data. Advances in neural information processing systems , 28, 2015.

Marco Fraccaro, Søren Kaae Sønderby, Ulrich Paquet, and Ole Winther. Sequential neural models with stochastic layers. Advances in Neural Information Processing Systems , 29, 2016.

Siddhartha Chib, Yasuhiro Omori, and Manabu Asai. Multivariate stochastic volatility. In Handbook of Financial Time Series , pages 365-400. Springer, 2009.

Alan L. Hodgkin and Andrew F. Huxley. A quantitative description of membrane current and its application to conduction and excitation in nerve. The Journal of Physiology , 117(4):500, 1952.

Arnaud Doucet and Adam M. Johansen. A tutorial on particle filtering and smoothing: Fifteen years later. In Dan Crisan and Boris Rozovsky, editors, The Oxford Handbook of Nonlinear Filtering , pages 656-704. Oxford University Press, 2011.

David Blei, Alp Kucukelbir, and Jon D. McAuliffe. Variational inference: A review for statisticians. Journal of the American Statistical Association , 112(518):859-877, 2017.

Martin J. Wainwright and Michael I. Jordan. Graphical models, exponential families, and variational inference. Foundations and Trends® in Machine Learning , 1(1-2):1-305, 2008.

Geoffrey E Hinton, Peter Dayan, Brendan J Frey, and Radford M Neal. The 'wake-sleep' algorithm for unsupervised neural networks. Science , 268(5214):1158-1161, 1995.

Jörg Bornschein and Yoshua Bengio. Reweighted Wake-Sleep. arXiv preprint arXiv:1406.2751 , 2014.

Tuan Anh Le, Maximilian Igl, Tom Rainforth, Tom Jin, and Frank Wood. Auto-encoding Sequential Monte Carlo. In 6th International Conference on Learning Representations , 2018.

Christian A. Naesseth, Scott Linderman, Rajesh Ranganath, and David Blei. Variational Sequential Monte Carlo. In International Conference on Artificial Intelligence and Statistics , pages 968-977. PMLR, 2018.

Chris J. Maddison, Dieterich Lawson, George Tucker, Nicolas Heess, Mohammad Norouzi, Andriy Mnih, Arnaud Doucet, and Yee Whye Teh. Filtering variational objectives. Advances in Neural Information Processing Systems , 30, 2017.

Dieterich Lawson, George Tucker, Christian A. Naesseth, Chris Maddison, Ryan P. Adams, and Yee Whye Teh. Twisted Variational Sequential Monte Carlo. In Third workshop on Bayesian Deep Learning (NeurIPS) , 2018.

Dieterich Lawson, Allan Raventós, Andrew Warrington, and Scott Linderman. SIXO: Smoothing Inference with Twisted Objectives. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022.

Shixiang Shane Gu, Zoubin Ghahramani, and Richard E. Turner. Neural Adaptive Sequential Monte Carlo. Advances in Neural Information Processing Systems , 28, 2015.

James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax .

Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Largescale machine learning on heterogeneous systems, 2015. URL https://www.tensorflow. org/ . Software available from tensorflow.org.

Art B Owen. Monte Carlo Theory, Methods and Examples . Stanford, 2013.

Christian A. Naesseth, Fredrik Lindsten, Thomas B. Schön, et al. Elements of Sequential Monte Carlo. Foundations and Trends® in Machine Learning , 12(3):307-392, 2019.

Pierre Del Moral. Feynman-Kac formulae: genealogical and interacting particle systems with applications , volume 88. Springer, 2004.

Nick Whiteley and Anthony Lee. Twisted particle filters. The Annals of Statistics , 42(1): 115-141, 2014.

Mark Briers, Arnaud Doucet, and Simon Maskell. Smoothing algorithms for state-space models. Annals of the Institute of Statistical Mathematics , 62(1):61-89, 2010.

Ming Lin, Rong Chen, and Jun S. Liu. Lookahead strategies for sequential Monte Carlo. Statistical Science , 28(1):69-94, 2013.

Vasileios Lioutas, Jonathan Wilder Lavington, Justice Sefas, Matthew Niedoba, Yunpeng Liu, Berend Zwartsenberg, Setareh Dabiri, Frank Wood, and Adam Scibior. Critic sequential monte carlo. arXiv preprint arXiv:2205.15460 , 2022.

Masashi Sugiyama, Taiji Suzuki, and Takafumi Kanamori. Density ratio estimation in machine learning . Cambridge University Press, 2012.

Rajesh Ranganath, Sean Gerrish, and David Blei. Black box variational inference. In International Conference on Artificial Intelligence and Statistics , pages 814-822. PMLR, 2014.

Matthew D. Hoffman, David Blei, Chong Wang, and John Paisley. Stochastic variational inference. Journal of Machine Learning Research , 2013.

Andriy Mnih and Danilo Rezende. Variational inference for Monte Carlo objectives. In International Conference on Machine Learning , pages 2188-2196. PMLR, 2016.

Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. In 4th International Conference on Learning Representations , 2016.

Christian Naesseth, Fredrik Lindsten, and Thomas Schon. Nested Sequential Monte Carlo Methods. In International Conference on Machine Learning , pages 1292-1301. PMLR, 2015.

Heiko Zimmermann, Hao Wu, Babak Esmaeili, and Jan-Willem van de Meent. Nested variational inference. Advances in Neural Information Processing Systems , 34:20423-20435, 2021.

Dieterich Lawson, George Tucker, Bo Dai, and Rajesh Ranganath. Energy-inspired models: Learning with sampler-induced distributions. Advances in Neural Information Processing Systems , 32, 2019.

Antonio Moretti, Zizhao Wang, Luhuan Wu, Iddo Drori, and Itsik Pe'er. Variational objectives for Markovian dynamics with backward simulation. In ECAI 2020 , pages 1371-1378. IOS Press, 2020.

Antonio Moretti, Zizhao Wang, Luhuan Wu, and Itsik Pe'er. Smoothing nonlinear variational objectives with sequential Monte Carlo. ICLR Workshop: Deep Generative Models for Highly Structured Data , 2019.

Emily Fox, Erik Sudderth, Michael Jordan, and Alan Willsky. Nonparametric Bayesian Learning of Switching Linear Dynamical systems. In D. Koller, D. Schuurmans, Y. Bengio, and L. Bottou, editors, Advances in Neural Information Processing Systems , volume 21. Curran Associates, Inc., 2008.

Scott W. Linderman, Matthew J. Johnson, Andrew C. Miller, Ryan P. Adams, David M. Blei, and Liam Paninski. Bayesian learning and inference in recurrent switching linear dynamical systems. In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics (AISTATS) , 2017.

David Zoltowski, Jonathan Pillow, and Scott Linderman. A general recurrent state space framework for modeling neural dynamics during decision-making. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pages 11680-11691. PMLR, 13-18 Jul 2020.

Peter Dayan and Laurence F. Abbott. Theoretical Neuroscience: Computational and Mathematical Modeling of Neural Systems . MIT press, 2005.

David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams. Learning internal representations by error propagation. Technical report, California Univ San Diego La Jolla Inst for Cognitive Science, 1985.

Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation , 9 (8):1735-1780, 1997.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is All You Need. Advances in Neural Information Processing Systems , 30, 2017.

Zhengdao Chen, Baranidharan Raman, and Ari Stern. Structure-preserving numerical integrators for Hodgkin-Huxley-type systems. SIAM Journal on Scientific Computing , 42(1):B273-B298, 2020.

Quanxin Wang, Song-Lin Ding, Yang Li, Josh Royall, David Feng, Phil Lesnar, Nile Graddis, Maitham Naeemi, Benjamin Facer, Anh Ho, et al. The Allen mouse brain common coordinate framework: a 3d reference atlas. Cell , 181(4):936-953, 2020.

Michael I Jordan. Serial order: A parallel distributed processing approach. In Advances in psychology , volume 121, pages 471-495. Elsevier, 1997.

Diederik Kingma, Jimmy Ba, Yoshua Bengio, and Yann LeCun. Adam: A method for stochastic optimization. In 3rd International Conference on Learning Representations , 2015.

Tom Rainforth, Adam Kosiorek, Tuan Anh Le, Chris Maddison, Maximilian Igl, Frank Wood, and Yee Whye Teh. Tighter variational bounds are not necessarily better. In International Conference on Machine Learning , pages 4277-4285. PMLR, 2018.

AIBS. Biophysical modeling - perisomatic. Technical report, Allen Institute for Brain Science, 10 2017. URL http://help.brain-map.org/display/celltypes/Documentation .

Costa M Colbert and Enhui Pan. Ion channel properties underlying axonal action potential initiation in pyramidal neurons. Nature neuroscience , 5(6):533-538, 2002.

Jacopo Magistretti and Angel Alonso. Biophysical properties and slow voltage-dependent inactivation of a sustained sodium current in entorhinal cortex layer-ii principal neurons: a whole-cell and single-channel study. The Journal of general physiology , 114(4):491-509, 1999.

Maarten HP Kole, Stefan Hallermann, and Greg J Stuart. Single ih channels in pyramidal neuron dendrites: properties, distribution, and impact on action potential output. Journal of Neuroscience , 26(6):1677-1687, 2006.

I Reuveni, A Friedman, Y Amitai, and Michael J Gutnick. Stepwise repolarization from ca2+ plateaus in neocortical pyramidal cells: evidence for nonhomogeneous distribution of hva ca2+ channels in dendrites. Journal of Neuroscience , 13(11):4609-4621, 1993.

Robert B Avery and Daniel Johnston. Multiple channel types contribute to the low-voltageactivated calcium current in hippocampal ca3 pyramidal neurons. Journal of Neuroscience , 16 (18):5567-5582, 1996.

AD Randall and RW Tsien. Contrasting biophysical and pharmacological properties of t-type and r-type calcium channels. Neuropharmacology , 36(7):879-893, 1997.

PR Adams, DA Brown, and A Constanti. M-currents and other potassium currents in bullfrog sympathetic neurones. The Journal of Physiology , 330(1):537-572, 1982.

Alon Korngreen and Bert Sakmann. Voltage-gated k+ channels in layer 5 neocortical pyramidal neurones from young rats: subtypes and gradients. The Journal of physiology , 525(3):621-639, 2000.

M Köhler, B Hirschberg, CT Bond, John Mark Kinzie, NV Marrion, James Maylie, and JP Adelman. Small-conductance, calcium-activated potassium channels from mammalian brain. Science , 273(5282):1709-1714, 1996.

Alastair J Walker. New fast method for generating discrete random numbers with arbitrary frequency distributions. Electronics Letters , 8(10):127-128, 1974.

Richard A Kronmal and Arthur V Peterson Jr. On the alias method for generating random variables from a discrete distribution. The American Statistician , 33(4):214-218, 1979.

7 Theoretical Analyses

Proposition 1. Consistency of NAS-X's gradient estimates. Suppose the twists are optimal so that 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 . Let ˆ ∇ θ log p θ ( y 1: T ) be NAS-X's weighted particle approximation to the true gradient of the log marginal likelihood ∇ θ log p θ ( y 1: T ) . Then ˆ ∇ θ log p θ ( y 1: T ) a.s. → ∇ θ log p θ ( y 1: T ) as N →∞ .

Proof. This is a direct application of Theorem 7.4.3 from Del Moral [25]. SMC methods provide strongly consistent estimates of expectations of test functions with respect to a normalized target distribution. That is, consider some test function h and let SMC's particle weights be denoted as w t i . We have that ∑ K i =1 w k t h ( x k 1: t ) a.s. → ∫ h ( x 1: t ) π t ( x 1: t ) as K →∞ where π t ( x 1: t ) is the normalized target distribution. NAS-X sets π t ( x 1: t ) ∝ γ t ( x 1: t ) = p θ ( x 1: t , y 1: t ) r ψ ∗ ( y t +1: T , x t ) , which by assumption is proportional to p θ ( x 1: t , y 1: t ) p θ ( y t +1: T | x t ) = p θ ( x 1: t , y 1: T ) . The desired result follows immediately.

Proposition 2. Unbiasedness of NAS-X's gradient estimates. Assume that proposal distribution q ϕ ( x t | x 1: t -1 , y 1: T ) is optimal so that q ϕ ( x t | x 1: t -1 , y 1: T ) = p ( x t | x 1: t -1 , y 1: T ) for t = 1 , . . . , T , and the twists r ψ ( y t +1: T , x t ) are optimal so that 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 . Let ˆ ∇ θ log p θ ( y 1: T ) be NAS-X's weighted particle approximation to the true gradient of the log marginal likelihood, ∇ θ log p θ ( y 1: T ) . Then, for any number of particles, E [ ˆ ∇ θ log p θ ( y 1: T )] = log p θ ( y 1: T ) .

Proof. We first provide a proof sketch then give a more detailed derivation, adapting the proof of a similar theorem in Lawson et al. [19]. We will prove that particles produced from running SMC with the smoothing targets and optimal proposal are exact samples from the posterior distribution of interest. The claim then follows immediately. Under the stated assumptions, both NAS-X and NASMC propose particles from the true posterior. However, the different intermediate target distributions will affect how these particles are distributed after reweighting. For NAS-X, the particles will have equal weight since they are reweighted using the smoothing targets. Thus, after reweighting, the particles are still samples from the true posterior. In contrast, in NASMC, the samples are reweighted by filtering targets and will not be distributed as samples from the true posterior.

We first consider NAS-X which uses the smoothing targets. We will show that the particles drawn at each timestep are sampled from the posterior distribution. This will follow from the fact that the particle weights at each timestep in the SMC sweep are equal; that is, w k t = 1 or w k t = p ( y 1: T ) for k = 1 , . . . , K , depending on whether resampling occurred. We proceed by induction on t , the timestep in the SMC sweep.

For t = 1 note that

This implies that the incremental weight α k 1 is

which does not depend on k . Because w k 0 ≜ 1 , we have that w k 1 = w k 0 α k 1 = p ( y 1: T ) for all k .

Note that, since the proposal is optimal, prior to SMC's reweighting step, the particles were distributed as follows x k 1 ∼ p ( x 1 | y 1: T ) . Since the incremental weights are equal, the distribution of particles is unchanged.

For the induction step, assume that w 1: K t -1 equals 1 or p ( y 1: T ) and that x k 1: t -1 ∼ p ( x 1: t -1 | y 1: T ) . The particles are distributed as follows, x k t ∼ p ( x t | x 1: t -1 , y 1: T ) . This implies that ( x k 1: t -1 , x k t ) ∼ p ( x t | x 1: t -1 , y 1: T ) p ( x 1: t -1 | y 1: T ) = p ( x 1: t | y 1: T ) .

We now show that the incremental particle weights α k t are equal using the following identities/assumptions

γ t ( x k 1: t ) = p ( x k 1: t , y 1: t ) p ( y t +1: T | x k t ) ,

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 .

and q t ( x k t ) = p ( x k t | x k 1: t -1 , y 1: T )

This shows that α k t is given by

for k = 1 , . . . , K .

There are two cases depending on the value of the weights at the previous timestep. If w 1: K t -1 = 1 , then w k t = w k t -1 α k t = 1 for all k . On the other hand, if w 1: K t -1 = p ( y 1: T ) then w k t = p ( y 1: T ) for all k . Therefore, even after reweighting, the particles are still drawn from the true posterior distribution. If resampling occurs, since the weights are equal in both cases, the distribution of the particles remains unchanged.

To conclude, we note that the incremental particle weights are, in general, not the same for NASMC. To see this, consider NASMC's incremental weights at timestep 1 .

After reweighting, the particles will be distributed according to the filtering distributions. The distribution of particles after reweighting is proportional to p ( x k 1 | y 1 ) . This is because the distribution of the reweighted particles is proportional to the incremental weights times the optimal proposal distribution. Therefore, the term in the denominator corresponding to the proposal cancels out with the proposal distribution term.

Interestingly, the particle weights will be the same at each iteration under a certain dependency structure for the model p ( x 1: t -1 | y 1: t ) = p ( x 1: t -1 | y 1: t -1 ) that was identified in Maddison et al. [17]. However, this dependency is not satisfied in general and therefore NASMC's gradients are not unbiased.

8 Derivations

8.1 Gradient of the Marginal Likelihood

We derive the gradients for the marginal likelihood. This identity is known as Fisher's identity.

The key steps were the log-derivative trick and Bayes rule.

8.2 Gradient of Inclusive KL Divergence

Below, we derive the gradient of the inclusive KL divergence for a generic Markovian model. In this derivation, we assume there are no shared parameters between the proposal and model.

We use the assumption that there are no shared parameters in the second equality.

8.3 Density Ratio Estimation via Classification

Here we briefly summarize density ratio estimation (DRE) via classification. For a full treatment, see Sugiyama et al. [30].

Let a ( x ) and b ( x ) be two distributions defined over the same space X , and consider a classifier g : X → R that accepts a specific x ∈ X and classifies it as either being sampled from a or b . We will train this classifier to predict whether a given x was sampled from a ( x ) or b ( x ) . The raw outputs (logits) of this classifier will approximate log( a ( x ) /b ( x )) up to a constant that does not depend on x .

To see this, define an expanded generative model where we first sample z ∈ { 0 , 1 } from a Bernoulli random variable with probability 0 < ρ < 1 , and then sample x from a ( x ) if z = 1 , and sample x

from b ( x ) if z = 0 . This defines the joint distribution

where p ( x z = 1) = a ( x ) and p ( x z = 0) = b ( x ) .

Let g : X → R be a function that accepts x ∈ X and produces the logit for Bernoulli distribution over z . This function will parameterize a classifier via the sigmoid function, meaning that the classifier's Bernoulli conditional distribution over z is defined as

where σ is the sigmoid function and σ -1 is its inverse, the logit function

The optimal function g ∗ will be selected by solving the maximum likelihood problem

The solution to this problem is the true p ( z | x ) . Because we have not restricted g , this solution can be obtained. Thus,

Thus, the optimal solution to the classification problem, g ∗ , is proportional to log( a ( x ) /b ( x )) up to a constant that does not depend on x . In practice we observe empirically that as long as a sufficiently flexible parametric family for g is selected, g ∗ will closely approximate the desired density ratio.

In the case of learning the ratio required for smoothing SMC,

Lawson et al. [19] instead learn the equivalent ratio

This in turn implies that

0

100

200

300

400

500

600

700

Proposal training steps (1000s)

10

-

1

10

0

10

1

10

2

log

p

(

y

1:

T

)

-L

128

NAS-X

NASMC

ELBO

IWAE

RWS

SIXO

FIVO

0

100

200

300

400

500

600

700

Proposal training steps (1000s)

10

-

2

10

-

1

10

0

Proposal relative error

0

10

20

30

40

50

Model timesteps

-

4

-

2

0

2

4

6

µ

value

As per the previous derivation, it suffices to train a classifier to distinguish between samples from the numerator and denominator of Eq. 35. To accomplish this, Lawson et al. [19] draw paired and unpaired samples from the model that are distributed marginally according to the desired densities. Specifically, consider drawing

and note that any x t , y t +1: T selected from the sample will be distributed marginally according to p θ ( x t , y t +1: T ) . Similarly, any ˜ x t , y t +1: T will be distributed marginally as p θ ( x t ) p θ ( y t +1: T ) . In this way, T -1 positive and negative training examples for the DRE classifier can be drawn using a single set of samples as in Eq. (36).

The twist training process is summarized in Algorithm 2.

9 LGSSM

Model Details We consider a one-dimensional linear Gaussian state space model with joint distribution

In our experiments we set the dynamics variance σ 2 x = 1 . 0 and the observation variance σ 2 y = 1 . 0 .

Proposal Parameterization For both NAS-X and NASMC, we use a mean-field Gaussian proposal factored over time

with parameters µ 1: T and σ 2 1: T corresponding to the means and variances at each timestep. In total, we learn 2 T proposal parameters.

Twist Parametrization We parameterize the twist as a quadratic function in x t whose coefficients are functions of the observations and time step and are learned via the density ratio estimation procedure described in [19]. We chose this form to match the analytic log density ratio for the model defined in Eq 10. Given that p ( x 1: T , y 1: T ) is a multivariate Gaussian, we know that p ( x t | y t +1: T ) and p ( x t ) are both marginally Gaussian. Let

Then,

where Z ( σ ) = 1 σ √ 2 π , so log Z ( σ ) = -log( σ √ 2 π ) .

Collecting terms gives:

So we'll define

We'll explicitly model log σ 2 1 , log σ 2 2 and µ 1 . Both log σ 2 1 and log σ 2 2 are only functions of t , not of y t +1: T , so those can be vectors of shape T initialized at 0 . µ 1 is a linear function of y t +1: T and t , so that can be parameterized by a set of T × T weights, initialized to 1 /T and T biases initialized to 0.

Training Details We use a batch size of 32 for the density ratio estimation step. Since we do not perform model learning, we do not repeatedly alternate between twist training and proposal training for NAS-X. Instead, we first train the twist for 3,000,000 iterations with a batch size of 32 using samples from the model. We then train the proposal for 750 , 000 iterations. For the twist, we used Adam with a learning rate schedule that starts with a constant learning rate of 1 e -3 , decays the learning by 0 . 3 and 0 . 33 at 100 , 000 and 300 , 000 iterations. For the proposal, we used Adam with a constant learning rate of 1 e -3 . For NASMC, we only train the proposal.

Evaluation In the right panel of Figure 1, we compare the bound gaps of NAS-X and NASMC averaged across 20 different samples from the generative model. To obtain the bound gap for NAS-X, we run SMC 16 times with 128 particles and with the learned proposal and twists. We then record the average log marginal likelihood. For NASMC, we run SMC with the current learned proposal (without any twists).

10 rSLDS

Model details The generative model is as follows. At each time t , there is a discrete latent state z t ∈ { 1 , . . . , 4 } as well as a two-dimensional continuous latent state x t ∈ R 2 . The discrete state

0

100

200

300

400

500

Proposal training steps (100s)

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

1

.

2

1

.

4

z

t

absolute error

NAS-X,

σ

2

O

= 0

.

1

NASMC,

σ

2

O

= 0

.

1

0

100

200

300

400

500

Proposal training steps (100s)

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

x

t

absolute error

NAS-X,

σ

2

O

= 0

.

1

NASMC,

σ

2

O

= 0

.

1

transition probabilities are given by

Here R i and r i are weights for the discrete state z i .

These discrete latent states dictates two-dimensional latent state x t ∈ R 2 which evolves according to linear Gaussian dynamics.

Here A k , Q k ∈ R 2 x 2 and b k ∈ R 2 . Importantly, from Equations 40 and 39 we see that the dynamics of the continuous latent states and discrete latents are coupled. The discrete latent states index into specific linear dynamics and the discrete transition probabilities depend on the continuous latent state.

The observations y t ∈ R 10 are linear projections of the continuous latent state x t with some additive Gaussian noise.

Here C, S ∈ R 10 x 10 and d ∈ R 10 .

Proposal Parameterization We use a mean-field proposal distribution factorized over the discrete and continuous latent variables (i.e. q ( z 1: T , x 1: T ) = q ( z 1: T ) q ( x 1: T ) ). For the continuous states, q ( x 1: T ) is a Gaussian factorized over time with parameters µ 1: T and σ 2 1: T . For the discrete states, q ( z 1: T ) is a Categorical distribution over K categories factorized over time with parameters p 1: K 1: T . In total, we learn 2 T + TK proposal parameters.

Twist Parameterization We parameterize the twists using a recurrent neural network (RNN) that is trained using density ratio estimation. To produce the twist values at each timestep, we first run a RNN backwards over the observations y 1: T to produce a sequence of encodings e 1: T -1 . We then concatenate the encodings of x t and z t into a single vector and pass that vector into an MLP which outputs the twist values at each timestep. The RNN has one layer with 128 hidden units. The MLP has 131 hidden units and ReLU activations.

Model Parameter Evaluation We closely follow the parameter initialization strategy employed by Linderman et al. [41]. First, we use PCA to obtain a set of continuous latent states and initialize the matrices C and d . We then fit an autoregressive HMM to the estimated continuous latent states in order to initialize the dynamics matrices { A k , b k } . Importantly, we do not initialize the proposal with the continuous latent states described above.

Training Details We use a batch size of 32 for the density ratio estimation step. We alternate between 100 steps of twist training and 100 steps of proposal training for a total of 50,000 training steps in total. We used Adam and considered a grid search over the model, proposal, and twist learning rates. In particular, we considered learning rates of 1 e -4 , 1 e -3 , 1 e -2 for the model, proposal, and twist.

Bootstrap Bound Evaluation To obtain the log marginal likelihood bounds and standard deviations in Table 5.2, we ran a bootstrapped particle filter (BPF) with the learned model parameters for all three methods (NAS-X, NASMC, Laplace EM) using 1024 particles. We repeat this across 30 random seeds. Initialization of the latent states was important for a fair comparison. To initialize the latent

states, for NAS-X and NASMC, we simply sampled from the learned proposal at time t = 0 . To initialize the latent state for Laplace EM, we sampled from a Gaussian distribution with the learned dynamics variance at t = 0 .

Inference Comparison In the top panel of Figure 6, we compare NAS-X and NASMC on inference in the SLDS model. We report (average) posterior parameter recovery for the continuous and discrete latent states across 5 random samples from the generative model. NAS-X systematically recovers better estimates of both the discrete and continuous latent states.

11 Inference in Squid Giant Axon Model

11.1 HHModel Definition

For the inference experiments (Section 5.3.1) we used a probabilistic version of the squid giant axon model [9, 43]. Our experimental setup was constructed to broadly match [19], and used a single-compartment model with dynamics defined by

where C m is the membrane capacitance; v is the potential difference across the membrane; I ext is the external current; ¯ g Na , ¯ g K , and ¯ g leak are the maximum conductances for sodium, potassium, and leak channels; E Na , E K , and E leak are the reversal potentials for the sodium, potassium, and leak channels; m and h are subunit states for the sodium channels and n is the subunit state for the potassium channels. The functions α and β that define the dynamics for n , m , and h are defined as

This system of ordinary differential equations defines a nonlinear dynamical system with a fourdimensional state space: the instantaneous membrane potential v and the ion gate subunit states n , m , and h .

As in [19], we use a probabilistic version of the original HH model that adds zero-mean Gaussian noise to both the membrane voltage v and the 'unconstrained' subunit states. The observations are produced by adding Gaussian noise with variance σ 2 y to the membrane potential v .

Specifically, let x t be the state vector of the system at time t containing ( v t , m t , h t , n t ) , and let φ dt ( x ) be a function that integrates the system of ODEs defined above for a step of length dt . Then the probabilistic HH model can be written as

where the 4-D state distributions p ( x 1 ) and p ( x t | φ dt ( x t -1 )) are defined as

4

8

16

32

64

128

256

-148

-197

-262

-347

-461

-612

-812

Log-likelihood lower bound

Inference in Hodgkin-Huxley

50ms trace, 1 obs/ms,

σ

2

y

= 20

FIVO-BS

FIVO

NASMC

SIXO

NAS-X

4

8

16

32

64

128

256

-167

-174

-181

-188

-195

-203

-211

Log-likelihood lower bound

SIXO vs NAS-X

50ms trace, 1 obs/ms,

σ

2

y

= 20

SIXO

NAS-X

Num Particles

Num Particles

In words, we add Gaussian noise to the voltage ( x t, 1 ) and logit-normal noise to the gate states n, m, and h . The logit-normal is defined as the distribution of a random variable whose logit has a Gaussian distribution, or equivalently it is a Gaussian transformed by the sigmoid function and renormalized. We chose the logit-normal because its values are bounded between 0 and 1, which is necessary for the gate states.

Problem Setting For the inference experiments we sampled 10,000 noisy voltage traces from a fixed model and used each method to train proposals (and possibly twists) to compute the marginal likelihood assigned to the data under the true model.

As in [19], we sampled trajectories of length 50 milliseconds, with a single noisy voltage observation every millisecond. The stability of our ODE integrator allowed us to integrate at dt = 0 . 1 ms, meaning that there were 10 latent states per observation.

Proposal and Twist Details Each proposal was parameterized using the combination of a bidirectional recurrent neural network (RNN) that conditioned on all observed noisy voltages as well as a dense network that conditioned on the RNN hidden state and the previous latent state x t -1 [45, 49]. The twists for SIXO and NAS-X were parameterized using an RNN run in reverse over the observations combined with a dense network that conditioned on the reverse RNN hidden state and the latent being 'twisted', x t . Both the proposal and twists were learned in an amortized manner, i.e. they were shared across all trajectories. All RNNs had a single hidden layer of size 64, as did the dense networks. All models were fit with ADAM [50] with proposal learning rate of 10 -4 and twist learning rate of 10 -3 .

A crucial aspect of fitting the proposals was defining them in terms of a 'residual' from the prior, a technique known as Res q [7]. In our setting, we defined the true proposal density as proportional to the product of a unit-variance Gaussian centered at φ ( x t ) and a Gaussian with parameters output from the RNN proposal.

11.2 Experimental Results

In Figure 7 we plot the performance of proposals and twists trained with 4 particles and evaluated across a range of particle numbers. All methods except FIVO perform roughly the same when evaluated with 256 particles, but with lower numbers of evaluation particles the smoothing methods emerge as more particle-efficient than the filtering methods. To achieve NAS-X's inference perfor-

4

8

16

32

64

128

256

-169

-179

-189

-200

Log-likelihood lower bound

Train num particles = 4

4

8

16

32

64

128

256

Train num particles = 8

4

8

16

32

64

128

256

Train num particles = 16

IWAE

FIVO

NASMC

SIXO

NAS-X

0

5

10

15

20

25

30

35

40

45

50

80

60

40

20

0

20

40

Voltage (mV)

True Voltage

Inferred Voltage

Resampling Event

NASMC

-

50

0

50

Voltage (mV)

SIXO

True Voltage

Inferred Voltage

Resampling Event

0

5

10

15

20

25

30

35

40

45

50

Time (ms)

-

50

0

50

Voltage (mV)

NAS-X

mance with 4 particles, NASMC would need 256 particles, a 64-times increase. NAS-X is also more particle-efficient than SIXO, achieving on average a 2x particle efficiency improvement. We show the effect of changing the number of training particles in Figure 8.

The FIVO method with a parametric proposal drastically underperformed all smoothing methods as well as NASMC, indicating that the combination of filtering SMC and the exclusive KL divergence leads to problems optimizing the proposal parameters. To compensate, we also evaluated the performance of 'FIVO-BS", a filtering method that uses a bootstrap proposal. This method is identical to a bootstrap particle filter, i.e. it proposes from the model and has no trainable parameters. FIVO-BS far outperforms standard FIVO, and is only marginally worse than NASMC in this setting.

In Figure 9 we investigate these results qualitatively by examining the inferred voltage traces of each method. We see that NASMC struggles to produce accurate spike timings and generates many spurious spikes, likely because it is unable to incorporate future information into its proposal or resampling method. SIXO performs better than NASMC, accurately inferring the timing of most spikes but resampling at a high rate. High numbers of resampling events can lead to particle degeneracy and poor inferences. NAS-X is able to correctly infer the voltage across the whole trace with no suprious or mistimed spikes. Furthermore NAS-X rarely resamples, indicating it has learned a high-quality proposal that does not generate low-quality particles that must be resampled away. These qualitative results seem to support the quantitative results in Figure 7 - SIXO's high resampling rate and NASMC's filtering approach lead to lower bound values.

L

-

-

12 Model Learning in Mouse Pyramidal Neuron Model

12.1 Model Definition

For the model learning experiments in Section 5.3.2 we used a generalization of the Hodgkin-Huxley model developed for modeling mouse visual cortex neurons by the Allen Institute for Brain Science [48, 52]. Specifically we used the perisomatic model with ID 482657528 developed to model cell 480169178. The model is detailed in the whitepaper [52] and the accompanying code, but we reproduce the details here to ensure our work is self-contained.

Similar to the squid giant axon model, the mouse visual cortex model is composed of ion channels that affect the current flowing in and out of the cell. Let I be the set of ions { Na + , Ca 2+ , K + } . Each ion has associated with it

A set of channels that transport that ion, denoted C i for i ∈ I .

A reversal potential, E i .

An instantaneous current density, I i , which is computed by summing the current density flowing through each channel that transports that ion.

Correspondingly, let C be the set of all ion channels so that C = ⋃ i ∈I C i . Each c ∈ C has associated with it

A maximum conductance density, g c .

A set of subunit states, referred to collectively as the vector λ c . Let λ c ∈ [0 , 1] d c , i.e. λ c is a d c -dimensional vector of values in the interval [0 , 1] .

A function g c that combines the gate values to produce a number in [0 , 1] that weights the maximum conductance density, g c · g c ( λ c ) .

Functions A c ( · ) and b c ( · ) which compute the matrix and vector used in the ODE describing λ c dynamics. A c and b c are functions of both the current membrane voltage v and calcium concentration inside the cell [Ca 2+ ] i . If the number of subunits (i.e. the dimensionality of λ c ) is d c , then the output of A c ( v, [Ca 2+ ] i ) is a d c × d c diagonal matrix and the output of b c ( v, [Ca 2+ ] i ) is a d c -dimensional vector.

With this notation we can write the system of ODEs

Most symbols are as described earlier, SA is the membrane surface area of the neuron, [Ca 2+ ] i is the calcium concentration inside the cell, [Ca 2+ ] min is the minimum interior calcium concentration with a value of 1 nanomolar, τ is the rate of removal of calcium with a value of 80 milliseconds, and k and is a constant with value

where 10000 is a dimensional constant, γ is the percent of unbuffered free calcium, F is Faraday's constant, and depth is the depth of the calcium buffer with a value of 0.1 microns.

Because the concentration of calcium changes over time, this model calculates the reversal potential for calcium E Ca 2+ using the Nernst equation

where G is the gas constant, T is the temperature in Kelvin (308.15 ◦ ), F is Faraday's constant, and [Ca 2+ ] o is the extracellular calcium ion concentration which was set to 2 millimolar.

Probabilistic Model The probabilistic version of the deterministic ODEs was constructed similarly to the probabilistic squid giant axon model - Gaussian noise was added to the voltage and unconstrained gate states. One difference is that the system state now includes [Ca 2+ ] i which is constrained to be greater than 0. To noise [Ca 2+ ] i we added Gaussian noise in the log space, analagous to the logit-space noise for the gate states.

Model Size The 38 learnable parameters of the model include:

Conductances g for all ion channels (10 parameters).

Reversal potentials of sodium, potassium, and the non-specific cation: E K + , E Na + , and E NSC + .

The membrane surface area and specific capacitance.

Leak channel reversal potential and max conductance density.

The calcium decay rate and free calcium percent.

Gaussian noise variances for the voltage v and interior calcium concentration [Ca 2+ ] i .

Gaussian noise variances for all subunit states (16 parameters).

Observation noise variance.

The 18-dimensional state includes:

Voltage v

Interior calcium concentration [Ca 2+ ] i

All subunit states (16 dimensions)

12.2 Channel Definitions

In this section we provide a list of all ion channels used in the model. In the following equations we often use the function exprel which is defined as

A numerically stable implementation of this function was critical to training our models.

Additionally, many of the channel equations below contain a 'temperature correction' q t that adjusts for the fact that the original experiments and Allen Institute experiments were not done at the same temperature. In those equations, T is the temperature in Celsius which was 35 ◦ .

12.2.1 Transient Na +

From Colbert and Pan [53].

12.2.2 Persistent Na +

From Magistretti and Alonso [54].

12.2.3 Hyperpolarization-activated cation conductance

From Kole et al. [55]. This channel uses a 'nonspecific cation current' meaning it can transport any cation. In practice, this is modeled by giving it its own special ion NSC + with resting potential

E NSC + .

12.2.4 High-voltage-activated Ca 2+ conductance

From Reuveni et al. [56]

12.2.5 Low-voltage-activated Ca 2+ conductance

From Avery and Johnston [57], Randall and Tsien [58].

12.2.6 M-type (Kv7) K + conductance

12.2.7 Kv3-like K + conductance

12.2.8 Fast inactivating (transient, Kv4-like) K + conductance

From Korngreen and Sakmann [60].

12.2.9 Slow inactivating (persistent) K + conductance

From Korngreen and Sakmann [60].

12.2.10 SK-type calcium-activated K + conductance

From Köhler et al. [61]. Note this is the only calcium-gated ion channel in the model.

12.3 Training Details

Dataset The dataset used to fit the model was a subset of the stimulus/response pairs available from the Allen Institute. First, all stimuli and responses were downloaded for cell 480169178. Then, sections of length 200 milliseconds were extracted from a subset of the stimuli types. The stimuli types and sections were chosen so that the neuron was at rest and unstimulated at the beginning of the trace. We list the exclusion criteria below.

Any 'Hold' stimuli: Excluded because these traces were collected under voltage clamp conditions which we did not model.

Test: Excluded because the stimulus is 0 mV for the entire trace.

Ramp/Ramp to Rheobase: Excluded because the cell is only at rest at the very beginning of the trace.

Short Square: 250 ms to 450 ms.

Short Square - Triple: 1250 ms to 1450 ms.

Noise 1 and Noise 2: 1250 ms to 1450 ms, 9250 ms to 9450 ms, 17250 ms to 17450 ms.

Long Square: 250 ms to 450 ms.

Square - 0.5ms Subthreshold: The entire trace.

Square - 2s Suprathreshold: 250 ms to 450 ms.

All others: Excluded.

For cell 480169178, the criteria above selected 95 stimulus/response pairs of 200 milliseconds each. Each trace pair was then downsampled to 1 ms (from the original 0.005 ms per step) and corrupted with mean-zero Gaussian noise of variance 20 mV 2 to simulate voltage imaging conditions. Finally, the 95 traces were randomly split into 72 training traces and 23 test traces.

Proposal and Twist The proposal and twist hyperparameters were broadly similar to the squid axon experiments, with the proposal being parameterized by a bidirectional RNN with a single hidden layer of size 64 and an MLP with a single hidden layer of size 64. The RNN was conditioned on the observed response and stimulus voltages at each timestep, and the MLP accepted the RNN hidden state, the previous latent state, and a transformer positional encoding of the number of steps since the last voltage response observation. The twist was similarly parameterized using an RNN run in reverse across the stimulus and response, combined with an MLP that accepted the RNN hidden state, the latent state being evaluated, and a transformer positional encoding of the number of steps elapsed since the last voltage response observation. The positional encodings were used to inform the twist and proposal of the number of steps elapsed since the last observation because the model was integrated with a stepsize of 0.1ms while observations were received once every millisecond.

Hyperparameter Sweeps To evaluate the methods we swept across the parameters

Initial observation variance: e 2 , e 3 , e 5

Initial voltage dynamics variance: e, e 2 , e 3

Bias added to scales produced by the proposal: e 2 , e 5

We also evaluated the models across three different data noise variances (20, 10, and 5) but the results were similar for all values, so we reported only the results for variance 20. This amounted to 3 · 3 · 3 · 2 different hyperparameter settings, and 5 seeds were run for each setting yielding a total of 270 runs.

When computing final performance, a hyperparameter setting was only evaluated if it had at least 3 runs that achieved 250,000 steps without NaN-ing out. For each hyperparameter setting selected for evaluation, all successful seeds were evaluated using early stopping on the train 4-particle log likelihood lower bound.

13 Strang Splitting for Hodgkin-Huxley Models

Because the Hodgkin-Huxley model is a stiff ODE, integrating it can be a challenge, especially at large step sizes. The traditional solution is to use an implicit integration scheme with varying step size, allowing the algorithm to take large steps when the voltage is not spiking. However, because our model adds noise to the ODE state at each timestep adaptive step-size methods are not viable as the different stepsizes would change the noise distribution.

Instead, we sought an explicit, fixed step-size method that could be stably integrated at relatively large stepsizes. Inspired by Chen et al. [47], we developed a splitting approach that exploits the conditional linearity of the system. The system of ODEs describing the model can be split into two subsystems

0

500

1000

1500

2000

2500

3000

Twist training steps

10

-

1

10

0

10

1

Twist parameter error

0

500

1000

1500

2000

2500

3000

Twist training steps

0

.

65

0

.

70

0

.

75

0

.

80

Twist accuracy

of linear first-order ODEs when conditioned on the state of the other subsystem. Specifically, the dynamics of the channel subunit states { λ c | c ∈ C} is a system of linear first-order ODEs when conditioned on the voltage v and interior calcium concentration [Ca 2+ ] i . Similarly, the dynamics for v and [Ca 2+ ] i is a system of linear first-order ODEs when conditioned on the subunit states.

Because the conditional dynamics of each subsystem are linear first-order ODEs, an exact solution to each subsystem is possible under the assumption that the states being conditioned on are constant for the duration of the step. Our integration approach uses these exact updates in an alternating fashion, first performing an exact update to the voltage and interior calcium concentration while holding the subunit states constant, and then performing an exact update to the subunit states while holding the voltage and interior calcium concentration constant. For details on Strang and other splitting methods applied to Hodgkin-Huxley type ODEs, see [47].

14 Robustness of Twist Learning

NAS-X uses SIXO's twist learning framework to approximate the smoothing distributions. The twist learning approaches involves density ratio estimation. In brief, the density ratio estimation procedure involves training a classifier to distinguish between samples from p θ ( x t | y t +1: T ) and p θ ( x t ) . These samples can be obtained from the generative model. For details see Section 2.2. In principle, incorporating twists complicates the overall learning problem and traditional methods for twist learning can indeed be challenging. However, in practice, twist learning using the SIXO framework is robust and easy. In Figure 10, we present twist parameter recovery and classification accuracy for the Gaussian SSM experiments (Section 5.1); in this setting, the optimal twists have a known parametric form. The optimal twist parameters are recovered quickly, the classification accuracy is high, and training is stable. This suggests that, with an appropriate twist parameterization, twist learning via density ratio estimation is tractable and straightforward.

15 Computational Complexity and Wall-clock Time

Theoretically, all multi-particle methods considered (NAS-X, SIXO, FIVO, NASMC, RWS, IWAE) have O ( KT ) time complexity, where K is the number of particles and T is the number of time steps. Once concern is that the resampling operation in SMC could require super-linear time in the number of particles, but drawing K samples from a K -category discrete distribution can be accomplished in O ( K ) time using the alias method [62, 63]. Additionally, for NAS-X, evaluating the twists is amortized across timesteps as in Lawson et al. [19], giving time linear in T .

NAS-X and SIXO have similar wall-clock speeds but are slower than FIVO and NASMC, primarily because of twist training, see Table 2. Even if FIVO and NASMC were run with more particles to equalize wall-clock times, they would still far underperform NAS-X in log marginal likelihood lower bounds, see Figure 7.

Specifically, SIXO and NAS-X take 3.5x longer per step than NASMC and FIVO and 2.5x longer per step than RWS and IWAE. However, Figure 7 shows that FIVO, NASMC, IWAE, and RWS cannot match NAS-X's performance even with 64 times more computation (4 vs. 256 particles). SIXO only matches NAS-X's performance with 4x as many particles (4 vs. 16 particles). Therefore, NAS-X uses computational resources much more effectively than other methods.

±

±

±

0

100

200

300

Proposal training steps (1000s)

10

-

3

10

-

1

10

1

10

3

Gradient variances

0

100

200

300

Proposal training steps (1000s)

10

0

10

1

Gradient bias

0

100

200

300

Proposal training steps (1000s)

-

10

4

-

10

3

LML lower bound

NAS-X

NASMC

SIXO

RWS

FIVO

IWAE

16 Empirical analysis of bias and variance of gradients

In Figure 11, we analyze the gradient variance and bias for the Hodgkin-Huxley experiments, supplementing our theoretical analyses in Section 3.1. Figure 11 (left) shows NAS-X attains lower variance gradient estimates than IWAE, FIVO, and SIXO with comparable variance to RWS. We also studied the bias (middle) by approximating the true gradient by running a bootstrap particle filter with 256 particles using the best proposal from the inference experiments. NAS-X's gradients are lower bias than all methods but FIVO, but FIVO's gradients are also the highest variance. We hypothesize that FIVO's gradients appear less biased because its parameters are pushed towards degenerate values where gradient estimation is 'easier'. We illustrate this in the right panel, where we plot log-marginal likelihood bounds.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes NAS-X, a novel approach to estimate the posterior expectations in Reweighted Wake Sleep (RWS) architectures. Instead of a traditional estimation using self-normalized importance sampling (SNIS), and similar to Neural Adaptive Sequential Monte-Carlo (NASMC), the proposed method employs a Sequential Monte-Carlo (SMC) approach to estimate the necessary expectations. In contrast to NASMC, NAS-X uses smoothing distributions instead of filtering distributions as targets, which improves particle efficiency and reduces the variance of the estimates. This requires the estimation of twist sequences, which the paper does using the density-ratio approach SIXO.\n\nSTRENGTHS:\n- Efficient particle-based estimation of posteriors in sequence models is an important topic and has a long history in the machine learning community.\n\n- The paper is well-written and relatively easy to follow. I was particularly impressed by the paper’s natural motivation: starting with a concise recap of RWS, the paper clearly explains the need for smoothing SMC and its estimation via twists.\n\n- The main technical contribution (Eq.(8) and Eq.(9)) is SIXO-based smoothing SMC for posterior inference in RWS architectures. From a technical point of view this contribution is relatively simple, because all necessary pieces (RWS, NASMC, and SIXO) were already available, but the insight that the RWS updates are compatible with SIXO estimation is noteworthy and must be appreciated.\n\n- The experiments validate the proposed approach in settings with known ground truth and discrete latent variables, as well as challenging inference in Hodgkin-Huxley models. While the experiments with Gaussian-linear models and Switching Linear Dynamical Systems are important sanity checks and confirm the advantages of smoothing SMC over filtering SMC, I particularly enjoyed the real-world experiments on voltage dynamics in neural membranes. They not only demonstrate that NAS-X is more particle-efficient than its competitors but also that the proposed model could be useful beyond synthetic environments.\n\nWEAKNESSES:\n- My main concern with this paper is its potential impact: SMC-based inference in RWS architectures is already a niche topic and there is nothing fundamentally wrong with filtering SMC. Smoothing SMC introduces the additional complication of twist estimation, but even that has already been addressed with SIXO. My worry is that the remaining contribution likely has a relatively small audience.\n\n- Another point of concern are NAS-X’s multi-level approximations: SMC is already approximate in nature, but now, in addition to the quality and efficiency of the particles, accurate estimation of the twist sequence becomes a separate challenge. There is a trade-off between the additional information contained in future observations and a decrease in robustness due to a more complex learning problem, and I would have liked to see experiments that directly evaluate the quality of the learned twists.\n\n- The paper claims that smoothing SMC leads to lower-variance estimates compared to filtering SMC (l.32f, l.92f), but the experiments do not investigate this claim directly. Some indirect evidence is provided in Figure 1 and Figure 3, but I would have appreciated additional insights.\n\n- The description of the experimental setup could be better structured. I did find most of the information I was looking for, but it required constant jumping between different sections of the paper and between the paper and the supplemental material. The information contained in the supplemental material is also not referenced well enough in the main paper and I would encourage the authors to include more pointers to the supplemental material.\n\nQUESTIONS:\n- The paper generally does a good job giving credit to prior work, but the relationship between SIXO and NAS-X is not clear enough. If SIXO already performs smoothing SMC, what is the main difference between the two? Is NAS-X more than SIXO applied to RWS? What architecture is used for SIXO *without* NAS-X?\n\n- The effect of the number of training particles on the reported performance metrics remains a bit of a mystery. I would appreciate a reproduction of Figure 3(a) with 32 and 128 training particles.\n\n- The meaning of “inclusive KL”, in contrast to “exclusive KL”, is not clear enough. I found the definitions in the literature, but they should be mentioned in the paper. \n\nTypos: l.48.5 (“,”), l.52 (“goals”).\n\nLIMITATIONS:\n- The paper does not discuss the limitations of the proposed model. Completely absent in this work is a runtime analysis. Does NAS-X's higher particle efficiency translate to faster (absolute wall time) inference compared to NASMC/SIXO?\n\n- The paper does not discuss the potential negative societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper tackles the problem of learning sequential latent variable models using a new method called NAS-X. It combines two previous research: 1. The smoothing with twisting that’s applied in variational inference (SIXO in particular). It approximates smoothing distribution as targets for proposal learning, instead of using filtering distribution, to avoid sample degeneracy; 2. The Reweighted wake-sleep method, which can handle discrete latent variables. They used experiments to illustrate that NAS-X can learn proposals that match the true posterior marginals, and can be applied to discrete variables and handle complex tasks.\n\nSTRENGTHS:\nThe paper is well written. The authors have done extensive experiments to show the effectiveness of their work. The proposed method is somewhat novel, as it combines two existing ideas and outperforms them.\n\nWEAKNESSES:\nThe paper lacks discussions of future work and limitations of the proposed method.\n\nQUESTIONS:\n1. Section 2.1: I find the structure there a bit unclear. you first described a two-step coordinate ascent method in the first paragraph, but then it is unclear whether your formular afterwards are talking about the first or the second step, it might make sense to be clearer there.\n\n2. Section 2.2 SMC description: the “repeats three steps” paragraph doesn’t make it clear that it is repeatedly sampling the next time point (if I understand correctly). Maybe it is better to just be clear that each step is for a new t.\n\n3. Section 5.2: it is unclear to me that the qualitative comparison in figure 2 brings that much insight, but maybe it at least shows how the data is generated?\n\n4. Section 5.3.1: for figure 3(a) you showed that is It more particle efficient, I wonder if it would be helpful to also comment on the computation efficiency there?\n\n5. Section 5.3.1 last paragraph: The discussion about RWS vs VI seems intriguing but I wish it can be discussed a bit more clearly and thoroughly and provide more context. Also if we believe this to be an important enough point, maybe it is worth highlighting this point in the intro or conclusion?\n\n6. Figure 4: Why would we expect NAS-X to outperform in certain metrics but underperform in other metrics? \n\n7. Section 5.3.2: I wonder why NASMC gets dropped from the results in this section?\n\nLIMITATIONS:\nI wonder if the authors could add more discussions on the limitations/future work of their proposed method?", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAn algorithm based on reweighted wake-sleep (RWS) is applied to sequential latent variable models $p_\\theta(x_{1:T}, y_{1:T})$. Model parameters are fit by maximizing the evidence $p_\\theta(y_{1:T})$ in $\\theta$, and posterior inference is performed by minimizing the forward KL divergence $KL(p_\\theta(x_{1:T} \\mid y_{1:T}) \\mid \\mid q_\\phi(x_{1:T} \\mid \\mid y_{1:T}))$ in variational parameters $\\phi$. Gradients for each of these tasks are estimated by smoothing SMC, which can be viewed as a lower-variance alternative to self-normalized importance sampling (SNIS) in the sequential setting. This submission extends the filtering SMC approach of previous work. However, the smoothing distributions are not available and must be approximated by learning twists, which is accomplished by training a discriminator. This results in an estimate of an appropriate density ratio that allows for twisting to be used.\n\nSTRENGTHS:\n* The exposition is clear the and potential advantages with respect to forward KL minimization and smoothing are intuitive.\n* The method is general and and appears to be applicable in any learning/inference setting with sequential latent variable models. The only significant design choice for the user appears to be the design of the variational distribution, and the additional computational overhead from the training of the classifier for estimating the twists is not prohibitive.\n* The method is novel to the best of my knowledge, and combining a forward KL minimization with SIXO seems like a good idea. The forward KL objective makes it easy to handle discrete latent variables, as the second experiment shows, whereas other methods (see below) could not do so as easily.\n* The authors show the applicability of the method to three experiments of increasing complexity, with an emphasis on dynamical systems.\n\nWEAKNESSES:\nThe biggest weaknesses of the paper are:\n* Lack of comparison to existing methods. In addition to NASMC, FIVO (Maddison et al., 2017) and VSMC (Naesseth et al., 2018), and AESMC (Le et al.,  2018) are closely related methods that instead target the reverse KL divergence. While these are cited in related work, they should be compared to when possible (perhaps not in discrete latent variable models, though) as they are natural competitors for these sequential latent variable models. \n* Lack of discussion of, and comparison to, other SMC variants. The main intuition provided in the paper for using smoothing SMC is that both SNIS and filtering SMC can result in either high-variance estimates and/or particle degeneracy. Metrics (e.g. effective sample size) and figures for validating this intuition would be helpful, as this intuition is the central motivation for the use of smoothing SMC.\n\nAblation studies incorporating these competitors/alternatives would make more clear the importance of the contribution of this work. At present, it is difficult to tell whether both the smoothing approach or the forward KL formulation contribute to the success of the proposed method, or whether in some cases just one of these does.\n\nQUESTIONS:\n* Was vanilla RWS compared to? Naive implementations wouldn’t use sequential structure, and would rely only on SNIS, so I’d expect all of NASMC, SIXO, and NAS-X to beat RWS easily. Nevertheless, it would be a nice baseline.\n* Can there be more discussion on the classifier being used to approximate the density ratio? Some precise results from the GAN literature or the literature on likelihood free inference by ratio estimation (Thomas et al., 2016) stating formally that the optimal classification rule yields the likelihood ratio in some form would make this part of the exposition more clear; at least some citations should be added. \n* To what extent are the advantages of the proposed method from 1) use of the forward KL divergence and 2) use of smoothing SMC? \n Comparison with FIVO, VSMC, AESMC along with naive RWS or ELBO could really help make this clear.\n* Is the $\\chi$ in the title on purpose? Or should there be an X as in the rest of the body? It should be consistent.\n* On Open Review, the “smothing” in the title should corrected if possible.\n\nLIMITATIONS:\nThe authors have made clear the class of probabilistic models that this method is designed for.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: SUMMARY:
The authors propose to use SIXO particle approximation to calculate the expectations in reweighted wake-sleep algorithm for state space models.

STRENGTHS:
The paper is written in a clear way and all required background is well-explained (but description of  SIXO and twists is too short).

The idea of replacing the variational expectations with best possible practical approximation of the posterior is interesting.

The empirical findings are well-presented.

WEAKNESSES:
I would lengthen the description of SIXO and twists.

More empirical evidence and discussion why NAS-X actually warrants better performance (as it introduces biases) would strengthen the paper.

Some theoretical analysis would strengthen the paper.

QUESTIONS:
What do you think are main limitations of NAS-X?

When the algorithm is expected to perform poorly (or be outperformed by Laplace EM)?

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07