SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We start with  $\nu = 10$  and decrease it every 10 iterations by factor 10. In each iteration, the standard CG method is employed to iteratively solve the regularized linear system until the residual norm drops below a threshold of  $10^{-7}$  or a maximum of 200 inner CG iterations is reached. This inner loop is repeated until the norm of the residuals
$$
\mathbf {r} ^ {(i)} = \left(\left(\nu \mathbf {I} + \mathbf {A}\right) \mathbf {x} ^ {(i)}\right) - \left(\nu \mathbf {x} ^ {(i)} + \mathbf {b}\right) \tag {36}
$$
drops below a threshold of  $10^{-7}$  or a maximum of 100 outer iterations is reached. The initial guess  $\mathbf{x}^{(0)}$  is given by the mean of the variational distribution  $q_{\phi}(\mathbf{x})$ .
In this section, we present additional results regarding the scalability of our approach (Section C.1), and provide more detailed results for the experiments in Section 4 of the main paper (Section C.2 and C.3). Finally, in Section C.4 we provide estimates of the total computation time required for our experiments.
C.1 Scalability
To empirically demonstrate the scalability of our method, we generate additional advection-diffusion datasets with varying lattice size and compare the runtime of our ST-DGMRF approach to a naive Kalman smoother (KS) [47] approach. To this end, we consider a model with advection-diffusion transition matrix using  $L_{\mathrm{temporal}} = 2$  temporal and  $L_{\mathrm{spatial}} = 2$  spatial layers. To avoid additional matrix inversions in the KS approach, we set the spatial and temporal bias terms  $\mathbf{b}_s$ ,  $\mathbf{b}_f$  to zero, resulting in  $\mu_0 = 0$  and  $c_k = 0$ . We train the model for 1000 iterations and measure the average wall clock time per iteration. In addition, we measure the wall clock time needed to perform inference with the trained model.
ST-DGMRF For the ST-DGMRF approach, we proceed as before using a variational distribution with one temporal diffusion layer during training. For better comparability, we employed the standard (non-regularized) CG method (with a tolerance of  $10^{-7}$ ) to compute the posterior mean and marginal variances (based on 100 CG samples) and provide measurements of the average time per CG iteration instead of the total time needed to perform inference. Multiplying this with the average number of CG iterations needed until convergence results in an estimate of the average total time for inference.
Kalman smoother For the KS approach, instead of approximating the true posterior with a variational distribution and estimating the ELBO based on Monte-Carlo samples, we use the KS to obtain exact marginal posterior estimates, which are used to compute expectations in closed form. The marginal covariance and transition matrices required for the KS equations are extracted from the ST-DGMRF model in every iteration. The associated parameters are then optimized via a form of Generalized EM-algorithm [16], where in each iteration a single gradient ascent step is taken.
Results Figure 7 shows how ST-DGMRF and KS training and inference scale as the number of nodes  $N$  in the system increases. Clearly, the time per training iteration increases super-linearly when the KS is used to obtain exact marginal posterior distributions, while the variational ST-DGMRF training time increases only marginally and remains below the fastest KS iteration for all tested  $N$ . In addition, KS memory requirements (due to storing  $K$  dense  $N \times N$  covariance matrices) exceeded the available GPU memory for  $N > 1024$ , making this approach infeasible for larger systems. In contrast, the ST-DGMRF exploits the sparsity of spatial and temporal graph-structured layers and thereby avoids storing dense matrices, remaining feasible for  $N \gg 1024$ .
For the tested systems, ST-DGMRF posterior inference with the CG method is slower than exact KS inference. However, the memory requirements of the KS approach again limit its feasibility to  $N \leq 1024$ , while the CG method only requires storing vectors of size  $\mathcal{O}(N)$  making it feasible for  $N \gg 1024$ . Moreover, Figure 7 confirms that computations per CG iteration scale linearly in  $N$ . And since the number of CG iterations required for convergence remains approximately constant, the
Figure 7: Comparison of ST-DGMRF and KS computation time in seconds for training (per iteration) and inference. For ST-DGMRF inference, the time per CG iteration is plotted together with the average number of CG iterations needed to converge (in gray). All quantities are plotted as mean  $\pm$  std based on 5 runs with different random seeds. The vertical dotted lines indicate the maximum  $N$  for which the KS approach was applicable.
total computation time for CG inference also scales linearly in  $N$ . In contrast, KS inference again scales super-linearly. This means that even if the KS approach would remain feasible in terms of memory requirements, its computation time will quickly approach, and eventually exceed, the time needed for CG inference.
C.2 Advection-diffusion experiments
Table 4 summarizes all results for the advection-diffusion dataset with mask size  $w = 9$ , including standard deviations for all metrics based on 5 runs with different random seeds. As discussed in the main paper, the ST-DGMRF variants provide more accurate posterior estimates than the baselines relying on simplified spatiotemporal dependency structures.
Ablation results Table 4 contains additional results for the ST-DGMRF variants using different settings for the variational distribution (see Section B.1.2) For this dataset, we do not find a significant effect of introducing temporal dependencies in the variational distribution. Further, Figure 8 shows additional results for the ST-DGMRF variants when varying the number of temporal layers  $L_{\mathrm{temporal}}$ . For all metrics, the performance improves significantly as we start adding temporal layers and stabilizes around  $L_{\mathrm{temporal}} = 3$ . Note that around the same point, both ST-DGMRF variants converge towards the EnKS using the true data-generating dynamics, in terms of the  $\mathrm{RMSE}_{\mu}$ , and even drop below it in terms of the CRPS. Only in terms of  $\mathrm{RMSE}_{\sigma}$ , the ST-DGMRF models remain inferior to both EnKS variants. We hypothesize that increasing the expressivity (i.e.  $L_{\mathrm{spatial}}$ ) of the noise terms can further reduce this gap.
Table 4: Model performance for the advection-diffusion dataset with  $w = 9$ , reported as mean ± std over 5 runs with different random seeds. All ST-DGMRF variants use  $L_{\mathrm{spatial}} = 2$  and  $L_{\mathrm{temporal}} = 4$ .
VI dynamics	RMSEμ↓	RMSEσ↓	CRPS↓
ARMA	-	2.3054 -	0.6812 -	1.7064 -
ST-AR	-	1.4595±0.0098	1.9216±1.0392	0.9707±0.0163
DGMRF	-	0.5901±0.0037	0.3808±0.0010	0.3495±0.0022
EnKS
true dynamics	-	0.0661±0.0030	0.0046±0.0000	0.1027±0.0035
estimated dynamics	-	0.1654±0.2031	0.0039±0.0005	0.1434±0.0902
ST-DGMRF (ours)
advection-diffusion	none	0.0526±0.0001	0.1148±0.0003	0.0726±0.0001
advection-diffusion	diffusion	0.0526±0.0001	0.1146±0.0005	0.0726±0.0000
neural network	none	0.0839±0.0022	0.1334±0.0089	0.0833±0.0008
neural network	diffusion	0.0854±0.0027	0.1402±0.0061	0.0839±0.0008
Figure 8:  $\mathrm{RMSE}_{\mu}$ ,  $\mathrm{RMSE}_{\sigma}$  and CRPS as a function of the number of temporal layers  $L_{\mathrm{temporal}}$  for the advection-diffusion dataset with  $w = 9$ , plotted as mean  $\pm$  std over 5 runs with different random seeds. Both ST-DGMRF variants are trained with a variational distribution using one temporal diffusion layer. Note that  $L_{\mathrm{temporal}} = 0$  corresponds to the spatial-only DGMRF baseline.
C.3 Air quality experiments
Table 5 summarizes all results for the air quality dataset. It contains additional results for the ST-DGMRF variants using different settings for the variational distribution (see Section B.2.2), and provides standard deviations for all metrics based on 5 runs with different random seeds.
Ablation results We find that, in contrast to our experiments on the advection-diffusion data, accounting for temporal dependencies in the variational distribution is clearly beneficial in the real world setting. Especially for the ST-DGMRF with neural network based transitions, adding the temporal diffusion layer results in significantly improved posterior estimates, and at the same time reduces the variability across different runs. Further, we find that at least two temporal layers are needed to achieve good posterior estimates that improve on the baselines (see Figure 9).
Example model outputs Figure 10 shows state estimates and associated uncertainties together with sensor measurements for two example sensors within the masked out area of the network, for ST-DGMRF, DGMRF and ARMA respectively. For all three models, state estimates are obtained by conditioning on the input data points (used for training), resulting in low errors and uncertainties for observed time points and higher errors and uncertainties for masked out time points. Moreover, for both ST-DGMRF and DGMRF, higher uncertainties coincide with larger errors and larger fluctuations in the measurements (top), while more accurate state estimates come with smaller uncertainties (bottom). Finally, Figure 11 visualizes how spatial and temporal ST-DGMRF layers transform samples from the estimated posterior over system states into (approximately) independent Gaussian noise, as derived in Section 3.1.2 and visualized in Figure 1. Clearly, temporal layers remove daily patterns and overall trends, while spatial layers remove dependencies between close-by sensors and increase temporal fluctuations.
Figure 9: RMSE and CRPS for increasing  $L_{\mathrm{temporal}}$  (mean ± std over 5 runs). Both models use  $p = 2.5$ . As before,  $L_{\mathrm{temporal}} = 0$  corresponds to the spatial-only DGMRF baseline.
Figure 10: Model outputs for two air quality sensors falling within the masked area. Solid lines represent posterior mean estimates, while shaded areas represent posterior std estimates. Time points that have either missing data or fall within the masked time window are shaded in gray.
Table 5: Model performance for the air quality dataset, reported as mean ± std over 5 runs with different random seeds. All ST-DGMRF variants use  $L_{\mathrm{spatial}} = 2$  and  $L_{\mathrm{temporal}} = 4$ .
p	VI dynamics	RMSE ↓	CRPS ↓
ARMA	-	-	0.6820 -	0.3625 -
ST-AR	-	-	0.7350±0.0006	0.4261±0.0003
DGMRF	-	-	0.7368±0.0135	0.3966±0.0032
MLP	-	-	0.8038±0.0245	-
ST-DGMRF (ours)
diffusion	1	none	0.6147±0.0082	0.3239±0.0058
diffusion	1	diffusion	0.6190±0.0073	0.3258±0.0043
diffusion	2	none	0.6020±0.0112	0.3214±0.0051
diffusion	2	diffusion	0.5928±0.0119	0.3161±0.0054
neural network	1	none	0.5995±0.0887	0.3147±0.0494
neural network	1	diffusion	0.5853±0.0457	0.3092±0.0257
neural network	2	none	0.5825±0.0626	0.3062±0.0353
neural network	2	diffusion	0.5565±0.0184	0.2925±0.0097
C.4 Total compute
Most computations were performed on a Nvidia Titan X GPU. On top of the final experiments, we performed hyperparameter sweeps and additional test runs. Here, we provide estimates of the total compute time grouped by experiment:
Advection-diffusion dataset:
- Performance comparison & ablations: ca. 50 GPU hours
- Varying mask size: ca. 30 GPU hours per  $w$ , resulting in ca. 210 GPU hours in total
- Scalability: ca. 15 GPU hours
Air quality dataset:
Performance comparison & ablations: ca. 100 GPU hours
Figure 11: Effects of applying temporal  $(\mathbf{f}_{\theta})$  and spatial  $(s_{\theta})$  ST-DGMRF layers to the predictive posterior. Left column: The dark yellow line (top row) shows the estimated posterior mean for an air quality sensor falling within the masked area. Light yellow lines represent corresponding posterior samples. Similarly, green lines (center row) represent states after applying the temporal transformation  $(\mathbf{f}_{\theta})$ , and gray lines (bottom row) represent states after applying both temporal and spatial layers  $(s_{\theta} \circ \mathbf{f}_{\theta})$ . As a reference, we also plot ground truth log-transformed and normalized PM2.5 measurements (dashed black lines). Right column: corresponding (transformed) states for all sensors at time  $k = 100$ .