Efficient Inference for Coupled Hidden Markov Models in Continuous Time and Discrete Space

Abstract

Systems of interacting continuous-time Markov chains are a powerful model class, but inference is typically intractable in high-dimensional settings. Auxiliary information, such as noisy observations, is typically only available at discrete times, and incorporating it via a Doob’s $h$-transform gives rise to an intractable posterior process that requires approximation. We introduce Latent Interacting Particle Systems, a model class parameterizing the generator of each Markov chain in the system. Our inference method involves estimating look-ahead functions (twist potentials) that anticipate future information, for which we introduce an efficient parameterization. We incorporate this approximation in a twisted Sequential Monte Carlo sampling scheme. We demonstrate the effectiveness of our approach on a challenging posterior inference task for a latent SIRS model on a graph, and on a neural model for wildfire spread dynamics trained on real data.

1. INTRODUCTION

Many real-world phenomena, from epidemics to wildfires, can be modeled as systems of interacting components evolving in continuous time, where the underlying dynamics are governed by discrete latent states (Lanchier, 2024). This modeling approach builds upon concepts from continuous-time hidden Markov models (Baum and Petrie, 1966; Kouemou, 2011) and extends them to spatially-structured, high-dimensional processes. Interacting particle systems (IPSs) (Liggett, 1985; Lanchier, 2024) provide a powerful mathematical framework for describing local propagation dynamics in discrete state spaces and continuous time, and are an important subset of the broader class of continuous-time Markov chains (CTMCs). We formulate our goal as performing probabilistic inference on systems whose latent dynamics follow an IPS, given only incomplete or indirect information.

We aim to endow the IPS with a flexible parameterization (e.g., a neural network), yielding a discrete analog of latent/neural stochastic differential equations (Movellan et al., 2002; Tzen and Raginsky, 2019; Liu et al., 2020; Li et al., 2020; Bartosh et al., 2025). While neural-based inference methods for continuous-time discrete-state processes have been explored (Seifner and Sánchez, 2023; Berghaus et al., 2024), they have primarily demonstrated efficacy in low-dimensional settings.

The core technical challenge for inference with latent IPSs lies in sampling from a smoothed path measure over latent trajectories. To address this problem, we make the following contributions:

• We propose a twisted sequential Monte Carlo (tSMC) scheme (Guarniero et al., 2017; Heng et al., 2020) tailored to latent IPSs. We learn a twist function that can approximate the likelihood of future observations, and directly incorporate it into the rate matrix of an approximate posterior process, lifting the need to learn a separate proposal distribution (Lawson et al., 2018, 2022, 2023).

• We learn the twist function using a mass-covering Kullback-Leibler divergence loss in an amortized fashion (Zhao et al., 2024). Moreover, we design an efficient parameterization with favorable inductive biases.

• We demonstrate this on (i) a spatial susceptible–infected–recovered–susceptible (SIRS) model on graphs with up to 256 nodes, and (ii) a neural wildfire-spread model on 64×64 grids using the WildFireSpreadTS dataset (Gerard et al., 2023).

2. BACKGROUND

2.1. Continuous-time Markov chains

Consider a Markov process $Z_{[0,T]}≔{\left\{Z_t\right\}}_{t\in [0,T]}$ taking values in a discrete state space $𝒵=\{1,\dots ,V\}$ over a finite time horizon $[0,T]$. This system, known as a CTMC, is characterized by an initial distribution $p_0\in 𝒫(𝒵)$ and a measurable family of rate matrices ${\left[R_t(z,\tilde{z})\right]}_{z,\tilde{z}\in 𝒵},t\in [0,T]$, where each entry satisfies

$$R_t\left(z,\tilde{z}\right)≔\underset{\Delta t\to 0}{\text{lim}}\frac{1}{\Delta t}\mathbb{P}\left(Z_{t+\Delta t}=\tilde{z}\mid Z_t=z\right),$$

for $\tilde{z}\ne z$, and $R_t(z,z)≔-\sum _{\tilde{z}\ne z}{R}_t(z,\tilde{z})$ on the diagonal. Assuming ${\text{sup}}_{t,z}-R_t(z,z)<\infty$ (non-explosion), sample paths are càdlàg, piecewise-constant, with finitely many jumps on $[0,T]$.

We refer to the induced distribution on the Skorokhod space $D([0,T],𝒵)$ as the path measure. Writing $0<t_1<\cdots <t_N<T$ for the jump times and $z_t$ for the state at a time $t\in [0,T]$, the density of this measure with respect to counting measure on states and Lebesgue measure on jump times is

$$p_0\left(z_0\right)\left[\prod _{n=1}^N{R}_{t_n}\left(z_{t_{n-1}},z_{t_n}\right)\right]\text{exp}\left({\int }_0^T{R}_t\left(z_t,z_t\right)dt\right).$$

A detailed introduction can be found in Norris (1998); Del Moral and Penev (2017). We provide an overview of inference with CTMCs in Appendix A.

2.2. Interacting particle systems

While CTMCs can be extended to high-dimensional systems, the size of the associated rate matrix increases exponentially with the number of dimensions, making exact inference intractable. We restrict attention to state spaces of the form $𝒵={𝒱}^d$, where $𝒱=\{1,\dots ,V\}$ is a fixed vocabulary. We assume that transitions affect only a single coordinate (i.e., dimension) at a time, and that the rate for updating a coordinate $i$ may depend on the current global state $z\in 𝒵$. The dependence of the rate on $z$ is often specified through a graph $𝒢=(ℐ,ℰ)$, where $ℐ=\{1,\dots ,d\}$ indexes coordinates and $ℰ$ encodes neighborhood structure (e.g., spatial dependence) (Liggett, 1985).

To describe local updates, we use the shorthand

$$z^{i\to v}=\left(z^1,\dots ,z^{i-1},v,z^{i+1},\dots ,z^d\right)$$

