SlowGuess/NeurIPS_2023_OCR huggingface dataset:
NAS-  $\mathcal{X}$ : Neural Adaptive Smoothing via Twisting
Dieterich Lawson*†
dieterichl@google.com
Michael Y. Li*
michaelyli@stanford.edu
Scott W. Linderman
scott.linderman@stanford.edu
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.
Sequential latent variable models (SLVMs) are a foundational model class in statistics and machine learning, propelled by the success of hidden Markov models [1] and linear dynamical systems [2]. To model more complex data, SLVMs have incorporated nonlinear conditional dependencies, resulting in models such as sequential variational autoencoders [3-7], financial volatility models [8], and biophysical models of neural activity [9]. While these nonlinear dependencies make SLVMs more flexible, they also frustrate inference and model learning, motivating the search for approximate methods.
One popular method for inference in nonlinear SLVMs is sequential Monte Carlo (SMC), which provides a weighted particle approximation to the true posterior and an unbiased estimator of the marginal likelihood. For most SLVMs, SMC is a significant improvement over standard importance sampling, providing an estimator of the marginal likelihood with variance that grows linearly in the length of the sequence rather than exponentially [10]. SMC's performance, however, depends on having a suitable proposal. The optimal proposal is often intractable, so in practice, proposal parameters are learned from data. Broadly, there are two approaches to proposal learning: variational inference [11, 12] and reweighted wake-sleep [13, 14].
* Equal contribution. † Work performed while at Stanford University.
Variational inference (VI) methods for SLVMs optimize the model and proposal parameters by ascending a lower bound on the log marginal likelihood that can be estimated with SMC, an approach called variational SMC [15-17]. Recent advances in variational SMC have extended it to work with smoothing SMC [18, 19], a crucial development for models where future observations are strongly related to previous model states. Without smoothing SMC, particle degeneracy can cause high variance of the lower bound and its gradients, resulting in unstable learning [18, 10].
Reweighted wake-sleep (RWS) methods instead use SMC's posterior approximation to directly estimate gradients of the log marginal likelihood [14, 20]. This approach can be interpreted as descending the inclusive KL divergence from the true posterior to the proposal. Notably, RWS methods work with discrete latent variables, providing a compelling advantage over VI methods that typically resort to high-variance score function estimates for discrete latent variables.
In this work, we combine recent advances in smoothing variational SMC with the benefits of reweighted wake-sleep. To this end, we introduce neural adaptive smoothing via twisting (NAS-X), a method for inference and model learning in nonlinear SLVMs that uses smoothing SMC to approximate intractable posterior expectations. The result is a versatile, low-bias and low-variance estimator of the gradients of the log marginal likelihood suitable for fitting proposal and model parameters in both continuous and discrete SLVMs.
After introducing NAS-X, we present two theoretical results that highlight the advantages of NAS-X over other RWS-based methods. We also demonstrate NAS-X's performance empirically in model learning and inference in linear Gaussian state-space models, discrete latent variable models, and high-dimensional ODE-based mechanistic models of neural dynamics. In all experiments, we find that NAS-X substantially outperforms several VI and RWS alternatives including ELBO, IWAE, FIVO, SIXO. Furthermore, we empirically show that our method enjoys lower bias and lower variance gradients, requires minimal additional computational overhead, and is robust and easy to train.
This work considers model learning and inference in nonlinear sequential latent variable models with Markovian structure; i.e. models that factor as
$$
p _ {\boldsymbol {\theta}} \left(\mathbf {x} _ {1: T}, \mathbf {y} _ {1: T}\right) = p _ {\boldsymbol {\theta}} \left(\mathbf {x} _ {1}\right) p _ {\boldsymbol {\theta}} \left(\mathbf {y} _ {1} \mid \mathbf {x} _ {1}\right) \prod_ {t = 2} ^ {T} p _ {\boldsymbol {\theta}} \left(\mathbf {x} _ {t} \mid \mathbf {x} _ {t - 1}\right) p _ {\boldsymbol {\theta}} \left(\mathbf {y} _ {t} \mid \mathbf {x} _ {t}\right), \tag {1}
$$
with latent variables  $\mathbf{x}_{1:T} \in \mathcal{X}^T$ , observations  $\mathbf{y}_{1:T} \in \mathcal{Y}^T$ , and global parameters  $\pmb{\theta} \in \Theta$ . By nonlinear, we mean latent variable models where the parameters of the conditional distributions  $p_{\pmb{\theta}}(\mathbf{x}_t \mid \mathbf{x}_{t-1})$  and  $p_{\pmb{\theta}}(\mathbf{y}_t \mid \mathbf{x}_t)$  depend nonlinearly on  $\mathbf{x}_{t-1}$  and  $\mathbf{x}_t$ , respectively.
Estimating the marginal likelihood  $p_{\theta}(\mathbf{y}_{1:T})$  and posterior  $p_{\theta}(\mathbf{x}_{1:T} \mid \mathbf{y}_{1:T})$  for this model class is difficult because it requires computing an intractable integral over the latents,
$$
p _ {\boldsymbol {\theta}} (\mathbf {y} _ {1: T}) = \int_ {\mathcal {X} ^ {T}} p _ {\boldsymbol {\theta}} (\mathbf {y} _ {1: T}, \mathbf {x} _ {1: T}) d \mathbf {x} _ {1: T},
$$
We begin by introducing two algorithms, reweighted wake-sleep [14, 13] and smoothing sequential Monte Carlo [19], that are crucial for understanding our approach.
2.1 Reweighted Wake-Sleep
Reweighted wake-sleep (RWS, [13, 14]) is a method for maximum marginal likelihood in LVMs that estimates the gradients of the marginal likelihood using self-normalized importance sampling. This is motivated by Fisher's identity, which allows us to write the gradients of the marginal likelihood as a posterior expectation,
$$
\nabla_ {\theta} \log p _ {\theta} (\mathbf {y} _ {1: T}) = \mathbb {E} _ {p _ {\theta} \left(\mathbf {x} _ {1: T} \mid \mathbf {y} _ {1: T}\right)} \left[ \nabla_ {\theta} \log p _ {\theta} \left(\mathbf {x} _ {1: T}, \mathbf {y} _ {1: T}\right) \right], \tag {2}
$$
as proved in Appendix 8.1. The term inside the expectation is computable with modern automatic differentiation tools [21, 22], but the posterior  $p_{\theta}(\mathbf{x}_{1:T} \mid \mathbf{y}_{1:T})$  is unavailable. Thus, SNIS is used to form a biased but consistent Monte Carlo estimate of Eq. (2) [23]. Specifically, SNIS draws  $N$  IID
samples from a proposal distribution,  $q_{\phi}(\mathbf{x}_{1:T} \mid \mathbf{y}_{1:T})$  and weights them to form the estimator
$$
\sum_ {i = 1} ^ {N} \bar {w} ^ {(i)} \nabla_ {\theta} \log p _ {\theta} \left(\mathbf {x} _ {1: T} ^ {(i)}, \mathbf {y} _ {1: T}\right), \quad \mathbf {x} _ {1: T} ^ {(i)} \sim q _ {\phi} \left(\mathbf {x} _ {1: T} \mid \mathbf {y} _ {1: T}\right), \quad \bar {w} ^ {(i)} \propto \frac {p _ {\theta} \left(\mathbf {x} _ {1 : T} ^ {(i)} , \mathbf {y} _ {1 : T}\right)}{q _ {\phi} \left(\mathbf {x} _ {1 : T} ^ {(i)} \mid \mathbf {y} _ {1 : T}\right)} \tag {3}
$$
where  $\overline{w}^{(i)}$  are normalized weights, i.e.  $\sum_{i=1}^{N} \overline{w}^{(i)} = 1$ .
The variance of this estimator is reduced as  $q_{\phi}$  approaches the posterior [23], so RWS also updates  $q_{\phi}$  by minimizing the inclusive Kullback-Leibler (KL) divergence from the posterior to the proposal. Crucially, the gradient for this step can also be written as the posterior expectation
$$
\nabla_ {\phi} \mathrm {K L} \left(p _ {\theta} \left(\mathbf {x} _ {1: T} \mid \mathbf {y} _ {1: T}\right) \mid q _ {\phi} \left(\mathbf {x} _ {1: T} \mid \mathbf {y} _ {1: T}\right)\right) = - \mathbb {E} _ {p _ {\theta} \left(\mathbf {x} _ {1: T} \mid \mathbf {y} _ {1: T}\right)} \left[ \nabla_ {\phi} \log q _ {\phi} \left(\mathbf {x} _ {1: T} \mid \mathbf {y} _ {1: T}\right) \right], \tag {4}
$$
as derived in Appendix 8.2. This allows RWS to estimate Eq. (4) using SNIS with the same set of samples and weights as Eq. (3),
$$
- \mathbb {E} _ {p _ {\theta} (\mathbf {x} _ {1: T} | \mathbf {y} _ {1: T})} \left[ \nabla_ {\phi} \log q _ {\phi} \left(\mathbf {x} _ {1: T} \mid \mathbf {y} _ {1: T}\right) \right] \approx - \sum_ {i = 1} ^ {N} \bar {w} ^ {(i)} \nabla_ {\phi} \log q _ {\phi} \left(\mathbf {x} _ {1: T} ^ {(i)} \mid \mathbf {y} _ {1: T}\right). \tag {5}
$$
Importantly, any method that provides estimates of expectations w.r.t. the posterior can be used for gradient estimation within the RWS framework, as we will see in the next section.
2.2 Estimating Posterior Expectations with Smoothing Sequential Monte Carlo
As we saw in Eqs. (2) and (4), key quantities in RWS can be expressed as expectations under the posterior. Standard RWS uses SNIS to approximate these expectations, but in sequence models the variance of the SNIS estimator can scale exponentially in the sequence length. In this section, we review sequential Monte Carlo (SMC) [10, 24], an inference algorithm that can produce estimators of posterior expectations with linear or even sub-linear variance scaling.
SMC approximates the posterior  $p_{\theta}(\mathbf{x}_{1:T} \mid \mathbf{y}_{1:T})$  with a set of  $N$  weighted particles  $\mathbf{x}_{1:T}^{1:N}$  constructed by sampling from a sequence of target distributions  $\{\pi_t(\mathbf{x}_{1:t})\}_{t=1}^T$ . Since these intermediate targets are often only known up to some unknown normalizing constant  $Z_t$ , SMC uses the unnormalized targets  $\{\gamma_t(\mathbf{x}_{1:t})\}_{t=1}^T$ , where  $\pi_t(\mathbf{x}_{1:t}) = \gamma_t(\mathbf{x}_{1:t}) / Z_t$ . Provided mild technical conditions are met and  $\gamma_T(\mathbf{x}_{1:T}) \propto p_{\theta}(\mathbf{x}_{1:T}, \mathbf{y}_{1:T})$ , SMC returns weighted particles that approximate the posterior  $p_{\theta}(\mathbf{x}_{1:T} \mid \mathbf{y}_{1:T})$  [10, 24]. These weighted particles can be used to compute biased but consistent estimates of expectations under the posterior, similar to SNIS.
SMC repeats the following steps for each time  $t$ :
1. Sample latents  $\mathbf{x}_{1:t}^{1:N}$  from a proposal distribution  $q_{\phi}(\mathbf{x}_{1:t} \mid \mathbf{y}_{1:T})$ .
2. Weight each particle using the unnormalized target  $\gamma_{t}$  to form an empirical approximation  $\hat{\pi}_{t}$  to the normalized target distribution  $\pi_{t}$ .
3. Draw new particles  $\mathbf{x}_{1:t}^{1:N}$  from the approximation  $\hat{\pi}_t$  (the resampling step).
By resampling away latent trajectories with low weight and focusing on promising particles, SMC can produce lower variance estimates than SNIS. For a thorough review of SMC, see Doucet and Johansen [10], Naesseth et al. [24], and Del Moral [25].
Filtering vs. Smoothing The most common choice of unnormalized targets  $\gamma_{t}$  are the filtering distributions  $p_{\theta}(\mathbf{x}_{1:t},\mathbf{y}_{1:t})$ , resulting in the algorithm known as filtering SMC or a particle filter. Filtering SMC has been used to estimate posterior expectations within the RWS framework in neural adaptive sequential Monte Carlo (NASMC) [20], but a major disadvantage of filtering SMC is that it ignores future observations  $\mathbf{y}_{t + 1:T}$ . Ignoring future observations can lead to particle degeneracy and high-variance estimates, which in turn causes poor model learning and inference [17, 18, 26, 27].
We could avoid these issues by using the smoothing distributions as unnormalized targets, choosing  $\gamma_{t}(\mathbf{x}_{1:t}) = p_{\boldsymbol{\theta}}(\mathbf{x}_{1:t},\mathbf{y}_{1:T})$ , but unfortunately the smoothing distributions are not readily available from the model. We can approximate them, however, by observing that  $p_{\boldsymbol{\theta}}(\mathbf{x}_{1:t},\mathbf{y}_{1:T})$  is proportional to the filtering distributions  $p_{\boldsymbol{\theta}}(\mathbf{x}_{1:t},\mathbf{y}_{1:t})$  times the lookahead distributions  $p_{\boldsymbol{\theta}}(\mathbf{y}_{t+1:T} \mid \mathbf{x}_t)$ . If the lookahead distributions are well-approximated by a sequence of twists  $\{r_{\psi}(\mathbf{y}_{t+1:T},\mathbf{x}_t)\}_{t=1}^T$ , then running SMC with targets  $\gamma_{t}(\mathbf{x}_{1:t}) = p_{\boldsymbol{\theta}}(\mathbf{x}_{1:t},\mathbf{y}_{1:t}) r_{\psi}(\mathbf{y}_{t+1:T},\mathbf{x}_t)$  approximates smoothing SMC [26].
Twist Learning We have reduced the challenge of obtaining the smoothing distributions to learning twists that approximate the lookahead distributions. Previous twist-learning approaches include maximum likelihood training on samples from the model [18, 28] and Bellman-type losses motivated by writing the twist at time  $t$  recursively in terms of the twist at time  $t + 1$  [18, 29]. For NAS-X we use density ratio estimation (DRE) via classification to learn the twists, as introduced in Lawson et al. [19]. This method is motivated by observing that the lookahead distribution is proportional to a ratio of densities up to a constant independent of  $x_{t}$ ,
$$
p _ {\boldsymbol {\theta}} \left(\mathbf {y} _ {t + 1: T} \mid \mathbf {x} _ {t}\right) = \frac {p _ {\boldsymbol {\theta}} \left(\mathbf {x} _ {t} \mid \mathbf {y} _ {t + 1 : T}\right) p _ {\boldsymbol {\theta}} \left(\mathbf {y} _ {t + 1 : T}\right)}{p _ {\boldsymbol {\theta}} \left(\mathbf {x} _ {t}\right)} \propto \frac {p _ {\boldsymbol {\theta}} \left(\mathbf {x} _ {t} \mid \mathbf {y} _ {t + 1 : T}\right)}{p _ {\boldsymbol {\theta}} \left(\mathbf {x} _ {t}\right)}. \tag {6}
$$
Results from the DRE via classification literature [30] provide a way to approximate this density ratio: train a classifier to distinguish between samples from the numerator  $p_{\theta}(\mathbf{x}_t \mid \mathbf{y}_{t+1:T})$  and denominator  $p_{\theta}(\mathbf{x}_t)$ . Then, the pre-sigmoid output of the classifier will approximate the log of the ratio in Eq. (6). For an intuitive argument for this fact see Appendix 8.3, and for a full proof see Sugiyama et al. [30].
In practice, it is not possible to sample directly from  $p_{\theta}(\mathbf{x}_t \mid \mathbf{y}_{t+1:T})$ . Instead, Lawson et al. [19] sample full trajectories from the model's joint distribution, i.e., draw  $\mathbf{x}_{1:T}, \mathbf{y}_{1:T} \sim p_{\theta}(\mathbf{x}_{1:T}, \mathbf{y}_{1:T})$ , and discard unneeded timesteps, leaving only  $\mathbf{x}_t$  and  $\mathbf{y}_{t+1:T}$  which are distributed marginally as  $p_{\theta}(\mathbf{x}_t, \mathbf{y}_{t+1:T})$ . Training the DRE classifier on data sampled in this manner will approximate the ratio  $p_{\theta}(\mathbf{x}_t, \mathbf{y}_{t+1:T}) / p_{\theta}(\mathbf{x}_t) p_{\theta}(\mathbf{y}_{t+1:T})$ , which is equivalent to Eq. (6), see Appendix 8.3.
3 NAS-X: Neural Adaptive Smoothing via Twisting
The goal of NAS-X is to combine recent advances in smoothing SMC with the advantages of reweighted wake-sleep. Because SMC is a self-normalized importance sampling algorithm, it can be used to estimate posterior expectations and therefore the model and proposal gradients within a reweighted wake-sleep framework. In particular, NAS-X repeats the following steps:
1. Draw a set of  $N$  trajectories  $\mathbf{x}_{1:T}^{(1:N)}$  and weights  $\overline{w}_{1:T}^{(1:N)}$  from a smoothing SMC run with model  $p_{\theta}$ , proposal  $q_{\phi}$ , and twist  $r_{\psi}$ .
2. Use those trajectories and weights to form estimates of gradients for the model  $p_{\theta}$  and proposal  $q_{\phi}$ , as in reweighted wake-sleep. Specifically, NAS-X computes the gradients of the inclusive KL divergence for learning the proposal  $q_{\phi}$  as
$$
- \sum_ {t = 1} ^ {T} \sum_ {i = 1} ^ {N} \bar {w} _ {t} ^ {(i)} \nabla_ {\phi} \log q _ {\phi} \left(\mathbf {x} _ {t} ^ {(i)} \mid \mathbf {x} _ {t - 1} ^ {(i)}, \mathbf {y} _ {t: T}\right) \tag {7}
$$
and computes the gradients of the model  $p_{\theta}$  as
$$
\sum_ {t = 1} ^ {T} \sum_ {i = 1} ^ {N} \bar {w} _ {t} ^ {(i)} \nabla_ {\theta} \log p _ {\theta} \left(\mathbf {x} _ {t} ^ {(i)}, \mathbf {y} _ {t} \mid \mathbf {x} _ {t - 1} ^ {(i)}\right). \tag {8}
$$
3. Update the twists  $r_{\psi}$  using density ratio estimation via classification.
A full description is available in Algorithms 1 and 2.
A key design decision in NAS-X is the specific form of the gradient estimators. Smoothing SMC provides two ways to estimate expectations of test functions with respect to the posterior: both the timestep- $t$  and timestep- $T$  approximations of the target distribution could be used, in the latter case by discarding timesteps after  $t$ . Specifically,
$$
p _ {\theta} \left(\mathbf {x} _ {1: t} \mid \mathbf {y} _ {1: T}\right) \approx \sum_ {i = 1} ^ {N} \bar {w} _ {t} ^ {(i)} \delta \left(\mathbf {x} _ {1: t}; \mathbf {x} _ {1: t} ^ {(i)}\right) \approx \sum_ {i = 1} ^ {N} \bar {w} _ {T} ^ {(i)} \delta \left(\mathbf {x} _ {1: t}; \left(\mathbf {x} _ {1: T} ^ {(i)}\right) _ {1: t}\right) \tag {9}
$$
where  $\delta(a; b)$  is a Dirac delta of  $a$  located at  $b$  and  $(\mathbf{x}_{1:T})_{1:t}$  denotes selecting the first  $t$  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, a key advantage that we explore theoretically in Section 3.1 and empirically in Section 5.
Algorithm 1: NAS-X
Procedure NAS-X  $(\theta_0,\phi_0,\psi_0,\mathbf{y}_{1:T})$ $\theta \leftarrow \theta_{0},\quad \phi \leftarrow \phi_{0},\quad \psi \leftarrow \psi_{0}$    
while not converged do   
 $\begin{array}{r}\mathbf{x}_{1:T}^{1:N},\overline{\boldsymbol{w}}_{1:T}^{1:N}\gets \mathrm{SMC}(\{p_{\theta}(\mathbf{x}_{1:t},\mathbf{y}_{1:t}),q_{\phi}(\mathbf{x}_t|\mathbf{x}_{t - 1},\mathbf{y}_{t:T}),r_{\psi}(\mathbf{x}_t,\mathbf{y}_{t + 1:T})\}_{t = 1}^T)\\ \Delta \theta = \sum_{t = 1}^{T}\sum_{i = 1}^{N}\overline{w}_{t}^{(i)}\nabla_{\theta}\log p_{\theta}(\mathbf{x}_{t}^{(i)},\mathbf{y}_{t}|\mathbf{x}_{t - 1}^{(i)})\\ \Delta \phi = -\sum_{t = 1}^{T}\sum_{i = 1}^{N}\overline{w}_{t}^{(i)}\nabla_{\phi}\log q_{\phi}(\mathbf{x}_{t}^{(i)}|\mathbf{x}_{t - 1}^{(i)},\mathbf{y}_{t:T})\\ \theta \leftarrow \mathrm{grad - step}(\theta ,\Delta \theta)\\ \phi \leftarrow \mathrm{grad - step}(\phi ,\Delta \phi)\\ \psi \leftarrow \mathrm{twist - training}(\theta ,\psi) \end{array}$    
