SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Proof. We first provide a proof sketch then give a more detailed derivation, adapting the proof of a similar theorem in Lawson et al. [19]. We will prove that particles produced from running SMC with the smoothing targets and optimal proposal are exact samples from the posterior distribution of interest. The claim then follows immediately. Under the stated assumptions, both NAS-X and NASMC propose particles from the true posterior. However, the different intermediate target distributions will affect how these particles are distributed after reweighting. For NAS-X, the particles will have equal weight since they are reweighted using the smoothing targets. Thus, after reweighting, the particles are still samples from the true posterior. In contrast, in NASMC, the samples are reweighted by filtering targets and will not be distributed as samples from the true posterior.
We first consider NAS-X which uses the smoothing targets. We will show that the particles drawn at each timestep are sampled from the posterior distribution. This will follow from the fact that the particle weights at each timestep in the SMC sweep are equal; that is,  $w_{t}^{k} = 1$  or  $w_{t}^{k} = p(\mathbf{y}_{1:T})$  for  $k = 1,\dots ,K$ , depending on whether resampling occurred. We proceed by induction on  $t$ , the timestep in the SMC sweep.
For  $t = 1$  note that
1.  $\gamma_{1}(\mathbf{x}_{1}^{k}) = p(\mathbf{x}_{1}^{k},\mathbf{y}_{1})p(\mathbf{y}_{2:T}\mid \mathbf{x}_{1}^{k}),$
2.  $\gamma_0\triangleq 1$
3. and  $q_{1}(\mathbf{x}_{1}^{k}) = p(\mathbf{x}_{1}^{k}\mid \mathbf{y}_{1:T})$
This implies that the incremental weight  $\alpha_{1}^{k}$  is
$$
\alpha_ {1} ^ {k} = \frac {p \left(\mathbf {x} _ {1} ^ {k} , \mathbf {y} _ {1}\right) p \left(\mathbf {y} _ {2 : T} \mid \mathbf {x} _ {1} ^ {k}\right)}{p \left(\mathbf {x} _ {1} ^ {k} \mid \mathbf {y} _ {1 : T}\right)} = \frac {p \left(\mathbf {x} _ {1} ^ {k} , \mathbf {y} _ {1 : T}\right)}{p \left(\mathbf {x} _ {1} ^ {k} \mid \mathbf {y} _ {1 : T}\right)} = p \left(\mathbf {y} _ {1: T}\right) \tag {13}
$$
which does not depend on  $k$ . Because  $w_0^k \triangleq 1$ , we have that  $w_1^k = w_0^k \alpha_1^k = p(\mathbf{y}_{1:T})$  for all  $k$ .
Note that, since the proposal is optimal, prior to SMC's reweighting step, the particles were distributed as follows  $\mathbf{x}_1^k\sim p(\mathbf{x}_1\mid \mathbf{y}_{1:T})$  . Since the incremental weights are equal, the distribution of particles is unchanged.
For the induction step, assume that  $w_{t-1}^{1:K}$  equals 1 or  $p(\mathbf{y}_{1:T})$  and that  $\mathbf{x}_{1:t-1}^k \sim p(\mathbf{x}_{1:t-1} \mid \mathbf{y}_{1:T})$ . The particles are distributed as follows,  $\mathbf{x}_t^k \sim p(\mathbf{x}_t \mid \mathbf{x}_{1:t-1}, \mathbf{y}_{1:T})$ . This implies that  $(\mathbf{x}_{1:t-1}^k, \mathbf{x}_t^k) \sim p(\mathbf{x}_t \mid \mathbf{x}_{1:t-1}, \mathbf{y}_{1:T}) p(\mathbf{x}_{1:t-1} \mid \mathbf{y}_{1:T}) = p(\mathbf{x}_{1:t} \mid \mathbf{y}_{1:T})$ .
We now show that the incremental particle weights  $\alpha_{t}^{k}$  are equal using the following identities/assumptions
1.  $\gamma_{t}(\mathbf{x}_{1:t}^{k}) = p(\mathbf{x}_{1:t}^{k},\mathbf{y}_{1:t})p(\mathbf{y}_{t + 1:T}\mid \mathbf{x}_{t}^{k}),$
2.  $\gamma_{t - 1}(\mathbf{x}_{1:t - 1}^k) = p(\mathbf{x}_{1:t - 1}^k,\mathbf{y}_{1:t - 1})p(\mathbf{y}_{t:T}|\mathbf{x}_{t - 1}^k),$
3.  $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,\dots ,T - 1$
4. and  $q_{t}(\mathbf{x}_{t}^{k}) = p(\mathbf{x}_{t}^{k} \mid \mathbf{x}_{1:t-1}^{k}, \mathbf{y}_{1:T})$
This shows that  $\alpha_{t}^{k}$  is given by
$$
\begin{array}{l} \alpha_ {t} ^ {k} = \frac {p \left(\mathbf {x} _ {1 : t} ^ {k} , \mathbf {y} _ {1 : t}\right) p \left(\mathbf {y} _ {t + 1 : T} \mid \mathbf {x} _ {t} ^ {k}\right)}{p \left(\mathbf {x} _ {1 : t - 1} ^ {k} , \mathbf {y} _ {1 : t - 1}\right) p \left(\mathbf {y} _ {t : T} \mid \mathbf {x} _ {t - 1} ^ {k}\right) p \left(\mathbf {x} _ {t} ^ {k} \mid \mathbf {x} _ {1 : t - 1} ^ {k} , \mathbf {y} _ {1 : T}\right)} (14) \\ = \frac {p \left(\mathbf {x} _ {1 : t} ^ {k} , \mathbf {y} _ {1 : T}\right)}{p \left(\mathbf {x} _ {1 : t - 1} ^ {k} , \mathbf {y} _ {1 : T}\right) p \left(\mathbf {x} _ {t} ^ {k} \mid \mathbf {x} _ {1 : t - 1} ^ {k} , \mathbf {y} _ {1 : T}\right)} (15) \\ = \frac {p \left(\mathbf {x} _ {1 : t - 1} ^ {k} , \mathbf {y} _ {1 : T}\right) p \left(\mathbf {x} _ {t} ^ {k} \mid \mathbf {x} _ {1 : t - 1} , \mathbf {y} _ {1 : T}\right)}{p \left(\mathbf {x} _ {1 : t - 1} ^ {k} , \mathbf {y} _ {1 : T}\right) p \left(\mathbf {x} _ {t} ^ {k} \mid \mathbf {x} _ {1 : t - 1} ^ {k} , \mathbf {y} _ {1 : T}\right)} (16) \\ = 1 (17) \\ \end{array}
$$
for  $k = 1,\ldots ,K$
There are two cases depending on the value of the weights at the previous timestep. If  $w_{t-1}^{1:K} = 1$ , then  $w_t^k = w_{t-1}^k\alpha_t^k = 1$  for all  $k$ . On the other hand, if  $w_{t-1}^{1:K} = p(\mathbf{y}_{1:T})$  then  $w_t^k = p(\mathbf{y}_{1:T})$  for all  $k$ . Therefore, even after reweighting, the particles are still drawn from the true posterior distribution. If resampling occurs, since the weights are equal in both cases, the distribution of the particles remains unchanged.
To conclude, we note that the incremental particle weights are, in general, not the same for NASMC. To see this, consider NASMC's incremental weights at timestep 1.
$$
\alpha_ {1} ^ {k} = \frac {p \left(\mathbf {x} _ {1} ^ {k} , \mathbf {y} _ {1}\right)}{p \left(\mathbf {x} _ {t} ^ {k} \mid \mathbf {y} _ {1 : T}\right)} \tag {18}
$$
After reweighting, the particles will be distributed according to the filtering distributions. The distribution of particles after reweighting is proportional to  $p(\mathbf{x}_1^k \mid \mathbf{y}_1)$ . This is because the distribution of the reweighted particles is proportional to the incremental weights times the optimal proposal distribution. Therefore, the term in the denominator corresponding to the proposal cancels out with the proposal distribution term.
Interestingly, the particle weights will be the same at each iteration under a certain dependency structure for the model  $p(\mathbf{x}_{1:t-1}|\mathbf{y}_{1:t}) = p(\mathbf{x}_{1:t-1}|\mathbf{y}_{1:t-1})$  that was identified in Maddison et al. [17]. However, this dependency is not satisfied in general and therefore NASMC's gradients are not unbiased.
8 Derivations
8.1 Gradient of the Marginal Likelihood
We derive the gradients for the marginal likelihood. This identity is known as Fisher's identity.
$$
\begin{array}{l} \nabla_ {\theta} \log p (\mathbf {y} _ {1: T}) = \nabla_ {\theta} \log \int p _ {\theta} (\mathbf {x} _ {1: T}, \mathbf {y} _ {1: T}) d \mathbf {y} _ {1: T} \\ = \frac {1}{p _ {\theta} (\mathbf {y} _ {1 : T})} \nabla_ {\theta} \int p _ {\theta} (\mathbf {x} _ {1: T}, \mathbf {y} _ {1: T}) d \mathbf {x} _ {1: T} \\ = \frac {1}{p _ {\theta} (\mathbf {y} _ {1 : T})} \int \nabla_ {\theta} p _ {\theta} (\mathbf {x} _ {1: T}, \mathbf {y} _ {1: T}) d \mathbf {x} _ {1: T} \\ = \frac {1}{p _ {\theta} (\mathbf {y} _ {1 : T})} \int p _ {\theta} (\mathbf {x} _ {1: T}, \mathbf {y} _ {1: T}) \nabla_ {\theta} \log p _ {\theta} (\mathbf {x} _ {1: T}, \mathbf {y} _ {1: T}) d \mathbf {x} _ {1: T} \\ = \int p _ {\theta} (\mathbf {x} _ {1: T} | \mathbf {y} _ {1: T}) \nabla_ {\theta} \log p _ {\theta} (\mathbf {x} _ {1: T}, \mathbf {y} _ {1: T}) d \mathbf {x} _ {1: T} \\ = \int p _ {\theta} (\mathbf {x} _ {1: T} | \mathbf {y} _ {1: T}) \nabla_ {\theta} \sum_ {t} \log p _ {\theta} (\mathbf {y} _ {t}, \mathbf {x} _ {t} | \mathbf {x} _ {t - 1}) d \mathbf {x} _ {1: T} \\ = \sum_ {t} \int p _ {\theta} (\mathbf {x} _ {1: T} | \mathbf {y} _ {1: T}) \nabla_ {\theta} \log p _ {\theta} (\mathbf {y} _ {t}, \mathbf {x} _ {t} | \mathbf {x} _ {t - 1}) d \mathbf {x} _ {1: T} \\ = \sum_ {t} \mathbb {E} _ {p _ {\theta} \left(\mathbf {x} _ {1: T} \mid \mathbf {y} _ {1: T}\right)} \left[ \nabla_ {\theta} \log p _ {\theta} \left(\mathbf {y} _ {t}, \mathbf {x} _ {t} \mid \mathbf {x} _ {t - 1}\right) \right] \\ \end{array}
$$
The key steps were the log-derivative trick and Bayes rule.
8.2 Gradient of Inclusive KL Divergence
Below, we derive the gradient of the inclusive KL divergence for a generic Markovian model. In this derivation, we assume there are no shared parameters between the proposal and model.
$$
\begin{array}{l} - \nabla_ {\phi} \mathrm {K L} (p _ {\theta} | | q _ {\phi}) = \nabla_ {\phi} \int p _ {\theta} (\mathbf {x} _ {1: T} | \mathbf {y} _ {1: T}) \log q _ {\phi} (\mathbf {x} _ {1: T} | \mathbf {y} _ {1: T}) d \mathbf {x} _ {1: T} \\ = \int p _ {\theta} (\mathbf {x} _ {1: T} | \mathbf {y} _ {1: T}) \nabla_ {\phi} \log q _ {\phi} (\mathbf {x} _ {1: T} | \mathbf {y} _ {1: T}) d \mathbf {x} _ {1: T} \\ = \int p _ {\theta} (\mathbf {x} _ {1: T} | \mathbf {y} _ {1: T}) \nabla_ {\phi} \left(\sum_ {t} \log q _ {\phi} (\mathbf {x} _ {t} | \mathbf {x} _ {t - 1}, \mathbf {y} _ {t: T})\right) d \mathbf {x} _ {1: T} \\ = \sum_ {t} \int p _ {\theta} (\mathbf {x} _ {1: T} | \mathbf {y} _ {1: T}) \nabla_ {\phi} \log q _ {\phi} (\mathbf {x} _ {t} | \mathbf {x} _ {t - 1}, \mathbf {y} _ {t: T}) d \mathbf {x} _ {1: T} \\ = \sum_ {t} \mathbb {E} _ {p _ {\theta} (\mathbf {x} _ {1: T} | \mathbf {y} _ {1: T})} \left[ \nabla_ {\phi} \log q _ {\phi} (\mathbf {x} _ {t} | \mathbf {x} _ {t - 1}, \mathbf {y} _ {t: T}) \right] \\ \end{array}
$$
We use the assumption that there are no shared parameters in the second equality.
8.3 Density Ratio Estimation via Classification
Here we briefly summarize density ratio estimation (DRE) via classification. For a full treatment, see Sugiyama et al. [30].
Let  $a(x)$  and  $b(x)$  be two distributions defined over the same space  $\mathcal{X}$ , and consider a classifier  $g: \mathcal{X} \to \mathbb{R}$  that accepts a specific  $x \in \mathcal{X}$  and classifies it as either being sampled from  $a$  or  $b$ . We will train this classifier to predict whether a given  $x$  was sampled from  $a(x)$  or  $b(x)$ . The raw outputs (logits) of this classifier will approximate  $\log(a(x) / b(x))$  up to a constant that does not depend on  $x$ .
To see this, define an expanded generative model where we first sample  $z \in \{0,1\}$  from a Bernoulli random variable with probability  $0 < \rho < 1$ , and then sample  $x$  from  $a(x)$  if  $z = 1$ , and sample  $x$
from  $b(x)$  if  $z = 0$ . This defines the joint distribution
$$
p (x, z) = p (z) p (x | z) = \operatorname {B e r n o u l l i} (z; \rho) a (x) ^ {z} b (x) ^ {(1 - z)}, \tag {19}
$$
where  $p(x\mid z = 1) = a(x)$  and  $p(x\mid z = 0) = b(x)$ .
Let  $g: \mathcal{X} \to \mathbb{R}$  be a function that accepts  $x \in \mathcal{X}$  and produces the logit for Bernoulli distribution over  $z$ . This function will parameterize a classifier via the sigmoid function, meaning that the classifier's Bernoulli conditional distribution over  $z$  is defined as
$$
p _ {g} (z | x) \triangleq \sigma (g (x)) ^ {z} (1 - \sigma (g (x))) ^ {1 - z}, \quad z \in \{1, 0 \} \tag {20}
$$
where  $\sigma$  is the sigmoid function and  $\sigma^{-1}$  is its inverse, the logit function
$$
\sigma (\ell) = \frac {1}{1 + e ^ {- \ell}}, \quad \sigma^ {- 1} (p) = \log \left(\frac {p}{1 - p}\right). \tag {21}
$$
The optimal function  $g^{*}$  will be selected by solving the maximum likelihood problem
$$
g ^ {*} \triangleq \arg \max  _ {g} \mathbb {E} _ {p (x, z)} [ p _ {g} (z \mid x) ]. \tag {22}
$$
The solution to this problem is the true  $p(z \mid x)$ . Because we have not restricted  $g$ , this solution can be obtained. Thus,
$$
\begin{array}{l} p (z = 1 \mid x) = p _ {g ^ {*}} (z = 1 \mid x) (23) \\ = \sigma \left(g ^ {*} (x)\right) ^ {1} \left(1 - \sigma \left(g ^ {*} (x)\right)\right) ^ {1 - 1} (24) \\ = \sigma \left(g ^ {*} (x)\right). (25) \\ \end{array}
$$
This in turn implies that
$$
\begin{array}{l} g ^ {*} (x) = \sigma^ {- 1} (p (z = 1 \mid x)) (26) \\ = \log \left(\frac {p (z = 1 \mid x)}{1 - p (z = 1 \mid x)}\right) (27) \\ = \log \left(\frac {p (z = 1 \mid x)}{p (z = 0 \mid x)}\right) (28) \\ = \log \left(\frac {p (z = 1 \mid x) p (x)}{p (z = 0 \mid x) p (x)}\right) (29) \\ = \log \left(\frac {p (z = 1 , x)}{p (z = 0 , x)}\right) (30) \\ = \log \left(\frac {p (x \mid z = 1) p (z = 1)}{p (x \mid z = 0) p (z = 0)}\right) (31) \\ = \log \left(\frac {p (x \mid z = 1)}{p (x \mid z = 0)}\right) + \log \left(\frac {p (z = 1)}{p (z = 0)}\right) (32) \\ = \log \left(\frac {a (x)}{b (x)}\right) + \log \left(\frac {\rho}{1 - \rho}\right) (33) \\ \end{array}
$$
Thus, the optimal solution to the classification problem,  $g^{*}$ , is proportional to  $\log (a(x) / b(x))$  up to a constant that does not depend on  $x$ . In practice we observe empirically that as long as a sufficiently flexible parametric family for  $g$  is selected,  $g^{*}$  will closely approximate the desired density ratio.
In the case of learning the ratio required for smoothing SMC,
$$
\frac {p _ {\boldsymbol {\theta}} \left(\mathbf {x} _ {t} \mid \mathbf {y} _ {t + 1 : T}\right)}{p _ {\boldsymbol {\theta}} \left(\mathbf {x} _ {t}\right)}, \tag {34}
$$
Lawson et al. [19] instead learn the equivalent ratio
$$
\frac {p _ {\boldsymbol {\theta}} \left(\mathbf {x} _ {t} , \mathbf {y} _ {t + 1 : T}\right)}{p _ {\boldsymbol {\theta}} \left(\mathbf {x} _ {t}\right) p _ {\boldsymbol {\theta}} \left(\mathbf {y} _ {t + 1 : T}\right)}. \tag {35}
$$
Figure 5: Comparison of NAS-X vs 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 margins.
As per the previous derivation, it suffices to train a classifier to distinguish between samples from the numerator and denominator of Eq. 35. To accomplish this, Lawson et al. [19] draw paired and unpaired samples from the model that are distributed marginally according to the desired densities. Specifically, consider drawing
$$
\mathbf {x} _ {1: T}, \mathbf {y} _ {1: T} \sim p _ {\theta} \left(\mathbf {x} _ {1: T}, \mathbf {y} _ {1: T}\right) \quad \tilde {\mathbf {x}} _ {1: T} \sim p _ {\theta} \left(\mathbf {x} _ {1: T}\right) \tag {36}
$$
and note that any  $\mathbf{x}_t, \mathbf{y}_{t+1:T}$  selected from the sample will be distributed marginally according to  $p_\theta(\mathbf{x}_t, \mathbf{y}_{t+1:T})$ . Similarly, any  $\tilde{\mathbf{x}}_t, \mathbf{y}_{t+1:T}$  will be distributed marginally as  $p_\theta(\mathbf{x}_t)p_\theta(\mathbf{y}_{t+1:T})$ . In this way,  $T - 1$  positive and negative training examples for the DRE classifier can be drawn using a single set of samples as in Eq. (36).
The twist training process is summarized in Algorithm 2.
Algorithm 2: Twist Training
Procedure twist-training  $(\theta ,\psi_0)$ $\psi \gets \psi_{0}$    
while not converged do   
 $\begin{array}{rl} & {\mathbf{x}_{1:T},\mathbf{y}_{1:T}\sim p_{\boldsymbol{\theta}}(\mathbf{x}_{1:T},\mathbf{y}_{1:T})}\\ & {\tilde{\mathbf{x}}_{1:T}\sim p_{\boldsymbol{\theta}}(\mathbf{x}_{1:T})}\\ & {\mathcal{L}(\psi) = \frac{1}{T - 1}\sum_{t = 1}^{T - 1}\log \sigma (r_{\psi}(\mathbf{x}_t,\mathbf{y}_{t + 1:T})) + \log (1 - \sigma (r_{\psi}(\tilde{\mathbf{x}}_t,\mathbf{y}_{t + 1:T})))}\\ & {\psi \leftarrow \mathrm{grad - step}(\psi ,\nabla_{\psi}\mathcal{L}(\psi))} \end{array}$    
