SlowGuess/NeurIPS_2023_OCR huggingface dataset:
13 Strang Splitting for Hodgkin-Huxley Models
Because the Hodgkin-Huxley model is a stiff ODE, integrating it can be a challenge, especially at large step sizes. The traditional solution is to use an implicit integration scheme with varying step size, allowing the algorithm to take large steps when the voltage is not spiking. However, because our model adds noise to the ODE state at each timestep adaptive step-size methods are not viable as the different step sizes would change the noise distribution.
Instead, we sought an explicit, fixed step-size method that could be stably integrated at relatively large stepsizes. Inspired by Chen et al. [47], we developed a splitting approach that exploits the conditional linearity of the system. The system of ODEs describing the model can be split into two subsystems
Figure 10: Twist learning in LG-SSM. (left) Twist parameter error relative to optimal twist parameters 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  $\{\lambda_c\mid c\in \mathcal{C}\}$  is a system of linear first-order ODEs when conditioned on the voltage  $v$  and interior calcium concentration  $[\mathrm{Ca}^{2 + }]_i$ . Similarly, the dynamics for  $v$  and  $[\mathrm{Ca}^{2 + }]_i$  is a system of linear first-order ODEs when conditioned on the subunit states.
Because the conditional dynamics of each subsystem are linear first-order ODEs, an exact solution to each subsystem is possible under the assumption that the states being conditioned on are constant for the duration of the step. Our integration approach uses these exact updates in an alternating fashion, first performing an exact update to the voltage and interior calcium concentration while holding the subunit states constant, and then performing an exact update to the subunit states while holding the voltage and interior calcium concentration constant. For details on Strang and other splitting methods applied to Hodgkin-Huxley type ODEs, see [47].
14 Robustness of Twist Learning
NAS-X uses SIXO's twist learning framework to approximate the smoothing distributions. The twist learning approaches involves density ratio estimation. In brief, the density ratio estimation procedure involves training a classifier to distinguish between samples from  $p_{\theta}(\mathbf{x}_t \mid \mathbf{y}_{t+1:T})$  and  $p_{\theta}(\mathbf{x}_t)$ . These samples can be obtained from the generative model. For details see Section 2.2. In principle, incorporating twists complicates the overall learning problem and traditional methods for twist learning can indeed be challenging. However, in practice, twist learning using the SIXO framework is robust and easy. In Figure 10, we present twist parameter recovery and classification accuracy for the Gaussian SSM experiments (Section 5.1); in this setting, the optimal twists have a known parametric form. The optimal twist parameters are recovered quickly, the classification accuracy is high, and training is stable. This suggests that, with an appropriate twist parameterization, twist learning via density ratio estimation is tractable and straightforward.
15 Computational Complexity and Wall-clock Time
Theoretically, all multi-particle methods considered (NAS-X, SIXO, FIVO, NASMC, RWS, IWAE) have  $O(KT)$  time complexity, where  $K$  is the number of particles and  $T$  is the number of time steps. Once concern is that the resampling operation in SMC could require super-linear time in the number of particles, but drawing  $K$  samples from a  $K$ -category discrete distribution can be accomplished in  $O(K)$  time using the alias method [62, 63]. Additionally, for NAS-X, evaluating the twists is amortized across timesteps as in Lawson et al. [19], giving time linear in  $T$ .
NAS-X and SIXO have similar wall-clock speeds but are slower than FIFO and NASMC, primarily because of twist training, see Table 2. Even if FIFO 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.5\mathrm{x}$  longer per step than NASMC and FIVO and  $2.5\mathrm{x}$  longer per step than RWS and IWAE. However, Figure 7 shows that FIVO, NASMC, IWAE, and RWS cannot match NAS-X's performance even with 64 times more computation (4 vs. 256 particles). SIXO only matches NAS-X's performance with 4x as many particles (4 vs. 16 particles). Therefore, NAS-X uses computational resources much more effectively than other methods.
Method	ms / global step	ms / proposal step	ms / twist step
IWAE	70.3 ± 15.9	70.3 ± 15.9	N/A
RWS	71.6 ± 8.2	71.6 ± 8.2	N/A
NASMC	53.9 ± 6.6	53.9 ± 6.6	N/A
FIVO	51.4 ± 13.3	51.4 ± 13.3	N/A
NAS-X	163.2 ± 39.8	73.5 ± 18.5	89.7 ± 21.3
SIXO	175.3 ± 42.4	85.6 ± 21.1	89.7 ± 21.3
Table 2: Wall-clock speeds of various methods during HH inference.
Figure 11: Hodgkin-Huxley gradient variances (left) gradient bias (middle), and log-marginal likelihood lower bounds (right) over training.
16 Empirical analysis of bias and variance of gradients
In Figure 11, we analyze the gradient variance and bias for the Hodgkin-Huxley experiments, supplementing our theoretical analyses in Section 3.1. Figure 11 (left) shows NAS-X attains lower variance gradient estimates than IWAE, FIVO, and SIXO with comparable variance to RWS. We also studied the bias (middle) by approximating the true gradient by running a bootstrap particle filter with 256 particles using the best proposal from the inference experiments. NAS-X's gradients are lower bias than all methods but FIVO, but FIVO's gradients are also the highest variance. We hypothesize that FIVO's gradients appear less biased because its parameters are pushed towards degenerate values where gradient estimation is "easier". We illustrate this in the right panel, where we plot log-marginal likelihood bounds.