SlowGuess/NeurIPS_2023_OCR huggingface dataset:
- By using the approximate smoothing distributions as targets for proposal learning, NAS-X can learn proposals that match the true posterior margins, while NASMC and other baseline methods cannot. We illustrate this in Section 5.1, in a setting where the true posterior is tractable. We illustrate the practical benefits on inference in nonlinear mechanistic models in Section 5.3.
- By optimizing the proposal within the RWS framework (e.g., descending the inclusive KL), NAS-X can perform learning and inference in discrete latent variable models, which SIXO cannot. We explore this in Section 5.2.
- We explore the practical benefits of this combination in a challenging setting in Section 5.3, where we show NAS-X can fit ODE-based mechanistic models of neural dynamics with 38 model parameters and an 18-dimensional latent state.
In addition to these experiments, we analyze the computational complexity and wall-clock speed of each method and the bias and variance of the gradient estimates in Sections 15 and 16 of the Appendix.
5.1 Linear Gaussian State Space Model
We first consider a one-dimensional linear Gaussian state space model with joint distribution
$$
p \left(\mathbf {x} _ {1: T}, \mathbf {y} _ {1: T}\right) = \mathcal {N} \left(\mathbf {x} _ {1}; 0, \sigma_ {x} ^ {2}\right) \prod_ {t = 2} ^ {T} \mathcal {N} \left(\mathbf {x} _ {t + 1}; \mathbf {x} _ {t}, \sigma_ {x} ^ {2}\right) \prod_ {t = 1} ^ {T} \mathcal {N} \left(\mathbf {y} _ {t}; \mathbf {x} _ {t}, \sigma_ {y} ^ {2}\right). \tag {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,  $q(\mathbf{x}_{1:T}) = \prod_{t=1}^{T} q_t(\mathbf{x}_t) = \prod_{t=1}^{T} \mathcal{N}(\mathbf{x}_t; \mu_t, \sigma_t^2)$ , with parameters  $\mu_{1:T}$  and  $\sigma_{1:T}^2$  corresponding to the means and variances at each time-step. For twist-based methods, we parameterize the twist as a quadratic function in  $\mathbf{x}_t$  whose coefficients are functions of the observations and time step. We chose this form to match the functional form of the analytic log density ratio. For details, see Section 9 in the Appendix. NAS-X outperforms all baseline methods, achieving a tighter lower bound on the log-marginal likelihood and lower parameter error.
In the right panel of Figure 1, we compare the learned proposal variances against the true posterior variance, which can be computed in closed form. See Section 9 for comparison of proposal means; we do not report this comparison in the main text since all methods recover the posterior mean. This comparison gives insight into NAS-X's better performance. NASMC's learned proposal overestimates the posterior variance and fails to capture the true posterior distribution, because it employs a filtering
approximation to the gradients of the proposal distribution. On the other hand, by using the twisted targets, which approximate the smoothing distributions, to estimate proposal gradients, NAS-X recovers the true posterior.
5.2 Switching Linear Dynamical Systems
To explore NAS-X's ability to handle discrete latent variables, we consider a switching linear dynamical system (SLDS) model [40, 41]. Specifically, we adapt the recurrent SLDS example from Linderman et al. [41] in which the latent dynamics trace ovals in a manner that resembles cars racing on a NASCAR track. There are two coupled sets of latent variables: a discrete state  $\mathbf{z}_t$ , with  $K = 4$  possible values, and a two-dimensional continuous state  $\mathbf{x}_t$  that follows linear dynamics that depend on  $\mathbf{z}_t$ . The observations are a noisy projection of  $\mathbf{x}_t$  into a ten-dimensional observation space. There are 1000 observations in total. For the proposal, we factor  $q$  over both time and the continuous and discrete states. The continuous distributions are parameterized by Gaussians, and categorical distributions are used for the discrete latent variables. For additional details on the proposal and twist, see Section 10 in the Appendix and for details on the generative model see Linderman et al. [41].
(a) Comparing learned dynamics of NAS-X, NASMC, and Laplace EM on rSLDS.
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.
(b) Train  $\mathcal{L}_{\mathrm{BPF}}^{1024}$  for rSLDS.
Method	σ2O=0.001	σ2O=0.01	σ2O=0.1
NAS-X	19.837 ± 0.0234	8.63 ± 0.0015	-2.79 ± 0.0009
NASMC	19.834 ± 0.0018	8.53 ± 0.001	-2.874 ± 0.0007
Laplace EM	19.154 ± 0.057	8.54 ± 0.0039	-2.765 ± 0.0012
RWS	17.148 ± 0.087	6.314 ± 0.023	-5.78 ± 0.0026
We present qualitative results from model learning and inference in the top panel of Figure 2. We compare the learned dynamics for NAS-X, NASMC, and a Laplace-EM algorithm designed specifically for recurrent state space models [42]. In each panel, we plot the vector field of the learned dynamics and the posterior means, with different colors corresponding to the four discrete states. NAS-X recovers the true dynamics accurately. In the Table in Figure 2, we quantitatively compare the model learning performances across these three approaches by running a bootstrap proposal with the learned models and the true dynamics and observation variances. We normalize the bounds by the sequence length. NAS-X outperforms or performs on par with both NASMC and Laplace EM across the different observation noises  $\sigma_O^2$ . See Section 10, for additional results on inference.
5.3 Biophysical Models of Neuronal Dynamics
For our final set of experiments we consider inference and parameter learning in Hodgkin-Huxley (HH) models [9, 43] — mechanistic models of voltage dynamics in neurons. These models use systems of coupled nonlinear differential equations to describe the evolution of the voltage difference across a neuronal membrane as it changes in response to external stimuli such as injected current. Un
derstanding how voltage propagates throughout a cell is central to understanding electrical signaling and computation in the brain.
Voltage dynamics are governed by the flow of charged ions across the cell membrane, which is in turn mediated by the opening and closing of ion channels and pumps. HH models capture the states of these ion channels as well as the concentrations of ions and the overall voltage, resulting in a complex dynamical system with many free parameters and a high dimensional latent state space. Model learning and inference in this setting can be extremely challenging due to the dimensionality, noisy data, and expensive and brittle simulators that fail for many parameter settings.
Model Description We give a brief introduction to the models used in this section and defer a full description to the appendix. We are concerned with modeling the potential difference across a neuronal cell membrane,  $v$ , which changes in response to currents flowing through a set of ion channels,  $c \in \mathcal{C}$ . Each ion channel  $c$  has an activation which represents a percentage of the maximum current that can flow through the channel and is computed as a nonlinear function  $g_{c}$  of the channel state  $\lambda_{c}$ , with  $g_{c}(\lambda_{c}) \in [0,1]$ . This activation specifies the time-varying conductance of the channel as a fraction of the maximum conductance of the channel,  $\overline{g}_{c}$ . Altogether, the dynamics for the voltage  $v$  can be written as
$$
c _ {m} \frac {d v}{d t} = \frac {I _ {\mathrm {e x t}}}{S} - \sum_ {c \in \mathcal {C}} \bar {g} _ {c} g _ {c} \left(\lambda_ {c}\right) \left(v - E _ {c _ {\text {i o n}}}\right) \tag {11}
$$
where  $c_{m}$  is the specific membrane capacitance,  $I_{\mathrm{ext}}$  is the external current applied to the cell,  $S$  is the cell membrane surface area,  $c_{\mathrm{ion}}$  is the ion transported by channel  $c$ , and  $E_{c_{\mathrm{ion}}}$  is that ion's reversal potential. In addition to the voltage dynamics, the ion channel states  $\{\lambda_c\}_{c\in \mathcal{C}}$  evolve as
$$
\frac {d \lambda_ {c}}{d t} = A (v) \lambda_ {c} + b (v) \quad \forall c \in \mathcal {C} \tag {12}
$$
where  $A(v)$  and  $b(v)$  are nonlinear functions of the membrane potential that produce matrices and vectors, respectively. Together, equations (11) and (12) define a conditionally linear system of first-order ordinary differential equations (ODEs), meaning that the voltage dynamics are linear if the channel states are known and vice-versa.
Following Lawson et al. [19], we augment the deterministic dynamics with zero-mean, additive Gaussian noise to the voltage  $v$  and unconstrained gate states  $\mathrm{logit}(\lambda_c)$  at each integration time-step. The observations are produced by adding Gaussian noise to the instantaneous membrane potential.
**Proposal and Twist Parameterization** For all models in this section we amortize proposal and twist learning across datapoints. SIXO and NAS-X proposals used bidirectional recurrent neural networks (RNNs) [44, 45] with a hidden size of 64 units to process the raw observations and external current stimuli, and then fed the processed observations, previous latent state, and a transformer positional encoding [46] into a 64-unit single-layer MLP that produced the parameters of an isotropic Gaussian distribution over the current latent state. Twists were similarly parameterized with an RNN run in reverse across observations, combined with an MLP that accepts the RNN outputs and latent state and produces the twist values.
Integration via Strang Splitting The HH ODEs are stiff, meaning they are challenging to integrate at large step sizes because their state can change rapidly. While small step sizes can ensure numerical stability, they also make methods prohibitively slow. For example, many voltage dynamics of interest unfold over hundreds of milliseconds, which could take upwards of 40,000 integration steps at the standard 0.005 milliseconds per step. Because running our models for even 10,000 steps would be too costly, we developed new numerical integration techniques based on an explicit Strang splitting approach that allowed us to stably integrate at 0.1 milliseconds per step, a 20-time speedup [47]. For details, see Section 13 in the Appendix.
5.3.1 Hodgkin-Huxley Inference Results
First, we evaluated NAS-X, NASMC, and SIXO in their ability to infer underlying voltages and channel states from noisy voltage observations. For this task we sampled 10,000 noisy voltage traces from a probabilistic Hodgkin-Huxley model of the squid giant axon [9], and used each method to
(a) HH inference performance.
(b) Inferred voltage traces for SIXO and NAS-X.
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], we sampled trajectories of length 50 milliseconds, with a single noisy voltage observation every millisecond. The stability of the Strang splitting based ODE integrator allowed us to integrate at  $dt = 0.1\mathrm{ms}$ , meaning there were 10 latent states per observation.
In Figure (3a) we plot the performance of proposals and twists trained with 4 particles and evaluated across a range of particle numbers. All methods perform roughly the same when evaluated with 256 particles, but with lower numbers of evaluation particles the smoothing methods emerge as more particle-efficient than the filtering methods. To achieve NAS-X's inference performance with 4 particles, NASMC would need 256 particles, a 64x increase, and NAS-X is also on average 2x more particle-efficient than SIXO.
In Figure 3b we further investigate these results by examining the inferred voltage traces of NAS-X and SIXO. SIXO accurately infers the timing of most spikes but resamples at a high rate, which can lead to particle degeneracy and poor bound performance. NAS-X correctly infers the voltage across the whole trace with no spurious or mistimed spikes and almost no resampling events, indicating it has learned a high-quality proposal that does not generate poor particles that must be resampled away. These qualitative results support the quantitative results in Figure 3a: SIXO's high resampling rate and NASMC's filtering approach lead to lower bound values.
These results highlight a benefit of RWS-based methods over VI methods: when the model is correctly specified, it can be beneficial to have a more deterministic proposal. Empirically, we find that maximizing the variational lower bound encourages the proposal to have high entropy, which in this case resulted in SIXO's poorer performance relative to NAS-X. In the next section, we explore the implications of this on model learning.
5.3.2 Hodgkin-Huxley Model Learning Results
In this section, we assess NAS-X and SIXO's ability to fit model parameters in a more complex, biophysically realistic model of a pyramidal neuron from the mouse visual cortex. This model was taken from the Allen Institute Brain Atlas [48] and includes 9 different voltage-gated ion channels as well as a calcium pump/buffer subsystem and a calcium-gated potassium ion channel. In total, the model had 38 free parameters and an 18-dimensional latent state space, in contrast to the 1 free parameter and 4-dimensional state space of the model considered by Lawson et al. [19]. For full details of the models, see Appendix Section 12.
We fit these models to voltage traces gathered from a real mouse neuron by the Allen Institute, but downsampled and noised the data to simulate a more common voltage imaging setting. We ran a hyperparameter sweep over learning rates and initial values of the voltage and observation noise variances (270 hyperparameter settings in all), and selected the best performing model via early stopping on the train log marginal likelihood lower bound. Each hyperparameter setting was run for
Figure 4: Model learning in HH model of a mouse pyramidal neuron (top) Samples drawn from learned models when stimulated with a 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-trained models. All errors are absolute errors.
Method	\( \mathcal{L}_{\text{BPF}}^{32} \)	# Spikes Err.	Rest Voltage Err.	Cross-Corr.	% Runs Failed
NAS-X	-686.4 ± 6.8	0.76 ± 0.15	2.74 ± 0.1	6258 ± 11	18.9
SIXO	-660.6 ± 4.4	1.88 ± 0.41	1.8 ± 0.2	6055 ± 22	25.2
5 seeds, and each seed was run for 2 days on a single CPU core with  $7\mathrm{Gb}$  of memory. Because of the inherent instability of these models, many seeds failed, and we discarded hyperparameter settings with more than 2 failed runs.
In Figure 4 (bottom), we compare NAS-X and SIXO-trained models with respect to test set log-likelihood lower bounds as well as biophysically relevant metrics. To compute the biophysical metrics, we sampled 32 voltage traces for each input stimulus trace in the test set, and averaged the feature errors over the samples and test set. NAS-X better captures the number of spikes, an important feature of the traces, and attains a higher cross correlation. Both methods capture the resting voltage well, although SIXO attains a slightly lower error and outperforms NAS-X in terms of log-likelihood lower bound.
Training instability is a significant practical challenge when fitting mechanistic models. Therefore, we also include the percentage of runs that failed for each method. SIXO's more entropic proposals more frequently generate biophysically implausible latent states, causing the ODE integrator to return NaNs. In contrast, fewer of NAS-X's runs suffer from numerical instability issues, a great advantage when working with mechanistic models.
In this work we presented NAS-X, a new method for model learning and inference in sequential latent variable models, that combines reweighted wake-sleep framework and approximate smoothing SMC. Our approach involves learning twist functions to use in smoothing SMC, and then running smoothing SMC to approximate gradients of the log marginal likelihood with respect to the model parameters and gradients of the inclusive KL divergence with respect to the proposal parameters. We validated our approach in experiments including model learning and inference for discrete latent variable models and mechanistic models of neural dynamics, demonstrating that NAS-X offers compelling advantages in many settings.
This work was supported by the Simons Collaboration on the Global Brain (SCGB 697092), NIH (U19NS113201, R01NS113119, and R01NS130789), NSF (2223827), Sloan Foundation, McKnight Foundation, the Stanford Center for Human and Artificial Intelligence, and the Stanford Data Science Institute.
[1] Leonard E Baum and Ted Petrie. Statistical inference for probabilistic functions of finite state markov chains. The annals of mathematical statistics, 37(6):1554-1563, 1966.
[2] Rudolph Emil Kalman. A new approach to linear filtering and prediction problems. ASME Journal of Basic Engineering, 82(1):35-45, 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 on Learning Representations, 2014.
[5] Rahul G Krishnan, Uri Shalit, and David Sontag. Deep Kalman filters. arXiv preprint arXiv:1511.05121, 2015.
[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 to 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. Variational inference: A review for statisticians. Journal of the American Statistical Association, 112(518):859-877, 2017.
[12] Martin J. Wainwright and Michael I. Jordan. Graphical models, exponential families, and variational inference. Foundations and Trends® in Machine Learning, 1(1-2):1-305, 2008.
[13] Geoffrey E Hinton, Peter Dayan, Brendan J Frey, and Radford M Neal. The "wake-sleep" algorithm for unsupervised neural networks. Science, 268(5214):1158-1161, 1995.
[14] Jörg Bornschein and Yoshua Bengio. Reweighted Wake-Sleep. arXiv preprint arXiv:1406.2751, 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.
[52] AIBS. Biophysical modeling — perisomatic. Technical report, Allen Institute for Brain Science, 10 2017. URL http://help.brain-map.org/display/celltypes/Documentation.
[53] Costa M, Colbert and Enhui Pan. Ion channel properties underlying axonal action potential initiation in pyramidal neurons. Nature neuroscience, 5(6):533-538, 2002.
[54] Jacopo Magistretti and Angel Alonso. Biophysical properties and slow voltage-dependent inactivation of a sustained sodium current in entorhinal cortex layer-ii principal neurons: a whole-cell and single-channel study. The Journal of general physiology, 114(4):491-509, 1999.
[55] Maarten HP Kole, Stefan Hallermann, and Greg J Stuart. Single in channels in pyramidal neuron dendrites: properties, distribution, and impact on action potential output. Journal of Neuroscience, 26(6):1677-1687, 2006.
[56] I Reuveni, A Friedman, Y Amitai, and Michael J Gutnick. Stepwise repolarization from ca2+ plateaus in neocortical pyramidal cells: evidence for nonhomogeneous distribution of hva ca2+ channels in dendrites. Journal of Neuroscience, 13(11):4609-4621, 1993.
[57] Robert B 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. Neuropharmacology, 36(7):879-893, 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  $\mathbf{k}+$  channels in layer 5 neocortical pyramidal neurones from young rats: subtypes and gradients. The Journal of physiology, 525(3):621-639, 2000.
[61] M Kohler, B Hirschberg, CT Bond, John Mark Kinzie, NV Marrion, James Maylie, and JP Adelman. Small-conductance, calcium-activated potassium channels from mammalian brain. Science, 273(5282):1709-1714, 1996.
[62] Alastair J 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 a discrete distribution. The American Statistician, 33(4):214-218, 1979.
7 Theoretical Analyses
Proof. This is a direct application of Theorem 7.4.3 from Del Moral [25]. SMC methods provide strongly consistent estimates of expectations of test functions with respect to a normalized target distribution. That is, consider some test function  $h$  and let SMC's particle weights be denoted as  $w_{i}^{t}$ . We have that  $\sum_{i=1}^{K} w_{t}^{k} h(\mathbf{x}_{1:t}^{k}) \stackrel{a.s.}{\to} \int h(\mathbf{x}_{1:t}) \pi_{t}(\mathbf{x}_{1:t})$  as  $K \to \infty$  where  $\pi_{t}(\mathbf{x}_{1:t})$  is the normalized target distribution. NAS-X sets  $\pi_{t}(\mathbf{x}_{1:t}) \propto \gamma_{t}(\mathbf{x}_{1:t}) = p_{\theta}(\mathbf{x}_{1:t}, \mathbf{y}_{1:t}) r_{\psi^{*}}(\mathbf{y}_{t+1:T}, \mathbf{x}_{t})$ , which by assumption is proportional to  $p_{\theta}(\mathbf{x}_{1:t}, \mathbf{y}_{1:t}) p_{\theta}(\mathbf{y}_{t+1:T} \mid \mathbf{x}_{t}) = p_{\theta}(\mathbf{x}_{1:t}, \mathbf{y}_{1:T})$ . The desired result follows immediately.