SlowGuess/NeurIPS_2023_OCR huggingface dataset:
To enable fast log-determinant computations for arbitrary graph structures, we follow [43] and parameterize  $\mathbf{s}_{\theta}$  in terms of  $K + 1$  independent DGMRFs  $\mathbf{z}_k = \mathbf{S}_k\mathbf{h}_k + (\mathbf{b}_s)_k$ , each defining a GMRF w.r.t  $\mathcal{G}_{\mathrm{spatial}} = (\mathcal{V},\mathcal{E}_{\mathrm{spatial}})$ . Note that due to the multi-layer construction (see Section 2.4),  $\mathcal{G}_{\mathrm{spatial}}$  is implicitly defined by the base graph  $\bar{\mathcal{G}}_{\mathrm{spatial}}$  and the number of layers  $L$ , with special case  $\mathcal{G}_{\mathrm{spatial}} = \bar{\mathcal{G}}_{\mathrm{spatial}}$  if  $L_{\mathrm{spatial}} = 1$ .
3.2.2 Temporal layer(s)
Compared to  $\mathbf{s}_{\theta}$ , the definition of  $\mathbf{f}_{\theta}$  is much less constrained as it does not affect the log-determinant computation (see Section 3.3.1), and should hence incorporate as much domain knowledge into transition matrices  $\mathbf{F}_k$  as possible. This could be, for example, based on conservation laws or knowledge about relevant covariates. To increase the flexibility in cases where the dynamics are unknown or involve long-range dependencies, each  $\mathbf{F}_k$  can again be decomposed into simpler layers  $\mathbf{F}_k^{(L_{\mathrm{temporal}})}\cdots\mathbf{F}_k^{(1)}$ , each defined according to a temporal base graph  $\bar{\mathcal{G}}_{\mathrm{temporal}}$ . Table 1 provides several example layers to give an idea of what is possible.
Note that any function can be used to define the entries of  $\mathbf{F}_k^{(l)}$ , including neural networks taking available covariates or node/edge features as inputs [46]. The associated parameters can simply be included in  $\theta$ . Similarly, non-linear dynamics can be approximated either through linearization akin to the Extended Kalman Filter [37], or through neural linearization [19, 7]. We, however, recommend sharing parameters where appropriate to avoid overparameterization.
Table 1: Examples of linear transition layers  $\mathbf{F}_k^{(l)}$ . The adjacency matrix  $\mathbf{A}$  can be symmetric or asymmetric, weighted or unweighted.
layer definition	properties of Gtemporal
AR process	F(kl) = λk,lI	self-edges only
Diffusion	F(kl) = λk,lI + ωk,l(A - D)	bidirected (symmetric A)
Directed flow	F(kl) = λk,lI + ωk,l(A - Dout) + ζk,l(A^T - D_in)	directed
Advection	(Fk(l))ij = -1/2 w_ij n^T_{ij} v_l(Fk(l))_{ii} = 1 - ∑j∈n(i) (Fk(l))_{ij}	discretized Rd (e.g. triangulation), edge weights w_ij and normals n_ij
Neural network	(Fk(l))_{ij} = fNN(u_i, u_j, e_{ij})	node and edge features u_i, e_{ij}
Higher order Markov processes It is straightforward to extend  $\mathbf{f}_{\theta}$  to describe a  $p$ -th order Markov process
$$
\mathbf {x} _ {k} = \sum_ {\tau = 1} ^ {p} \mathbf {F} _ {k, \tau} \mathbf {x} _ {k - \tau} + \mathbf {c} _ {k} + \boldsymbol {\epsilon} _ {k}, \quad \boldsymbol {\epsilon} \sim \mathcal {N} (\mathbf {0}, \mathbf {Q} _ {k} ^ {- 1}), \quad \mathbf {x} _ {0} \sim \mathcal {N} (\boldsymbol {\mu} _ {0}, \mathbf {Q} _ {0} ^ {- 1}), \tag {12}
$$
where  $\mathbf{x}_{-\tau} = \mathbf{0}$  for  $\tau = 1,\dots ,p$ . This can be done by introducing edges  $\mathcal{E}_{\mathrm{temporal}(2)},\ldots ,\mathcal{E}_{\mathrm{temporal}(p)}$  and adding the corresponding higher-order transition matrices  $(\mathbf{F}_{\tau ,\tau},\dots ,\mathbf{F}_{K,\tau})$  to the  $\tau$ -th lower block diagonal of  $\mathbf{F}$  (see Appendix A.2 for derivations).
3.3 Learning and inference
Although the marginal likelihood  $p(\mathbf{y} \mid \boldsymbol{\theta})$  is available in closed form (see Section 2.3), maximum likelihood parameter estimation becomes computationally prohibitive in high dimensional settings
[50]. Instead, a variational approximation is used to learn parameters  $\hat{\pmb{\theta}}$  for large-scale DGMRFs. Then, the conjugate gradient method allows to efficiently compute the exact posterior mean and to draw samples from  $p(\mathbf{x} \mid \mathbf{y}, \hat{\pmb{\theta}})$ .
3.3.1 Scalable parameter estimation
Given a variational distribution  $q_{\phi}(\mathbf{x}) = \mathcal{N}(\pmb{\nu}_{\phi}, \pmb{\Lambda}_{\phi})$ , the parameters  $\{\pmb{\theta}, \pmb{\phi}\}$  are optimized jointly by maximizing the Evidence Lower Bound (ELBO)
$$
\begin{array}{l} \mathcal {L} \left(\mathbf {y} _ {0: K}, \boldsymbol {\theta}, \phi\right) = \mathbb {E} _ {q _ {\phi} (\mathbf {x})} \left[ \log p _ {\theta} (\mathbf {x}) + \sum_ {k = 0} ^ {K} \log p \left(\mathbf {y} _ {k} \mid \mathbf {x} _ {k}\right) \right] + H \left[ q _ {\phi} (\mathbf {x}) \right] (13) \\ = - \frac {1}{2} \mathbb {E} _ {q _ {\phi} (\mathbf {x})} \left[ \mathbf {g} _ {\theta} (\mathbf {x}) ^ {T} \mathbf {g} _ {\theta} (\mathbf {x}) + \sum_ {k = 0} ^ {K} \left(\mathbf {y} _ {k} - \mathbf {H} _ {k} \mathbf {x} _ {k}\right) ^ {T} \mathbf {R} _ {k} ^ {- 1} \left(\mathbf {y} _ {k} - \mathbf {H} _ {k} \mathbf {x} _ {k}\right) \right] (14) \\ + \log | \det (\mathbf {G} _ {\theta}) | + \frac {1}{2} \log | \det (\boldsymbol {\Lambda} _ {\phi}) | - \frac {1}{2} \sum_ {k = 0} ^ {K} \log | \det (\mathbf {R} _ {k}) | + \text {c o n s t}, (15) \\ \end{array}
$$
using stochastic gradient descent, where the expectation is replaced by a Monte-Carlo estimate based on samples  $\hat{\mathbf{x}}\sim q_{\phi}$
Variational distribution To facilitate efficient log-determinant computations and sampling via the reparameterization trick [33], we follow [43] and define  $q_{\phi}$  as an affine transformation  $\mathbf{x} = \mathbf{P}_{\phi}\mathbf{z} + \pmb{\nu}_{\phi}$  with  $\mathbf{z} \sim \mathcal{N}(\mathbf{0},\mathbf{I})$ , resulting in  $\Lambda_{\phi} = \mathbf{P}_{\phi}(\mathbf{P}_{\phi})^{T}$ . Then,  $\mathbf{P}_{\phi}$  can be defined as a block-diagonal matrix  $\mathrm{diag}(\mathbf{P}_0,\dots ,\mathbf{P}_K)$  with  $\mathbf{P}_k = \mathrm{diag}(\pmb {\rho}_k)\tilde{\mathbf{S}}_k\mathrm{diag}(\pmb {\psi}_k)$ , where  $\pmb {\rho}_k,\pmb {\psi}_k\in \mathbb{R}_+^N$  are variational parameters and  $\tilde{\mathbf{S}}_k$  is based on spatial DGMRF layers as discussed in Section 3.2.1. To relax the temporal independence assumptions in  $\Lambda_{\phi}$ , we propose an extension  $\mathbf{P}_{\phi} = \mathrm{diag}(\mathbf{P}_0,\dots ,\mathbf{P}_K)\tilde{\mathbf{F}}$  where  $\tilde{\mathbf{F}}$  has a similar structure to  $\mathbf{F}$  (see Eq. (8)). Again, the design of  $\tilde{\mathbf{F}}_k$  is highly flexible as it does not enter the log-determinant computation.
Log-determinant computations [50, 43] proposed specific lattice and graph layers for which the associated log-determinants  $\log |\operatorname{det}(\mathbf{G}_{\theta})|$  are computationally scalable. Conveniently, the spatiotemporal case does not require a new type of layer. Instead, we can build directly on top of existing spatial layers. Specifically, using Eq. (8)&(10), the log-determinant  $\log |\operatorname{det}(\mathbf{G}_{\theta})|$  simplifies to
$$
\log | \det (\mathbf {S F}) | \stackrel {(i)} {=} \log | \det (\mathbf {S}) | \stackrel {(i i)} {=} \sum_ {k = 0} ^ {K} \log | \det (\mathbf {S} _ {k}) |, \tag {16}
$$
where  $(i)$  follows from  $\mathbf{S}$  being block-diagonal and  $(ii)$  follows from  $\mathbf{F}$  being unit-lower triangular with  $\operatorname*{det}(\mathbf{F}) = 1$ . Note that each  $\mathbf{S}_k$  is defined by a spatial DGMRF of dimension  $N$  (see Section 3.2.1), for which efficient log-determinant methods have been developed. The same arguments hold for the variational distribution proposed above. Finally, due to the diagonality assumptions on  $\mathbf{R}_t$  (see Section 2.1),  $\frac{1}{2}\log |\operatorname*{det}(\mathbf{R}_k)|$  simplifies to  $\sum_{i=1}^{M_k}\log (\sigma_i)$ .
Computational complexity During training,  $\log |\operatorname{det}(\mathbf{G}_{\theta})|$  and  $\log |\operatorname{det}(\boldsymbol{\Lambda}_{\phi})|$  can be computed in  $\mathcal{O}(KNL_{\mathrm{spatial}})$  using Eq. (16) together with the methods proposed in [43]. Note that no complexity is added when introducing conditional dependencies between time steps. The necessary preprocessing steps, i.e. computing eigenvalues or traces, only need to be performed for the spatial base graph and are thus independent of the number of transitions  $K$ . Finally, assuming an average of  $d_{\mathrm{spatial}}$  and  $d_{\mathrm{temporal}}$  edges per node in  $\bar{\mathcal{E}}_{\mathrm{spatial}}$  and  $\bar{\mathcal{E}}_{\mathrm{temporal}}$ , the transformations  $\mathbf{s}_{\theta}$  and  $\mathbf{f}_{\theta}$  scale linearly with the number of nodes  $N$ . In particular,  $\mathbf{g}_{\theta}(\mathbf{x}) = (\mathbf{s}_{\theta} \circ \mathbf{f}_{\theta})(\mathbf{x})$  can be computed in  $\mathcal{O}(KNd_{\mathrm{spatial}}L_{\mathrm{spatial}} + KNd_{\mathrm{temporal}}L_{\mathrm{temporal}})$ , which dominates the computational complexity of the training loop. In addition, we can leverage massively parallel GPU computations to speed up this process even further.
3.3.2 Exact inference with conjugate gradients
As discussed in Section 2.3, the conjugate gradient method can be used to iteratively compute the posterior mean and obtain Monte Carlo estimates of marginal variances. Importantly, each iteration
is dominated by a single matrix-vector multiplication of the form
$$
\boldsymbol {\Omega} ^ {+} \mathbf {x} = \mathbf {F} ^ {T} \mathbf {S} ^ {T} \mathbf {S} \mathbf {F} \mathbf {x} + \mathbf {H} ^ {T} \mathbf {R} ^ {- 1} \mathbf {H} \mathbf {x}. \tag {17}
$$
This amounts to a series of sparse matrix-vector multiplications with total computational complexity  $\mathcal{O}(KNd_{\mathrm{spatial}}L_{\mathrm{spatial}} + KNd_{\mathrm{temporal}}L_{\mathrm{temporal}})$ , which is again linear in the number of nodes, time steps, and layers, and can be implemented in parallel on a GPU.
We implemented ST-DGMRF in Pytorch and Pytorch Geometric, and conducted experiments on a consumer-grade GPU, leveraging parallel computations in both spatial and temporal layers. In all experiments, we optimize parameters for 10 000 iterations using Adam [32] with learning rate 0.01, and draw 100 posterior samples to estimate marginal variances. Unless specified otherwise, we use  $L_{\mathrm{spatial}} = 2$ ,  $L_{\mathrm{temporal}} = 4$  and  $p = 1$ , and define the variational distribution based on one spatial layer and one temporal diffusion layer. Additional details on our experiments are provided in Appendix B.
4.1 Advection-diffusion process
We start with a synthetic dataset for which we have access to both the ground truth posterior distribution and transition matrix. The dataset consists of  $K = 20$  system states that are sampled from an ST-DGMRF with time-invariant transition matrix  $\mathbf{F}_{\mathrm{adv - diff}}$ . This matrix is defined according to the third-order Taylor approximation of an advection-diffusion process with constant velocity and diffusion parameters, discretized on a  $30 \times 30$  lattice with periodic boundary conditions. From the sampled state trajectory, we generate observations with varying amount of missing data by removing pixels within a square mask of width  $w$  for 10 consecutive time steps and adding noise with  $\sigma = 0.01$ . For all experiments, we use the masked pixels as test set, and  $10\%$  of the observed pixels as validation set for hyperparameter tuning.
Two ST-DGMRF variants are considered: (1) using advection-diffusion matrices  $\mathbf{F}_k^{(l)}$  defined based on the first-order Taylor approximation of the ground-truth dynamics, with trainable diffusion and velocity parameters, and (2) replacing parts of these advection-diffusion matrices with small neural networks, taking the edge unit vector  $\mathbf{n}_{ij}$  pointing from pixel  $i$  to  $j$  as input. In both cases,  $\bar{\mathcal{G}}_{\mathrm{temporal}}$  and  $\bar{\mathcal{G}}_{\mathrm{spatial}}$  are defined as the 4-nearest neighbor graph. Further, as the underlying dynamics are time-invariant, we enforce  $\mathbf{F}_k^{(l)} = \mathbf{F}_{k + 1}^{(l)}$  and  $\mathbf{S}_k^{(l)} = \mathbf{S}_{k + 1}^{(l)}\forall k\geq 1$
Baselines We compare our approach to a range of baselines accounting for varying degrees of spatial and/or temporal dependencies. In particular, we consider the original DGMRF applied to all time steps independently, an ARMA state-space model assuming spatial independence among the state variables, and a spatiotemporal AR state-space model (ST-AR) with spatially correlated error terms  $\epsilon_{k}$  for which an unconstrained covariance matrix  $\mathbf{Q}^{-1}$  is estimated using the EM algorithm. For both ARMA and ST-AR, we use the standard Kalman smoother [47] to obtain posterior estimates. Additionally, we consider two Ensemble Kalman Smoother (EnKS) variants, one with an advection-diffusion transition model matching the
Table 2: Performance on the advection-diffusion data with  $w = 9$ . We report the mean over 5 runs with different random seeds.
RMSEμ	RMSEσ	CRPS
ARMA	2.3054	0.6812	1.7064
ST-AR	1.4595	1.9216	0.9707
DGMRF	0.5901	0.3808	0.3495
EnKS
true dynamics	0.0661	0.0046	0.1027
estimated dynamics	0.1654	0.0039	0.1434
ST-DGMRF (ours)
advection-diffusion	0.0526	0.1146	0.0726
neural network	0.0854	0.1402	0.0839
true data-generating process, and one using state augmentation to estimate the velocity and diffusion parameters jointly with the system states. Note that, in contrast to the ST-DGMRF approach, we consider initial and transition noise parameters to be fixed in order to avoid divergence of the EnKS.
Performance evaluation We evaluate the estimated posterior mean and marginal standard deviations in terms of the root-mean-square-error  $(\mathrm{RMSE}_{\mu}$  and  $\mathrm{RMSE}_{\sigma})$  with respect to the ground truth posterior. In addition, we use the mean negative continuous ranked probability score (CRPS) [22] to evaluate the predictive distribution with respect to the masked out data. Table 2 shows that, by exploiting the spatiotemporal structure of the process, our ST-DGMRF variants provide much more
accurate estimates than the purely spatial DGMRF, the purely temporal ARMA model, and the ST-AR model with highly simplified transitions and unstructured noise terms. As expected, the two EnKS variants with fixed noise parameters provide the most accurate uncertainty estimates. However, the CRPS scores indicate that overall our ST-DGMRF approach results in better calibrated posterior distributions. More detailed results are reported in Appendix C.
DGMRF
ST-DGMRF
Figure 2: Left: snapshot of the advection-diffusion data at time  $k = 8$ , and reconstructions by the time-independent DGMRF and our ST-DGMRF with advection-diffusion  $\mathbf{F}_k$ . Center: corresponding time series for a single pixel (marked on the left). Shaded areas represent posterior mean ± std of a single run. Right:  $\mathrm{RMSE}_{\mu}$  as a function of the mask width (mean ± std over 5 runs).
Increasing mask size To evaluate the robustness to missing data, we analyze how posterior estimates change with increasing mask size. Figure 2 (right) shows the effect on  $\mathrm{RMSE}_{\mu}$  when varying  $w$  from 6 to 12. Clearly, the purely spatial DGMRF suffers the most from expanding the unobserved region. In contrast, the ST-DGMRF variants continue to provide accurate state estimates that are on par with the EnKS using ground truth dynamics. Note that the decreasing errors for ST-AR can be attributed to the changing set of pixels used for evaluation.
Figure 3: Stencil comparison. Center: absolute error between true and learned transition weights for the neural network based ST-DGMRF. Right: Pearson correlation for increasing temporal depth.
Evaluation of learned transition models Finally, we assess how well the temporal ST-DGMRF layers can approximate the true data-generating dynamics, given varying levels of complexity in the transition model. For this purpose, we extract the stencil, i.e. the weights assigned to nearby pixels, from the ground truth and learned transition matrices and compare them in terms of absolute error and Pearson correlation. We find that the learned transition weights converge rapidly towards the true weights as the number of temporal layers increases (see Figure 3). For small  $L_{\mathrm{temporal}}$ , the neural network based layers show a stronger agreement with the true dynamics, suggesting that some flexibility in the temporal layers helps to compensate for simplifications in the transition structure.
4.2 Air quality data
To test our method on a real world system exhibiting more complex dynamics and graph structures, we conduct experiments on an air quality dataset obtained from [59]. The dataset contains hourly PM2.5 measurements from 246 sensors distributed around Beijing, China, covering a time period of  $K = 400$  hours. To ensure that predicted PM2.5 concentrations are non-negative, we model both system states and observations in log-space. The spatial and temporal base graphs are defined
based on the Delaunay triangulation, with edge weights proportional to the inverse distance between sensors. To mimic a realistic scenario of local network failures, we randomly choose 10 time points  $t_k$  and mask out all data points within a spatial block containing  $50\%$  of all sensors, for time steps  $t_k, \ldots, t_k + 20$  (see Figure 4). As before, we use the masked data for model evaluation, and use  $10\%$  of the remaining data as validation set. The observation noise is assumed to be uniform with  $\sigma = 0.01$ .
The spatiotemporal distribution of particulate matter is strongly influenced by atmospheric processes that transport and diffuse emitted particles. To incorporate this knowledge into the transition model, we extract temperature and wind conditions for all sensors and time points from the ERA5 reanalysis dataset [25] and feed them together with static graph features into a set of neural networks which transform them into spatially and temporally varying bias and velocity parameters. The matrices  $\mathbf{F}_k^{(l)}$  are then formed in the same way as in Section 4.1. In addition, we consider a simplified diffusion transition model (see Table 1) which cannot capture any directional transport processes.
Results Next to the baselines introduced in Section 4.1, we consider a multi-layer perceptron (MLP) mapping local weather features to the corresponding log-transformed PM2.5 concentration. Following [50], we also include weather features into the spatial DGMRF model by adding linear effects to the measurement model. Table 3 reports the resulting CRPS and the RMSE with respect to the masked out data. Clearly, the neural network based ST-DGMRF, accounting for time-varying directional transport processes, provides the most accurate state estimates. However, even with highly simplified diffusion transitions our approach provides better estimates than the considered baselines. We attribute this to the expressive DGMRF noise terms which can capture complex error structures
Table 3: Performance on the air quality data. We report the mean over 5 runs with different random seeds.
p	RMSE	CRPS
ARMA	0.6820	0.3625
ST-AR	0.7350	0.4261
DGMRF	0.7456	0.4037
MLP	0.8038	-
ST-DGMRF (ours)
diffusion	1	0.6190	0.3258
diffusion	2	0.5928	0.3161
neural network	1	0.5853	0.3092
neural network	2	0.5565	0.2925
resulting from inaccuracies in the transition model. Increasing the Markov order from  $p = 1$  to 2 clearly improves the resulting posterior estimates for both ST-DGMRF variants, reflecting the complexity of the modeled process. A similar trend is observed as we increase  $L_{\mathrm{temporal}}$  (see Appendix C).
Figure 4: ST-DGMRF posterior estimates (neural network  $\mathbf{F}_k$  and  $p = 2$ ). Left: time series of log-transformed and normalized PM2.5 levels for a masked sensor. Light red areas represent posterior mean  $\pm$  std of a single run. Unobserved time points are marked with gray bars. Right: Marginal std for all sensors at  $k = 100$ . Masked out sensors (gray box) feature higher uncertainties.
We have presented ST-DGMRF, an extension to Deep Gaussian Markov Random Fields for inference in spatiotemporal dynamical systems with partial and noisy observations, (partially) unknown dynamics, and limited historical data. Our reformulation of graph-structured state-space models as multi-layer space-time GMRFs enables computationally efficient learning and inference even in high-dimensional settings, scaling linear w.r.t. the number of both time steps and state variables. Empirically, ST-DGMRF provides more accurate posterior estimates than other scalable approaches relying on simplifications of the dependency structure. While our approach relies on linear Gaussian assumptions, which can be restrictive for real systems, we find that expressive DGMRF noise terms can compensate (to a certain extent) for approximations in the transition model. In the future, non-linearities between the layers could be explored as discussed in [50]. Further, we see potential in defining more flexible time-varying transition matrices based on a hierarchy of latent variables that again follow a ST-DGMRF.
[51] A. Solin and S. Särkkä. Explicit link between periodic covariance functions and state space models. In International Conference on Artificial Intelligence and Statistics (AISTATS). PMLR, 2014.
[52] J. R. Stroud and T. Bengtsson. Sequential state and variance estimation within the ensemble Kalman filter. Monthly Weather Review, 135(9):3194-3208, 2007.
[53] I. Szunyogh, E. J. Kostelich, G. Gyarmati, E. Kalnay, B. R. Hunt, E. Ott, E. Satterfield, and J. A. Yorke. A local ensemble transform Kalman filter data assimilation system for the NCEP global model. Tellus A: Dynamic Meteorology and Oceanography, 60(1):113-130, 2008.
[54] H. Wang, A. Bhattacharya, D. Pati, and Y. Yang. Structured variational inference in Bayesian state-space models. In International Conference on Artificial Intelligence and Statistics (AISTATS). PMLR, 2022.
[55] J. S. Whitaker, T. M. Hamill, X. Wei, Y. Song, and Z. Toth. Ensemble data assimilation with the NCEP global forecast system. Monthly Weather Review, 136(2):463-482, 2008.
[56] C. K. Wikle and M. B. Hooten. A general science-based framework for dynamical spatio-temporal models. TEST, 19:417-451, 2010.
[57] C. K. Wikle and A. Zammit-Mangion. Statistical deep learning for spatial and spatiotemporal data. Annual Review of Statistics and Its Application, 10:247-270, 2023.
[58] W. Wilkinson, A. Solin, and V. Adam. Sparse algorithms for Markovian Gaussian processes. In International Conference on Artificial Intelligence and Statistics (AISTATS). PMLR, 2021.
[59] Y. Zheng, X. Yi, M. Li, R. Li, Z. Shan, E. Chang, and T. Li. Forecasting fine-grained air quality based on big data. In SIGKDD Conference on Knowledge Discovery and Data Mining, 2015.
Appendix A provides derivations supporting Section 3 in the main paper. In Appendix B, we explain our experimental setup, including dataset preparation and model implementation, in more detail. Finally, Appendix C provides additional results supporting our claims regarding the scalability of our method, together with additional results from the experiments presented in Section 4.
A ST-DGMRF derivations
In this section we provide detailed derivations of the ST-DGMRF joint distribution, for both first-order transition models (Section A.1) and higher-order transition models (Section A.2).
A.1 Joint distribution
The LDS (see Section 2.2 and 3.1 in the main paper) defines a joint distribution over system states  $\mathbf{x}_{0:K}$  that factorizes as
$$
p \left(\mathbf {x} _ {0: K}\right) = \mathcal {N} \left(\mathbf {x} _ {0} \mid \boldsymbol {\mu} _ {0}, \mathbf {Q} _ {0} ^ {- 1}\right) \prod_ {k = 1} ^ {K} \mathcal {N} \left(\mathbf {x} _ {k} \mid \mathbf {F} _ {k} \mathbf {x} _ {k - 1} + \mathbf {c} _ {k}, \mathbf {Q} _ {k} ^ {- 1}\right), \tag {18}
$$