NAS-X: Neural Adaptive Smoothing via Twisting
Authors: Dieterich Lawson
Michael Y. Li
Scott Linderman
Venue: NeurIPS 2023 poster (2023)
Keywords: sequence models
probabilistic inference
reweighted wake-sleep
sequential monte carlo
smoothing
mechanistic models
TL;DR: We introduce a method for fitting sequential latent variable models that combines the benefits of reweighted wake-sleep and smoothing sequential Monte Carlo.

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-X: Neural Adaptive Smoothing via Twisting
Dieterich Lawson *
Michael Y. Li"
Google Research Meterrich@@googee.com
Stanford University
mi whwehhallllae ciiceelyl@sstanford.cco
Scott W. Linderman
Stanford University
scott www..iinnermnnttaaooroooeoo edu
Abstract
Sequential latent variable models (SLVMs) are essential tools in statisties and machine learning, with applications ranging from healthcare to neuroscience. As their flexibility increases, analytic inference and model learning can become chal- lenging, 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 X to provide low-bias and low-variance gradient estimates, and both discrete and continuous latent variable models. We illustrate the thcoretical advantages of NAS-X over previous methods and explore these advantages em- pirically 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-bbsed methods in inference and model learning, achieving lower parameter error and tighter likelihood bounds.
Introduction
Sequential latent variable models (SLVMs) are 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 weighted particle approximation to the true posterior and an unbiased estimator of the marginal likelihood. For most SLVMs. SMC is 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 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].
Equal contribution. WWrk performed while at Stanford University.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Variational inference (VI) methods for SLVMs optimize the model and proposal parameters by ascending 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 compelling advantage over VI methods that typically resort to high-vane variance score function estimates for discrete latent variables.
In this work, combine recent advances in smoothing variational SMC with the benefits of reweighted wake-sleep. To this end, we introduce neural adaptive smoothing via twisting (NAS- X), 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 X, 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 nonlinear sequential latent variable models with Markovian structure; i.e. models that factor
po (I.T, pe (po)) x1) II po(x; xt- (y: xt),
with latent variables X1:7 E XT= observations y1.r E yT. and global parameters O O. By nonlinear, we mean latent variable models where the parameters of the conditional distribu- tions (x x and py, xx) depend nonlinearly x;-1 and x,, respectively.
Estimating the marginal likelihood i:TT) and posterior po(x yur) for this model class is difficult because it requires computing an intractable integral over the latents,
( = / po(yI.T X1:T) dx
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 method for maximum marginal likelihood in LVMs that estimates the gradients of the marginal likelihood using eerf-normllized importance sampling. This is motivated by Fisher's identity, which allows us write the gradients of the marginal likelihood as posterior expectation,
Vo log Po (y = Ep. ir brarr)
V
(2)
as proved in Appendix 8.1. The term inside the expectation is computable with modern automatic differentiation tools [21, 22], but the posterior po (xir) Yı.r) is unavailable. Thus, SNIS used to form biased but consistent Monte Carlo estimate of Eq. (2) [23]. Specifically, SNIS draws N IID

---

samples from proposal distribution, (x1:T V1:T) and weights them to form the estimator
Va log po (x VIrT),
i XiT (o(er yirr),
w
Pe (x (2)
1: MI.T)
go 1:7 yI.T)
where w(() normalized weights, ie. E w(()
The variance of this estimator is reduced as le approaches the posterior [23], so RWS also updates le by minimizing the inclusive dullback-Leible (KL) divergence from the posterior to the proposal. Crucially, the gradient for this step can also be written as the posterior expectation
(pKL(pg (x1:T YI.T) 4o (x1TT Yi-)) po lyarr) [V log qo (xL:T yirr),, (4)
as derived in Appendix 8.2. This allows RWS estimate Eq. (4) using SNIS with the same set of samples and weights as Eq. (3),
N
-E Epo(xur(yier) V log (X1TT Yy:r)] =
Σ
( V log go (x YIT).
(5)
Importantly, any method that provides estimates of expectations w.r.t. the posterior can used for gradient estimation within the RWS framework, 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 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 (x Yı-r) with a set of N weighted particles x con- structed by sampling from a sequence of target distributions (xx(x)))){== Since these intermediate targets are often only known up some unknown normalizing constant Z,, SMC uses the unnor- malized targets (xx:)) ) where Ty (x14) = % (x1:)/ /Z, Provided mild technical conditions met and (>1:2) pø (X1:T Yi.T). N1:T.Yi:T) SMC returns weighted particles that approximate the poste-
rior po (X yI.T) [10, 24]. These weighted particles can be used to compute biased but consistent
estimates of expectations under the posterior, similar SNIS.
SMC repeats the following steps for each time t::
1. Sample latents x1:N from a proposal distribution (x Yıer).
2. Weight each particle using the unnormalized target 7 form an empirical approximation
R to the normalized target distribution Ne.
3. Draw new particles x iit from the approximation t, (the resampling step).
By resampling away latent trajectories with low weight and focusing on promising particles, SMC produce lower variance estimates than SNIS. For thorough review of SMC, Doucet and Johansen [10], Naesseth et al. [24], and Del Moral [25].
Filtering vs. Smoothing The most common choice of unnormalized targets "r are the filtering distributions pa (Xi:t; V1:t), resulting in the algorithm known as filtering SMC or 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 ignores future observations Yt+1:7. 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 (x1:)) = Pe X1 )1:r), but unfortunately the smoothing distributions are not readily available from the model. We can approximate them, however, by observing that (X1t. V1.r) is proportional to the filtering distributions po(X1:ts y1.)) times the lookahead distributions p xx).
the lookahead distributions are well-approximated by sequence of twists (F+4r,, x=)) then running SMC with targets 7 (x(x;))) = Po po (x Y1.) re (Y++IT, xx) approximates smoothing SMC [26].
3

---

Twist Learning We have reduced the challenge of obtaining the smoothing distributions to learning twists that approximate the lookahead distributions. Previous twist-leariigg approaches include maximum likelihood training on samples from the model [18, 28] and Bellman-type losses motivated by writing the twist at time recursively in terms of the twist at time / + 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 motivated by observing that the lookahead distribution is proportional to a ratio of densities up to a constant independent of 2t,
Po(x; Y++1r)) po(y)1o)))
Po y++1T) (x
Po 1 Xx)
(6)
po (x(x))
Results from the DRE via classification literature [30] provide a way to approximate this density ratio:
train classifier distinguish between samples from the numerator po Vr+1:7) and denominator
(x.). 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 full proof Sugiyama et al. [30].
In practice, it is not possible to sample directly from px(x: Y(+1.). Instead, Lawson et al. [19] sample full trajectories from the model's joint distribution, i.c. draw X1:7. VI-T (o(Xi:T.Y1 and discard unneeded timesteps, leaving only x, and Yt+1:T which are distributed marginally as
po po(X,, Xt, Y++1:r). Training the DRE classifier data sampled in this manner will approximate the ratio po Xx, Y++1:1)////)) P/Po (xe)po (++1:r), which is equivalent to Eq. (6), see Appendix 8.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 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 X1:7 and weights w..r from a smoothing SMC run with model pe. proposal 4o, and twist Tyg
2. Use those trajectories and weights to form estimates of gradients for the model pe and proposal qe. as in reweighted wake-sleep. Specifically, NAS-X computes the gradients of the inclusive KI. divergence for learning the proposal 4o as
T
N
a log (x)
xx2 (x(2)
(7)
(=1
and computes the gradients of the model pa
ΣΣ
w
Vo
( yv x xx2)).
(8)
3. Update the twists Ty using density ratio estimation via classification.
full description is available in Algorithms 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 timestep-t and timestep-T approximations of the target distribution could be used, in the latter case by discarding timesteps after t. Specifically,
N
N
pa YI.) R
Σ W 6(X
x ~
Σ
( W7
0(x (x
1=1
1=1
where 8(a; b) is a Dirac delta of located at b and (x1:)):: 1: denotes selecting the first i timesteps of a timestep-T particle; due to SMC's ancestral resampling step these are not in general equivalent. For NAS-X we choose the time-t 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, key advantage that explore theoretically in Section 3.1 and empirically in Section 5.

---

Algorithm 1: NAS-X
Procedure NAS-X (fo, on, to yI.T) 0 <- 0o, Φ 6 00;
1 + 100
while not converged do
x xI.
1:1
+ SMC( (x Ylah P
AO
Σ Ei= == wi 10)
Vo lx, yy xx_1)
ΑΦ - Σ t= ΣN
1= mi
(
(x (eT))
0 steddssep(() A0)
Φ + grad-step Aω)
3 f- twist training v))
end
return 0, Φ. vb
Procedure twist- twwisttntriinnggnae..n (0, vo)) See Algorithm 2 in Appendix 8.3.
3.1 Theoretical Analysis of NAS-X
In this section, state two theoretical results that illustrate NAS-X's advantages over NASMC, with proofs given in Appendix 7.
Proposition Consistency of NAS-X's gradient estimates. Suppose the twists are optimal so that (Y xx) 1(7
x,) up constant independent of x, for t= 1,...,T - Vo log po Yr.r) NAS-X's weighted particle approximation the true gradient of the log marginal likelihood V log.elyi.)). Then Ve log poly py(T).)) ": Va logloooo pa (y(i)) as N -0.
Let
Proposition 2. Unbiasedness of NAS-X's gradient estimates. Assume that proposal distribution (x, X1:1-1-Y::)) is optimal that (x X1:1-1, V1.T) p(x: X1.-1.y1:1) for t 1,..., T, and the twists (y X+) are optimal so that Tw (Y Xx) )) xx p(yt+1eT Xx) up to constant independent of Xt for t t= I....,T approximation to the true gradient of the log marginal likelihood, Va log pe (YIT)) Then, for any
Let Va log po(T).)) be NAS-X's weighted particle
number of particles, E log bogpo(s (y(Y))) = log (y
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 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 Hinton et al. [13] as way to train deep directed graphical models. Bomnschein and Bengio [14] interpreted the wake-sleep algorithm as self-normalizer importance sampling and proposed reweighted wake- sleep, 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 upproximate 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 lower bound on the log marginal likelihood. 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
5

---

011
NASK
Progroe polatio unne
1.
013
ELBO
025
RIWS SIXO
Maaaethhhaada.. Addde M
603-
0.43
0.3
025
110
XC
Propuoal Inarin (1004)
Prepocal
(1000g
Madd lannegs
Figure 1: Comparison of NAS-X and baseline methods on inference in LG-SSM. (left) Compari- son of -orrgiall likelihood bounds (lower is better), (middle) proposal parameter enror (lower better), and (right) learned proposal variances. NAS-X outperforms all baseline methods and recovers the true posterior marginals.
addition these empirical advantages, NAS-X can fit discrete latent variable models while SIXO would require high-variance score function estimates.
Experiments
We empirically validate the following advantages of NAS-X:
By using the approximate smoothing distributions 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 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 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 show NAS-X can ODE-based mechanistic models of neural dynamics with 38 model parameters and an 18 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 one-dimensional linear Gaussian state space model with joint distribution
p(x.T.YL.) (x 0, of)
W of) II (NYx,,
(10)
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, ((117)) = IT== de (x.) IN= Axx;; Me; of). with parameters μ1.7 and of corresponding to the means and variances at each time-step. For twist-based methods, we parameterize the twist as a quadratic function in x, whose coefficients functions of the observations and time step. We chose this form to match the functional form the analytic log density ratio. For details, see Section in the Appendix. NAS-X outperforms all baseline methods, achieving tighter lower bound on the log-marging 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 computed closed form. See Section for comparison 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 upproximate the smoothing distributions, 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, 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 coupled sets of latent variables: a discrete state 24, with K = possible values, and a two-dimensional continuous state X that follows linear dynamics that depend on Z₁. The observations are noisy projection of x; into a cen-dimensional observation space. There are 1000 observations in total. For the proposal, we factor 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 al. [41]-
(a) Comparing learned dynamics of NAS-X, NASMC, and Laplace EM rSLDS.
Ground truth
NAS-X
NASMC
Laplace EM
201
10
60-
00
-05
-03
-45
-15-
13
-25
-15
21
21
J1
(b) Train 1024 for rSLDS.
Method
03 = 0.001
o = 0.01
03 = 0.1
NAS-X
19.837 0.0234
8.63 + 0.0015
-2.79 = 0.0009
NASMC
19.834 = 0.0018
-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
Figure 2: Inference and model learning in switching linear dynamical systems (SLDS). (top) Comparison of learned dynamics and inferred latent states in model learning. Laplace EM sometimes learns incorrect segmentations, as seen in the rightmost panel. (bottom) Quantitative comparison of log marginal likelihood lower bounds obtained from running bootstrap particle filter (BPF) with learned models.
We present qualitative results from model learning and inference in the top panel of Figure We compare the learned dynamics for NAS-X, NASMC, and Laplace-EM algorithm designed specifically for recurrent state space models [42]. each panel, we plot the vector field the learned dynamics and the posterior means, with different colors corresponding to the four discrete states. NAS-X 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 par with both NASMC and Laplace EM across the different observation noises og- See Section 10, for additional results on inference.
5.3 Biophysical Models of Neuronal Dynamics
For our final set 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 neuronal membrane it changes in response to external stimuli such injected current. Un-
7

---

derstanding how voltage propagates throughout 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 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 brief introduction to the models used in this section and defer full deseription to the appendix. We are concerned with modeling the potential difference across neuronal cell membrane, v, which changes in response to currents flowing through a set of ion channels, c C. Each ion channel has an activation which represents percentage of the maximum current that can flow through the channel and is computed
nonlinear function Je of the channel state Aes with go (Ae) E [0, l]. This activation specifies the time-varying conductance the channel fraction of the maximum conductance of the channel, Te- Altogether, the dynamics for the voltage can be written as
du Cm
Text
EEaae(de)(r EcCOo) S
(11)
CEC
where Cn is the specific membrane capacitance, Lext is the external current applied to the cell, S is the cell membrane surface area, Gon the ion transported by channel c, and Eeg. is that ion's reversal potential. In addition to the voltage dynamics, the ion channel states {deLeEc evolve as
dde dt
A(v)Ae + b(u) E
(12)
where A(v) and b(u) are nonlinear functions of the membrane potential that produce matrices and vectors, respectively. Together, equations (11) and (12) define conditionally linear system of first-order ordinary differential equations (ODEs), meaning that the voltage dynamics are linear if the channel states are known and vece-versa.
Following Lawson et al. [19]. we augment the deterministic dynamics with zero-mean, additive Gaussian noise to the voltage and unconstrained gate states (Ag) 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 hidden size of 64 units to process the raw observations and external current stimuli, and then fed the processed observations, previous latent state, and transformer positional encoding [46] into 64-unit single-layer MLP that produced the parameters 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 large step sizes because their state 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 60,000 integration steps the standard 0.005 milliseconds per step. Because running our models for even 10,000 steps would be 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, 20-time speedup [47]. For details, 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 probabilistic Hodgkin Huxley model of the squid giant axon [9], and used each method to
8

---

STXO
-109
Vollogge (wi))
Trae Valtea Infernd
Posongeing Fivel
-181
Log-Wlehhoo
-195
Vollage (aN))
FIVO
NASMO
SINO
NASX
2005
32
128
Eval Num Particles
(ms)
(a) HH inference performance. Proposals were trained with 4 particles
(b) Inferred voltage traces for SIXO and NAS-X.
(top) SIXO generates high number resampling events Icading to particle degeneracy and single mistimed spike, (bottom) NAS- X perfectly infers the latent voltage with no mistimed spikes, and resamples infrequently. See Fig. 9 in Appendix for NASMC's traces.
and evaluated across a range of particle
numbers. RWS performed too poorly
be included.
Figure 3: Inference in Mechanistic HH Model
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], sampled trajectories of length 50 milliseconds, with single noisy voltage observation every millisecond. The stability of the Strang splitting based ODE integrator allowed us to integrate dt 0. 0.1ms, 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 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-efficiet than the filtering methods. To achieve NAS-X's s inference performance with particles, NASMC would need 256 particles, a 64x increase, and NAS-X also on average 2x more particle-efficien 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 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 beneficial 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 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 model parameters 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-gates ion channels 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, contrast to the 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 real mouse neuron by the Allen Institute, but downsampled and noised the data to simulate a more common voltage imaging setting. We ran hyperparameter sweep over learning rates and initial values of the voltage and observation noise variances (270 hyperparameter settings all), and selected the best performing model via early stopping the train log marginal likelihood lower bound. Each hyperparameter setting run for

---

NAS-X Sumplees
SIXO Sampless
Thee Voltune
200
Time (ms)
Method
LBPF
# Spikes Err.
Rest Voltage Err.
Cross-Core
% Runs Failed
NAS-X
686.4 + 6.8
0.76 0.15
2.74 + 0.1
6258
18.9 25.2
SIXO -660.6 + 4.4
1.88 = 0.41
1.8 + 0.2
6055 + 22
Figure 4: Model learning in HH model of mouse pyramidal neuron (top) Samples drawn from learned models when stimulated with square pulse of 250 picoamps beginning at 20 milliseconds (vertical grey dashed line). NAS-X's samples are noisier than SIXO's, but spike more consistently. (bottom) A comparison of NAS-X- and SIXO-trained models along various evaluation metrics. SIXO's models achieve higher bounds, but are less stable and capture overall spike count more poorly than NAS-X MAS-X-trained models. All errors are absolute errors.
5 seeds, and each seed was run for days on a single CPU core with 7 Gb of memory. Because the inherent instability of these models, many seeds failed, and discarded hyperparameter settings with more than failed runs.
In Figure 4 (bottom), we compare NAS-X and SIXO-trained models with respect to test set log- likelihood lower bounds as well as biophysically relevant metrics. To compute the biophsyical metrics, 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 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 '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 's runs suffer from numerical instability issues, great advantage when working with mechanistic models.
Conclusion
In this work presented NAS-X. X, 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 KI. divergence with respect 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.
10

---