end   
return  $\theta ,\phi ,\psi$    
Procedure twist-training  $(\theta ,\psi_0)$    
See Algorithm 2 in Appendix 8.3.
3.1 Theoretical Analysis of NAS-X
In this section, we state two theoretical results that illustrate NAS-X's advantages over NASMC, with proofs given in Appendix 7.
Proposition 1. Consistency of NAS-X's gradient estimates. Suppose the twists are optimal so that  $r_{\psi}(\mathbf{y}_{t + 1:T},\mathbf{x}_t)\propto p(\mathbf{y}_{t + 1:T}\mid \mathbf{x}_t)$  up to a constant independent of  $\mathbf{x}_t$  for  $t = 1,\ldots ,T - 1$ . Let  $\hat{\nabla}_{\theta}\log p_{\theta}(\mathbf{y}_{1:T})$  be NAS-X's weighted particle approximation to the true gradient of the log marginal likelihood  $\nabla_{\theta}\log p_{\theta}(\mathbf{y}_{1:T})$ . Then  $\hat{\nabla}_{\theta}\log p_{\theta}(\mathbf{y}_{1:T})\stackrel {a.s.}{\to}\nabla_{\theta}\log p_{\theta}(\mathbf{y}_{1:T})$  as  $N\rightarrow \infty$ .
Proposition 2. Unbiasedness of NAS-X's gradient estimates. Assume that proposal distribution  $q_{\phi}(\mathbf{x}_t \mid \mathbf{x}_{1:t-1}, \mathbf{y}_{1:T})$  is optimal so that  $q_{\phi}(\mathbf{x}_t \mid \mathbf{x}_{1:t-1}, \mathbf{y}_{1:T}) = p(\mathbf{x}_t \mid \mathbf{x}_{1:t-1}, \mathbf{y}_{1:T})$  for  $t = 1, \ldots, T$ , and the twists  $r_{\psi}(\mathbf{y}_{t+1:T}, \mathbf{x}_t)$  are optimal so that  $r_{\psi}(\mathbf{y}_{t+1:T}, \mathbf{x}_t) \propto p(\mathbf{y}_{t+1:T} \mid \mathbf{x}_t)$  up to a constant independent of  $\mathbf{x}_t$  for  $t = 1, \ldots, T-1$ . Let  $\hat{\nabla}_{\theta} \log p_{\theta}(\mathbf{y}_{1:T})$  be NAS-X's weighted particle approximation to the true gradient of the log marginal likelihood,  $\nabla_{\theta} \log p_{\theta}(\mathbf{y}_{1:T})$ . Then, for any number of particles,  $\mathbb{E}[\hat{\nabla}_{\theta} \log p_{\theta}(\mathbf{y}_{1:T})] = \log p_{\theta}(\mathbf{y}_{1:T})$ .
VI Methods There is a large literature on model learning via stochastic gradient ascent on an evidence lower bound (ELBO) [31, 32, 4, 33]. Subsequent works have considered ELBOs defined by the normalizing constant estimates from multiple importance sampling [34], nested importance sampling, [35, 36], rejection sampling, and Hamiltonian Monte Carlo [37]. Most relevant to our work is the literature that uses SMC's estimates of the normalizing constant as a surrogate objective. There are a number of VI methods based on filtering [17, 16, 15] and smoothing SMC [18, 38, 39, 19, 27], but filtering SMC approaches can suffer from particle degeneracy and high variance [18, 19].
Reweighted Wake-Sleep Methods The wake-sleep algorithm was introduced in Hinton et al. [13] as a way to train deep directed graphical models. Bornschein and Bengio [14] interpreted the wake-sleep algorithm as self-normalized importance sampling and proposed reweighted 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 approximate posterior expectations instead of SNIS [20]. To combat particle degeneracy, NASMC approximates the posterior with the filtering distributions, which introduces bias.
NAS-X vs. SIXO Both NAS-X and SIXO [19] leverage smoothing SMC with DRE-learned twists, but NAS-X uses smoothing SMC to estimate gradients in an RWS-like framework while SIXO uses it within a VI-like framework. Thus, NAS-X follows biased but consistent estimates of the log marginal likelihood while SIXO follows unbiased estimates of a lower bound on the log marginal likelihood. It is not clear a-priori which approach would perform better, but we provide empirical evidence in Section 5 that shows NAS-X is more stable than SIXO and learns better models and proposals. In
Figure 1: Comparison of NAS-X and 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 variances. NAS-X outperforms all baseline methods and recovers the true posterior margins.
addition to these empirical advantages, NAS-X can fit discrete latent variable models while SIXO would require high-variance score function estimates.
We empirically validate the following advantages of NAS-X: