SlowGuess/NeurIPS_2023_OCR huggingface dataset:
with  $\mathbf{x}_k, \pmb{\mu}_0, \mathbf{c}_k \in \mathbb{R}^N$  and  $\mathbf{F}_k, \mathbf{Q}_k \in \mathbb{R}^{N \times N}$ . As a product of Gaussian distributions,  $p(\mathbf{x}_{0:K})$  can be written as a joint Gaussian  $\mathcal{N}(\pmb{\mu}, \pmb{\Omega}^{-1})$  with mean  $\pmb{\mu} \in \mathbb{R}^{(K + 1)N}$  and precision (inverse covariance) matrix  $\Omega \in \mathbb{R}^{(K + 1)N \times (K + 1)N}$ . Here, we derive expressions for  $\pmb{\mu}$  and  $\pmb{\Omega}$  in terms of  $\pmb{\mu}_0, \mathbf{c}_k, \mathbf{F}_k, \mathbf{Q}_k$ .
First, note that Eq. (18) can be written as a set of linear equations
$$
\mathbf {x} _ {0} = \boldsymbol {\mu} _ {0} + \boldsymbol {\epsilon} _ {0} \quad \boldsymbol {\epsilon} _ {0} \sim \mathcal {N} (\mathbf {0}, \mathbf {Q} _ {0} ^ {- 1})
$$
$$
\mathbf {x} _ {1} = \mathbf {F} _ {1} \mathbf {x} _ {0} + \mathbf {c} _ {1} + \boldsymbol {\epsilon} _ {1} \quad \boldsymbol {\epsilon} _ {1} \sim \mathcal {N} (\mathbf {0}, \mathbf {Q} _ {1} ^ {- 1})
$$
$$
\mathbf {x} _ {2} = \mathbf {F} _ {2} \mathbf {x} _ {1} + \mathbf {c} _ {2} + \epsilon_ {2} \quad \epsilon_ {2} \sim \mathcal {N} (\mathbf {0}, \mathbf {Q} _ {2} ^ {- 1})
$$
$$
\dots
$$
$$
\mathbf {x} _ {K} = \mathbf {F} _ {K} \mathbf {x} _ {K - 1} + \mathbf {c} _ {K} + \boldsymbol {\epsilon} _ {K} \quad \boldsymbol {\epsilon} _ {K} \sim \mathcal {N} (\mathbf {0}, \mathbf {Q} _ {K} ^ {- 1}).
$$
Moving all  $\mathbf{x}_k$ -terms to the left-hand side, we can rewrite this as a matrix-vector multiplication
$$
\underbrace {\left[ \begin{array}{c c c c c} \mathbf {I} & & & & \\ - \mathbf {F} _ {1} & \mathbf {I} & & & \\ & - \mathbf {F} _ {2} & \mathbf {I} & & \\ & & \dots & \dots & \\ & & & - \mathbf {F} _ {K} & \mathbf {I} \end{array} \right]} _ {= \mathbf {F}} \cdot \underbrace {\left[ \begin{array}{l} \mathbf {x} _ {0} \\ \mathbf {x} _ {1} \\ \mathbf {x} _ {2} \\ \vdots \\ \mathbf {x} _ {K} \end{array} \right]} _ {= \mathbf {x}} = \underbrace {\left[ \begin{array}{l} \boldsymbol {\mu} _ {0} \\ \mathbf {c} _ {1} \\ \mathbf {c} _ {2} \\ \vdots \\ \mathbf {c} _ {K} \end{array} \right]} _ {= \mathbf {c}} + \underbrace {\left[ \begin{array}{l} \boldsymbol {\epsilon} _ {0} \\ \boldsymbol {\epsilon} _ {1} \\ \boldsymbol {\epsilon} _ {2} \\ \vdots \\ \boldsymbol {\epsilon} _ {K} \end{array} \right]} _ {= \epsilon}, \tag {19}
$$
with block-matrix  $\mathbf{F} \in \mathbb{R}^{(K + 1)N \times (K + 1)N}$  and vectorized  $\mathbf{x} = \mathrm{vec}(\mathbf{x}_0, \ldots, \mathbf{x}_K) \in \mathbb{R}^{(K + 1)N}$ ,  $\mathbf{c} = \mathrm{vec}(\pmb{\mu}_0, \mathbf{c}_1, \ldots, \mathbf{c}_K) \in \mathbb{R}^{(K + 1)N}$  and  $\pmb{\epsilon} = \mathrm{vec}(\pmb{\epsilon}_0, \ldots, \pmb{\epsilon}_K) \in \mathbb{R}^{(K + 1)N}$ . Empty positions in  $\mathbf{F}$  represent zero-blocks.
Now, we can express  $\mathbf{x}$  as an affine transformation of  $\epsilon$
$$
\mathbf {x} = \mathbf {F} ^ {- 1} \mathbf {c} + \mathbf {F} ^ {- 1} \boldsymbol {\epsilon}, \tag {20}
$$
where  $\mathbf{F}^{-1}$  exists because  $\operatorname*{det}(\mathbf{F}) = 1$ . Since  $\epsilon$  is distributed as  $\epsilon \sim \mathcal{N}(\mathbf{0},\mathbf{Q}^{-1})$  with  $\mathbf{Q} = \mathrm{diag}(\mathbf{Q}_0,\mathbf{Q}_1,\dots ,\mathbf{Q}_K)$ , and  $\mathbf{c}$  is deterministic, we can use the affine property of Gaussian distributions to obtain the joint distribution
$$
\mathbf {x} \sim \mathcal {N} \left(\mathbf {F} ^ {- 1} \mathbf {c}, \mathbf {F} ^ {- 1} \mathbf {Q} ^ {- 1} \mathbf {F} ^ {- T}\right). \tag {21}
$$
Thus, the joint precision matrix  $\Omega$  factorizes as
$$
\boldsymbol {\Omega} = \left(\mathbf {F} ^ {- 1} \mathbf {Q} ^ {- 1} \mathbf {F} ^ {- T}\right) ^ {- 1} = \mathbf {F} ^ {T} \mathbf {Q} \mathbf {F} \tag {22}
$$
and has a block-tridiagonal structure
$$
\boldsymbol {\Omega} = \left[ \begin{array}{c c c c c} \mathbf {Q} _ {0} + \mathbf {F} _ {1} ^ {T} \mathbf {Q} _ {1} \mathbf {F} _ {1} & - \mathbf {F} _ {1} ^ {T} \mathbf {Q} _ {1} & & & \\ - \mathbf {Q} _ {1} \mathbf {F} _ {1} & \mathbf {Q} _ {1} + \mathbf {F} _ {2} ^ {T} \mathbf {Q} _ {2} \mathbf {F} _ {2} & - \mathbf {F} _ {2} ^ {T} \mathbf {Q} _ {2} & & \\ & \dots & \dots & \dots & \\ & & - \mathbf {Q} _ {K - 1} \mathbf {F} _ {K - 1} & \mathbf {Q} _ {K - 1} + \mathbf {F} _ {K} ^ {T} \mathbf {Q} _ {K} \mathbf {F} _ {T} & - \mathbf {F} _ {K} ^ {T} \mathbf {Q} _ {K} \\ & & & - \mathbf {Q} _ {K} \mathbf {F} _ {K} & \mathbf {Q} _ {K} \end{array} \right]. \tag {23}
$$
Note that for matrix-vector multiplications of the form  $\Omega \mathbf{x}$ , the sparse structure of  $\mathbf{F}$  and  $\mathbf{Q}$  can be leveraged by performing three consecutive matrix-vector multiplications, instead of first forming the full precision matrix and then computing the matrix vector product. This reduces both computations and memory requirements.
To compute the joint mean  $\pmb{\mu} = \mathbf{F}^{-1}\mathbf{c}$  without expensive matrix inversion, the components  $\pmb{\mu}_k$  have to be computed iteratively as
$$
\boldsymbol {\mu} _ {k} = \mathbf {F} _ {k} \boldsymbol {\mu} _ {k - 1} + \mathbf {c} _ {k}. \tag {24}
$$
In contrast, the information vector  $\eta = \Omega \mu$  can be expressed compactly as
$$
\boldsymbol {\eta} = \mathbf {F} ^ {T} \mathbf {Q} \mathbf {F} \mathbf {F} ^ {- 1} \mathbf {c} = \mathbf {F} ^ {T} \mathbf {Q} \mathbf {c}, \tag {25}
$$
which can be computed efficiently using sparse and parallel matrix-vector multiplications on a GPU. We make use of this property in the DGMRF formulation and in the conjugate gradient method.
A.2 Extension to higher-order Markov processes
We can easily adjust the joint distribution to accommodate higher-order processes with dependencies on multiple past time steps.
For a  $p$ -th order Markov process, the dynamics are defined by equations
$$
\mathbf {x} _ {1} = \mathbf {F} _ {1, 1} \mathbf {x} _ {0} + \mathbf {c} _ {1} + \boldsymbol {\epsilon} _ {1} \quad \boldsymbol {\epsilon} _ {1} \sim \mathcal {N} (\mathbf {0}, \mathbf {Q} _ {1} ^ {- 1})
$$
$$
\mathbf {x} _ {k} = \mathbf {F} _ {k, 1} \mathbf {x} _ {k - 1} + \mathbf {F} _ {k, 2} \mathbf {x} _ {k - 2} + \dots + \mathbf {F} _ {k, p} \mathbf {x} _ {k - p} + \mathbf {c} _ {k} + \boldsymbol {\epsilon} _ {k} \qquad \qquad \boldsymbol {\epsilon} _ {k} \sim \mathcal {N} (\mathbf {0}, \mathbf {Q} _ {k} ^ {- 1})
$$
$$
\mathbf {x} _ {K} = \mathbf {F} _ {K, 1} \mathbf {x} _ {K - 1} + \mathbf {F} _ {K, 2} \mathbf {x} _ {K - 2} + \dots + \mathbf {F} _ {K, p} \mathbf {x} _ {K - p} + \mathbf {c} _ {K} + \boldsymbol {\epsilon} _ {K} \quad \boldsymbol {\epsilon} _ {K} \sim \mathcal {N} (\mathbf {0}, \mathbf {Q} _ {K} ^ {- 1}).
$$
Following the same steps as before, this results in a linear system
$$
\underbrace {\left[ \begin{array}{c c c c c c} \mathbf {I} & & & & & \\ - \mathbf {F} _ {1 , 1} & \mathbf {I} & & & & \\ - \mathbf {F} _ {2 , 2} & - \mathbf {F} _ {2 , 1} & \mathbf {I} & & & \\ \dots & \dots & \dots & \dots & & \\ - \mathbf {F} _ {p , p} & - \mathbf {F} _ {p , p - 1} & \dots & - \mathbf {F} _ {p , 1} & \mathbf {I} & \\ & \dots & \dots & \dots & \dots & \dots \\ & & - \mathbf {F} _ {K, p} & \dots & - \mathbf {F} _ {K, 2} & - \mathbf {F} _ {K, 1} & \mathbf {I} \end{array} \right]} _ {= \mathbf {F}} \cdot \underbrace {\left[ \begin{array}{l} \mathbf {x} _ {0} \\ \mathbf {x} _ {1} \\ \mathbf {x} _ {2} \\ \vdots \\ \mathbf {x} _ {K} \end{array} \right]} _ {= \mathbf {x}} = \underbrace {\left[ \begin{array}{l} \boldsymbol {\mu} _ {0} \\ \mathbf {c} _ {1} \\ \mathbf {c} _ {2} \\ \vdots \\ \mathbf {c} _ {K} \end{array} \right]} _ {= \mathbf {c}} + \underbrace {\left[ \begin{array}{l} \boldsymbol {\epsilon} _ {0} \\ \boldsymbol {\epsilon} _ {1} \\ \boldsymbol {\epsilon} _ {2} \\ \vdots \\ \boldsymbol {\epsilon} _ {K} \end{array} \right]} _ {= \epsilon}. \tag {26}
$$
This means that the matrix  $\mathbf{F}$  is extended by adding the higher-order transition matrices  $(\mathbf{F}_{\tau,\tau},\ldots,\mathbf{F}_{K,\tau})$  to the  $\tau$ -th lower block diagonal of  $\mathbf{F}$  for all  $\tau = 1,\dots,p$ . The expressions for  $\Omega, \mu$  and  $\eta$  remain the same (using the extended  $\mathbf{F}$ ), resulting in a block  $p$ -diagonal precision matrix.
B Experimental details
B.1 Advection-diffusion process
The advection-diffusion dataset is a random sample from a ST-DGMRF for which the transition matrices are defined according to an advection-diffusion process
$$
\frac {\partial \rho (t , s)}{\partial t} = D \nabla^ {2} \rho (t, s) - \nabla \cdot (\mathbf {v} \rho (t, s)) \tag {27}
$$
with constant diffusion coefficient  $D$  and velocity vector  $\mathbf{v} = [u, v]^T$ . The process is discretized on a  $30 \times 30$  lattice with grid cell size  $\Delta x = \Delta y = 1$  and periodic boundary conditions. The spatial discretization results in a system of ordinary differential equations
$$
\frac {\partial \boldsymbol {\rho} (t)}{\partial t} = \mathbf {M} \boldsymbol {\rho} (t), \tag {28}
$$
where  $\rho(t)$  is a vector containing the system states of all grid cells. Using a finite difference discretization, matrix  $\mathbf{M}$  is defined as
$$
\mathbf {M} _ {i j} = \left\{ \begin{array}{l l} D - \frac {1}{2} \mathbf {n} _ {i j} ^ {T} \mathbf {v} & \text {i f} d (i, j) = 1 \\ - 4 D & \text {i f} i = j \\ 0 & \text {o t h e r w i s e ,} \end{array} \right. \tag {29}
$$
where  $d(i,j)$  denotes the distance between cell  $i$  and  $j$ , and  $\mathbf{n}_{ij}$  denotes the unit vector pointing from lattice cell  $i$  to its neighbor  $j$ . For example, for cell  $i = (s_x,s_y)$  and cell  $j = (s_x,s_y - 1)$  it is  $\mathbf{n}_{ij} = [0, -1]^T$ , and thus  $\mathbf{n}_{ij}^T\mathbf{v} = -v$ .
Eq. 28 is converted into a discrete-time dynamical system by approximating
$$
\boldsymbol {\rho} _ {t + \Delta t} = \exp (\Delta t \cdot \mathbf {M}) \boldsymbol {\rho} _ {t} \approx \left(\sum_ {k = 0} ^ {3} \frac {1}{k !} (\Delta t) ^ {k} (\mathbf {M}) ^ {k}\right) \boldsymbol {\rho} _ {t} = \mathbf {F} _ {\text {a d v - d i f f}} \boldsymbol {\rho} _ {t} \tag {30}
$$
using a third-order Taylor series expansion. For simplicity, we use time resolution  $\Delta t = 1$  resulting in
$$
\mathbf {F} _ {\text {a d v - d i f f}} = \mathbf {I} + \mathbf {M} + \frac {1}{2} \mathbf {M} ^ {2} + \frac {1}{6} \mathbf {M} ^ {3}. \tag {31}
$$
B.1.1 Process simulation
We sample the initial state  $\rho_0$  from a GMRF with  $\mu_0 = 0$  and precision matrix  $\mathbf{Q}_0 = \mathbf{S}_0^T\mathbf{S}_0$  with  $\mathbf{S}_0 = (\mathbf{D} - \mathbf{A})$ , where  $\mathbf{A}$  is the adjacency matrix of the 4-nearest neighbour graph  $\mathcal{G}_{\mathrm{lattice}}$ , and  $\mathbf{D} = 4\cdot \mathbf{I}$  is the corresponding degree matrix. This corresponds to a 1-layer DGMRF with parameters  $\alpha = 1, \beta = -1, \gamma = 1$  and  $b = 0$ .
Starting from  $\rho_0$ , we iteratively sample the next system state according to
$$
\rho_ {k} = \mathbf {F} _ {k} \rho_ {k - 1} + \epsilon_ {k} \quad \epsilon_ {t} \sim \mathcal {N} (\mathbf {0}, \mathbf {Q} _ {k} ^ {- 1}) \tag {32}
$$
with time-invariant transition matrix  $\mathbf{F}_k = (\mathbf{F}_{\mathrm{adv - diff}})^4$  , where  $\mathbf{F}_{\mathrm{adv - diff}}$  is defined according to Eq. (31). This effectively aggregates four simulation steps into one, i.e.  $\Delta t_{k} = (t_{k + 1} - t_{k}) = 4$  , resulting in larger differences between consecutive system states. For the noise terms  $\epsilon_{k}$  , we use a time-invariant precision matrix  $\mathbf{Q}_k = \mathbf{S}_k^T\mathbf{S}_k$  where  $\mathbf{S}_k = (10\cdot \mathbf{I} - \mathbf{A})$  . This corresponds to a 1-layer DGMRF with parameters  $\alpha = \frac{10}{4},\beta = -1,\gamma = 1$  and  $b = 0$
We simulate for  $K = 20$  time steps, using  $D = 0.01$  and  $\mathbf{v} = [-0.3, 0.3]^T$ , and generate observations by masking out grid cells within a square of width  $w \in \{6, \dots, 12\}$  for 10 consecutive time steps and applying white noise with standard deviation  $\sigma = 0.01$ .
Figure 5: Advection-diffusion dataset with ground truth system states (bottom) and corresponding observations using masks of width  $w = 9$  (top).
B.1.2 ST-DGMRF parameterization
We consider two ST-DGMRF variants that capture different amounts of prior knowledge. In both cases, spatial and temporal layers are defined based on  $\mathcal{G}_{\mathrm{lattice}}$ , and temporal bias terms are, similar to spatial bias terms, defined as  $\mathbf{b}_f^{(l)} = b_f^{(l)}\mathbf{1}$ .
Variant 1 If prior knowledge is available in the form of a parameterized transition model, the ST-DGMRF transition matrices can be parameterized accordingly. Here, we consider temporal layers  $\mathbf{F}_k^{(l)}$  that are a simplified first-order approximation to the true transition matrix  $\mathbf{F}_{\mathrm{adv - diff}}$  used to generate the data (see Eq. (31)), i.e.  $\mathbf{F}_k^{(l)} = \mathbf{I} + \mathbf{M}^{(l)}$ , with time-invariant learnable diffusion coefficients  $D^{(l)}$  and velocity vectors  $\mathbf{v}^{(l)}$ . To ensure that the diffusion coefficient is non-negative, we model it as  $D^{(l)} = (d^{(l)})^2$ . This leaves us with four learnable parameters  $d^{(l)}, u^{(l)}, v^{(l)}$  and  $b_{f}^{(l)}$  per temporal layer.
Variant 2 If only partial knowledge about the underlying dynamics is available, the unknown parts can, for example, be replaced by a small neural network. Here, we consider temporal layers of the form  $\mathbf{F}_k^{(l)} = \mathbf{I} + \mathbf{M}^{(l)}$  with
$$
\mathbf {M} _ {i j} ^ {(l)} = \left\{ \begin{array}{l l} (d ^ {(l)}) ^ {2} + \phi_ {i j, 1} ^ {(l)} & \text {i f} j \in n (i) \\ - 4 (d ^ {(l)}) ^ {2} + \sum_ {j \in n (i)} \phi_ {i j, 2} ^ {(l)} & \text {i f} i = j \\ 0 & \text {o t h e r w i s e ,} \end{array} \right. \tag {33}
$$
where we define  $\phi_{ij,1}^{(l)},\phi_{ij,2}^{(l)} = f_{MLP}^{(l)}(\mathbf{n}_{ij})$  where  $f_{MLP}^{(l)}:\mathbb{R}^2\to \mathbb{R}^2$  is a multilayer perceptron (MLP) with one hidden layer of width 16 with ReLU non-linearity, and Tanh output non-linearity. Again, we define the transition model to be time-invariant and share MLP parameters across time and space. This amounts to 83 learnable parameters per temporal layer.
Log-determinant computations In our experiments, the spatial base graph  $\mathcal{G}_{\mathrm{lattice}}$  is small enough to pre-compute eigenvalues exactly and use the eigenvalue method for log-determinant computations proposed in [43].
Variational distribution For the variational distribution, we also consider two variants, one without temporal dependencies (equivalent to the DGMRF baseline) and one with a single temporal layer with time-invariant diffusion transition matrices  $\tilde{\mathbf{F}}_k = \lambda \mathbf{I} + \omega (\mathbf{A} - \mathbf{D})$ . Note that for  $\omega = 0$ , this reduces to a simple auto-regressive process.
Observation model All ST-DGMRF variants assume a temporally and spatially invariant observation noise level of  $\sigma = 0.01$ . The observation matrices  $\mathbf{H}_k$  are defined as selection matrices matching the training masks during the learning phase and the training plus validation masks during the testing phase.
B.2 Air quality data
The air quality dataset is based on hourly PM2.5 measurements obtained from [59]. We consider 246 sensors within the metropolitan area of Beijing, China, for which we extracted time series of  $K = 400$  hours between 13 March 2015 at 12pm and 30 March 2015 at 3am. Relevant weather covariates (surface temperature, as well as u and v wind components at 10 meters above ground level) were extracted from the ERA5 reanalysis dataset [25].
B.2.1 Data preprocessing
We define both the spatial and the temporal base graph based on the Delaunay triangulation of sensor locations,  $\mathcal{G}_{\text{Delaunay}}$ , where we disregard edges between sensors that are more than 160 kilometers apart. Edge weights are defined as the inverse distance between sensors, normalized to range between 0 and 1. The raw PM2.5 measurements are log-transformed and standardized to zero mean and unit variance. Finally, we remove clear outliers where the transformed values jump up and down by more than a threshold of  $\delta = 2.0$  within three consecutive time steps. The ERA5 covariates are normalized to range between -1 and 1.
To mimic a realistic setting of repeatedly occurring partial network failures, we define our test set by masking out all measurements within a predefined spatial block (containing  $50\%$  of all sensors) within 10 randomly placed windows of 20 time steps (see Figure 6). Note that these windows may overlap, resulting in fewer periods of missing data with variable length. The masked out measurements are used for the final model evaluation.
Figure 6: Left: air quality sensor network. Ca.  $50\%$  of the nodes are masked out (purple nodes within the gray box) during 10 randomly placed (partially overlapping) windows of 20 time steps. Right: associated log-transformed and normalized PM2.5 measurements for two sensors falling within the masked area. Time points that have either missing data or fall within a masked time window are shaded in gray.
B.2.2 ST-DGMRF parameterization
As with the advection-diffusion dataset, we consider two ST-DGMRF variants with different types of temporal layers. In both cases, spatial and temporal layers are defined based on the Delaunay triangulation described in Section B.2.1, and temporal bias terms  $\mathbf{b}_f$  are defined in terms of a neural network mapping local weather covariates to temporally and spatially varying biases. We use a simple MLP with one hidden layer of width 16 with ReLU activations and no output non-linearity. The MLP parameters are shared over both space and time.
Variant 1 This variant accounts for directional transport processes, adopting a transition model similar to the neural network model used in the advection-diffusion experiments. In particular, we consider temporal layers of the form  $\mathbf{F}_k^{(l)} = \mathbf{I} + \mathbf{M}_k^{(l)}$  with
$$
\left(\mathbf {M} _ {k} ^ {(l)}\right) _ {i j} = \left\{ \begin{array}{l l} (d ^ {(l)}) ^ {2} + \phi_ {k, i j} ^ {(l)} & \text {i f} j \in n (i) \\ - 4 (d ^ {(l)}) ^ {2} + \sum_ {j \in n (i)} \psi_ {k, i j} ^ {(l)} & \text {i f} i = j \\ 0 & \text {o t h e r w i s e ,} \end{array} \right. \tag {34}
$$
where we define  $\phi_{k,ij}^{(l)},\psi_{k,ij}^{(l)} = f_{MLP}^{(l)}(\mathbf{n}_{ij},w_{ij},(\mathbf{u}_k)_i)$  where  $f_{MLP}^{(l)}:\mathbb{R}^6\to \mathbb{R}^2$  is a MLP with one hidden layer of width 16 with ReLU activations, and Tanh output non-linearity.  $w_{ij}$  are the edge weights of the base graph (see Section B.2.1), and  $(\mathbf{u}_k)_i$  is the vector of weather covariates for node  $i$  at time  $k$ . Since we use these time-dependent covariates as input to the MLP, the resulting transition model is not time-invariant anymore. However, the parameters of the MLP remain shared across time and space. As before, diffusion parameter  $d^{(l)}$  is assumed to be spatially and temporally invariant.
Variant 2 The second variant uses highly simplified diffusion temporal layers of the form  $\mathbf{F}_k^{(l)} = \lambda^{(l)}\mathbf{I} + \omega^{(l)}(\mathbf{A} - \mathbf{D})$  with spatially and temporally invariant parameters  $\lambda^{(l)}$  and  $\omega^{(l)}$ .
Log-determinant computations Again, the spatial base graph  $\mathcal{G}_{\mathrm{Delaunay}}$  is small enough to precompute eigenvalues exactly and use the eigenvalue method for log-determinant computations [43].
Variational distribution As with the advection-diffusion dataset, we consider two variants for the variational distribution, one without temporal dependencies and one with a single temporal diffusion layer.
B.3 Baseline models
B.3.1 DGMRF
We apply the DGMRF for general graphs introduced by [43] to each time frame of the time series, not accounting for temporal dependencies. The DGMRF parameters are not shared across time, allowing for dynamically changing spatial covariance patterns. We use one spatial layer in the variational distribution, as proposed in [43], and run a hyperparameter search over  $L_{\mathrm{spatial}} \in \{1,2,3\}$  with  $L_{\mathrm{spatial}} = 2$  performing best.
Including covariates In our experiments on the air quality dataset, for which we have access to relevant covariates, we follow [50] and add linear effects to the measurement model. Note that the vector of coefficients if shared across both space and time.
B.3.2 ARMA
We implemented ARMA  $(p,q)$  models with  $p = 1$  and  $q = 1$  for the advection-diffusion data, and with  $p = 2$  and  $q = 2$  for the air quality data, using the Python statsmodels package. For each node in the test set, maximum likelihood parameter estimation is performed based on the observed time points. Given the estimated model coefficients, we obtain posterior mean and variance estimates using the standard Kalman smoother [47]. As the maximum likelihood estimates are deterministic, we do not provide standard deviations of the evaluation metrics for these models.
B.3.3 ST-AR
The spatiotemporal autoregressive (ST-AR) model takes the form  $\mathbf{x}_k = \alpha \cdot \mathbf{x}_{k - 1} + \epsilon_k$ , with initial state  $\mathbf{x}_0\sim \mathcal{N}(\pmb {\mu}_0,\pmb {\Sigma}_0)$  and unconstrained spatial error terms  $\epsilon_{k}\sim \mathcal{N}(\mathbf{0},\mathbf{Q}^{-1})$ . We fix  $\boldsymbol{\Sigma}_0 = 10\cdot \mathbf{I}$  to encode high uncertainty about the initial state  $\mathbf{x}_0$ , and fit  $\alpha ,\pmb{\mu}_0$  and  $\mathbf{Q}^{-1}$  to the data using closed-form EM updates. The EM algorithm is initialized with  $\alpha = 1,\pmb {\mu}_0 = \mathbf{0}$  and  $\mathbf{Q}^{-1} = \mathrm{diag}(\mathbf{q})$  where elements  $\mathbf{q}_i$  are drawn randomly from the interval [5, 6]. After convergence of the EM-algorithm, the final state estimates are obtained with the Kalman smoother [47].
B.3.4 EnKS
We consider an Ensemble Kalman Smoother (EnKS) variant for which the transition model matches the true data-generating process of the advection-diffusion dataset, as well as an EnKS variant for which we use a state augmentation approach to estimate unknown parameters  $\mathbf{v}$  and  $d = \sqrt{D}$  jointly with the system states. For both variants, we use  $10^{4}$  ensemble members (the maximum feasible on our machine). We fix the initial state distribution to  $\mathbf{x}_0\sim \mathcal{N}(\mathbf{0},10\cdot \mathbf{I})$ , and sample transition noise terms as  $\epsilon_{k}\sim \mathcal{N}(\mathbf{0},0.1\cdot \mathbf{I})$
For the state augmentation approach, we define the initial distribution over velocities  $\mathbf{v}$  as  $\mathcal{N}(\pmb{\mu}_v, 0.1 \cdot \mathbf{I})$ , where  $\pmb{\mu}_v$  is randomly drawn from  $[-1, 1]$  for each repeated run of the EnKS. Similarly, the initial distribution for diffusion parameter  $d$  is defined as  $\mathcal{N}(\mu_d, 0.01)$  where  $\mu_d$  is randomly drawn from  $[0, 0.2]$  for each repeated run of the EnKS. Finally, the transition noise terms for parameters  $\mathbf{v}$  and  $d$  are sampled from  $\mathcal{N}(\mathbf{0}, 0.01 \cdot \mathbf{I})$ .
B.3.5 MLP
For the air quality dataset, the MLP baseline maps local weather covariates  $(\mathbf{u}_k)_i \in \mathbb{R}^3$  to log-transformed PM2.5 measurements. We use one hidden layer of width 16 with ReLU activations and no output non-linearity. The MLP parameters are shared over both space and time.
B.4 Regularized Conjugate Gradients
We use a regularized variant of the conjugate gradient (CG) method [4] to avoid slow convergence in the case of ill-conditioned matrices. Instead of directly solving a potentially ill-conditioned linear system  $\mathbf{A}\mathbf{x} = \mathbf{b}$ , the idea is to iteratively solve a sequence of regularized (i.e. well conditioned) linear systems
$$
(\nu \mathbf {I} + \mathbf {A}) \mathbf {x} = \nu \mathbf {x} ^ {(i)} + \mathbf {b}. \tag {35}
$$
At each iteration, the solution from the previous iteration  $\mathbf{x}^{(i)}$  is used to obtain the next solution  $\mathbf{x}^{(i + 1)}$ . Eventually, this sequence will converge towards the true solution  $\mathbf{x}^*$  of the original system  $\mathbf{A}\mathbf{x} = \mathbf{b}$ .