for the configuration obtained from $z$ by replacing the $i$-th coordinate with $v\in 𝒱$. The local rate $r_{i,t}(v\mid z)\ge 0$ is defined as the instantaneous intensity of transitioning from $z$ to $z^{i\to v}$, i.e.

$$R_t\left(z,z^{i\to v}\right)=r_{i,t}\left(v∣z\right),v\ne z^i.$$

Under these assumptions, the global generator $R_t$ decomposes as a sum of local generators:

$$R_t\left(z,\tilde{z}\right)=\sum _{i\in ℐ}{r}_{i,t}\left({\tilde{z}}^i\mid z\right)\prod _{j\ne i}{\delta }_{z^j}\left({\tilde{z}}^j\right).$$ (1)

reflecting the assumption that only one coordinate may change at a time. Typically the local rates $r_{i,t}(v\mid z)$ depend only on the neighborhood ${\left(z^j\right)}_{j\in 𝒩(i)}$ of $i$ in the graph $𝒢$. CTMCs with generators of this form are known as interacting particle systems (IPS) on a finite state space and finite graph (Liggett, 1985; Lanchier, 2017, 2024). We are interested in system that may exhibit time dependence, i.e. non-homogeneous processes (Norris, 1998).

Example (SIR model):

In an SIR model on a network of $d$ nodes, the state space is $\{S,I,R{\}}^d$ and the rate of infection of each node can be a function of the states of its neighbors. For an illustration, see the process $Z_{[0,T]}$ in Figure 1.

Figure 1:

Example of a latent IPS as a state space model. Here, the latent trajectory is sampled from a continuous-time SIR model, and potentials are emission distributions of partially masked observations.

Simulation.

Exact simulation from IPSs is feasible for models like the contact process and the voting process (Lanchier, 2017), leveraging a construction from independent Poisson processes (Lanchier, 2017). However, simulation can be challenging for non-homogeneous processes. For this reason, in this work we employ a simple first-order Euler discretization (Sun et al., 2023). We assume conditional independence among coordinates in intervals of width ${\Delta }_t$, and let the transition probability be:

$$q_{t,{\Delta }_t}(\tilde{z}\mid z)≔\prod _{i\in ℐ}\left({\delta }_{{\tilde{z}}^i,z^i}+{\Delta }_t{r}_{i,t}\left({\tilde{z}}^i\mid z\right)\right),$$ (2)

where $r_{i,t}\left(z^i\mid z\right)≔-\sum _{v\ne z^i}{r}_{i,t}(v\mid z)$. Note that $q_{t,{\Delta }_t}$ may assign mass to multi-site flips for any finite step size. Despite its simplicity, this simulation scheme has shown good performance in generative modeling applications, such as simulating the reverse process of discrete diffusion and discrete flow matching models (Campbell et al., 2024; Gat et al., 2024).

In the following Proposition, we show that the error of this kernel can be bounded in total variation. Let ${\lambda }_t(z)≔\sum _{i\in ℐ}\sum _{v\ne z^i}{r}_{i,t}(v\mid z)$. We assume that the following properties hold for each $t\in [0,T]$:

• (A1)Bounded total rate: ${\text{sup}}_{t,z}{\lambda }_t(z)\le \bar{\lambda }<\infty$.
• (A2)Small time interval: we bound the step size as ${\Delta }_t\le 1/{\text{sup}}_{t,z,i}\sum _{v\ne z^i}{r}_{i,t}(v\mid z)$.
• (A3)Local Lipschitz in time: for all $u\in \left[0,{\Delta }_t\right],\left|r_{i,t+u}(v\mid z)-r_{i,t}(v\mid z)\right|\le Lu$.

Since $𝒵$ is finite, we identify the transition kernel of an IPS $P_{t,{\Delta }_t}\left(z_t,dz\right)≔\mathbb{P}\left(Z_{t+{\Delta }_t}\in dz\mid Z_t=z_t\right)$ with its pmf $p_{t,{\Delta }_t}(\tilde{z}\mid z)≔P_{t,{\Delta }_t}(z,\{\tilde{z}\})$.

Proposition 1. Under (A1)–(A3), there exists $C<\infty$ such that

$${‖q_{t,{\Delta }_t}(\cdot \mid z)-p_{t,{\Delta }_t}(\cdot \mid z)‖}_{\text{TV}}\le C{\Delta }_t^2,$$

uniformly in $t\in [0,T]$ and $z\in 𝒵$.

A proof is provided in Appendix B.1.

2.3. Sequential Monte Carlo

In our approximate posterior inference scheme for IPSs, we employ sequential Monte Carlo (SMC) methods. SMC approximates a terminal target ${\pi }_T$ by evolving a population of weighted samples through a sequence of intermediate targets ${\left\{{\pi }_t\right\}}_{t\in [0,T]}$ (Doucet and Johansen, 2009; Naesseth et al., 2019; Chopin and Papaspiliopoulos, 2020). We assume ${\pi }_0$ is easy to sample from, and that we can evaluate unnormalized densities ${\left\{{\gamma }_t\right\}}_{t\in [0,T]}$ with ${\pi }_t\propto {\gamma }_t$. In our setting, it is natural to let ${\pi }_t$ be a distribution on path prefixes $Z_{[0,t]}\in D([0,t],𝒵)$. Given particles ${\left\{Z_{\left[0,t-{\Delta }_t\right]}^{(s)},w_{t-{\Delta }_t}^{(s)}\right\}}_{s=1}^S$ targeting ${\pi }_{t-{\Delta }_t}$ and a proposal kernel $q_t\left(\cdot \mid Z_{t-{\Delta }_t}^{(s)}\right)$ that advances each trajectory by ${\Delta }_t$, we draw a segment $Z_{\left(t-{\Delta }_t,t\right]}^{(s)}~q_t\left(\cdot \mid Z_{t-{\Delta }_t}^{(s)}\right)$ and update unnormalized incremental weights as

$$\begin{gathered} {\tilde{w}}_t^{(s)}=\frac{{\gamma }_t\left(d{Z}_{[0,t]}^{(s)}\right)}{{\gamma }_{t-{\Delta }_t}\left(d{Z}_{\left[0,t-{\Delta }_t\right]}^{(s)}\right)q_t\left(Z_{\left(t-{\Delta }_t,t\right]}^{(s)}\mid Z_{t-{\Delta }_t}^{(s)}\right)}, \\ {w}_t^{(s)}\propto w_{t-{\Delta }_t}^{(s)}{\tilde{w}}_t^{(s)},{\bar{w}}_t^{(s)}=\frac{w_t^{(s)}}{\sum _{j=1}^S{w}_t^{(j)}}. \end{gathered}$$ (3)

A distinctive feature of SMC is resampling, where samples are selected and propagated based on their importance weights. This is implemented by drawing ancestry indices according to ${\bar{w}}_t^{(1:S)}$, adaptively when the effective sample size ${\text{ESS}}_t={\left(\sum _{s=1}^S{\left({\bar{w}}_t^{(s)}\right)}^2\right)}^{-1}$ falls below a threshold (Naesseth et al., 2019). Two design choices are central to performance: (i) intermediate targets $\left\{{\pi }_t\right\}$, which can be crafted to reduce variance of incremental weights; and (ii) the proposal distribution $q_t$, which should ensure proposed segments land in high-probability regions of ${\pi }_t$. Poor choices of either can trigger rapid weight collapse, especially in high-dimensional, sparse-observation regimes (Naesseth et al., 2019; Chopin et al., 2023).

3. EFFICIENT INFERENCE FOR LATENT INTERACTING PARTICLE SYSTEMS

3.1. Posterior path measure for an IPS

Let $Z_{[0,T]}$ be an Interacting Particle System (IPS) with path measure $P$, initial distribution $p_0$, and local transition rates $r_{i,t}(v\mid z)$. We are interested in conditioning this process on auxiliary information available at discrete times ${\tau }_1,\dots ,{\tau }_K\in [0,T]$. We represent this information by nonnegative potential functions $G_{{\tau }_k}:𝒵\to {\mathbb{R}}_{+}$ that can be evaluated pointwise. The resulting posterior path measure is

$$P^{\star }\left(d{Z}_{[0,T]}\right)\propto \left(\prod _{k=1}^K{G}_{{\tau }_k}\left(Z_{{\tau }_k}\right)\right)P\left(d{Z}_{[0,T]}\right),$$ (4)

a special case of a Feynman–Kac (FK) path measure (Chopin and Papaspiliopoulos, 2020; Lu and Wang, 2024; Park et al., 2025). A key property is that the process $Z_{[0,T]}^{\star }$ governed by $P^{\star }$ is itself an IPS.

This follows from Doob’s $h$-transform applied to the original dynamics (Del Moral and Penev, 2017; Corstanje et al., 2023; Corstanje and Meulen, 2025). Define the look-ahead function

$$h_t^{\star }(z)≔{\mathbb{E}}_P\left[\prod _{k:{\tau }_k>t}{G}_{{\tau }_k}\left(Z_{{\tau }_k}\right)\mid Z_t=z\right],h_T^{\star }(z)=1,$$ (5)

which is right-continuous with left limits (càdlàg) with reset conditions at potential times (Eich et al., 2025):

$$h_{{\tau }_k^{-}}^{\star }\left(z\right)=G_{{\tau }_k}\left(z\right)h_{{\tau }_k}^{\star }\left(z\right).$$ (6)

Note that we can marginalize equation 4 to any time $t\in [0,T]$, and get

$$P^{\star }\left(d{Z}_{[0,t]}\right)\propto P\left(d{Z}_{[0,t]}\right)h_t^{\star }\left(Z_t\right)\prod _{k:{\tau }_k\le t}{G}_{{\tau }_k}\left(Z_{{\tau }_k}\right).$$ (7)

The initial distribution and local rates of the posterior IPS are characterized in the following Proposition.

Proposition 2 (Doob’s $h$-transform for IPS). Under $P^{\star }$, the process $\left\{Z_t^{\star }\right\}$ is an IPS with initial distribution

$$p_0^{\star }\left(z\right)\propto p_0\left(z\right)h_0^{\star }\left(z\right),$$

and off-diagonal local transition rates

$$r_{i,t}^{\star }(v\mid z)=r_{i,t}(v\mid z)s_{i,t}^{\star }(v,z),v\ne z^i,$$ (8)

where $s_{i,t}^{\star }(v,z)≔h_t^{\star }\left(z^{i\to v}\right)/h_t^{\star }(z)$. Diagonal terms are defined as usual by $r_{i,t}^{\star }\left(z^i\mid z\right)=-\sum _{v\ne z^i}{r}_{i,t}^{\star }(v\mid z)$.

Note that this proposition follows from standard arguments on the conditioning of CTMCs, see (Huang et al., 2016; Corstanje et al., 2023; Corstanje and Meulen, 2025). We provide a complete proof of Proposition 2 in Appendix B.2. This formulation generalizes the posterior distribution used in reward-guided fine-tuning for discrete diffusion models (Li et al., 2025; Wang et al., 2025; Cheuk-Kit et al., 2025) to multiple potential times. Simulating directly from these transformed rates $r_{i,t}^{\star }$ is generally intractable because computing the look-ahead function $h_t^{\star }(z)$ requires solving a high-dimensional integral over future trajectories under $P$.

3.2. Latent IPS

Let $\left\{P_{\theta }:\theta \in \Theta \right\}$ denote a family of prior path measures on $D([0,T],𝒵)$, where each $P_{\theta }$ corresponds to an IPS with initial distribution $p_0^{\theta }$ and local transition rates $r_{i,t}^{\theta }(v\mid z)$.

Auxiliary information about the latent trajectory is available in the form of discrete-time observations

$$Y=\left(Y_{{\tau }_1},\dots ,Y_{{\tau }_K}\right),0<{\tau }_1<\cdots <{\tau }_K\le T.$$