Acknowledgements
This work was supported by the Simons Collaboration on the Global Brain (SCGB 697092), NIH (U19NS113001, ROINS113119, and R0INS130789) NSF (2223827), Sloan Foundation, McKnight Foundation, the Stanford Center for Human and Artificial Intelligence, and the Stanford Data Science Institute.
References
[1] Leonard E Baum and Ted Petrie. Statistical inference for probabilistic functions of finite state markov chains. The annals mathematical statistics, 37(6))1554 1563, 1966.
[2] Rudolph Emil Kalman. A approach linear filtering and prediction problems. ASME Journal of Basic Engineering, 82(1):35 March 1960.
[3] 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.
[4] Diederik Kingma and Max Welling, Auto-encoding variational Bayes. In 2nd International Conference Learning Representations, 2014.
[5] Rahul G Krishnan, Uri Shalit, and David Sontag. Deep Kalman filters. 111:01151,.0512,, 2015.
arXiv preprint
[6] 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.
[7] 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.
[8] Siddhartha Chib, Yasuhiro Omori, and Manabu Asai. Multivariate stochastic volatility. In Handbook of Financial Time Series, pages 365-400. Springer, 2009.
[9] Alan L. Hodgkin and Andrew F. Huxley. A quantitative description of membrane current and its application conduction and excitation in nerve. The Journal of Physiology, 117(4):500, 1952.
[10] 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.
[11] David Blei, Alp Kucukelbir, and Jon D. McAuliffe. statisticians. Journal of the American Statistical Association, 1
Variational inference: A review for 877, 2017.
[12] Martin J. Wainwright and Michael I. Jordan. Graphical models, exponential families, and variational inference. Foundations and Trends@ Machine Learning, ((1-2):1-055 2008
[13] Geoffrey E Hinton, Peter Dayan, Brendan Frey, and Radford M Neal. The "wake-sleep" algorithm for unsupervised neural networks. Science, 268(5214): 1158-111111110 1995.
[14] Jörg Bornschein and Yoshua Bengio. Reweighted Wake-Sleep Sleep. arXiv preprint arXiv:1406.221 2014.
[15] 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.
[16] 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.
[17] 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.

---

[18] 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.
[19] Dieterich Lawson, Allan Raventós, Andrew Warrington, and Scott Linderman. SIXO: Smooth- ing Inference with Twisted Objectives. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[20]
Shixiang Shane Gu, Zoubin Ghahramani, and Richard E. Turner. Neural Adaptive Sequential Monte Carlo. Advances in Neural Information Processing Systems, 28, 2015.
[21] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman--Min Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL com/ monn/goole/jax.
http://giie
[22] Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrry 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: Large- scale machine learning on heterogeneous systems, 2015. URL https: htsttps://ww.... /www///wwwww// tensorflow. org/ Software available from tensorflow.org.
[23] Art B Owen. Monte Carlo Theory, Methods and Examples. Stanford, 2013.
[24]
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.
[25] Pierre Del Moral. Feynman-Kan Kac formulae: genealogical and interacting particle systems with applications, volume 88. Springer, 2004.
[26] Nick Whiteley and Anthony Lee. Twisted particle filters. The Annals of Statistics, 42(1): 115-141, 2014.
[27] Mark Briers, Amnad Doucet, and Simon Maskell. Smoothing algorithms for state space models. Annals of the Institute of Statistical Mathematics, 62(1):61 2010.
[28] Ming Lin, Rong Chen, and Jun S. Liu. Lookahead strategies for sequential Monte Carlo. Statistical Science, 28(1):69-94, 2013.
[29] 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: /5460 2022.
[30] Masashi Sugiyama, Taiji Suzuki, and Takafumi Kanamori. Density ratio estimation in machine learning. Cambridge University Press, 2012.
[31] Rajesh Ranganath, Sean Gerrish, and David Blei. Black box variational inference. In Interna- tional Conference on Artificial Intelligence and Statistics, pages 814-822. PMLR, 2014.
[32] Matthew D. Hoffman, David Blei, Chong Wang, and John Paisley. Stochastic variational inference. Journal of Machine Learning Research, 2013.
[33] Andriy Mnih and Danilo Rezende. Variational inference for Monte Carlo objectives. In International Conference on Machine Learning, pages 2188-2196. PMLR, 2016.
[34] Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders, In 4th International Conference on Learning Representations, 2016.
[35] Christian Naesseth, Fredrik Lindsten, and Thomas Schon. Nested Sequential Monte Carlo Methods. In International Conference on Machine Learning, pages 1292-1301. PMLR, 2015.
12

---

[36] Heiko Zimmermann, Hao Wu, Babak Esmacili, and Jan- Van-Willem van de Meent. Nested vari- ational inference. Advances in Neural Information Processing Systems, 34:20423-20435. 2021.
[37] Dieterich Lawson, George Tucker, Bo Dai, and Rajesh Ranganath. Energy-inspircd models: Learning with sampler-induced distributions. Advances in Neural Information Processing Systems, 32, 2019.
[38] 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.
[39] 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.
[40] 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
[41] 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.
[42] 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 11880111911 PMLR, 13-18 Jul 2020.
[43] Peter Dayan and Laurence F. Abbott. Theoretical Neuroscience: Computational and Mathemat- ical Modeling of Neural Systems. MIT press, 2005.
[44] David E Rumelhart, Geoffrey E Hinton, and Ronald Williams. Learning internal represen- tations by propagation. Technical report, California Univ San Diego La Jolla Inst for Cognitive Science, 1985.
[45] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 9 (8):1735-1780, 1997.
[46] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is All You Need. Advances in Neural Information Processing Systems, 30, 2017.
[47] Zhengdao Chen, Baranidharan Raman, and Ari Stemn Structure- presteriig numerical integrators for Mux-kiin-Hully-tyye systems. SIAM Journal on Scientific Computing, 42(1):B273-B3) 2020.
[48] 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)9936-953, 2020.
[49] Michael Jordan. Serial order: A parallel distributed processing approach. In Advances in psychology, volume 121, pages 471-495. Elsevier, 1997.
[50] Diederik Kingma, Jimmy Ba, Yoshua Bengio, and Yann LeCun. Adam: method for stochastic optimization. In 3rd International Conference on Learning Representations, 2015.
[51] Tom Rainforth, Adam Kosiorek, Tuan Anh Le, Chris Maddison, Maximilian Igl, Frank Wood, and Yee Whye Teh. Tighter variational bounds are not necessarily better. International Conference Machine Learning, pages 4277-4285. PMLR, 2018.
13

---

[52] AIBS. Biophysical modeling perisomatic. Technical report, Allen Institute for Brain Science, Domupppppp//oommtttattiooo
10 2017. URL htth://hhhe brain-
[53] Costa M Colbert and Enhui Pan. Ion channel properties underlying axonal action potential initiation in pyramidal neurons. Nature neuroscience, 5(6):533 2002.
[54] Jacopo Magistretti and Angel Alonso. Biophysical properties and slow voltage-dependent inactivation of sustained sodium current entorhinal cortex layer-ii principal neurons: whole-cell and single-channel study. The Journal of general physiology, 114(4):491-50) 1999.
[55] Maarten HP Kole, Stefan Hallermann, and Greg Stuart. Single ih channels in pyramidal neuron dendrites: properties, distribution, and impact on action potential output. Journal of Neuroscience, 26(6):1677-167) 2006.
[56] 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-4611 1993.
[57] Robert Avery and Daniel Johnston. Multiple channel types contribute to the low-voltage- activated calcium current in hippocampal ca3 pyramidal neurons. Journal of Neuroscience, 16 (18):5567-5582, 1996.
[58] AD Randall and RW Tsien. Contrasting biophysical and pharmacological properties of t-type and r-type calcium channels. Neuropharmacoloy. 36(7):879-89) 1997.
[59] 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.
[60] Alon Korngreen and Bert Sakmann. Voltage-gated k+ channels layer neocortical pyramidal neurones from young rats: subtypes and gradients. The Journal of physiology. 325(3):621-63)) 2000.
[61] M Köhler, Hirschberg, CT Bond, John Mark Kinzie, NV Marrion, James Maylie, and JP Adelman. Small-conductance, calciun-activated potassium channels from mammalian brain. Science, 27366282) 1996.
[62] Alastair Walker. New fast method for generating discrete random numbers with arbitrary frequency distributions. Electronics Letters, 8(10):127-128, 1974.
[63] Richard A Kronmal and Arthur V Peterson Jr. On the alias method for generating random variables from discrete distribution. The American Statistician, 33(4):214-218, 1979.
14

---