end   
return  $\psi$
9 LGSSM
Model Details We 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 {37}
$$
In our experiments we set the dynamics variance  $\sigma_x^2 = 1.0$  and the observation variance  $\sigma_y^2 = 1.0$ .
Proposal Parameterization For both NAS-X and NASMC, we use a mean-field Gaussian proposal factored over time
$$
q \left(x _ {1: T}\right) = \prod_ {t = 1} ^ {T} q _ {t} \left(x _ {t}\right) = \prod_ {t = 1} ^ {T} \mathcal {N} \left(x _ {t}; \mu_ {t}, \sigma_ {t} ^ {2}\right), \tag {38}
$$
with parameters  $\mu_{1:T}$  and  $\sigma_{1:T}^2$  corresponding to the means and variances at each timestep. In total, we learn  $2T$  proposal parameters.
Twist Parametrization We parameterize the twist as a quadratic function in  $x_{t}$  whose coefficients are functions of the observations and time step and are learned via the density ratio estimation procedure described in [19]. We chose this form to match the analytic log density ratio for the model defined in Eq 10. Given that  $p(x_{1:T}, y_{1:T})$  is a multivariate Gaussian, we know that  $p(x_{t} \mid y_{t+1:T})$  and  $p(x_{t})$  are both marginally Gaussian. Let
$$
p (x _ {t} \mid y _ {t + 1: T}) \triangleq \mathcal {N} (\mu_ {1}, \sigma_ {1} ^ {2})
$$
$$
p (x _ {t}) \triangleq \mathcal {N} (0, \sigma_ {1} ^ {2})
$$
$$
\begin{array}{l} \log \left(\frac {p (x _ {t} \mid y _ {t + 1 : T})}{p (x _ {t})}\right) = \log \mathcal {N} (x _ {t}; \mu_ {1}, \sigma_ {1} ^ {2}) - \log \mathcal {N} (x _ {t}; 0, \sigma_ {2} ^ {2}) \\ = \log Z (\sigma_ {1}) - \frac {1}{2 \sigma_ {1} ^ {2}} x _ {t} ^ {2} + \frac {\mu_ {1}}{\sigma_ {1} ^ {2}} x _ {t} - \frac {\mu_ {1} ^ {2}}{2 \sigma_ {1} ^ {2}} - \log Z (\sigma_ {2}) + \frac {1}{2 \sigma^ {2}} x _ {t} ^ {2} \\ \end{array}
$$
where  $Z(\sigma) = \frac{1}{\sigma\sqrt{2\pi}}$ , so  $\log Z(\sigma) = -\log (\sigma \sqrt{2\pi})$ .
Collecting terms gives:
$$
\begin{array}{l} - \log (\sigma_ {1} \sqrt {2 \pi}) + \log (\sigma_ {2} \sqrt {2 \pi}) \\ - \frac {1}{2} \left(\frac {1}{\sigma_ {1} ^ {2}} - \frac {1}{\sigma_ {2} ^ {2}}\right) x _ {t} ^ {2} \\ + \frac {\mu_ {1}}{\sigma_ {1} ^ {2}} x _ {t} \\ - \frac {\mu_ {1} ^ {2}}{2 \sigma_ {1} ^ {2}} \\ \end{array}
$$
So we'll define
$$
a \triangleq - \frac {1}{2} \left(\frac {1}{\sigma_ {1} ^ {2}} - \frac {1}{\sigma_ {2} ^ {2}}\right)
$$
$$
b \triangleq \frac {\mu_ {1}}{\sigma_ {1} ^ {2}}
$$
$$
c \triangleq - \frac {\mu_ {1} ^ {2}}{2 \sigma_ {1} ^ {2}} - \log (\sigma_ {1} \sqrt {2 \pi}) + \log (\sigma_ {2} \sqrt {2 \pi})
$$
We'll explicitly model  $\log \sigma_1^2$ ,  $\log \sigma_2^2$  and  $\mu_{1}$ . Both  $\log \sigma_1^2$  and  $\log \sigma_2^2$  are only functions of  $t$ , not of  $y_{t + 1:T}$ , so those can be vectors of shape  $T$  initialized at 0.  $\mu_{1}$  is a linear function of  $y_{t + 1:T}$  and  $t$ , so that can be parameterized by a set of  $T\times T$  weights, initialized to  $1 / T$  and  $T$  biases initialized to 0.
Training Details We use a batch size of 32 for the density ratio estimation step. Since we do not perform model learning, we do not repeatedly alternate between twist training and proposal training for NAS-X. Instead, we first train the twist for 3,000,000 iterations with a batch size of 32 using samples from the model. We then train the proposal for 750,000 iterations. For the twist, we used Adam with a learning rate schedule that starts with a constant learning rate of  $1e - 3$ , decays the learning by 0.3 and 0.33 at 100,000 and 300,000 iterations. For the proposal, we used Adam with a constant learning rate of  $1e - 3$ . For NASMC, we only train the proposal.
Evaluation In the right panel of Figure 1, we compare the bound gaps of NAS-X and NASMC averaged across 20 different samples from the generative model. To obtain the bound gap for NAS-X, we run SMC 16 times with 128 particles and with the learned proposal and twists. We then record the average log marginal likelihood. For NASMC, we run SMC with the current learned proposal (without any twists).
10 rSLDS
Model details The generative model is as follows. At each time  $t$ , there is a discrete latent state  $z_{t} \in \{1, \ldots, 4\}$  as well as a two-dimensional continuous latent state  $x_{t} \in \mathbb{R}^{2}$ . The discrete state
Figure 6: Inference in NASCAR experiments.
transition probabilities are given by
$$
p \left(z _ {t + 1} = i \mid z _ {t} = j, x _ {t}\right) \propto \exp \left(r _ {i} + R _ {i} ^ {T} x _ {t - 1}\right) \tag {39}
$$
Here  $R_{i}$  and  $r_i$  are weights for the discrete state  $z_{i}$ .
These discrete latent states dictates two-dimensional latent state  $x_{t} \in \mathbb{R}^{2}$  which evolves according to linear Gaussian dynamics.
$$
x _ {t + 1} = A _ {z _ {t + 1}} x _ {t} + b _ {z _ {t + 1}} + v _ {t}, \quad v _ {t} \sim^ {\mathrm {i d}} \mathcal {N} \left(0, Q _ {z _ {t + 1}}\right) \tag {40}
$$
Here  $A_{k}, Q_{k} \in \mathbb{R}^{2x^{2}}$  and  $b_{k} \in \mathbb{R}^{2}$ . Importantly, from Equations 40 and 39 we see that the dynamics of the continuous latent states and discrete latents are coupled. The discrete latent states index into specific linear dynamics and the discrete transition probabilities depend on the continuous latent state.
The observations  $y_{t} \in \mathbb{R}^{10}$  are linear projections of the continuous latent state  $x_{t}$  with some additive Gaussian noise.
$$
y _ {t} = C x _ {t} + d + w _ {t}, \quad v _ {t} \sim^ {\mathrm {i d}} \mathcal {N} (0, S) \tag {41}
$$
Here  $C, S \in \mathbb{R}^{10x10}$  and  $d \in \mathbb{R}^{10}$ .
**Proposal Parameterization** We use a mean-field proposal distribution factorized over the discrete and continuous latent variables (i.e.  $q(\mathbf{z}_{1:T}, \mathbf{x}_{1:T}) = q(\mathbf{z}_{1:T})q(\mathbf{x}_{1:T})$ ). For the continuous states,  $q(\mathbf{x}_{1:T})$  is a Gaussian factorized over time with parameters  $\mu_{1:T}$  and  $\sigma_{1:T}^2$ . For the discrete states,  $q(\mathbf{z}_{1:T})$  is a Categorical distribution over  $K$  categories factorized over time with parameters  $p_{1:T}^{1:K}$ . In total, we learn  $2T + TK$  proposal parameters.
Twist Parameterization We parameterize the twists using a recurrent neural network (RNN) that is trained using density ratio estimation. To produce the twist values at each timestep, we first run a RNN backwards over the observations  $\mathbf{y}_{1:T}$  to produce a sequence of encodings  $\mathbf{e}_{1:T-1}$ . We then concatenate the encodings of  $\mathbf{x}_t$  and  $\mathbf{z}_t$  into a single vector and pass that vector into an MLP which outputs the twist values at each timestep. The RNN has one layer with 128 hidden units. The MLP has 131 hidden units and ReLU activations.
Model Parameter Evaluation We closely follow the parameter initialization strategy employed by Linderman et al. [41]. First, we use PCA to obtain a set of continuous latent states and initialize the matrices  $C$  and  $d$ . We then fit an autoregressive HMM to the estimated continuous latent states in order to initialize the dynamics matrices  $\{A_k, b_k\}$ . Importantly, we do not initialize the proposal with the continuous latent states described above.
**Training Details** We use a batch size of 32 for the density ratio estimation step. We alternate between 100 steps of twist training and 100 steps of proposal training for a total of 50,000 training steps in total. We used Adam and considered a grid search over the model, proposal, and twist learning rates. In particular, we considered learning rates of  $1e - 4$ ,  $1e - 3$ ,  $1e - 2$  for the model, proposal, and twist.
Bootstrap Bound Evaluation To obtain the log marginal likelihood bounds and standard deviations in Table 5.2, we ran a bootstrapped particle filter (BPF) with the learned model parameters for all three methods (NAS-X, NASMC, Laplace EM) using 1024 particles. We repeat this across 30 random seeds. Initialization of the latent states was important for a fair comparison. To initialize the latent
states, for NAS-X and NASMC, we simply sampled from the learned proposal at time  $t = 0$ . To initialize the latent state for Laplace EM, we sampled from a Gaussian distribution with the learned dynamics variance at  $t = 0$ .
Inference Comparison In the top panel of Figure 6, we compare NAS-X and NASMC on inference in the SLDS model. We report (average) posterior parameter recovery for the continuous and discrete latent states across 5 random samples from the generative model. NAS-X systematically recovers better estimates of both the discrete and continuous latent states.
11 Inference in Squid Giant Axon Model