Here and throughout, we use $\theta$ to denote the parameters of the full generative model, including both the prior IPS dynamics and the observation model. For an observed trajectory $y_{1:K}$, potential functions in equation 4 correspond to emission distributions

$$G_{{\tau }_k,\theta }(z)=p_{\theta }\left(y_k\mid Z_{{\tau }_k}=z\right),k=1,\dots ,K,$$

where we omit $Y_{{\tau }_k}$ to simplify the notation. We denote the corresponding posterior path measure by $P_{\theta }^{\star }$, which, by Proposition 2, is again an IPS. We refer to this model as a Latent Interacting Particle System (Latent IPS). Figure 1 illustrates this construction for a latent susceptible-infected-recovered (SIR) model on a graph.

Maximum likelihood.

The logarithm of the marginal likelihood of the observations is

$$ℒ\left(\theta \right)≔\text{log}{\mathbb{E}}_{P_{\theta }}\left[\prod _{k=1}^K{G}_{{\tau }_k,\theta }\left(Z_{{\tau }_k}\right)\right].$$ (9)

This objective can be optimized via gradient steps using Fisher’s identity (see Appendix B.3), rewriting ${\nabla }_{\theta }ℒ(\theta )$ as

$${\mathbb{E}}_{P_{\theta }^{\star }}\left[\sum _{k=1}^K{\nabla }_{\theta }\text{log}{G}_{{\tau }_k,\theta }\left(Z_{{\tau }_k}\right)+{\nabla }_{\theta }\text{log}{P}_{\theta }\left(Z_{[0,T]}\right)\right].$$ (10)

These updates are not directly tractable, as they require samples from $P_{\theta }^{\star }$. Viable options include optimizing an evidence lower bound (Hinton et al., 1995), importance sampling (Bornschein and Bengio, 2015), and SMC (Lawson et al., 2023; McNamara et al., 2024). Due to the sequential nature of latent IPSs, we opt for an SMC approach.

3.3. Twisted SMC for posterior inference in latent IPS

Classic design choices for SMC algorithms, such as the bootstrap particle filter (BPF) (Doucet and Johansen, 2009), notoriously display poor performance in continuous-time problems with sparse potentials, as weights are uniform in between potential times, leading to particle degeneracy (Chopin et al., 2023).

Twisted targets.

Twisted SMC (tSMC) addresses this limitation by specifying a favorable choice for the intermediate target distributions of SMC (Guarniero et al., 2017; Heng et al., 2020). The core idea is to introduce a learnable function $h_t^{\psi }:𝒵\to {\mathbb{R}}_{+}$, the twist function, approximating the look-ahead function $h_t^{\star }(z)$ from equation 5. The tSMC algorithm targets a sequence of twisted distributions

$$P_{\theta }^{\psi }\left(d{Z}_{[0,t]}\right)\propto P_{\theta }\left(d{Z}_{\left[0,t\right]}\right)h_t^{\psi }\left(Z_t\right)\prod _{k:{\tau }_k\le t}{G}_{{\tau }_k,\theta }\left(Z_{{\tau }_k}\right),$$ (11)

for $t\in [0,T]$. Note that $h_t^{\psi }$ can be a function of future observations $y_{>t}$, and we omit this dependence for notational simplicity.

Twist-induced proposal.

We can approximately sample from this IPS using an approximation of a first-order Euler discretization as in 2.2, resulting in the twist-induced proposal distribution

$$q_{t,{\Delta }_t}^{\theta ,\psi }\left(z\mid z_t\right)=\prod _{i\in ℐ}\left({\delta }_{z_t^i,z^i}+{\Delta }_t{r}_{i,t}^{\theta ,\psi }\left(z^i\mid z_t\right)\right),$$ (12)