7 Theoretical Analyses
Proposition Consistency of NAS-X's gradient estimates. Suppose the twists are optimal so that
rs (Yr-1:T x²) P(y++1:T xx) up to a constant independent of X₁ for t= 1...... .T -
Let
Vo log po (y:r) be NAS-X's weighted particle approximation the true gradient of the log marginal likelihood Vg log po(r).) Then Ve polgo.r))) "3 Va logpo (YLe)) LT) as N +00
Proof. This is 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 normalized target distribution. That is, consider some test function and let SMC's particle weights be denoted as wf. We have that Ei= w h(x(:)) " ) (x1.e) as K + oc where Te (x1:)) is the normalized target distribution. NAS-X sets (x(a))) or (x1a) (x yia)rg MI+1-T xz). which by assumption is proportional to Po V1a)Po xx) = Po Vi.)). The desired result follows immediately.
Proposition 2. Unbiasedness of NAS-X's gradient estimates, Assume that proposal distribution V1.T) (1.T)
(x( X1--11, Y1.T) is optimal so that q X1:t-1, I...., T, and the twists T (V++1:1; X2) are optimal so that T (Y+IIT, x₂) x p(y++1:T (x)) up to a
p(x: X1a-1, for t
constant independent of xt for t= ==...... - 1. Ler V log (yi:)) be NAS-X's weighted particle approximation to the true gradient of the log marginal likelihood, Vo log po (y Then, for any number of particles, V loglllopyeee (y(rr)) = log po (y1:)).
Proof. We first provide a proof sketch then give a more detailed derivation, adapting the proof of 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 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 equal; that is, w = 1 w p(y1:)) for h = 1..... K, depending on whether resampling occurred. We proceed by induction on f, the timestep in the SMC sweep.
For = note that
1. 71 (xt) v(xxxr))p(yer xf).
2. 70 - 1,
3. and qı (x+)) = p(xt VLT)).
This implies that the incremental weight af
p(xt,y 1))ppp2(2) x²)
p(x),
at
= p(V1:))
(13)
p(x Yı:r)
p(x+
yI.T)
which does not depend on k. Because w A 1, we have that w = w = p(y)er) for all k.
Note that, since the proposal is optimal, prior to SMC's reweighting step, the particles were distributed as follows xf ~ p(x1 y..r). Since the incremental weights are equal, the distribution of particles is unchanged.
For the induction step, assume that wi'k equals or p(yer) and that x ~ p(x (1:r). The particles are distributed as follows, x ~ p(x, Nirr). This implies that (xf xf)
p(x; X1:t-1, Mic-1.y1.T// N1.T) (1.T)p(xx--1111 y1.r) p(x Vi.T).
15

---

We now show that the incremental particle weights af are equal using the following identi- ties/assumptions
1. (x(x).)) - V1x) / xxx),
2. - -xx((x(--1)) p(x+ 1. 11 xx_)),
3. (y xx) oc p(y++1.7 x,) up to constant independent of x, for I,....? T
4. and (x+) p(x x
This shows that at is given by
p(xp.r- V V1:e)((y+tt) xx)
p(x 1 Y1:-1 PY x )p(x Nirr) a
(14)
p(xf VI.T)
(15)
11x X1:1-1;
(16)
p(x p(xf (11r-1:V1.))
(17)
for k = 1..... K.
There cases depending the value of the weights the previous timestep. If wttk then wt a for all k. On the other hand, if wi:k (()).)) then wt p(yI.T) for all 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, 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.
p(xt,y1) y1) af
(18)
p(x² V1.T
After reweighting, the particles will be distributed according to the filtering distributions. The distribution of particles after reweighting is proportional to p(x{ y1). 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 same each iteration under a certain dependency structure for the model p(xi:--Ilyss) - that was identified in Maddison et al. [17]. However, this dependency not satisfied general and therefore NASMC's gradients are not unbiased.
16

---

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.
Ve log p(y::)) = Ve log (o(XI:T. Vurl)yy
/ Po P V
(o(y Vo
/
Vi. Yur))
po(y1:))
po Vo log Vir))
po(y1:)) po(xi:T|yI:T) Vo (ir
/ Vo Σ
log xx
Σ/ Do
Σ Beotarrtrr) Ve log
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 generic Markovian model. In this derivation, we assume there are no shared parameters between the proposal and model.
- -VoKL(pollg (lp)) V
pe pe(xirlyy:r) log qa (X1:T 1yrryy.
/ Do po(x1Tlyi log to (
/ po(xirlyyTT) log (x (xx dxir
Σ/ po(xi:l/11e V
X1 (x )
Σ E Epo
V log g0 (x(
We use the assumption that there are no shared parameters 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 al. [30].
Let (x)) and b(z) be two distributions defined over the same space X' and consider a classifier g X' - RR that accepts specific E X' and classifies it as either being sampled from b. We will train this classifier to predict whether given x was sampled from a(x) or b(z). The raw outputs (logits) of this classifier will approximate log(a (og(a())/b(a))) up constant that does not depend on
To see this, define an expanded generative model where we first sample E {0, 1} from a Bernoulli random variable with probability 0 p 1, and then sample x from a(x) if 1, and sample
17

---

from b(x) if = O. This defines the joint distribution
p 2) - p(=)p((())))))))))()) 2) (ernoall((x:))))((]]])))))
(19)
where p( 2= = 1) a(z) and p( 2 = 0) b(z).
Let g: X + IR be a function that accepts x E X' and produces the logit for Bernoulli distribution over This function will parameterize a classifier via the sigmoid function, meaning that the classifier's s Bernoulli conditional distribution over is defined as
P=(2|>) e (f ( = = (1,0)
(20)
where a is the sigmoid function and o is its inverse, the logit function
(() I-e-t o (p) = log (1 P
(21)
p
The optimal function g" will be selected by solving the maximum likelihood problem
arg maxxEp(px (py(2 x)]
(22)
9
The solution to this problem is the true p(= x). Because we have not restricted g, this solution can be obtained. Thus,
p(z xx) Pp (a = | x))
(23)
=o (x)) o(gg (x)))
(24)
(0(gg(x)))
(25)
This in turn implies that
g" g'(r) =0 =0 = 1 x:))
(26)
p(2 = x)
=log
(27)
1 - p(z = x))
p( = 1 x)
log
(28)
p( = 0 x)
p( = 1 x #))(x)
log
(29)
p( = 0 x)p(T)
p(: = 1, x)
log
(30)
p(z = 0, x)
p(x = = 1) 1)p(= = 1)
log
(31)
p(x = 0)p(z 0)
p(x 2 1)
p( 1)
log
+ log
(32)
p(x 2 0)
p(z 0)
(axx) Mx))
p
log
+ log
(33)
Thus, the optimal solution to the classification problem, g", is proportional to log(a(x)) ((x))))()))))))))))))))))))))))) up to constant that does not depend on x. practice we observe empirically that long as sufficiently flexible parametric family for g is selected, 9 will closely approximate the desired density ratio.
In the case of learning the ratio required for smoothing SMC,
px Y(+IT)
(34)
po (x4)
Lawson et al. [19] instead learn the equivalent ratio
pe (e41.T)
(35)
po (x po (:T
18

---

MASI ELBO
Prepooa polation seru
1.
RIWS
Jlgaethhhaadd... A Addd A
SIXO
110
Propuoal (1004)
Prepocal (1300g)
Madd Rianoeg
Figure 5: Comparison of NAS-X baseline methods on Inference in LG-SSM. (left) Comparison of log- marginal likelihood bounds (lower is better), (middle) proposal parameter error (lower is better), and (right) learned proposal means. NAS-X outperforms several baseline methods and recovers the true posterior marginals.
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 [19] draw paired and unpaired samples from the model that distributed marginally according to the desired densities. Specifically, consider drawing
XI:T, & (x YI.T) X1:7 ~ (X(T)))
(36)
and note that any X,, Y1+1: selected from the sample will be distributed marginally according to po(x,, Vt+1:7) Similarly, any X,, Vt+1.T will be distributed marginally as po (X)) (y++ir)) In this way. T - positive and negative training examples for the DRE classifier can be drawn using single set of samples as Eq. (36).
The twist training process summarized in Algorithm 2.
Algorithm 2: Twist Training
Procedure twist- trainint (8, 1/o)
v <- who
while not converged do
X1.7, NI.T ~ po N1.T)
&1.T po (XITT)
[(b) - Ei= log
log(1 (++1.))
1 - grad-step(4) ((b))
end
return v
LGSSM
Model Details
We consider a one-dimensional linear Gaussian state space model with joint distri-
bution
T Wx 2 N(yxx,
P(XIT = (xxx;
II
(37)
1=2
our experiments we set dynamics variance o² 1.0 and the observation variance o = 1.0.
Proposal Parameterization For both NAS-X and NASMC, we use mean-field Gaussian proposal factored over time
T
T
((11T) = II q: (x)) = II N(z:
(38)
with parameters M1. and of corresponding the means and variances at each timestep. In total, we learn 2I' proposal parameters.
19

---

Twist Parametrization We parameterize the twist as quadratic function in 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(x1:T, V1.T) a multivariate Gaussian, we know that p(rt (t+1:r) and p(xz)) are both marginally Gaussian. Let
pl V++ITT) _ NN(1,, of) p(xx) _ (N0,
Then,
p(ax ((tr:) p(())
log
Log p1, m1, oî) 1oz;
log (
207
1
+
of of
M1
/¹ 20
log Z(oz) +
20² 2
where Z(r) = ov2z. so log Z(r) - loog(ov
Collecting terms gives:
log(o v2xx + =oog002/22
1
1
0
02
f11
of 24
20
So we'll Il define
1 2
1 of
1
A/ a
o3
b b
11 of
log(on 2m) log( (o2 22
207
We'll explicitly model log of. log of and fu, Both log of and log of are only functions of t, not of Ve+1.T, so those can be vectors of shape T initialized at O. μ4 is linear function of Vt+1:7 and f, so that can be parameterized by a set of T weights, initialized to I/T and T biases initialized
Training Details We use a batch size of 32 for the density ratio estimation step. Since we do not perform model learning, 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 le - 3, decays the learning by 0.3 and 0.33 100, 000 and 300, 000 iterations. For the proposal, we used Adam with constant learning rate of le 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, 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 I, there is a discrete latent state {I...., 4) as well two-dimensional continuous latent state 2; E 122 The discrete state
20

---

NASXX, 60
% -01
NASMC, 53 -0.1
NASMC, 06 o6-
Propossl stege (100s)
Propose
stope (100+)
Figure 6: Inference in NASCAR experiments.
transition probabilities are given by
p(z+11 = i 1 = x.))) x exp (ra + R xz-1)
(39)
Here R, and weights for the discrete state 24-
These discrete latent states dictates two-dimensional latent state Iy E R.2 which evolves according to linear Gaussian dynamics.
Le+1 = T + bect1 + Ur. Up iid NO0 Q2111)
(40)
Here 4g. Qs E 222 and bg E TR2. 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 Ye R linear projections of the continuous latent state Ia with some additive Gaussian noise.
y/ Cz; + d Wt,
Dt rid (O,
(41)
Here C, S E R 10z10 and d E 10
Proposal Parameterization We use a mean-field proposal distribution factorized over the discrete and continuous latent variables (i.e. q(Z1.T,X x1.r) = (xx:T))() For the continuous states, q(X1.)) a Gaussian factorized over time with parameters P1.T and oirr. For the discrete states, ((Z1:)) is Categorical distribution over K categories factorized over time with parameters p In total, we learn 27 TK proposal parameters.
Twist Parameterization We parameterize the twists using a recurrent neural network (RNN) that trained using density ratio estimation. To produce the twist values at each timestep, we first run RNN backwards over the observations Y1:r to produce sequence of encodings 01:T-1- We then concatenate the encodings of x; and 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 set of continuous latent states and initialize the matrices C and We then fit autoregressive HMM to the estimated continuous latent states in order to initialize the dynamics matrices (Ax, bb 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 twist training and 100 steps of proposal training for a total 50,000 training steps in total. We used Adam and considered grid search over the model, proposal, and twist learning rates. In particular, we considered learning rates of le le 3, le for the model, proposal, and twist.
Bootstrap Bound Evaluation 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 random seeds. Initialization of the latent states was important for a fair comparison. To initialize the latent
21

---

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 / 0.
Inference Comparison In the top panel of Figure 6, we compare NAS-X X and NASMC on inference in the SLDS model. We report (average) posterior parameter recovery for the continuous and discrete latent states across 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 HH Model Definition
For the inference experiments (Section 5.3.1) we used a probabilistic version of the squid giant axon model 19, 43]. Our experimental setup constructed to broadly match [19], and used single-compartment model with dynamics defined by
du dt. Text INam h(v ENa) JKn (v Ek) - greak (v Eteak)
Cm
(42)
dm am (v)(1 m) Bm (m
(43)
dh
di an (v)(1 ht) - (u)h
(44)
dn (1) n) (n
(45)
dt
where Om is the membrane capacitance; is the potential difference across the membrane; Lext is the external current; INa. Ik, and Ueeak are the maximum conductances for sodium, potassium, and leak channels; ENa: Ek, and Ejeak 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 a and that define the dynamics for n, m, and h are defined as
4 r/10
am(u) exp( -4 - v/10) -1
Bm(()) = 4. expeep(((()) 65 v)/18)
(46)
1
(v) 0.07 exp((-63 u)/200,
(n(v) exp( -3.5 101 +1
(47)
5.5 v/10
an(v)= exp(-5.8 1/10) 1 Bm(()) = 0.125. exp(( (pp((665 65 v))8
(48)
This system of ordinary differential equations defines nonlinear dynamical system with four- dimensional state space: the instantaneous membrane potential and the ion gate subunit states n, m, and h.
As in [19], noise to both the membrane voltage and the "unconstrained" subunit states. The observations are produced by adding Gaussian noise with variance of the membrane potential
probabilistic version of the original HH model that adds zero-mean Gaussian
Specifically, let be the state vector of the system at time containing (ve, m., h4, ne), and let spar(x (x) be a function that integrates the system of ODEs defined above for step of length dr. Then the probabilistic HH model can be written as
p(XE:1 NI.T) pxx1) I p(x, Pai (x--1) 1)) II N (
(49)
where the 4-D state distributions p(x1) and p(x; pai (x are defined as
p(x( (xx(xx)))) = (
LogitNormmalxx.: (x o2.). (50)
22

---

Inference Hodeken Hundey trace, oba/mus, og
SINO NAS X treee, obs//ts
-189
-207
197
-171
Logg Logittellooo brwer bornd
202
-282
Logg iieellnood bovee
347
2.4
-00
115-
FIVO-RS
NASMC NAS-X
NAS-X X
155
Partickos
Nom Porticlis
Figure 7: HH inference performance across different numbers of particles.
(left) Log-likelihowe lower bounds for proposals trained with 4 particles and evaluated across a range of particle numbers. NAS-X's inference performance decays only minimally the number of particles is decreased, while other methods experience significant performance degradation. (right) A comparison of SIXO and NAS-X containing the same values as the left panel, but zoomed in. NAS-X is roughly twice as particle efficient as SIXO, and outperforms SIXO by roughly 34 nats at 4 particles.
In words, add Gaussian noise the voltage (x,,1) and logit-normal noise the gate states ni, m, and h. The logit-normat is defined as the distribution of a random variable whose logit has Gaussian distribution, or equivalently it Gaussian transformed by the sigmoid function and renormalized. We chose the logit-normal because its values are bounded between 0 and I, which is necessary for the gate states.
Problem Setting For the inference experiments sampled 10,000 noisy voltage traces from fixed model and used each method to train proposals (and possibly twists) to compute the marginal likelihood assigned 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 integrate dt = 0.1ms, meaning that there were 10 latent states per observation.
Proposal and Twist Details Each proposal was parameterized using the combination of a bidi- rectional recurrent neural network (RNN) that conditioned on all observed noisy voltages well as a dense network that conditioned on the RNN hidden state and the previous latent state X,-1 I [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', Xy, Both the proposal and twists were learned in an amortized manner, i.c. 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 'residual' from the prior, technique known as Res, [7]. our setting, we defined the true proposal density as proportional to the product of unit-variance Gaussian centered at (xx.) and 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 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-
23

---

Triin num particies
Triin num particles
Train mam particles 16
-100
ibbilhood bound
-179
STXO
NAS-X
32 250
111 1.28
16
1.08
250
Figure 8: Training HH proposals with increasing numbers of particles. HH proposal performance plots similar to Figure 7, but trained with varying numbers of particles. Increasing the number of particles at training time has a negligible effect on NAS-X performance in this setting, but caused many VI-based methods perform worse. This could be due to signal-to-noiss issues in proposal gradients, as discussed in Rainforth et al. [51].
NASMC
Vallage
Inferred Valage
20
Resamoling Svent
10
20
25
30
35
45
SIXO
50
True Voltage
Voltage (mV)
Inferredd Voltage
Resempling Event
50
NAS-X
50
Voltage (m))
-50
10
15
20
25
30
35
40
45
50
Time (ms)
Figure 9: Inferred voltage traces for NASMC, SIXO, and NAS-X.
(top) NASMC exhibits poor performance, incorrectly inferring the timing of most spikes. (middle) SIXO's inferred voltage traces are more accurate than NASMC's with only single mistimed spike, but SIXO generates high number of resampling events leading to particle degeneracy. (bottom) NAS-X perfectly infers the latent voltage with no mistimed spikes, and resamples very infrequently.
24

---

mance with 4 particles, NASMC would need 256 particles, a 64-66m4-m timemmes increase. NAS-X is also more particle-efficiet than SIXO, achieving on average 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 problems optimizing the proposal parameters. To compensate, we also evaluated the performance of "FIVO-BS", filtering method that uses a bootstrap proposal. This method is identical 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 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 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 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 has learned 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 SIXO's s high resampling rate and NASMC's s filtering approach lead to lower bound values.
Table 1: Train Bound comparison
Metric
NAS-X
SIXO
Cu
660.7003
636.2579
Li
664.3528
668.6865
662.8712
653.6352
train
662.0753
644.8764
L32
661.5387
639.5388
train
prain
660.8040
636.5131
C
660.5102
6633.7875
train
660.3423
1372.1377
12
Model Learning in Mouse Pyramidal Neuron Model
12.1
Model Definition
For the model learning experiments Section 5.3.2 we used a generalization the Bodgkin-Huxled 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 model cell 187169178. The model is detailed in the whitepaper [52] and the accompanying code, but we reproduce the details here to ensure our work 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 the set of ions [Na* Ca 2+ _K+). Each ion has associated with it
A set of channels that transport that ion, denoted C, for E I.
2. A reversal potential, E..
3. An instantaneous current density, 7, 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' = Utez C.. Each e e C' has associated with it
1. A maximum conductance density, Je.
2. A set of subunit states, referred to collectively as the vector de. Let d E 1jde i.e. da is de dimeennsion vector of values in the interval [0, 1].
25

---

3. A function 9. that combines the gate values to produce a number in [0, 1] that weights the maximum conductance density, 9. ge(Ae).
4. Functions A-(-) and (-(-) which compute the matrix and vector used in the ODE describing Ae dynamics. A. and be are functions of both the current membrane voltage v and calcium concentration inside the cell [Ca*+], the number of subunits (i.e. the dimensionality of
A₂) is des then the output (u, [Ca²*1;) +1)) is de x de diagonal matrix and the output of bu( (Ca?+]] is a de-dimensional vector.
With this notation we can write the system of ODEs
Cm
du dl
Text
Deakk (v Bleak) - 4
(51)
SA
iEions
4
Σ Jele (No)(- Ei)
(52)
eEC:
dee
Ac(v: (Ce, (Ca2+11)) li)d + be(u, [Ca2+];)
(53)
dt
d[Caa 2+
(ta2+ [Ca ICa?+Imin
dt = klca²
(54)
Most symbols are as described earlier, SA is the membrane surface area of the neuron, [Ca²-]; the calcium concentration inside the cell, 2+ min is the minimum interior calcium concentration with a value of nanomolar, is the rate of removal of calcium with a value of 80 milliseconds, and & and a constant with value
k 10000
(55)
2 F depth
where 10000 is a dimensional constant, y is the percent unbuffered free calcium, F Faraday's constant, and depth is the depth of the calcium buffer with value of 0.1 microns.
Because the concentration of calcium changes time, this model calculates the reversal potential for calcium E using the Nernst equation
G log CCa (Ca2-]0) Caa 2+]
FCC2- 2 F
(56)
where G the gas constant, T' is the temperature in Kelvin (308.15%), Fis Faraday's 's constant, and [Ca? 2aa *I is the extracellular calcium ion concentration which was set millimolar.
Probabilistic Model The probabilistic version of the deterministic ODEs was constructed similarly the probabilistic squid giant axon model Gaussian noise was added to the voltage and uncon- strained gate states. One difference is that the system state now includes [Ca?+]; which is constrained to be greater than O. To noise [Ca2*] added Gaussian noise in the log space, analagous the logit-space noise for the gate states.
Model Size
The 38 learnable parameters of the model include:
1. Conductances for all ion channels (10 parameters).
2. Reversal potentials of sodium, potassium, and the non-specific cation: Ex+, Exat, and E'scI.
3. The membrane surface area and specific capacitance.
4. Leak channel reversal potential and max conductance density.
5. The calcium decay rate and free calcium percent.
6. Gaussian noise variances for the voltage and interior calcium concentration [Ca2-]
7. Gaussian noise variances for all subunit states (16 parameters).
8. Observation noise variance.
26

---

The 18-dimensional state includes:
1. Voltage
2. Interior calcium concentration [Ca?*]
3. 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
if x
exprel(x)
exp(x) -
(57)
otherwise
numerically stable implementation of this function was critical to training our models.
Additionally, many of the channel equations below contain a 'temperature correction 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].
d (m, h), 9c(Ac) mm
1 dm
am (v) - -) - 3m (m)m
q dt
1 dh
an(v)(1 h) 3n(w)h
qi dt
= 2.3 (-22)
0.182
0.124
Om (m(U) (v) exprel (-(v + 40)/6)
Bm(w) exprel + 40)/6)
0.015 6
0.015 6
(n(v) exprel((v 66)/6)
Ba(v)= exprel(-(t 66)/6)
12.2.2 Persistent Na*
From Magistretti and Alonso [54].
d- h ge(de) mooh
mac
It 1+exp(--0 (v + 52.6)/4.6)
(exp(-(v-
dh
an(v) h) (r(w)
4 dt
4: 2.3( (2_211
2.88 10 4463 an(v) exprel((v (v + 17.013)/4.63)
6.94 6 2.63 (n(v) exprel(-(() (u + 64.4)/2.63)
12.2.3 Hpperpolarzzatton-etiiteedd cation conductance
From Kole al. [55]. This channel uses a 'nonspecific cation current' meaning can transport any cation. In practice, this is modeled by giving its own special ion NSC* with resting potential
27

---

+
A = m,
Do(Ao) = m
ENSC 45.0
dm di
am am (1) m) - (mm
0.001 6.43 11.9 exprel((v 154.9)/11.9) 8m(u) - 0.001 - 193 exp(e/33.1)
Om (v)
12.2.4
High-vollage-actiiate Ca conductance
From Reuveni et al. [56]
A- = (m, h), go(No) = m²²
dm di am (v) m) -33(v)m (v)m
dh
dt (n( h) --,(v)h
0.055 3.8
O' (v) exprel (-(r 27)/3.8)
Bm(()) 0.94 exp(-(v 75)
0.0065
On (n (v) = exp04 + 13) /50), (n(v)) exp( (xp 15))/88) +
12.2.5 Low-voltage-activaee Ca conductance
From Avery and Johnston [57], Randall and Tsien [58].
A (m, h), 9e (c()e) m²h
dm
mcx
m
m q: dt
1100 - dh
df
h
223 -221/10
20
meo exp( (r + 40)/6)'
m4 = 5 +
+ exp( 35)/5)
50
1o0 exp((v + 90)/6.4)
h 20 +
+ exp((1 + 50)/7)
12.2.6 M-type (Kv7) K* conductance
From Adams et al. [59].
Ao = m, go(dc) = m
1 dm
(t am (1)(u(( m) 3. (v)m
(t 2.3 T-
Om (u) = 0.0033 exp (0.10x (v+ + 35)), (m(u) = 0.0033 exp(-0..pp0 + 35))
12.2.7 Kv3-like K* conductance
A m, dm m dl m
ge()c) m
mx + exp(-(1 18.7)/9.7)
m,
+ exeeppp(re (v 46.56)/44.14)
28

---

Fast inactivating (transient, Kv4-like) K conductance From Korngreen and Sakmann [60].
12.2.8
A- (m, (c(Ac) = m²h
dm
max m m,
dt
hx h dh
(t dl
l1 Q: 2.3 (T-)/222)/200
1
0.92
m, = 0.34
1 expl (v 47)/29)
exp(((*+ 71)/59)3)
1
49
h
h, = 8 +
1+exp((1 + 66)/10)
exp(((v + 23)
g = = 10-5
12.2.9 Slow inactivating (persistent) K* conductance
From Korngreen and Sakmann [60].
A (m, hi), dm
9c(Acc) m²h
ma
my 9: dt
dh
h - he -2.3 -2--11//10
4: dt
(v exp(p
+
14.3)/14.6)
[1255 + 175.03 0.026r
if < 50
m.
1.25 13 0.026 1
if 50
hoc
+ exp((e + 54)/11)
24v 2690
h expl exp(((v + (7)0 75)/48)2) 48)2)
g = x 10-5
12.2.10
SK-type calcium-activated K* conductance
From Köhler al. [61]. Note this is the only calcium-gated ion channel in the model.
Ae 2,
(c()e) =2
dz
dt 24
2x0 1 + (0.00043/ (0.00043/(Ca" (Ca2
2, 2=
12.3 Training Details
Dataset The dataset used to fit the model was subset of the stimulus/response pairs available from Allen Institute. First, all stimuli and responses were downloaded for cell 480169178. Then, sections of length 200 milliseconds were extracted from subset of the stimuli types. The stimuli types and sections were chosen that the neuron was at rest and unstimulated at the beginning of the trace. We list the exclusion criteria below.
29

---

1. Any "Hold" stimuli: Excluded because these traces were collected under voltage clamp conditions which we did model.
2. Test: Excluded because the stimulus is O mV for the entire trace.
3. Ramp//amp to Rheobase: Excluded because the cell is only at rest at the very beginning of the trace.
4. Short Square: 250 ms to 450 ms.
5. Short Square Triple: 1250 ms to 1450 ms.
6. Noise and Noise 2: 1250 ms to 1450 ms, 9250 ms to 9450 ms, 17250 ms to 17450 ms.
7. Long Square: 250 ms to 450 ms.
8. Square 0.5ms Subthreshold: The entire trace.
9. Square 2s Suprathreshold: 250 ms 450 ms.
10. 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 I ms (from the original 0.005 ms per step) and corrupted with mean-zero Gaussian noise of variance 20 m to simulate voltage imaging conditions. Finally, the 95 traces 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 bidirectional RNN with 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 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 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 stepsize of O. 1ms while observations were received once every millisecond.
Hy Hyperparameter Sweeps To evaluate the methods we swept across the parameters
1. Initial observation variance: e,e2,
2. Initial voltage dynamics variance: e, 6,,
3. Bias added to scales produced by the proposal: e²,e5
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. 3-31 different hyperparameter settings, and 5 seeds were run for each setting yielding total of 270 runs.
When computing final performance, hyperparameter setting was only evaluated if it had 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- 4-particle log likelihood lower bound.
13 Strang Splitting for Hodgkin-Huxley Models
Because the Bodgkin-Huxley, model is a stiff ODE, integrating it can be a challenge, especially at large step sizes. The traditional solution is to use 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 not viable 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 [47], we developed splitting approach that exploits the conditional linearity of the system. The system of ODEs describing the model can be split into two subsystems
30

---

Twist parameter error
10
Twist securacy
0.50
10
0.75
0.70
10
0.65
300
1000
1500
2000
2500
3000
500
1000
1500
2000
2500
5000
Twist training steps
Twist training steps
Figure 10: Twist learning in LG-SSM. (left) Twist parameter error relative to optimal twist param- eters for LG-SSM task; (right) Classification accuracy of learned twist. With an appropriate twist parameterization, twist learning via density ratio estimation is robust.
of linear first-order ODEs when conditioned on the state of the other subsystem. Specifically, the dynamics of the channel subunit states {>e c E C} system of linear first-order ODEs when conditioned on the voltage and interior calcium concentration [Ca2+]; Similarly, the dynamies for and [Ca } ++1
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 the duration of the step. Our integration approach uses these exact updates in an alternating fashion, first performing an exact update 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 Strang and other splitting methods applied to Hodgkin-Huxley type ODEs, see [47].
constant for
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 distinguish between samples from (x Yc+1:)) and (x)).. 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 robust and easy. 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 appropriate twist parameterization. twist learning via density ratio estimation is tractable and straightforward.
15 Computational Complexity and Wall-clock Time
Theoretically, all have OKTT) time complexity, where K the number of particles and T is the number of time steps. Once concern is that the resampling operation in SMC could require super-line liper-linear time in the number of particles, but drawing K samples from K 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 Lawson et al. [19), giving time linear T.
methods considered (NAS-X, SIXO, FIVO, NASMC, RWS, IWAE)
NAS-X and SIXO have similar wall-clock speeds but slower than FIVO and NASMC, primarily because of twist training, see Table 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 vs. particles). Therefore, NAS-X uses computational resources much more effectively than other methods.
31

---

Method
ms ms/global step
propool step
twist step
TWAE
70.3 = 15.9
70.3
N/A
RWS
71.6 + 8.2
71.6
N/A
NASMC
53.9 + 6.6
N/A 51.4+
53.9
FIVO
51.4 13.3
N/A 89.7 + 21.3
NAS-X
163.2 + 39.8
73.5
SIXO
175.3 = 42.4
85.6 21.1
89.7 + 21.3
Table 2: Wall-clock speeds of various methods during HH inference.
Gradient variances
10P
Greclent hise
inf
INL lover honad
NAS-X
NASAC
101
-10F
SINO
FIVO
IANE
110
300
300
350
Propoool trannng stepe CEOUAN
Propert triting stape (1000))
Prigonal training stoge (10006)
Figure 11: Hodgkin-Huu Huxley gradient variances (left) gradient bias (middle), and log- gee-magiinal
likelihood lower bounds (right) over training.
16 Empirical analysis of bias and variance of gradients
In Figure 11, 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 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-marging likelihood bounds.
32


Reviews:
[{"review_id": "huwUxyPsPY", "reviewer": "Reviewer_Dwhv", "summary": "The authors propose to use SIXO particle approximation to calculate the expectations in reweighted wake-sleep algorithm for state space models.", "questions": "What do you think are main limitations of NAS-X?\n\nWhen the algorithm is expected to perform poorly (or be outperformed by Laplace EM)?", "limitations": "N/A", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "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).\n\nThe idea of replacing the variational expectations with best possible practical approximation of the posterior is interesting.\n\nThe empirical findings are well-presented.\n", "weaknesses": "I would lengthen the description of SIXO and twists.\n\nMore empirical evidence and discussion why NAS-X actually warrants better performance (as it introduces biases) would strengthen the paper.\n\nSome theoretical analysis would strengthen the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive comments and constructive feedback. Below we respond to the questions/concerns raised.\n\n**I would lengthen the description of SIXO and twists.**\n\nThank you for the suggestion. We have added a  detailed discussion of the technical details behind SIXO and twists in the revised version of the paper.\n \n**More empirical evidence and discussion why NAS-X actually warrants better performance.**\n\nIn new experiments (see PDF), we compare NAS-X to several more baselines, including FIVO, SIXO, RWS, IWAE, and ELBO. In short, NAS-X either outperforms or performs comparably to these baselines in all experiments. These additional results establish that 1) we outperform methods that employ reverse KL + smoothing while being able to handle discrete latents (e.g., SIXO) 2) we outperform both forward KL and reverse KL-based methods that use filtering SMC (e.g., NASMC, FIVO) and 3) we outperform several baselines that ignore sequential structure (RWS, IWAE, ELBO). Additionally, in new empirical results, we show that NAS-X has lower variance and lower bias gradient estimates than these methods (see general response).  \n\nAltogether, our new empirical results highlight the benefits of using smoothing SMC in conjunction with a forward-KL divergence objective for sequential latent variable models. \n\n**Some theoretical analysis would strengthen the paper.**\n\nWe will include two theoretical guarantees in the revised paper: 1) NAS-X’s gradient estimates are consistent when the twists are optimal 2) NAS-X produces unbiased gradient estimates even for a finite number of particles, provided that the twists and proposal are optimal. Please see the general response for a detailed discussion of this.\n\nImportantly, NASMC does not have these guarantees because it uses filtering SMC’s intermediate particle approximations to estimate expectations w.r.t. the posterior. Since these intermediate particles approximate the filtering distributions and NOT the smoothing distributions, this introduces additional bias into gradient estimates. Importantly, this bias persists even in the infinite particle limit. From this perspective, you can view the twists not as introducing bias but as compensating for the bias inherent in NASMC’s approach.\n\n**What do you think are main limitations of NAS-X?**\n\nOne limitation is the NAS-X does not have a unified objective (i.e., the proposal and model are updated using different objectives). In certain settings, this could lead to a divergence of model and proposal parameters.\n\nFIVO does have a unified objective for the model and proposal, namely SMC’s lower bound on the log-marginal likelihood.  SIXO shares this objective for the model and proposal but learns the twist using the separate density ratio estimation objective, as in NAS-X.\n\nWhile this had no practical impact on our experiments, the lack of a unified objective could be relevant in other applications. \n\nAnother limitation is that SMC’s gradient estimates are consistent but biased for finite numbers of particles in practice (when twists and proposals are imperfect). We think a thorough theoretical analysis of how the bias of the gradients depends on the quality of the twist approximation and the number of particles is an exciting topic for future work.\n\nA final limitation is that NAS-X only works in the offline setting where all observations are available. It cannot be used for streaming data.\n\n**When is the algorithm expected to perform poorly (or be outperformed by Laplace EM)?**\n\nBecause NAS-X is a sampling-based method, we would not expect it to perform as well as methods that analytically marginalize out discrete latent variables, such as Laplace EM [1]. However, analytic marginalization is only tractable when the number of latent variables is small. Our rSLDS experiments were designed to address this directly and show NAS-X matching or slightly outperforming Laplace EM even in a model with 4 discrete states. This shows that our sampling-based method performs reasonably even in settings where it is at a disadvantage.\n\n[1] Zoltowski, David, Jonathan Pillow, and Scott Linderman. \"A general recurrent state space framework for modeling neural dynamics during decision-making.\" Proceedings of the 37th International Conference on Machine Learning, vol. 119, PMLR, 2020\n"}, {"review_id": "bSVgNCTA08", "reviewer": "Reviewer_nTEg", "summary": "An 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", "questions": "* 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", "limitations": "The authors have made clear the class of probabilistic models that this method is designed for. \n", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* 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", "weaknesses": "The 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments and feedback. In addition to the small changes you suggested, we also respond to your main questions and concerns:\n\n**Lack of comparison to existing methods.**\n\nIn new experiments (see PDF), we compare NAS-X to several more baselines, including FIVO, SIXO, RWS, IWAE, and ELBO. In short, NAS-X outperforms or performs comparably to these baselines in all three experiments. These additional results establish that 1) we outperform methods that employ reverse KL + smoothing while being able to handle discrete latent variables (e.g., SIXO), 2) we outperform both forward KL and reverse KL-based methods that use filtering SMC (e.g., NASMC, FIVO) and 3) we outperform several baselines that ignore sequential structure (RWS, IWAE, ELBO). Additionally, in new empirical results, we show that NAS-X has lower variance and lower bias gradient estimates than these methods (see general response).  \n \nAltogether, our new empirical results highlight the benefits of using smoothing SMC in conjunction with a forward-KL divergence objective for sequential latent variable models. \n\nWe did not include IWAE, FIVO, or SIXO in the discrete latent experiment (rSLDS) because (as you imply) they rely on continuous reparameterization gradients. Instead, we compared to Laplace EM, a very strong baseline that analytically marginalizes out the latents.\n\n**Lack of discussion of, and comparison to, other SMC variants… motivation for the use of smoothing SMC**\n\nFirst, to address your concerns we included an analysis of variance and bias of gradient estimates for the Hodgkin-Huxley experiments. This includes comparisons to methods based on SNIS (e.g., IWAE, RWS) and filtering SMC (e.g., FIVO, NASMC). In short (see general response), NAS-X’s use of smoothing SMC leads to lower variance gradient estimates and lower bias.\n\nSecond, while this is certainly problem-dependent, the claim that filtering SMC’s estimates can be high variance is relatively well-established in the literature [1,2,3,4,5]. Performing reliable inference and model learning in problems where smoothing information is highly important (e.g. neural voltage data) was a primary motivation for our work.\n\nIn addition to the practical motivations, we wish to highlight theoretical motivations for using smoothing SMC in the context of RWS: 1) NAS-X’s gradient estimates are consistent when the twists are optimal, 2) NAS-X produces unbiased gradient estimates even for finite particles, provided that the twists and proposal are optimal. Please see the general response for a detailed discussion. Importantly, NASMC does not have these guarantees because it uses filtering SMC’s intermediate particle approximations to estimate expectations wrt the true posterior, to avoid particle degeneracy. Since these intermediate particles approximate the filtering distributions and NOT the true posterior, this introduces additional bias into gradient estimates; this bias persists even in the infinite particle limit. We illustrated this in Experiment 1, where NASMC did not recover the true posterior. In contrast, NAS-X’s gradients are unbiased, in theory, when the twists and proposal are optimal. We illustrated this advantage in the Gaussian state-space model experiments, and will provide a discussion/formal proof of this result in the revision to motivate our method better. \n\n[1] Lawson, Dieterich et al. “Twisted Variational Sequential Monte Carlo.” (2018).\n\n[2] Lawson, Dieterich et al. “SIXO: Smoothing Inference with Twisted Objectives.” (2022).\n\n[3] Naesseth, Christian Andersson et al. “Elements of Sequential Monte Carlo.” Found. Trends Mach. Learn. 12 (2019): 307-392.\n\n[4] Heng, Jeremey et al. “Controlled Sequential Monte Carlo” (2019).\n\n[5] Guarniero, Pieralberto et al. \"The iterated auxiliary particle filter.\" Journal of the American Statistical Association 112.520 (2017): 1636-1647.\n\n**Ablation studies incorporating these competitors [...] would make more clear the importance of the contribution of this work**\n\nThanks for this suggestion. We wholeheartedly agree. In new ablation studies (see general response) we observe, broadly speaking, that the RWS-based methods (RWS, NASMC, NAS-X) provide lower-variance gradients and learn less-entropic proposals than their variational inference counterparts (ELBO, FIVO, SIXO). Lower-variance gradients are an important advantage as long as they are also low-bias, and our experiments show that learning the twists allows NAS-X to provide the lowest-bias gradient estimates of all methods considered. Thus, both the KL direction and the twist learning contribute to NAS-X’s performance.\nIn addition to the performance benefits of NAS-X over its competitors, the forward-KL methods are also more versatile because they can accommodate discrete latent variables.\n\n**Was vanilla RWS compared to?**\n\nWe have included comparisons to RWS and other non-sequential baselines (e.g., ELBO, IWAE). We find that in all experiments, NAS-X significantly outperforms RWS. \n\n **Can there be more discussion on the classifier being used to approximate the density ratio?**\n\nWe will discuss this in more detail in the revised version. For experiment 1, we chose a parametric form that matches the analytic density ratio for the Gaussian state-space model. For experiments 2 and 3, we use RNNs that take the observations in as inputs and produce encodings for each time step. We feed these encodings and the latent states into an MLP classifier for each time step. We’ll include the citations below and expand our discussion of SIXO’s density ratio approach in the revision; please suggest any additional citations.\n  \n[1] Sugiyama, Masashi et al. Density ratio estimation in machine learning. Cambridge University Press, 2012.\n\n[2] Thomas, Owen et al. “Likelihood-Free Inference by Ratio Estimation.” Bayesian Analysis (2016).\n\n[3] Uehara, Masatoshi et al. “Generative Adversarial Nets from a Density Ratio Estimation Perspective.” arXiv (2016).\n"}, {"review_id": "iojga1Swdq", "reviewer": "Reviewer_DvQm", "summary": "The 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\n", "questions": "1. 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\n", "limitations": "I wonder if the authors could add more discussions on the limitations/future work of their proposed method?", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The 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. ", "weaknesses": "The paper lacks discussions of future work and limitations of the proposed method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments and constructive feedback. Below we respond to the questions/concerns raised.\n\n**The paper lacks discussions of future work and limitations of the proposed method.**\n\nOne limitation is the NAS-X does not have a unified objective (i.e., the proposal and model are updated using different objectives). On the other hand, both FIVO and SIXO use a unified objective, namely SMC’s lower bound on the log-marginal likelihood.  While this had no practical impact on our experiments, the lack of a unified objective could be relevant in other applications. \n\nAnother limitation is that SMC’s gradient estimates are consistent but biased; empirically, this doesn’t seem to prevent NAS-X from recovering proposal and model parameters in settings where we know those parameters. However, we think a thorough theoretical analysis of how the bias of the gradients depends on the quality of the twist approximation and number of particles is an exciting topic for future work.  \n\n**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\nSee general response\n\n**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**\n\nWe can expand upon this in more detail in the revision but elaborate here briefly. We find that VI methods such as SIXO tend to learn more entropic proposals. We saw this in the Hodgkin-Huxley experiments and the Gaussian state-space model experiments. \nWe think this property is related to alternative interpretations of the VI methods. These VI methods were originally motivated as using Monte Carlo algorithms to derive tighter lower bounds on the log-marginal likelihood, hopefully resulting in a better learning objective.\n\nHowever, an alternative interpretation is that these VI methods maximize a standard ELBO objective but with an “expanded” variational family. The expanded family is defined by *both* the chosen proposal family (q) and the Monte Carlo algorithm. As a concrete example, you can interpret the IWAE bound as maximizing the ELBO with a proposal defined by this procedure:\n\n1. Draw K samples x_1, … , x_k from q(x | y).\n2. Weight the K samples using the density ratio p(x_i, y) / q(x_i | y)\n3. Draw a sample from the set of weighted samples.\n\nThe IWAE bound can therefore be thought of as maximizing the ELBO but with a variational family that allows the underlying proposal (q) to propose K samples and choose the “best” one. This perspective is also explored in the paper \"Energy Inspired Models\" (Lawson et al. 2019).\n\nIntuitively, if the proposal (q) in IWAE has K chances to draw a high-quality sample, it is incentivized to be more entropic, since only the best sample is chosen. A similar argument holds for bounds based on more complex Monte Carlo algorithms like FIVO and SIXO. Samples from the underlying proposal (q) drive Monte Carlo algorithms that select only the best particles. This allows the proposal to generate many dubious samples provided it generates at least one good sample.\n\nCrucially, this interpretation does not hold for RWS-based algorithms (RWS, NASMC, NAS-X). RWS algorithms use self-normalized importance sampling variants like filtering and smoothing SMC to directly estimate the gradients of the log marginal likelihood. The Monte Carlo algorithms are only used to improve the gradient estimates, and cannot be interpreted as augmenting the proposal distribution. In practice, this results in proposals that are not as entropic as VI-trained proposals because they are not “aware” that their bad particles will be resampled away.\n\nWe hope this clarifies things and are happy to discuss more.\n\n**Figure 4: Why would we expect NAS-X to outperform in certain metrics but underperform in other metrics?**\n\nWe are not entirely sure why NAS-X would outperform on some metrics and underperform on others, but suspect it could be related to the less entropic proposals it learns. We provided both standard model learning metrics and biophysical metrics because different users may value different metrics. For example, our biophysical metrics were motivated by the ones used in [1]. \n\nThis is similar to other fields like image modeling, where it has become standard practice to evaluate model samples on perceptual metrics like Frechet inception distance (FID) as well as likelihoods. This is because many models that achieve the highest likelihoods did not generate good samples, and many models with the best samples did not provide tractable likelihood evaluation (e.g., GANS). \n\n[1] Lueckmann, J.M et al. Flexible statistical inference for mechanistic models of neural dynamics. NeurIPS (2017).\n\n**Section 5.3.2: I wonder why NASMC gets dropped from the results in this section?**\n\nWe found HH model learning under the IWAE, NASMC, RWS, and FIVO objectives to be highly unstable. Few runs survived to converge, instead NaN-ing out early. For example, 46 percent of all IWAE runs failed to converge, compared to 19 percent of NAS-X runs and 25 percent of SIXO runs. This prevented us from reliably evaluating the methods, and we concluded that they were ill-suited for the model learning task. We were able to successfully evaluate all of these methods for proposal learning (see new figures), however, and found them to be far inferior to SIXO and NAS-X.\n"}, {"review_id": "cSfWWJoIYS", "reviewer": "Reviewer_mkCV", "summary": "The 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.", "questions": "- 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”).", "limitations": "- 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.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- 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.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your comments and feedback. Just to be clear on terminology --- In your review you refer to “RWS architectures”, but we weren’t sure if you meant model architecture or a general framework for developing new methods. We take the later interpretation and refer to it as the “RWS approach” in our response.\n\n**My main concern with this paper is its potential impact**\n\nThank you for this feedback, proper positioning of the work is extremely important!\n\nFirst, we underscore the broad applicability of our work. NAS-X is a method for model learning and inference in Markovian state-space models (SSMs). Model learning and inference in SSMs is an important topic to NeurIPS and the broader machine learning community, with many applications (e.g., neuroscience, healthcare, deep generative models) [1,2,3,4,5,6,7,8]. \n\nIn addition, we believe NAS-X is of broader methodological interest than SIXO, as it outperforms SIXO on most tasks and can fit discrete latent variable models which SIXO cannot. NAS-X is not limited in its applicability by using RWS; instead, it leverages RWS- and SIXO-inspired methods to tackle broader problems than SIXO or RWS can on their own.\n\nIn many problems, we agree there is nothing wrong with filtering SMC. However, in some problems, smoothing is crucial (e.g., the Hodgkin-Huxley model). We showed that NAS-X performs better on the HH inference task with 4 particles than FIVO or NASMC does with 256 particles, a 64x improvement over filtering SMC. Our theoretical results also support this --- FIVO and NASMC’s gradient estimates are not consistent, whereas NAS-X’s gradient estimates are consistent. While there is nothing wrong with filtering SMC, in many settings, there can be strong gains from using smoothing techniques.\n\n[1] Krishnan, Rahul G. et al. “Structured Inference Networks for Nonlinear State Space Models.” AAAI (2016).\n\n[2] Fox, Emily B. et al. “Nonparametric Bayesian Learning of Switching Linear Dynamical Systems.” NeurIPS (2008).\n\n[3] Costacurta, Julia C., et al. \"Distinguishing Discrete and Continuous Behavioral Variability Using Warped Autoregressive HMMs.\" NeurIPS (2022).\n\n[4] Alaa, Ahmed M. et al. “Attentive State-Space Modeling of Disease Progression.” NeurIPS (2019).\n\n[5] Ghahramani, Zoubin et al. “Variational Learning for Switching State-Space Models.” Neural Computation 12 (2000).\n\n[6] Miller, Andrew C. et al. “Learning Insulin-Glucose Dynamics in the Wild.” MLHC (2020)\n\n[7] Chung, Junyoung et al. “A Recurrent Latent Variable Model for Sequential Data.” NeuIPS (2015).\n\n[8] Wu, Luhuan et al. “Practical and Asymptotically Exact Conditional Sampling in Diffusion Models.” (2023)\n\n**The relationship between SIXO and NAS-X is not clear enough**\n\nWe view the RWS approach as ascending estimates of the gradients of the log marginal likelihood (LML) of a latent variable model. These gradients are expectations wrt the posterior over latents and can be approximated using biased but consistent estimates from self-normalized importance sampling (SNIS).\n\nChanging the SNIS method gives new model-fitting techniques. Thus, NASMC uses filtering SMC (also an SNIS method) to estimate gradients, and NAS-X uses smoothing SMC. NAS-X approximates smoothing SMC using twists learned via density ratio estimation, as in SIXO.\n\nThis contrasts the VI approach, which ascends unbiased estimates of a lower bound on the LML. FIVO uses filtering SMC’s estimate of the marginal likelihood to define a lower bound on the LML, and SIXO uses smoothing SMC similarly.\n\nThe crucial difference is that RWS-based methods follow biased but consistent estimates of the gradients of the true LML, while VI-based methods follow unbiased estimates of lower bounds of the LML. These different approaches result in gradient estimators with different strengths and weaknesses. \n\nBoth NAS-X and SIXO use (approximate) smoothing SMC as part of their algorithms. The difference is that NAS-X uses the samples from smoothing SMC to estimate the gradients in a way that allows it to outperform SIXO and be more broadly applicable, as supported by our empirical results.\n\n**Another point of concern are NAS-X’s multi-level approximations**\n\nNAS-X's approximations provide important theoretical and practical advantages. For example, NASMC incurs bias in the gradient estimates because it approximates the posterior distribution with filtering distributions. We view introducing the twists not as a further approximation but as an attempt to fix NASMC's approximation. Our experiments and additional comparisons show NAS-X performs favorably against all baselines, providing evidence that this change is helpful in practice.\n\nIn principle, incorporating twists complicates the learning problem. In practice, twist learning is robust and easy. In the PDF, we present twist parameter recovery and classification accuracy for the Gaussian SSM experiments; 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. Twist learning was stable and successful in all other experiments as well, we'll include full results in the revision.\n\n**The paper claims that smoothing SMC leads to lower-variance estimates**\n\nIn new results, we show that NAS-X has lower variance and lower bias gradient estimates than filtering SMC-based methods (see general response).\n\nThis claim is well-established in the literature (e.g. FIVO). Furthermore, as discussed above, the advantage of smoothing SMC in the context of NASMC is also to reduce bias of gradient estimates. \n\n**Effect of number of training particles...**\n\nWe didn't have the resources to reproduce Figure 3a for 32/128 particles; see PDF for 8 and 16 particles.\n\n**We are constrained by space, please see general response for the rest of your concerns.**"}]

Author rebuttal:
We thank the reviewers for their detailed feedback. We respond to reviewers individually and provide a general response below. We have strengthened our submission with several new experiments (see figures in PDF) and theoretical analyses (see below). If the reviewers feel that the new experimental results, analyses, and clarifications resolve their concerns, we kindly ask them to consider updating their overall scores. We are happy to provide further clarifications.

### New Experimental Results
#### Additional Baselines
In new experiments, we show NAS-X outperforms alternative methods suggested by the reviewers (ELBO, IWAE, FIVO, SIXO, and RWS). NAS-X’s improvements over these methods highlight 1) the benefits of using smoothing SMC versus filtering SMC and self-normalized importance sampling and 2) the benefits of forward KL over reverse KL in the context of sequential latent variable models. Results include:  
* Figure 1. Gaussian State Space Model: NAS-X achieves a tighter lower bound and lower proposal parameter error than SIXO, FIVO, ELBO, IWAE, NASMC, and RWS. 
* Figure 3. r-SLDS: NAS-X attains a higher bootstrap particle filter bound than NASMC and RWS. SIXO and FIVO rely on continuous reparameterization gradients and so were not compared to.
* Figure 4. Hodgkin-Huxley Inference experiments: We include new comparisons between IWAE, FIVO, NASMC, SIXO, and NAS-X for proposals trained with 4, 8, and 16 particles. These show that NAS-X substantially outperforms all other methods even as the number of training particles changes. NAS-X achieves the same performance as SIXO with 4x fewer particles, and achieves the same performance as FIVO and NASMC with 64x fewer particles. RWS results were also obtained but were too poor to fit on the plots.

#### Robustness of twist learning 
* Figure 2 illustrates that twist learning is robust, plotting convergence to true twist parameters and twist classification accuracy in the Gaussian SSM.

### Computational Complexity/Wall Clock Time
Several reviewers asked about computational complexity. We provide a summary of our results here, and will include a full discussion in the final version of the paper.

Theoretically, all methods have O(KT) time complexity, where K is the number of particles and T is the number of time steps.
Practically, NAS-X and SIXO have similar wall-clock times but are slower than FIVO and NASMC, primarily because of twist training.
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.

In the accompanying PDF, we include a table (Figure 6) detailing the wall-clock speed of each method in milliseconds per step during HH inference training runs.

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 4 shows that FIVO, NASMC, IWAE, and RWS cannot match NAS-X’s performance even with 64 times more computation (256 particles). SIXO only matches NAS-X’s performance with 4x as many particles. Therefore, NAS-X uses computational resources much more effectively than other methods.

### Gradient variance and bias
Figure 5 (left) shows NAS-X attains lower variance gradient estimates than IWAE, FIVO, and SIXO with comparable variance to RWS. We also studied the bias in Figure 5 (middle) by approximating the true gradient using SMC with 256 particles and the best proposal found in the HH inference experiments. NAS-X's gradients are lower bias than all methods but FIVO, but FIVO's gradients are also the highest variance. We believe FIVO’s gradients appear less biased because its parameters are pushed towards degenerate values where gradient estimation is “easier”. We illustrate this in Figure 5 (right), where we plot LML bounds.

### New theoretical results 
We present two theoretical results illustrating NAS-X’s advantages in idealized settings. One concerns the consistency of NAS-X’s gradients and the other concerns unbiasedness.

Proposition I (informal): For any proposal and the optimal twists, NAS-X’s gradient estimate of the log marginal likelihood converges almost surely to the true gradient as the number of particles approaches infinity. This is not true for NASMC. 

Sketch: This follows from 1) SMC’s strongly consistent estimates of expectations of test functions with respect to a normalized target distribution (Theorem 7.4.3 Del Moral [2004]), 2) NAS-X’s use of smoothing distributions as the intermediate targets, and 3) the twist optimality. 

This is not true for NASMC since it uses filtering SMC’s intermediate particle approximations (to avoid particle degeneracy), which approximate the filtering distributions and not the true posterior.  

Proposition II (informal): For any finite number of particles, NAS-X’s gradient estimates are unbiased, provided that the proposal equals the true posterior and twists are optimal. This is not true for NASMC. 

Sketch: 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 be distributed according to the filtering distributions instead of the true posterior. 

In the revision of this paper, we will include these statements and formal proofs. 

### Limitations 
NAS-X lacks a unified objective; the proposal and model are updated using different objectives. FIVO and SIXO have a unified objective for model and proposal. This didn’t have a practical impact on our experiments but could be relevant in other applications. NAS-X also does not work in online settings where future data is unavailable.


Meta-review: Based on unanimous approval from all reviewers, this paper is accepted. The authors effectively addressed concerns raised during the review process and provided comprehensive responses during the rebuttal stage.

In their rebuttal, the authors present new results showcasing substantial advancements over various existing methods. Their introduced approach, NAS-X, exhibits superior performance compared to NASMC, ELBO, IWAE, FIVO, SIXO, and RWS on synthetic (Gaussian SSM) and real-world (Hodgkin-Huxley) problems. They further support their claims by offering empirical and theoretical evidence on gradient variance and bias, as well as the computational efficiency of their method.

Additionally, the authors emphasize the robustness and simplicity of twist learning in their experiments, which continues to be effective beyond just the Gaussian SSM task. 

Furthermore, the authors highlight the practical benefits of NAS-X. It can readily replace SIXO, requiring only minimal code modifications, while significantly outperforming SIXO. Moreover, NAS-X allows for the handling of discrete latent variables, which SIXO fails to accomplish.

Overall, the paper's substantial improvements, backed by empirical and theoretical results, warrant its acceptance.
Decision: Accept (poster)