$$\begin{array}{c}{r}_{i,t}^{\theta ,\psi }\left(z^i\mid z_t\right)≔\left\{\begin{array}{ll}{r}_{i,t}^{\theta }\left(z^i\mid z_t\right)s_{i,t}^{\psi }\left(z^i,z_t\right),& z^i\ne z_t^i,\\ -\sum _{v\ne z_t^i}{r}_{i,t}^{\theta }\left(v\mid z_t\right)s_{i,t}^{\psi }\left(v,z_t\right),& z^i=z_t^i,\end{array}\right.\end{array}$$ (13)

where $s_{i,t}^{\psi }\left(z^i,z_t\right)≔h_t^{\psi }\left(z_t^{i\to z^i}\right)/h_t^{\psi }\left(z_t\right)$ is the concrete score (Meng et al., 2022). Note that the twist function $h^{\psi }$ might not satisfy the reset conditions in equation 6. We can quantify the extent to which it is violated by introducing the reset residual

$${\rho }_t\left(z\right)≔\left(\text{log}{h}_{t^{-}}^{\psi }\left(z\right)-\text{log}{G}_t\left(z\right)-\text{log}{h}_t^{\psi }\left(z\right)\right){\mathbf{1}}_{t\in \left\{{\tau }_k\right\}}.$$

Let the incremental effective sample size (ESS) with respect to the target posterior be defined as

$${\text{ESS}}_{t,{\Delta }_t}\left(z_t\right)≔{\mathbb{E}}_{q_{t,{\Delta }_t}^{\theta ,\psi }\left(\cdot \mid z_t\right)}{\left[{\left(\frac{p_{t,{\Delta }_t}^{\theta ,\star }\left(Z_{t+{\Delta }_t}\mid z_t\right)}{q_{t,{\Delta }_t}^{\theta ,\psi }\left(Z_{t+{\Delta }_t}\mid z_t\right)}\right)}^2\right]}^{-1},$$ (14)

where

$$p_{t,{\Delta }_t}^{\theta ,\star }\left(z\mid z_t\right)\propto p_{t,{\Delta }_t}^{\theta }\left(z\mid z_t\right)h_{t+{\Delta }_t}^{\star }(z)G_t(z{)}^{{\mathbf{1}}_{t\in \left\{{\tau }_k\right\}}}.$$

In the following Theorem, we show that the incremental ESS at each step can be lower bounded by composing terms depending on the approximation error of the twist function, the reset residual, and the discretization error.

Theorem 3. For $t\in [0,T]$, bound the twist error and the reset residual by

$${\epsilon }_t≔\underset{z\in 𝒵}{\text{sup}}\left|\text{log}{h}_t^{\star }(z)-\text{log}{h}_t^{\psi }(z)\right|,{\delta }_t≔\underset{z\in 𝒵}{\text{sup}}\left|{\rho }_t(z)\right|.$$

Under (A1)-(A3) for the prior dynamics and assuming $h^{\psi }$ to be Lipschitz continuous between observation times, there exists a $C<\infty$ such that the incremental ESS in equation 14 satisfies, uniformly in $z_t$,

$${\text{ESS}}_{t,{\Delta }_t}\left(z_t\right)\ge \frac{\text{exp}\left(-4\left({\epsilon }_{t+{\Delta }_t}+{\delta }_{t+{\Delta }_t}\right)\right)}{1+C{\Delta }_t^2},$$

For a proof, see Appendix B.4.

3.4. Efficient twist parameterization.

Parameterizing the twist–induced proposal in equation 12 is challenging because, at each step and for a given state $z_t$, we need all ratios

$$s_{i,t}^{\psi }\left(v,z_t\right)=\frac{h_t^{\psi }\left(z_t^{i\to v}\right)}{h_t^{\psi }\left(z_t\right)},i\in ℐ,v\in 𝒱,$$

that is, the twist evaluated at every single–site modification $z_t^{i\to v}$ of $z_t$. A naïve implementation would call $h_t^{\psi }$ separately on each configuration $z_t^{i\to v}$, requiring $O(dV)$ forward passes per time step. To make this structure explicit, define for each configuration $z\in 𝒵$ the matrix

$$H_t^{\psi }(z)\in {\mathbb{R}}_{+}^{d\times V},H_t^{\psi }(z{)}_{\left[i,v\right]}≔h_t^{\psi }\left(z^{i\to v}\right).$$

In particular, $H_t^{\psi }(z{)}_{\left[i,z^i\right]}=h_t^{\psi }(z)$ for all $i\in ℐ$. By construction, replacing the $i$-th coordinate of $z$ by any label $u\in 𝒱$ leaves the row $H_t^{\psi }(z{)}_{[i,\cdot ]}$ unchanged:

$$H_t^{\psi }(z{)}_{\left[i,v\right]}=H_t^{\psi }{\left(z^{i\to u}\right)}_{\left[i,v\right]},u,v\in 𝒱,i\in ℐ.$$ (15)

This is the functional identity we would like the learned mapping $z\mapsto H_t^{\psi }(z)$ to satisfy: each entry $H_t^{\psi }(z{)}_{[i,v]}$ should depend on $z$ only through the context provided by the other coordinates and the choice of index–value pair $(i,v)$, not on the current value $z^i$ itself.

A direct implementation of $H_t^{\psi }(z)$ by explicitly forming all $z^{i\to v}$ would be prohibitively expensive. Instead, we directly exploit the product structure of the state space to build an efficient alternative. We first use a context encoder that can depend on any contextual information $c_t$ we wish to make available for the model at time $t$, except for the state $z_t$ itself:

$${\Phi }_t\left(c_t\right)\in {\mathbb{R}}^{d\times V\times m},$$ (16)

where $c_t$ can include, for instance, future observations and observation times $\left(Y_{\ge t},{\tau }_{\ge t}\right)$, covariates, and positional information if available. We highlight that the encoder in equation 16 is independent of the current state $z_t$, and produces a tensor of embeddings associated with setting each coordinate $i$ to a state $v$. Based on these embeddings, we can compose the output $h_t^{\psi }(z)$ by aggregating the representations at the positions indexed by our target state, via a map

$$\text{log}{h}_t^{\psi }(z)=\varphi \left(\sum _{i\in ℐ}{\Phi }_t{\left(c_t\right)}_{\left[i,z^i\right]}\right),$$ (17)

where $\varphi :{\mathbb{R}}^m\to \mathbb{R}$ can be a learned map (e.g., a small MLP), a construction mirroring the DeepSets model (Zaheer et al., 2017). Let $S_t(z)≔\sum _{j\in ℐ}{\Phi }_t{\left(c_t\right)}_{\left[j,z^j\right]}$ denote the pooled representation of the current configuration. For each single–site update $z^{i\to v}$, we can reuse $S_t(z)$ and adjust only the contribution of coordinate $i$:

$$\text{log}{h}_t^{\psi }\left(z^{i\to v}\right)=\varphi \left(S_t(z)+{\Phi }_t{\left(c_t\right)}_{[i,v]}-{\Phi }_t{\left(c_t\right)}_{\left[i,z^i\right]}\right),$$ (18)

an operation trivial to parallelize. Moreover, the identity in equation 15 holds exactly for this parameterization: replacing $z^i$ by any $u$ changes both $S_t(z)$ and the subtraction $-{\Phi }_t{\left(c_t\right)}_{\left[i,z^i\right]}$ in a way that cancels out, leaving the argument of $\varphi$ (and thus $H_t^{\psi }(\cdot {)}_{[i,v]}$) invariant for every $v$.

We refer to this architecture, which (i) amortizes the cost of encoding contextual information across all $(i,v)$ pairs and (ii) enforces the desired invariance of equation 15, as TwistNet. Figure 2 illustrates this construction: the context encoder computes ${\Phi }_t\left(c_t\right)$ once, the current state $z_t$ selects one embedding per coordinate to form $S_t\left(z_t\right)$, and each single-site modification $z_t^{i\to v}$ is evaluated by replacing only the $i$-th contribution in the pooled representation.

Figure 2:

TwistNet uses a context encoder to produce a table of embeddings for every coordinate and every possible state value, independently of the current state $z_t$. The state $z_t$ then selects one embedding per coordinate, and the selected embeddings are pooled into a single representation of the full state. To evaluate the twist at any single-site modification of $z_t$, TwistNet reuses this pooled representation and updates only the contribution of the modified coordinate, enabling parallel evaluation of all single-site neighboring states.

Twist learning.

Multiple objective functions have been proposed to learn the twist, from approaches based on consistency inspired by reinforcement learning (Heng et al., 2020; Lawson et al., 2018), to density ratio estimation (Lawson et al., 2022, 2023). Recently, Zhao et al. (2024) proposed to minimize the mass-covering forward KL divergence with respect to the true posterior in the context of autoregressive language models. Forward KL objectives had previously been proposed for learning the proposal distribution in SMC by Gu et al. (2015); Lawson et al. (2023). We adapt this approach to our setting, and learn an amortized twist function by optimizing $D_{\text{KL}}\left(P_{\theta }^{\star }\Vert P_{\theta }^{\psi }\right)$, which is proportional to

$$\begin{gathered} {\mathbb{E}}_{P_{\theta }^{\star }}\left[\sum _{i\in ℐ}\left({\int }_0^T{r}_{i,t}^{\theta ,\psi }\left(Z_t^i\mid Z_t\right)dt+\sum _{u:Z_u^i\ne Z_{u^{-}}^i}\text{log}{s}_{i,u}^{\theta ,\psi }\left(Z_u^i,Z_{u^{-}}\right)\right)\right] \\ +{\mathbb{E}}_{p_0^{\star }}\left[\text{log}{q}_0^{\psi }\left(Z_0\right)\right], \end{gathered}$$ (19)

where $s_{i,t}^{\theta ,\psi }(v,z)=h_t^{\psi }\left(z^{i\to v}\right)/h_t^{\psi }(z)$ is the concrete score, and $h_t^{\psi }$ may depend on future observations $y_{>t}$.

Wake-sleep.

In order to compute a tractable approximation to equation 19, we employ the strategy proposed in Zhao et al. (2024) and in wake-sleep algorithms (Hinton et al., 1995; Bornschein and Bengio, 2015; Le et al., 2020; McNamara et al., 2024), and perform ancestral sampling of the trajectory from the prior, and of the observations from the emission distribution. We obtain the time-discretized sleep objective

$$\begin{gathered} {ℒ}_{\text{s}}\left(\psi ;z_{t_0:t_M},y_{1:K},\theta \right)= \\ -\text{log}{q}_0^{\psi }\left(z_0\mid y_{1:K}\right)-\sum _{m=0}^{M-1}\sum _{i\in ℐ}\left[{\Delta }_{t_{m+1}}{r}_{i,t_m}^{\theta ,\psi }\left(z_{t_m}^i\mid z_{t_m}\right)\right.\left.+\mathbf{1}\left[z_{t_m}^i\ne z_{t_{m+1}}^i\right]\text{log}{s}_{i,t_m}^{\psi }\left(z_{t_{m+1}}^i,z_{t_m}\right)\right], \end{gathered}$$ (20, 21)

where $0=t_0<t_1<\cdots <t_M=T$ is a time grid containing ${\left\{{\tau }_k\right\}}_{k=1}^K$ with ${\Delta }_{t_m}=t_m-t_{m-1}$, and $q_0^{\psi }$ is an approximate posterior initial distribution. Note that, although omitted for notational simplicity, $s_{i,t_m}^{\psi }$ is conditioned on future observations $y_{>t_m}$. We optimize this objective with respect to the parameters $\psi$, while holding $\theta$ fixed. Note that optimizing equation 20 does not require backpropagating through the sampled trajectories, avoiding the need to use reparameterization (Jang et al., 2017) or REINFORCE (Williams, 1992).

We alternate $\psi$–updates on the objective in equation 20 with $\theta$–updates approximating equation 10 based on trajectories sampled using tSMC, similarly to Lawson et al. (2023) and McNamara et al. (2024). For the $\theta$-updates, we use the discretized log path density of the prior

$$\begin{gathered} \text{log}{\hat{P}}_{\theta }\left(z_{[0,T]}\right)= \\ \text{log}{p}_0^{\theta }\left(z_0\right)+\sum _{m=0}^{M-1}\sum _{i\in ℐ}\left[{\Delta }_{t_{m+1}}{r}_{i,t_m}^{\theta }\left(z_{t_m}^i\mid z_{t_m}\right)\right.\left.+\mathbf{1}\left[z_{t_m}^i\ne z_{t_{m+1}}^i\right]\text{log}{r}_{i,t_m}^{\theta }\left(z_{t_{m+1}}^i\mid z_{t_m}\right)\right], \end{gathered}$$ (22)

to form the wake gradient estimator

$$\widehat{{\nabla }_{\theta }{ℒ}_{\text{w}}}\left(\theta ;{\left\{z_{t_0:t_M}^{(s)},{\bar{w}}_T^{(s)}\right\}}_{s=1}^S,y_{1:K}\right)=$$ (23)

$$\sum _{s=1}^S{\bar{w}}_T^{(s)}\left[\sum _{k=1}^K{\nabla }_{\theta }\text{log}{G}_{{\tau }_k,\theta }\left(z_{{\tau }_k}^{(s)}\right)+{\nabla }_{\theta }\text{log}{\hat{P}}_{\theta }\left(z_{[0,T]}^{(s)}\right)\right],$$ (24)

where ${\bar{w}}_T^{(s)}$ are SMC weights as in equation 3. We note that the estimator in equation 23 is asymptotically consistent for the gradient of the maximum likelihood objective (Lawson et al., 2023). Simplified pseudocode is presented in Algorithm 1, and the detailed training algorithm and its complexity are reported in Appendix C.

Algorithm 1.

Wake–sleep with tSMC for latent IPS

1:	Inputs: Dataset $𝒟$; time grid $0=t_0<\cdots <t_M=T$; particles $S$; batch size $B$
2:	Initialize $\theta ,\psi$
3:	repeat
4:	Sleep phase: # update $\psi$, hold $\theta$ fixed
5:	for $b=1,\dots ,B$ do
6:	Simulate $z_{t_0:t_M}^{(b)}~P_{\theta }$ via Euler steps
7:	Simulate ${\tilde{y}}_k^{(b)}~p_{\theta }\left(\cdot \mid z_{{\tau }_k}^{(b)}\right)$, $k=1,\dots ,K$
8:	Compute ${\ell }_{\text{s}}^{(b)}={ℒ}_{\text{s}}\left(\psi ;z_{t_0:t_M}^{(b)},{\tilde{y}}_{1:K}^{(b)},\theta \right)$ via equation 20
9:	end for
10:	$\psi \leftarrow {\text{G}\text{RAD}\text{S}\text{TEP}}_{\psi }\left({\nabla }_{\psi }\frac{1}{B}\sum _{b=1}^B{\ell }_{\text{s}}^{(b)}\right)$
11:	Wake phase: # update $\theta$, hold $\psi$ fixed
12:	for $b=1,\dots ,B$ do
13:	Sample $\left(y_{1:K}^{(b)},{\tau }_{1:K}^{(b)}\right)~𝒟$
14:	Run ${\left\{z_{t_0:t_M}^{(b,s)},{\overline{w}}_T^{(b,s)}\right\}}_{s=1}^S\leftarrow \text{tSMC}\left(P_{\theta }^{\psi };y_{1:K}^{(b)},{\tau }_{1:K}^{(b)}\right)$
15:	Compute $g_{\text{w}}^{(b)}=\widehat{{\nabla }_{\theta }{ℒ}_{\text{w}}^{(b)}}\left(\theta ;{\left\{z_{t_0:t_M}^{(b,s)},{\overline{w}}_T^{(b,s)}\right\}}_{s=1}^S,y_{1:K}^{(b)}\right)$ via equation 23
16:	end for
17:	$\theta \leftarrow {\text{G}\text{RAD}\text{S}\text{TEP}}_{\theta }\left(\frac{1}{B}\sum _{b=1}^B{g}_{\text{w}}^{(b)}\right)$
18:	until convergence

4. EXPERIMENTS

4.1. Latent SIRS model

We study inference in a latent spatial SIRS IPS on graphs with state space $𝒵=\{S,I,R{\}}^d$ (Lanchier, 2024). We simulate a dataset of trajectories using Gillespie’s algorithm (Gillespie, 1977; Wilkinson, 2018) with the following local rates:

$$r_{i,t}(I\mid z)=\left({\alpha }_0+{\alpha }_1{\sum }_{j\ne i}{a}_{ij}\sigma \left(⟨{\xi }_i,{\xi }_j⟩\right){\delta }_{z^j,I}\right){\delta }_{z^i,S},r_{i,t}(R\mid z)=\beta {\delta }_{z^i,I},r_{i,t}(S\mid z)=\gamma {\delta }_{z^i,R},$$ (25)

where $a_{ij}$ is an adjacency matrix sampled from a graph with fixed expected degree, ${\xi }_i$ are normalized features acting as edge weights, and $\sigma$ is a logistic function ensuring non-negativity. We let observation times be sampled uniformly in $[0,T]$ at $K=10$ irregular snapshots ${\tau }_1<\cdots <{\tau }_K$, and we sample observations from an emission distribution masking each node (denoted $y_k=\varnothing$) independently with probability $p_{\text{mask}}=0.5$.

Task.

We conduct two experiments. First, we perform latent trajectory inference (Eich et al., 2025): given observations $y_{1:K}$ and parameters $\theta$ at their true values, we measure how accurately our approximate posterior samples match the ground truth trajectory. We focus on scalability with respect to graph dimensionality, quantifying performance via categorical cross-entropy loss averaged over a discretized time grid. Second, we estimate parameters from an arbitrary initialization using only 50 training trajectories. Following Section 3.4, we employ a wake-sleep algorithm and evaluate accuracy via individual parameter estimates and total relative parameter error $\sum _j\left|{\hat{\theta }}_j-{\theta }_j\right|/\left|{\theta }_j\right|$ over the four parameters in equation 25.

Methods.

For our TwistNet parameterization, we experiment learning the twist using both our forward KL loss from equation 20 and the DRE loss introduced by (Lawson et al., 2022), using a graph transformer as context encoder and a two-layer MLP as $\rho$ in equation 17. We consider NASMC (Gu et al., 2015) and NAS-X (Lawson et al., 2023), and parameterize the proposal distribution with a graph transformer as the context encoder used in TwistNet, taking the current state as input and predicting a concrete score field. The twist for NAS-X follows the same parameterization, with output pooled to a scalar. We experiment using TAG (Nisonoff et al., 2025) with DRE loss, and a regular BPF (Doucet and Johansen, 2009). We discuss further details in Appendix D.1 and D.2.

Results.

As illustrated in Figure 3, we find the combination of TwistNet and the forward KL loss in equation 19 to scale the best with respect to dimensions for latent trajectory inference. We report results on parameter estimation on a graph with 32 nodes in Table 1, and display convergence through wake steps in Figure 4. Additional results are reported in Appendix E.1.

Figure 3:

Latent trajectory reconstruction, measured by cross-entropy loss on the test set of ground truth trajectories with respect to the posterior approximations. Error bars correspond to two standard errors computed across trajectories.

Table 1:

Parameter estimates and relative parameter error (RPE) for the SIRS model with 32 nodes, mean ± 2 standard deviations across 10 random seeds.

	${\alpha }_0$	${\alpha }_1$	$\beta$	$\gamma$	RPE
Ground truth	0.1	1.0	0.4	0.05	–
TwistNet - KL	0.113 ± 0.03	0.922 ± 0.08	0.393 ± 0.03	0.046 ± 0.01	0.330 ± 0.35
TwistNet - DRE	0.115 ± 0.03	0.957 ± 0.15	0.384 ± 0.03	0.048 ± 0.01	0.360 ± 0.34
TAG - DRE	0.124 ± 0.05	0.673 ± 0.43	0.377 ± 0.14	0.055 ± 0.02	0.982 ± 0.68
NAS-X	0.136 ± 0.04	0.880 ± 0.10	0.405 ± 0.04	0.058 ± 0.01	0.673 ± 0.39
NASMC	0.076 ± 0.09	0.936 ± 0.26	0.444 ± 0.11	0.061 ± 0.04	0.962 ± 0.56
BPF	0.791 ± 1.90	1.354 ± 2.05	0.618 ± 0.24	0.001 ± 0.00	9.461 ± 17.88

Figure 4:

Evolution of the parameters and log relative parameter error (log RPE) through wake steps updating the parameters $\theta$, for a SIRS model on a graph with 32 nodes.

4.2. Wildfire trajectories

Forecasting wildfire spread is a problem of critical importance, given the impacts of wildfires and their increasing severity globally in recent years (Abatzoglou and Williams, 2016; Walker et al., 2019). Physics-based models have long existed (Rothermel, 1972) but are difficult to calibrate given real-world constraints, limiting forecasting utility. With the increasing availability of satellite data, ML approaches for spatio-temporal prediction have emerged (e.g., Coffield et al. (2019); Prapas et al. (2021); Apostolakis et al. (2022); Gerard et al. (2023)), yet they can lack physical insight (e.g., local propagation constraints) and have not yet seen widespread application in real-world scenarios. A third class of model in this context is based on the notion of local spreading of fire (across pixels) via stochastic cellular automata (SCAs), closely related to our latent IPS framework. Prior work in developing SCAs for wildfires has largely focused on relatively simple discrete-time models (Grieshop and Wikle, 2024), often hand-tuned (Clarke et al., 1994; Hargrove et al., 2000).

We model wildfires as latent IPSs, and consider a state space where each pixel in a 64×64 grid can be unburned $(U)$, active $(A)$, or burned $(B)$, i.e. $𝒵=\{U,A,B{\}}^{64\times 64}$. We let the graph be a lattice with Moore neighborhood. We parameterize outgoing, off-diagonal local rates by a single neural network $F_{\theta }$ mapping to ${\mathbb{R}}_{+}^{d\times (V-1)}$, using the UTAE architecture (Garnot and Landrieu, 2021). We constrain the dynamics by zeroing out the rate of flipping to $A$ whenever no neighbors are ignited, and do not allow transitions from $U$ to $B$ and vice versa.

Task.

Using the WildFireSpreadTS dataset (Gerard et al., 2023), we extract 156 week-long trajectories on 64 × 64 grids where at least one pixel is initially active with two subsequent days of activity, split 130-26 into a train and test set. We evaluate reconstruction of active fire maps at observation times (i.e., conditioning on the whole sequence) and prediction (initialized by encoding only the first observation), both measured via binary cross-entropy from the empirical distribution given by 16 particles per datapoint, averaged across time. We report results in two regimes: conditioning on VIIRS reflectance channels for the full week, and conditioning only on channels observed up to each timepoint.

Methods.

We use the UTAE architecture for the context encoder in TwistNet, and for proposal and twist in NAS-X and NASMC. We train initial distributions – both posterior and prior – using a smaller model with the same architecture. We let our emission distribution be factorized, and let $p\left(y_k^i=1\mid Z_{{\tau }_k}^i=A\right)=\sigma \left({\theta }_{\text{detect}}\right)$, where ${\theta }_{\text{detect}}$ is a scalar logit and $\sigma$ the logistic function. Due to their poor performance, we do not compare with TAG and BPF. For each method we also train an encoder for the first observation, used in the prediction task, by minimizing a forward KL loss in the sleep phase. We report experimental details in Appendix D.3.

Results.

Table 2 reports results both with full-week covariates and with covariates restricted to those available up to each timepoint. TwistNet with the KL loss achieves the lowest average error in both regimes.

Table 2:

Binary cross-entropy loss for reconstruction and prediction of active fire maps, mean ± 2 standard errors across test trajectories.

	Full-week covariates		Past covariates only
	Reconstruction	Prediction	Reconstruction	Prediction
TwistNet - KL	0.877 ± 0.36	1.149 ± 0.41	1.046 ± 0.389	1.424 ± 0.499
TwistNet - DRE	1.985 ± 0.59	1.541 ± 0.52	1.702 ± 0.532	1.649 ± 0.536
NAS-X	2.392 ± 0.72	1.755 ± 0.53	1.787 ± 0.531	1.576 ± 0.523
NASMC	2.008 ± 0.60	1.759 ± 0.57	3.350 ± 0.897	1.992 ± 0.563

By examining trajectories from the second best model (TwistNet with a DRE loss) in the full-week regime in Figure 5, we see a large gap in terms of how natural the trajectories look, and especially in how well the approximate posterior can bridge between observed states. Additional results are reported in Appendix E.2.

Figure 5:

Observed active fire maps from WildFireSpreadTS (Gerard et al., 2023) and snapshots of approximate posterior samples at observation times, for the full-week covariate regime. A pixel can only turn active if another one in its Moore neighborhood is, and transitions from unburned to burned are impossible. Observations are noisy, as satellite imagery for wildfires can be obscured by smoke and clouds (Schroeder et al., 2014).

5. CONCLUSIONS

We introduced an efficient posterior inference framework for systems of latent interacting CTMCs, and empirically demonstrated its effectiveness on challenging tasks on latent state inference and parameter learning for both simulated and real data. Our proposed approach outperforms the baselines, and in our experiments we saw this performance gap increase with the dimensionality of the problem under consideration. A major bottleneck in the training of our method is the cost of simulating trajectories with twisted SMC. Recent advances in latent SDEs have shown that simulation-free algorithms are possible (Bartosh et al., 2025; Kiyohara et al., 2025), and we believe adapting these methods to latent IPSs is an exciting avenue for future work.