Deep Gaussian Markov Random Fields for Graph-Structured Dynamical Systems
Authors: Fiona Lippert
Bart Kranstauber
E. Emiel van Loon
Patrick Forré
Venue: NeurIPS 2023 poster (2023)
Keywords: probabilistic inference
graphical models
spatiotemporal dynamical systems
state-space models
TL;DR: We propose an extension to Deep Gaussian Markov Random Fields for computationally efficient learning and inference in large spatiotemporal dynamical systems.

Abstract: Probabilistic inference in high-dimensional state-space models is computationally challenging. For many spatiotemporal systems, however, prior knowledge about the dependency structure of state variables is available. We leverage this structure to develop a computationally efficient approach to state estimation and learning in graph-structured state-space models with (partially) unknown dynamics and limited historical data. Building on recent methods that combine ideas from deep learning with principled inference in Gaussian Markov random fields (GMRF), we reformulate graph-structured state-space models as Deep GMRFs defined by simple spatial and temporal graph layers. This results in a flexible spatiotemporal prior that can be learned efficiently from a single time sequence via variational inference. Under linear Gaussian assumptions, we retain a closed-form posterior, which can be sampled efficiently using the conjugate gradient method, scaling favourably compared to classical Kalman filter based approaches.

Deep Gaussian Markov Random Fields for Graph-Structured Dynamical Systems
Fiona Lippert
Bart Kranstauber University of Amsterdam b. kranstauberrrua....
University of Amsterdam f. l.lppert@uva.n.
E. Emiel van Loon University of Amsterdam e e vanloon@uva...
Patrick Forré
University of Amsterdam p.d. fdrforeeuuva.nl
Abstract
Probabilistic inference in high-dimensional state-space models is computationally challenging. For many spatiotemporal systems, however, prior knowledge about the dependency structure of state variables available. We leverage this structure to develop computationally efficient approach to state estimation and learning in graph-structured state-space models with (partially) unknown dynamics and limited historical data. Building on recent methods that combine ideas from deep learning with principled inference in Gaussian Markov random fields (GMRF), we reformulate graph-structured state-space models as Deep GMRFs defined by simple spatial and temporal graph layers. This results in a flexible spatiotemporal prior that can be learned efficiently from single time sequence via variational inference. Under linear Gaussian assumptions, we retain a closed-form posterior, which can be sampled efficiently using the conjugate gradient method, scaling favourably compared to classical Kalman filter based approaches.
1 Introduction
Consider the problem of monitoring air pollution, leading environmental risk factor for mortality worldwide [1] (22 In the past years, sensor networks have been installed across major metropolitan areas, measuring concentration pollutants across time and space [27] 38). These measurements are, however, prone to noise might be missing completely due to hardware failures limited sensor coverage. To identify sources of pollution or travel routes with limited exposure, it essential to recover pollution levels and provide principled uncertainty estimates for unobserved times and locations. Similar problems occur also in the geosciences, ecology, neuroscience, epidemiology, transportation systems, where animal movements, the spread of diseases, traffic load need to be estimated from imperfect data to facilitate scientific discovery and decision matiiiionmmaiing.. From a probabilistic inference perspective, all these problems amount to estimating the posterior distribution over the latent states of spatiotemporal system given partial and noisy multivariate time series data.
To make inference in these systems feasible, it is common to assume state-space model, where the latent states evolve according to Markov chain. Then, the classical Kalman filter (KF) [29] or its variants [47 37 can be used for efficient inference, scaling linearly with the time series length. However, the complexity with respect to the state dimension remains cubic due to matrix- matrix multiplications and inversions. Thus, the spatial coverage of available measurements and thereby the dimensionality of the state variables increase, KF-based approaches quickly become computationally prohibitive. This is especially problematic in cases where the underlying dynamics are (partially) unknown and need to be learned from data, requiring repeated inference during the
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

optimization loop 21 40, 91 If. additionally, the access to historical training data is limited, it is
essential incorporate prior knowledge for learning and inference to remain both data-efficient and computationally feasible [56]
For spatiotemporal systems, prior knowledge is often available in the form of graphs, representing the dependency structure of state variables. For air pollution, this could be the city road network and knowledge about wind directions which influence the spread of particulate matter. Incorporating this knowledge into the initial state distribution and the transition model of a state-space model results
39 [6], an interpretable and data-efficiens graphical model for multivariate time series data. If the joint space- timeeettmmmeee graph is sparse and acyclic, belief propagation methods allow for fast and principled inference of marginal posteriors in DBNs, scaling linear with the number of edges [4S] 39]. However, when the underlying dynamics are complex, requiring denser graph structures and flexible error terms with a loopy spatial structure, the scalability and convergence of these methods is no longer guaranteed.
in a Dynamic Bayesian network (DBN) 12
Combining the favorable statistical properties of classical graphical models with the flexibility and scalability of deep learning offers promis- ing solutions in this regard [57]. Deep Gaus- sian Markov random fields (DGMRF), for exam- ple, integrate principled inference in Gaussian Markov random fields (GMRF) with convolu- tional 50 or graph layers 43 to define new flexible family of GMRFs that can be learned efficiently by maximizing a variational lower bound, while facilitating principled Bayesian in- ference that remains computationally feasible even for high-dimensional complex models with dense dependency structures. While in principle DGMRFs can be applied spatiotemporal sys- tems by treating each time step independently, they are by design limited to spatial structures.
Squrtial Tapres)
(omaporn lopes
A De
Ville
-illie Rel
?
In this paper we extend the DGMRF framework to spatiotemporal systems with graph-structured states, transitions and noise terms. To this end, we formulate dynamical prior based on a DBN that relates state variables over adjacent time steps, and DGMRFs that capture the spatial structure of unmodeled exogenous influences and errors in the transition model. The key to our approach is the insight that, using locally linear Gaussian transition models, the precision matrix of the joint space-time process preserves the struc- ture of transition and noise precision matrices. This is in contrast to the marginal distributions used in Kalman filter based approaches for which both precision and covariance matrices become dense over time. A convenient factorization of the joint precision matrix then leads to a spatiotemporal DGMRF formulation that is equivalent to the generative state-space formulation, but inherits the favorable properties of DMGRFs for learning and inference.
Figure 1: ST-DGMRF overview. We reconstruct the latent states of graph-structured dynamical system from partial and noisy observations (orange arrow with question mark). Temporal and spatial layers transform the state x to standard Gaussian, which implicitly defines space-time GMRF prior.
1.1 Related work
The simplest approach to scaling Kalman filtering and smoothing to high-dimensional state variables to assume a decomposition into subprocesses, that can be reconstructed independently [30]. For spatiotemporal processes, however, this assumption can be detrimental. Sample-based approaches, like the Ensemble KF [7 estimates, balancing computational complexity and accuracy. This is widely used in the geosciences where accurate models of the underlying process are available. If the dynamics are (partially) unknown, the parameters are typically estimated jointly with the system states via state augmentation 2] However, this becomes problematic if the transition model only a poor approximation and complex error covariances need to be estimated [15] 52 Alternatively, expectation maximization
18 26 BI. instead approximate state distributions with Monte Carlo
53 5

---

(EM) is used to estimate parameters iteratively 21. which requires repeated ensemble simulations and thus quickly becomes computationally demanding.
In machine learning, neural network-based approaches map high-dimensional data to a latent space that is either low-dimensional [19] or factorizes conveniently [7], such that standard KF inference remains feasible. To learn suitable maps, however, these methods require sufficiently large training data sets. Moreover, the projection to an abstract latent space hinders the incorporation of structural or functional prior knowledge. Alternatiely, variational inference allows for fast inference in (deep) state space models 20
54 by approximating the true posterior with computationally convenient variational distribution. While this approach is highly flexible and scalable, it is known to severely underestimate uncertainties [5], which can be detrimental for science and decision-
Gaussian processes (GP) are another widely used model class for spatiotemporal inference [35] In contrast KF-based approaches, GPs encode prior knowledge into their covariance kernel. To increase the expressivity beyond standard kernels and overcome computational limitations, GPs have been combined with deep learning [14 [3.. However, scaling GPs large-scale spatiotemporal processes with graph structure remains challenging. Scalable approaches for temporal data involve reformulating GPs as state-space models and applying KF-based recursions 24 49
58
23
which, as discussed before, remains computationally prohibitive for large state spaces, Extensions to graph-structure domains
41 rely on the relation between GPs and stochastic differential equations (SDE) [36] which makes it cumbersome to design new kernels or to incorporate prior knowledge that cannot be formulated as an SDE.
1.2 Our contributions
In contrast to most previous approaches, we focus on settings where the dynamics are (partially) unknown and need to be learned from data, while historical training data limited or even unavailable. Our contributions can be summarized as follows:
1. We propose ST-DGMRF, which extends the DGMRF framework to spatiotemporal systems, by reformulating graph-structured state-space models as joint space-time GMRFs that are implicitly defined by simple spatial and temporal graph layers.
2. We show that the multi-layer space-time formulation facilitates efficient learning and infer- ence, scaling linearly w.r.t. the number of layers, time series length, and state dimensionality.
3. In experiments on synthetic and real world data, we demonstrate that our approach provides accurate state and uncertainty estimates, comparing favorably other scalable approaches relying on ensembles or simplifications of the dependency structure.
2 Preliminaries
2.1 Problem formulation
taken at consecutive time points lg,..., lk. Each yı E MR represents partial and noisy observations (e.g. from spatially distributed sensors) generated according to a linear Gaussian noise model
Consider a single sequence of high-dimensional measurements Yu:k =
Y& H + En. Ex ~ (O,
(1)
from the latent state X E R of an underlying dynamical process of dimensionality N 2 M2 (c.g. air pollution along roads within an urban area). It is common assume either R& diag(o?, ........... R,M, representing independent sensors with varying noise levels, or R₂ = o I'I representing a fixed noise level across all sensors and time points [. Further, the observation matrices H, R M2 N are typically very sparse with O(M)) number of non-zero entries.
Assuming that all Hg are known, we aim at reconstructing the latent system states Xo:K (Xo. (xo....,Xk) [K+I) from data Yo.k. Since observations are partial and noisy, X cannot be recovered without ambiguity. It is therefore essential to treat the problem probabilistically and in- corporate available domain knowledge into the prior distribution p(Xo:K 0). The task then amounts to (i) specifying a suitable prior that captures the spatiotemporal structure of Xo:7, (ii) estimating unknown parameters, and (iii computing the posterior distribution p(Xo.K Yo.k, 0).
3

---

Notation We use vec(Xo:k) R(K+1)N and y = vec(yo:k) E R with M = Σ
k=0 Mr
denote vectorized states and measurements. Similarly, we use H diag (H. Hg.... Hx)
R (K+N)) and R diag R R&) RMM denote the joint spatiotemporal observation and covariance matrix.
In the following sections, we briefly introduce two graphical models, linear dynamical systems and Gaussian Markov random fields, which form the temporal and spatial backbone of our method. We then introduce the deep GMRF framework on which build in Section to formulate a flexible multi-layer spatiotemporal prior.
2.2 Linear dynamical system
discrete-time linear dynamical system (LDS) is state-space model that defines the prior over states Xu:K in terms of linear Gaussian transition models
X F + Ck + ck, 6 ~ ~N000, QT1) Ak E (1,..,,K
and an initial distribution Xo (u 2T²). The classical Kalman filter (29 and its variants [47 37 use recursive algorithms that exploit the temporal independence structure the prior and its conjugacy the observation model (see Eq. (T) to compute marginal posterior distributions p(xx Yo:k) in O K N3) The same formulas can be used within the EM-algorithm to learn unknown parameters Hog Qu: F.. Ck Q in (() for J iterations. While this scales favorably for long time series, it remains computationally prohibitive for high-dimensional systems with N > K.
2.3 Gaussian Markoy random fields
with respect to an undirected graph g = (V,E) if the edges & correspond to the non-zero entries of the precision matrix Ω, i.e. S2 * 0 <+ (i,j) E E. Given a linear Gaussian observation model ~ WHHx, R), the posterior can be
multivariate Gaussian ((,, 22 forms a GMRF 48
written in closed form p(x y, 0) =N(xt (-+)-1) with precision matrix Ω Ω H H and mean f1 ($2+ (2pp + H Hy)
However, naively computing 1at by matrix inversion requires O(K* computations and O(K memory, which is infeasible for high-dimensional systems. Instead, the conjugate gradi- ent method [6] can be used to iteratively solve the sparse linear system 2 1 2 = S2 + H R R- for ut The same approach allows generate samples x NNx "(t=*,05 ({2+)-1) and obtain Monte Carlo estimates of marginal variances [44]. In practice, however, it remains difficult to design suit- able precision matrices that are expressive but sparse enough to remain computationally feasible.
2.4 Deep Gaussian Markov random fields
A DGMRF 50 is defined by an affine transformation
z = go (x) = Gøx + bg with 2 ~~ N/
(3)
(1. o...og) go is a composition
simple linear layers with convenient computational properties. This implicitly defines a GMRF x ~ (H, 2211) with μ = G and Ω = G G The multi-layer construction facilitates fast and statistically sound posterior inference using the conjugate gradient method, as well as fast parameter learning using a variational approximation, even when the resulting precision matrix Ω becomes dense.
where ge = gg
While originally, (50) defined their layers only for lattice-structured data,
43 generalized this to arbitrary graphs. Central to their approach is that each layer (4) = G (1-1) + b, is defined on
sparse base graph & with adjacency matrix A and degree matrix D such that
( = a. -1- and by 10 b,1
with parameters 0, = (an, 31, n,
(n,by b1). This construction allows for fast log-determinant computations during learning, and GPU-accelerated computation of go gx(x) using existing software for graph neural networks. Note that for T. layers defined on G, the sparsity pattern of corresponds the sparsity pattern of (A + 2) And thus, the resulting model defines GMRF w.r.t the L- graph of G.

---

3 Spatiotemporal DGMRE's
To extend the DGMRF framework to spatiotemporal systems, first formulate dynamical prior in terms of a GMRF that encodes prior knowledge in the form of spatial and temporal independence assumptions. We then parameterize this prior using simple spatial and temporal layers, which results in flexible model architecture that facilitates efficient learning and principled Bayesian inference in high-dimensional dynamical systems, scaling favorably compared Kalman filter based methods.
3.1 Graph-structured dynamical prior
Consider dynamical system for which the state evolution is well described by Eq. 2 We say that this process is graph-structured if each dimension d, of the system state can be associated with node V in multigraph
(V, Espatial: Enemporal), where the set of undirected spatial edges Expatial C V x V defines the sparsity pattern of noise precision (inverse covariance) matrices
(Q #0
(Q*);:
= (j, Espatial
w E {O...., K), (5)
and the set of directed temporal edges Etemporal S V x V defines the sparsity pattern of the state transition matrices
(F #0
=
(j E Exemporal
Ak E (1....,,K) K}.
(6)
This defines DBN over Xo:K encoding conditional independencies of the form
Tke IL X (Meepepaa)) (i) and Thii 1L. XI.V Iimmpe (6) 4I 1 (7)
where Neemporal (i) denotes the set of neighbors j i for which (j, E Etemporal- Intuitively, Exemprral represent causal effects over time, while Esparial represent the structure of random effects that are not captured by the transition model. To define the graph structure, prior knowledge about, for example, the physical system structure or underlying causal mechanisms can be exploited.
3.1.1 Joint distribution
The graph-structured dynamical prior induces multivariate Gaussian prior ((, on x with sparse precision matrix Q2 e x (K+))N Importantly, S2 can be shown to factorize as F Q" with block diagonal matrix Q and unit lower block bidiagonal matrix F defined as
I
F1
Q = diag( Qo: QI...., Qk),
FK
where empty positions represent zero-blocks. Further, while the mean fa E (KK+1) needs to be computed iteratively as f4g F -1 Ck. the information vector 7 Pez can be expressed compactly as
n F QF FC
with [Po: C1: ..., CK] E RK+N See Appendix A. for detailed derivations.
3.1.2 DGMRF formulation
We now reformulate x (N as DGMRE Importantly, in contrast to 50 the graph- structured dynamical prior allows us to impose additional structure, reflecting the spatiotemporal nature the underlying system. In particular, since Qo Q....., QK are symmetric positive definite, factorize Q S with symmetric block- block-diagonal matrix S, and thus Ω FT SF. Together with Eq. g this results GMRF defined by
Go = SF and by == Gof SSc
(10)
Finally, we separate go into temporal map fo [R(K+1)N (K+1)) and spatial map So [(K+ + RR (K) defined as
(x) = Fx by h and 86 (h) = Sh b.
(11)
Note that this results in an overall bias by Sby + b. and thus for modelling long-range spatial dependencies in P4 and C cI...., Cx. The combined transformation (Sp fo) (x) essentially defines standard two-layer DGMRE, which describes a graph-structured dynamical system the parameters A (b,, by, F F1,..., ......., F Qo...., Qk) are subject to sparsity
(by + bb)),), which allows
constraints (5) and 6
5

---

3.2 Parameterization
Learning the unknown parameters directly from a single sequence Yo:k will result in a highly over- parameterized model that is unsuitable for the problem of interest. In addition, careful parameteriza- tion of ge can greatly reduce the computational complexity of the transformation and the associated log-determinant computations required during learning 50 43). To achieve good trade-off between data-efficiency, expressivity and scalability, we define both So and fo in terms of simple layers with few parameters and convenient computational properties.
3.2.1 Spatial layer(s)
To enable fast log dog-denerminant computations for arbitrary graph structures, we follow 45 and
parameterize in terms of K + independent DGMRFs
Syhı + (b,)k. each defining GMRF w.r.t Spparial = (V, Epatiaa). Note that due to the multi-layer construction (see Section 2.4).
Goaaial is implicitly defined by the base graph Gqpatial and the number of layers L, with special case Ispatial Esaaial spatial 1.
3.2.2 Temporal layer(s)
Compared to the definition of fg is much less constrained as it does not affect the log-determinant computation (see Section B.3.I. and should hence incorporate as much domain knowledge into transition matrices F& 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 F2 can again be decomposed into simpler layers F (Lomppraa)
F(U) each defined according to temporal base graph Grepporll Table provides several example layers to give an idea of what is possible.
Note that any function can be used define the entries of F including neural networks taking available covariates or node/edge features as inputs [46]. The associated parameters can simply be included in O. Similarly, non-linear dynamics to the Extended Kalman Filter (37), or through neural linearization [19 7 We, however, recommend sharing parameters where appropriate to avoid overparameterizztion.
be approximated either through linearization akin
Table 1: Examples of linear transition layers F(. The adjacency matrix A can be symmetric or asymmetric, weighted or unweighted.
layer definition
properties of Seeppoa
AR process
F A F. A + (A
self-edges only bidirected (symmetric A) directed
Diffusion
Directed flow
Fx = Ar, + WA, Dan) + CA.A Dn)
(p(1))
= --
discretized 3 (e.g. triangulation),
Advection
ig
(F(")
ExEn(o) (F(")
edge weights way
n.j
Neural network (F(9) = fuu(uu,
made and edge features u. Ciy
Higher order Markov processes process
It is straight forward to extend fo to describe p-th order Markov
X
F X Ca Ex: 6~ N00 Qi Σ
N Q²!),
(12)
where x for 1,...,p This can be done by introducing edges (jempora(()) ..... Eteppoal(p)
to the 7-th lower and adding the corresponding oigher-order transition matrices (F (F......., F k,,) block diagonal of F (see Appendix A.2 for derivations).
3.3 Learning and inference
Although the marginal likelihood ply 0) available in closed form (see Section 2.3 maximum likelihood parameter estimation becomes computationally prohibitive in high dimensional settings

---

Instead, a variational approximation is used to learn parameters for large-scale DGMRFs. Then, the conjugate gradient method allows efficiently compute the exact posterior mean and to draw samples from p(x yy 0).
50
3.3.1 Scalable parameter estimation
Given variational distribution (x(x) =M(ps, A&), the parameters {0, 3) are optimized jointly by maximizing the Evidence Lower Bound (ELBO)
L(yo:K, 0, [gy(x))
log (x))
Σ logp(ye xxs)
+ H1 [xa(x)]
(13)
K
= 2 (x)(
Bx(x)) gx(x) + E(yx (yk
H.X+) RR, R (y* H&x+)
(14)
k=
1
+ log (ee(GG)) + log det((e (
log |det [det(Rx)] const,
(15)
2
2 k=0
using stochastic gradient descent, where the expectation is replaced by a Monte-Care estimate based on samples X ~ 9á-
Variational distribution To facilitate efficient log-determinant computations and sampling via the reparameterization trick (33), we follow [43 and define 96 as an affine transformation x = P with z (O,,,), I), resulting in As P (P& pT) Then, Ps can be defined as a block-diagonal matrix diag(P Po...... Px) with Pk- = diag( Pv) dadiaa(bb))) where Pk: &b E R are variational parameters and & is based on spatial DGMRE layers as discussed in Section 3.2.1 To relax the temporal independence assumptions in As, we propose an extension P diag(Po: o. .... Px F, where has a similar structure to F (see Eq. (8)) Again, the design of & is highly flexible as does not enter the log-determinant computation.
50 43 proposed specific lattice and graph layers for which the associated log-determinants log ddet Go) are computationally scalable. Conveniently, the spatiotem- poral 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 jdet G&)] simplifies to
dogg-eeterminant computations
K
(88)
log (
log [det(S)]
log det(SS)).
(16)
where (1) follows from S being block-diagonal and (it) follows from F being unit-lower triangu- lar with det(F) = 1. Note that each S₁ is defined by spatial DGMRF of dimension N (see Sec- tion 3.2.T ). 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 Mr log(o,).
R2 (see Section 2.1 2 log [de((R.)) simplifies Σ
Computational complexity During training, log (det(Go) Go)l and log |det( (s)()))) can be computed in O(KN ((KNT.)aaa) sppatial using Eq. (16) together with the methods proposed in (43. Note that no complexity added when introducing conditional dependencies between time steps. The necessary preprocess- ing steps, i.c. computing eigenvalues or traces, only need to be performed for the spatial base graph the number of transitions K. Finally, assuming an average of dspatial and ciemporal edges per node in Espatial and Exemponat, the transformations Se and fa scale linearly with the number of nodes N. In particular, (x)x) (Sg (x) can be computed in O K N dospttial Laspttall K N diemporal L.emporat), which dominates the computational complexity of the training loop. In addi- tion, we can leverage massively parallel GPU computations to speed up this process even further.
and are thus independent
3.3.2 Exact inference with conjugate gradients
As discussed Section 2.3 the conjugate gradient method can used to iteratively compute the posterior mean and obtain Monte Carlo estimates of marginal variances. Importantly, each iteration
7

---

is dominated by single matrix-vectos multiplication of the form
Ω x = F SF'SSFX=11 H'R
(17)
This amounts to series of sparse matrix-vector multiplications with total computational complexity O(K N d'spatial Lsspttal + KNN N (heemprra/<cemprra)). which is again linear in the number of nodes, time steps, and layers, and can be implemented in parallel on a GPU.
Experiments
We implemented ST-DGMRF in Pytorch and Pytorch Geometric, and conducted experiments consumer-grade GPU, leveraging parallel computations both spatial and temporal layers. In all experiments, we optimize parameters for 10000 iterations using Adam [22 with learning rate 0.01, and draw 100 posterior samples to estimate marginal variances. Unless specified otherwise, we use spptiil 2. Liemporal 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
4.1 Advection-diffusion process
We start with synthetic dataset for which have access to both the ground truth posterior distribution and transition matrix. The dataset consists of = 20 system states that are sampled from an ST-DGMRF with time-invariant transition matrix Fadv This matrix is defined according to the third-order Taylor approximation of an advection-diffusin process with constant velocity and diffusion parameters, discretized on 30 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 for 10 consecutive time steps and adding noise with 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: (I) using advection-diffusioon matrices F 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 n.j pointing from pixel i to j as input. In both cases, Gumporall and Sapanial are defined as the 4-nearest neighbor graph. Further, as the underlying dynamics are time-
invariant, we enforce F(!) = (!))
and (()) s(1) k+1
Vk 2
k+1
Baselines We compare approach range of baselines accounting for varying degrees of spatial and/or temporal dependencies. In particular, sider the original DGMRF applied to all time steps in- dependently, 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 Eμ for which an uncon- strained covariance matrix Q is estimated using the EM algorithm. For both ARMA and ST-AR, we use the standard Kalman smoother [47] to obtain pos- terior estimates. Additionally, we consider two En- semble Kalman Smoother (EnKS) variants, one with an advection-din -diffusiffonn transition model matching the true data- matta-geneaatng 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 be fixed in order to avoid divergence of the EnKS.
Table 2: Performance the advection- diffusion data with w = 9. We report the mean over 5 runs with different random seeds.
con-
RMSF,
RMSE,
CRPS
ARMA ST-R
2,3054
0.6812
1.7064
1,4595
1,9216
0.9707
DOMRE
0.5901
0.3808
0.3495
EnKS
true dynamies
0.0661
0.0046
0.1027
estimated dynamics
0.1654
0.0039
0.1434
ST-DGMRE (ours)
advection advection-diffusiion
0.0526
0.1146
0.0726
neural network
0.0854
0.1402
0.0839
Performance evaluation We evaluate the estimated posterior mean and marginal standard devia- tions in terms of the noo--mean-square-error (RMSE,, and RMSE,) 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
8

---

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 better calibrated posterior distributions. More detailed results are reported in Appendix
system state
data
175
-34
Data
53-44
STAT
node (pixel) state
cynamics
DOM17
ST.0V1 idbection ditucine
srowel
DGMRF
ST-DGMRF
-37
0.00
0
10
12 14 16
11
12
time
mask width W
Figure 2: Left: snapshot of the addection-diffusion data at time k=8, and reconstructions by the time- independent DGMRF and our ST-DGMRF with advection-diffit series for single pixel (marked the left). Shaded areas represent posterior mean + std of single run. Right: RMSE,, as a function the mask width (mean + std over 5 runs).
Fx. Center: corresponding time
Increasing mask size To evaluate the robustness to missing data, analyze how posterior estimates change with increasing mask size. Figure (right) shows the effect on RMSE,, when varying 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 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.
true stencil
Lesmoorai
Leemoora 2
tempora
1.00
Pearson
0.85
STOOMRF veroitinn
0.80
-a- advection difludos
meura retwork
0.75
-0.25
0.00 2.25 transition weight
0.00
0.05
0.30
0.15
0.20
0.25
absolute ertor
Etenssoal
Figure 3: Stencil comparison. Center: absolute error between true and learned transition weights for the neural network based ST-DGMRE. 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 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 Lremporal the neural network based layers show 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 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 realistic scenario of local network failures, we randomly choose 10 time points la and mask out all data points within spatial block containing 50% of all sensors, for time steps tk.....44 l + 20 (see F 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 o 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 F()) are then formed in the same way as in Section 4.1 In addition, we consider a simplified diffusion transition model (see Table [ which cannot capture any directional transport processes.
Results Next to the baselines introduced in Section 4.1 we consider multi-layer perceptron (MLP) mapping lo- cal 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ßreports 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 trans- port processes, provides the most accurate state estimates. However, even with highly simplified diffusion transitions our approach provides better estimates than the consid- ered baselines. We attribute this the expressive DGMRF noise terms which can capture complex error structures resulting from inaccuracies in the transition model. Increasing the Markov order from p to
Table 3: Performance on the air quality data, We report the mean over 5 runs with different random seeds.
RMSE
CRPS
ARMA
0.3625
STAR
0.7350
0.4261
DGMRF
0.7456
0.4037
MLP
0,8038
ST-DGMRE (ours)
diffasion
0.6190
0.3258
diffusion
0.5928
0.3161
newral network
0,5853
0.3092
newaal network
0.5565
0.2925
clearly improves the resulting posterior estimates for both ST-DGMRE variants, reflecting the com-
plexity of the modeled process. A similar trend is observed as we increase Lnemporal (see Appendix C
input data
test data
ST DGMRG mean
ST DGMRF sample
0.6
logiPM2.5)
0.4
STDGGME
0.2
0.0
50
100
150
200
250 300
350
400
time step k
Figure 4: ST-DGMRF posterior estimates (neural network F5 and 2). Left: time series of log- transformed and normalized PM2.5 levels for masked sensor. Light red areas represent posterior mean + std of single run. Unobserved time points are marked with gray bars. Right: Marginal std for all sensors 100. Masked out sensors (gray box) feature higher uncertainties.
Conclusion
We have presented ST-DGMRF,, an extension to Deep Gaussian Markov Random Fields for infer- ence 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 simee-time GMRFs enables computationally efficient learning and inference even in nigh-ddmensiooal settings, scaling linear w.r.t. the number of both time steps and state variables. Empirically, ST-DGMRR provides more accurate posterior estimates than other scalable approaches relying on simplifications of the dependency structure. While our approach relies on linear Gaus- sian assumptions, which can be restrictive for real systems, we find that expressive DGMRF noise terms can compensate (to a certain extend) for approximations non-linearities between the layers could be explored as discussed in [50] Further, we see potential in defining more flexible vare-ng transition matrices based on a hierarchy of latent variables that again follow a ST-DGMRF.
the transition model. the future,
10

---

References
[1] V. Adam, S. Eleftheriadis, A. Artemev, N. Durrande, and J. Hensman. Doubly sparse variational gaussian processes. In International Conference on Artificia Intelligence and Statistics (AISTATS). PMI.R, 2020.
[2] J. L., Anderson, An ensemble adjustment Kalman filter for data assimilation. Monthly Weather Review, 129((2))2884-2905, 2001.
[3] E. Archer, M. Park, L. Buesing, Cunningham, and L. Paninski. Black box variational inference for state space models. arXiv preprint 1rXivv1511.07337. 2015.
[4] 7. 7., Bai and S. I.. Zhang. A regularized conjugate gradient method for symmetric positive definite system of linear equations. Journal of Computational Mathematics, pages 437-448, 2002.
[5] D. Barber, A. T. Cemgil, and S. Chiappa. Bayesian time series models. Cambridge University Press, 2011.
[6] R. Barrett, M. Berry, T. F. Chan, J. Demmel, J. Donato, Dongarra, V. Eijkhout, R. Pozo, C. Romine, and H. Van der Vorst. Templates for the salution of linear systems: building blocks for iterative methods. SIAM, 1994.
[7] P. Becker, H. Pandya, G. Gebhardt, C. Zhao, C. J. Taylor, and G. Neumann. Recurrent Kalman networks: Factorized inference in digg--dimessioall deep feature spaces. In International Conference on Muchine Learning (ICML). PMLR, 2019.
[8] V. Borovitskiy, I. Azangulov, A. Terenin, P. Mostowsky, M. Deisenroth, and N. Durrande. Matérn Gaussian processes on graphs. In International Conference on Artificial Intelligence and Statistics (AISTATS).
PMLR, 2021.
[9] J. Brajard, A. Carrassi, M. Bocquet, and I.. Bertino. Combining data assimilation and machine learning to emulate dynamical model from sparse and noisy ohservations: case study with the Lorenz 96 model. Journal of Computational Science, 44:101171, 2020.
[10] P. E. Chang, W. J. Wilkinson, M. E. Khan, and A. Solin. Fast variational learning in state-space Gaussian process models. In IEEE International Workshop Machine Learning for Signal Processing (MLSP), volume 30. IEEE, 2020.
[11] A. J. Cohen, M. Brauer, R. Burnett, H. R. Anderson, J. Frostad, K. Estep. K. Balakrishnan, B. Brunekreef, L. Dandona. R. Dandona, et al. Estimates and 25- trends of the global burden of disease attributable to ambient air pollution: an analysis data from the Global Burden Diseases Study 2015. The Lancet, 2017.
389(10082):1900-19999
[12] P. Dagum, A. Galper, and E. Horvitz. Dynamic network models for forecasting. In Uncertainty in Artificial Intelligence (UAI). Elsevier, 1992.
[13] Z. Dai, A. Damianou, J. González, and N. Lawrence. Variational auto-encoded deep Gaussian processes. arXiv preprint arXiv: /51 /, 106455, 2015.
[14] A. Damianou and N. D. Lawrence, Deep Gaussian processes. In International Conference on Artificial Intelligence and Statistics (AISTATS). PMLR, 2013.
[15] T. DelSole and X. Yang. State and parameter estimation in stochastic dynamical models. Physica D: Nonlinear Phenomena, 1778111177111188 2010.
[16] A. P. Dempster, N. M. Laird, and D. B. Rubin. Maximum likelihood from incomplete data via the EM algorithm. Journal of the royal statistical society: series B (methodological), 39(1):1-22, 1977.
[17] G. Evensen. Sequential data assimilation with a nonlinear guusi-ccostropphcc model using Monte Carlo methods to forecast error statistics. Journal of Geophysicul Research: Oceans, 1014)::0143310166 1994.
[18] G. Evensen and J. Van Leeuwen. An ensemble Kalman smoother for nonlinear dynamics. Monthly Weather Review, 18826)2782118 2000.
[19] M. Fraccaro, Kamronn, U. Paquet, and O. Winther. disentangled recognition and nonlinear dynamies model for unsupervised learning. In Advances Neural Information Processing Systems (NeurIPS), 2017.
[20] Z. Ghahramani and G. E. Hinton. Variational learning for switching state-space models. Neural Computa- tion, 12(4):831-864 2000.
[21] Z. Ghahramani and S. Roweis. Learning nonlinear dynamical systems using an EM algorithm. In Advances in Neural Information Processing Systems (NeurIPS), 1998.
[22] T. Gneiting and A. E. Raftery. Strietly proper scoring rules, prediction, and estimation. Journal of the American statistical Association, 102(477):359-378 2007.
[23] O. Hamelijnek, W. Wilkinson, N. Loppi. A. Solin, and T. Damoulas. Spatio-temporal variational Gaussian processes. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
[24] J. Hartikainen, J. Rühimaki, and S. Särkkä. Sparse ppatio-tempora Gaussian processes with general likelihoods. In International Conference on Artificial Neural Networks (ICANN), 2011.
11

---

[25] H. Hersbach, B. Bell, P. Berrisford, S. Hirahara, A. Horányi, J. Muñoz-Sabatie J. Nicolas, C. Peubey, R. Radu, D. Schepers, al. The ERA5 global reanalysis. Quarterly Journal of the Royal Meteorological Society, (16(730):1999-2000
2020.
[26] P. L.. Houtekamer and F. Zhang. Review of the ensemble Kalman filter for atmospheric data assimilation. Monthly Weather Review, 144(12):4489 4432, 2016.
[27] H.-P. Hsieh, S.-D. Lin, and Y. Zheng. Inferring air quality for station location recommendation based on urban big data. In SIGKDD International Conference on Knowledge Discovery and Data Mining, 2015.
[28] M. Johnson, K. Duvenaud, Wiltschko, R. P. Adams, and S. R. Datta. Composing graphical models with neural networks for structured representations and fast inference. In Advances in Neural Information Processing Systems (NeurlPS), 2016.
[29] R. E. Kalman. new approach to linear filtering and prediction problems. 1960.
[30] M. Katebi and M. Johnson. Predictive control design for large-scale systems. Automatica, 33(3):421-455,
1997.
[31] M. Katzfuss and N. Cressie Spatio-temporal smoothing and EM estimation for massive remote-sensing data sets, Journal of Time Series Analysis, 32(4):30-446, 2011.
[32] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint |41.v:1412.6980, 2014.
[33] D. P. Kingma and M. Welling, Auto-encoding variational bayes. In International Conference on Learning Representations (ICLR), 2014.
[34] R. Krishnan, U. Shalit, and D. Sontag. Structured inference networks for nonlinear state space models. In AAAI Conference on Artificial Intelligence, 2017.
[35] P. C. Kyriakidis and A. G. Journel. Geostatistical space-time models; a review. Mathematical geology, 31:6516684, 1999.
[36] F. Lindgren, H. Rue, and J. Lindström. An explicit link between Gaussian fields and Gaussian Markov random fields: the stochastic partial differential equation approach. Journal of the Royal Statistical Society: Series B (Statistical Methodology). 73(4):423-49) 498, 2011.
[37] L. Ljung. Asymptotic behavior of the extended Kalman filter parameter estimator for linear systems. IEEE Transactions on Automatic Control, 24(1):36-50 50, 1979.
[38] M. B. Marinov, 1. Topalov, E. Gieva, and G. Nikolov. Air quality monitoring in urban environments. In 39th International Spring Seminar on Electronics Technology (ISSE). IEEE, 2016.
[39] K. P. Murphy. Dynamic Bayesian networks: California, Berkeley, 2002.
representation, inference and learning. University of partial observations. arxiv preprint arXiv: 1903. 10335, 2019.
[40] D. Nguyen, S. Ouala, L. Drumetz, and R. Fablet. EM-like learning chaotic dynamics from noisy and
[41] A. V. Nikitin, S. John, A. Solin, and S. Kaski. Non-separable spatio-tempora graph kemelss via SPDEs. In International Conference on Artificial Intelligence and Statistics (AISTATS). PMLR, 2022.
[42] S. Ohlwein, R. Kappeler, M. Kutlar Joss, N. Künzli, and B. Hoffmann. Health effects of ultrafine particles: systematic literature review update of epidemiological evidence. International Journal of Public Health, 64:547-559, 2019.
[43] J. Oskarsson, P. Sidén, and F. Lindsten. Scalable deep Gaussian Markov random fields for general graphs. In International Conference on Muchine Leurning (ICML). PMLR, 2022.
[44] G. Papandreou and A. L. Yuille. Gaussian sampling by local perturbations. In Advances in Neural Information Processing Systems (NewrIPS), 2010.
[45] J. Pearl. Probabilistic reasoning in intelligent systems: networks of plausible inference. Morgan Kaufmann, 1988.
[46] S. S. Rangapuram, M. W. Seeger. J. Gasthaus, L. Stella, Y. Wang, and T. Januschowski. Deep state space models for time series forecasting. In Advances in Neural Information Processing Systems (NeurIPS), 2018.
[47] H. E. Rauch, F. Tung, and C. T. Striebel. Maximum likelihood estimates of linear dynamic systems. AIAA journal, 3(8):1445 -454-4 1965.
[48] H. Rue and L. Held. Gaussian Markov random fandomfields: theory and applications. CRC press, 2005.
[49] S. Sarkka, A. Solin, and J. Hartikainen. Spatiotemporal learning via infinite-dimensional Bayesian filtering and smoothing: A look Gaussian process regression through Kalman filtering. IEEE Signal Processing Magazine, 30(4):51-6) 16151166,,22 2013.
[50] P. Sidén and F. Lindsten. Deep Gaussian Markov random fields. In International Conference on Machine Learning (ICML). PMLR, 2020.
12

---

[51] A. Solin and S. Särkkä. Explicit link between periodic covariance functions and state space models. In International Conference 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):319) 194 3208, 2007.
[53] Szunyogh, J. Kostelich, Gyarmati, E. Kalnay. R. Hunt, Ou, E. Satterfield. and A. Yorke. local ensemble transform Kalman filter data assimilation system for the NCEP global model. Tellux A: Dynamic Meteorology and Oceanography, 600 60(1):113-130, 2008.
[54] H. Wang, A. Bhattacharya, D. Pati, and Yang. Structured variational inference Bayesian 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, 1336())463-4822 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. Wikle and 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 Interna- tional 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-grainet air quality based on big data. In SIGKDD Conference Knowledge Discovery and Data Mining, 2015.
13

---

Appendices
Appendix A provides derivations supporting Section the main paper. In Appendix B explain our experimental setup, including dataset preparation and model implementation, in more detail. Finally, Appendix provides additional results supporting our claims regarding the scalability of our method, together with additional results from the experiments presented in Section
A ST-DGMRR derivations
In this section we provide detailed derivations the ST-DGMRE joint distribution, for both first- order transition models (Section A.I and higher-order transition models (Section A.2
A.1 Joint distribution
The LDS (see Section 2.2 and 3.1 X that factorizes as
the main paper) defines a joint distribution over system states
K II k=1
p(xo:)) = -x Mo: Q
M FAXA-1 C Q
(18)
with Xk, Mo: Cj E R and FA; Qu E RNNN As product of Gaussian distributions, p(xp:k) be written joint Gaussian W 12 with mean μ R (K+1) and precision (inverse covariance)
matrix Q E IR (F(++)) Here, derive expressions for ze and Z in terms of Nox Cs. F k. Qk-
First, note that Eq. 18 can be written set of linear equations
X N4 + 6o
cn 1 NN(0, Q)
X1 = F,xo + C1 + 11
61 ~NO, Q
X2 = F2x1 + C2 + 62
62 ~ ~NO, o2
Xx FK CK + CK
cx ~ N(O Q
Moving all xy-terms to the left-hand side, we can rewrite this as matrix- multiplication
Xo
Pu
Eo
F
I
C₁
é1
F2
X2
C2
+
E2
(19)
K
I
[K_
CK
with block-matrix F E R(K+))N((K+),) and vectorized x vec .o. xx) E RKK+IJN
c = vec( C1, cI,...,ck represent zero-blocks.
,K) E (K+1) and e = vec( 6p. (x) E RKK+IJ Empty positions in
Now, we can express x as an affine transformation of
x=F-ce+Flle
(20)
where F-1 exists because dct(F) diag( Qo. Q....., Qk), and с is deterministic, we can use the affine property of Gaussian distribu- tions obtain the joint distribution
Since c is distributed as NNO.Q-1) with Q
(21)
Thus, the joint precision matrix Ω factorizes as
Ω = (F-1Q (F-
= FFQF
(22)
14

---

and has block- block-tridiagone tridiagonal structure
Qo
F
F
-FQQ Q
QIFi
Q + F+
-F2 Q2
Ω
QK-IFK- Qx-1 F Qx K F
-F QK QK
(23)
Note that for matrix-vecter multiplications of the form f2x, the sparse structure of F and 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 F- without expensive matrix inversion, the components μ4 have to be computed iteratively as
μb F + Ck
(24)
In contrast, the information vector 7 S2 can be expressed compactly
7 = F" DFFFF-" c = Qc.
(25)
which can be computed efficiently using sparse and parallel matrix-vector multiplications GPU. We make use 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
60 ~ ~O, Q7
Xg Po 60
X1 F1 + C1 +1
61 ~ ~N, Qi
X₁ F F +.. Fkp + ck + Ex
Ex ~NO, Q
XK =x=FRIIM -K- + K22XK-22 +... KK + cк + tK
6K ~ ((0, QT²)
Following the same steps as before, this results in linear system
F1,1
I
X9
Ho
En
-2222
2,1
Xi
Ci
El
X2
C2
+
E2
-F -pp
F
Ppp-1
F p.1
X
CK
6
FK2
-FK,
(26)
This means that the matrix F extended by adding the higher-order transition matrices (Fr......, F.FF K,r) to the 7-th lower block diagonal of F for all - = 1,...,p. The expressions for Q u and n remain the same (using the extended F), resulting in 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 www.ecttiondifffsiiioo process
Op(t, s) 2 p(t, s) V (vp(t,s))
(27)
Øt
15

---

with constant diffusion coefficient D and velocity vector v= [u,, a x 30 lattice with grid cell size Ar Ay = 1 and periodic boundary conditions. The spatial discretization results in system of ordinary differential equations
[u,"]" The process is discretized on
0p(t) Mp((),
(28)
where p(1) vector containing the system states of all grid cells. Using a finite difference discretization, matrix M is defined
D D- 4D
ifd =1 ifi=j
Mij
(29)
otherwise,
where d(i,j) denotes the distance between cell and j. and my denotes the unit vector pointing
from lattice cell to its neighbor j. For example, for cell i = (5):8y) and cell j = (S Sy - 1) it is
naj [0,-1)² and thus niy v -.
Eq. 28 is converted into discrete-time dynamical system by approximating
Σ Σe (At)* (M)) P: F aa
PEtA: exp(At M)P,
(30)
using third-order Taylor series expansion. For simplicity, use time resolution At resulting in
adv- I+ M
M² 2 M
(31)
B.1.1 Process simulation
We sample the initial state Po from a GMRF with #4 O and precision matrix Qu = ST So with
S. (D A), where A is the adjacency matrix of the 4-nearest neighbour graph Guanices and
D 4. is the corresponding degree matrix. This corresponds to a 1-layer DGMRF with parameters
a = 1 -- 11,7 and
Starting from Po. we iteratively sample the next system state according to
F Ek E, (0,
(32)
with time-invariant transition matrix F5- (F au) where Fadv- is defined according to Eq. 31 This effectively aggregates four simulation steps into one, i.e. At, = (f&+1 fx) 4, resulting in larger differences between consecutive system states. For the noise terms eg. we time-invariant precision matrix Q2 S S₁ where S = (10- I- A). This corresponds to I-layer DGMRF with parameters a = 4 3 = - 1,7 = and b = 0.
10
We simulate for K 20 time steps, using D 0.01 and 100.3,0.33 and generate observations by masking out grid cells within square of width w E [6..... 12} for consecutive time steps and applying white noise with standard deviation 0.01.
R-2
k-2
K-5
kod
CHS
E-G
k-0
K-D
1-11
k-
1-
1-15
1-17
1-10
Figure 5: Advection-difius dataset with ground truth system states (bottom) and corresponding observations using masks 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 Gunice and temporal bias terms are, similar to spatial bias terms, defined as b b/ 1.
(1)
16

---

Variant 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 F()) that are a simplified first-order approximation to the true transition matrix F adv--diff used to generate the data (see Eq. (31p). i.e. F I+ + M with time-invariant learnable diffusion coefficients D and velocity vectors v(O,, To ensure that the diffusion coefficient is non-negative, we model it as D )2 This leaves us with four learnable parameters d((), ( v( and by) per temporal layer.
If only partial knowledge about the underlying dynamics is available, the unknown parts can, for example, be replaced by small neural network. Here, we consider temporal layers of the form F( = + M with
Variant 2
(d(0)== -0 A(d11) + +Ejen(e P3j
if E n(i) if =j otherwise.
M
(33)
where we define 1, 13.2 (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.
[1]
(M.))(n.)) where /0]
MLP R2 + R22 is a multilayer perceptron
Log-determinant computations In our experiments, the spatial base graph Gtanice is small enough to pre-compute eigenvalues exactly and use the eigenvalue method for fog-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 with a single temporal layer with time-invariant diffusion transition matrices F;- AI + w A - D). Note that for = 0, this reduces to a simple auto-regressive process.
Observation model All ST-DGMRF variants assume a temporally and spatially invariant observa- tion noise level of a = 0.01. The observation matrices H2 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 datasct 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, well and v wind components 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, Delammay 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 -rggtransformed and standardized to zero mean and unit variance. Finally, we remove clear outliers where the transformed values jump up and down by more than threshold of 8 2.0 within three consecutive time steps. The ERAS covariates normalized to range between 11 and
To mimic a realistic setting of repeatedly occurring partial network failures, we define our test set by masking out all measurements within predefined spatial block (containing 50% all sensors) within 10 randomly placed windows of 20 time steps (see Figure (). 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.
17

---

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 masked time window are shaded in gray.
B.2.2 ST-DGMRF parameterization
As with the diveecion-diffusion dataset, we consider two ST-DGMRF variants with different types of temporal layers. In both cases, spatial and temporal layers defined based the Delaunay triangulation described in Section B.2.1] and temporal bias terms by are defined in terms of a neural network mapping local weather covariates to temporally and spatially varying biases. We use 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.
This variant accounts for directional transport processes, adopting transition model similar the neural network model used in the advection-diffusion experiments. particular, we consider temporal layers of the form F!! = =I M (1) with
Variant 1
(d 2 + ok.ij 4(d00)2 + Σjen Zjen() Vko0
if n(i) if = otherwise.
(M)))
(34)
where we define Kk k. = SMLP (nz Way. (ux);) where fMIP - - is MLP with one hidden layer of width 16 with ReLU activations, and Tanh output non-linearity. wy are the edge weights of the base graph (see Section B.2.1, and (u=); is the vector of weather covariates for node i at time Since we use these time-dependent covariates as input to the MLP, the resulting transition model is not time-invarian anymore. However, the parameters of the MLP remain shared across time and space. As before, diffusion parameter (f)) is assumed to be spatially and temporally invariant.
(1)
Variant The second variant uses highly simplified diffiusion temporal layers of the form F. x4) + (9) D) D) with spatially and temporally invariant parameters (!) and (!)
Log-determinant computations Again, the spatial base graph DDeaamay small enough pre- compute eigenvalues exactly and use the eigenvalue method for log-determinant computations (4
Variational distribution As with the advection-diffusion dataset, we consider two variants for the variational distribution, without temporal dependencies and one with single temporal diffusion layer.
Observation model All ST-DGMRF variants assume a temporally and spatially invariant observa- tion noise level of = 0.01. The observation matrices H₃, are defined as selection matrices matching training masks during the learning phase and the training plus validation masks during the testing phase.
18

---

B.3 Baseline models
B.3.1 DGMRF
We apply the DGMRF for general graphs introduced by 43
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 hyperparameter search over Espatial E {1,2,3 3} with Lspatial = 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 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 g 1 for the advection-diffusion data, and with p = 2 and
for the air quality data, using the Python statsmodels package. For each node in the test set, maximum likelihood parameter estimation 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 X a X Es, with initial state Xo ~ (MH. Σo) and unconstrained spatial eπror terms E₁ (O, Q ). We fix Eo = 10 encode high uncertainty about the initial state Xg, and fit a, Mg and Q to the data using closed- form EM updates. The EM algorithm is initialized with = 1, 0 and Q = diag(q) where elements qi 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-gen settaageneating process of the advection difffusion dataset, as well as an EnKS variant for which we use a state augmentation approach to estimate unknown parameters v and d = /D jointly with the system states. For both variants, we use 10 ensemble members (the maximum feasible on our machine). We fix the initial state distribution to ~ NO, 10 I), and sample transition noise terms as 6₁ (0, 0.1 I).
For the state augmentation approach, we define the initial distribution over velocities v as ((,,,, 0.1 I), where n. is randomly drawn from [-1, for each repeated run of the EnKS. Similarly, the initial distribution for diffusion parameter d' is defined as A (pa, 0.01) where μa is randomly drawn from [o, 0.2] for each repeated run of the EnKS. Finally, the transition noise terms for parameters and are sampled from (0, 0.01 -I.
B.3.5 MLP
For the air quality dataset, the MLP baseline maps local weather covariates (ug)s R log- transformed PM2.5 measurements. We use one hidden layer of width 16 with ReLU activations and no output non-linenn non-linearit. The MLP parameters are shared over both space and time.
B.4 Regularized Conjugate Gradients
We use a regularized variant the conjugate gradient (CG) method 4 to avoid slow convergence in the case of ill-conditioned matrices. Instead of directly solving potentially ill-conditioned linear system Ax = b. the idea to iteratively solve a sequence of regularized (i.e. well conditioned) linear systems
( ( +b.
(35)
19

---

At each iteration, the solution from the previous iteration x()) is used to obtain the next solution x (++1). Eventually, this sequence will converge towards the true solution x" of the original system Ax b.
We start with v = 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 maximum of 200 inner CG iterations is reached. This inner loop is repeated until the norm of the residuals
(NI + A)x(() (2x(1) +b)
(36)
drops below a threshold of 10-7 or a maximum of 100 outer iterations is reached. The initial guess (0) is given by the mean of the variational distribution qx(x).
C Additional results
In this section, we present additional results regarding the scalability of our approach (Section and provide more detailed results for the experiments in Section If the main paper (Section C2 and C.3). Finally, in Section 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 ddlvectioo-diffusiiii 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 model with transition matrix using Liemponl 2 temporal and Lspatial 2 spatial layers. To avoid additional matrix inversions the KS approach, we set the spatial and temporal bias terms b, by to zero, resulting in 140 0 and wall clock time per itcration. In addition, we measure the wall clock time needed to perform inference with the trained model.
O. We train the model for 1 000 iterations and measure the average
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, regulargu)) CG method (with a tolerance of 10 7 to compute the posterior mean and marginal variances (based 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 estimate of the average total time for inference.
employed the standard
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 [66], where in each iteration a single gradient ascent step is taken.
Results Figure 7 shows how ST-DGMRF and KS training and inference scale the number of nodes N 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-DGMRR training time increases only marginally and remains below the fastest KS iteration for all tested N. In addition, KS memory requirements (due storing K dense N x 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 > 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 < 1024, while the CG method only requires storing vectors of size ((N) making it feasible for N > 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
20

---

1.75
ST-DGMRF training
ST-DGMRF inference
200
training
1.50
2.0
CG fterations
tma (s)
1s]
1.5
100
0.75
1.0
0.50
50
0.25
0.5
0.00
0.0
0
250
500
750
1000
1250 1500
250
500
750
1000
1250
1500
number of
number nodes
Figure 7: Comparison 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 + 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, 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 9, including standard deviations for all metrics based on 5 runs with different random seeds. 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 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 significant effect of introducing temporal dependencies the variational distribution. Further, Figure & shows additional results for the ST-DGMRF variants when varying the number of temporal layers Lemppral- For all metrics, the performance improves significantly as we start adding temporal layers and stabilizes around Liemporal 3. Note that around the same point, both ST-DGMRF variants converge towards EnKS using the true sata-generating dynamics, terms of the RMSE, and even drop below it in terms of the CRPS. Only in terms of RMSE,, the ST-DGMRF models remain inferior to both EnKS variants. We hypothesize that increasing the expressivity (i.c. Lypatia) of the noise terms can further reduce this gap.
Table 4: Model performance for the advection-diffusion dataset with 9, reported mean + std over runs with different random seeds. All ST-DGMRF variants use Lygutial 2 and Lnemporrl = 4.
dynamies
RMSE,
RMSE,
CRPS J
ARMA ST-AR
2.3054
0.6812
1.7064
20.00m
1.9216#1.m92
0.9707
DGMRF
0.5901 +0,007
0.3808
0.0095
EnKS
true dynamics estimated dynamics
0.066 +-0.0000
0.0046
0.1027 1.1022 +.0000
0.1654-.0.001
0.0039110.005
0.14342:0.00
ST-DGMRF (ours)
advection-diffusion
0.0526+a..00
0.1148
0.0726+a..00
difffuion
diffusion
0.0526+40.00l
0.1146+00.0s
0.0726+0.000
neural network
none
0.0839-10.000
0.1334420.080
0.0833 0.0000
neural network
diffusion
0.0854410.007
0.1402.:0.061
0.0839120.008
21

---

EnkS (estimated dynamics)
EnKS (true dynamics)
ST-DGMRF (advection-difusionn
1.50
1.25
0.8
0.8
1.00
0.6
0.75
RMSE:
CAPS
0.4
0.4
0.50
0.25
0.2
0.00
0.0
Loorporu
Eipeoponal
Figure 8: RMSE,,, RMSE, and CRPS as a function of the number of temporal layers Lemporal for the advection sdteetion-diiffuioon dataset with w = 9, plotted as mean + std over 5 runs with different random seeds. Both ST-DGMRF variants are trained with variational distribution using one temporal diffusion layer. Note that Lemprral = 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 variational distribution (see Section [.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 -iffuuionnn 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 of the network, for ST-DGMRF, DGMRF and ARMA respectively. For all three models, state estimates are obtained by conditioning 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 [ patterns and overall trends, while spatial layers remove dependencies between close-by sensors and increase temporal fluctuations.
Clearly, temporal layers remove daily
0.425
0.75
ST-DGMRT trensition
- diffusion
0.70
0.375
CRPS
0.65
0.350
0.325
0.300
0.55
Lienponal
Laamporal
Figure 9: RMSE and CRPS for increasing Leemporal (mean + std over 5 runs). Both models use 2. As before, Lemporll = 0 corresponds to the spatial-only DGMRF baseline.
22

---

input
ST-DGMRF
DGMRF
log(PM2.5)
-2
log 5)
20
60
100 120
140
160
180
200
220
240
260
280
300
320
340
360
360
400
time step k
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 Spotial and Lemporal
P VI dynamics
RMSF
CRPS 4
ARMA
0,6820
0.3625
ST-AR
0.7350
0.4261 4:00003
DGMRF
0.7368
0.3966110002
MLP
0.8038100045
ST-DGMRF (ours)
diffusion
1
none
0.6147 +:00002
0.3239 t00058
diffuusion
1
diffusion 0.6190
0.3258 1:000
diffussion
2
0.6020100112
0.3214-100081
diffusson
2
diffusion 0.5928:10118
0.3161 1:00054
neural network
none 0.5995 0:00
0.31
000444 1:0494
neural network
1
diffusion 0.5853
0.3092+00257
neural network
2
none 0.5825
0.3062
neural network
2
diffusion
0.5565
C.4 Total compute
Most computations were performed on Nvidia Titan X GPU. On top of the final experiments, performed hyperparameter sweeps and additional test runs. Here, we provide estimates of the total compute time grouped by experiment:
Advection-diffusian dataset:
Performance comparison & ablations: 50 GPU hours
Varying mask size: ca. 30 GPU hours per resulting in ca. 210 GPU hours in total
Scalability: ca. 15 GPU hours
Air quality dataset:
Performance comparison & ablations: ca. 100 GPU hours
23

---

x
-2
-4
250
300
350
400
(g
30
100
150
200
250
300
350
400
(Sa fx)(x)
-2
-4
400
time step
Figure 11: Effects of applying temporal (fg) and spatial (sg) 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 (fg), and gray lines (bottom row) represent states after applying both temporal and spatial layers (sp fo). As reference, also plot ground truth grgtrraasformed and normalized PM2.5 measurements (dashed black lines). Right column: corresponding (transformed) states for all sensors at time & 100.
24


Reviews:
[{"review_id": "MLDoCto6PJ", "reviewer": "Reviewer_9F96", "summary": "This work extends DGMRF to account for temporal dependency in data with a SSM reformulation. With necessary assumptions, the proposed method produces accurate posterior, competitive performance to DGMRF and faster inference to Kalman smoother.", "questions": "- Were the observation parameters $H$ and $\\xi$ trainable or fixed at the true values?\n- Were the hyperparameters such as state noise and observation noise trainable or fixed at the true values?", "limitations": "The technical limitations were discussed in the last section.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The method extends DGMRF to account for temporal dependency in data.\n- It has been shown that the method has low computational complexity and good scalability.\n- The method produces higher accurate posterior comparing to DGMRF and other transitional TS models.", "weaknesses": "- Since the method is extending to dynamical systems, DGMRF is a bit of \"strawman\". The method should be also compared with other filter/smoother e.g. deep KL, ensemble KL, particle smoothing and etc.\n- Though the posterior was evaluated in the manuscript, the dynamical model (F) wasn't. It is very often that you got a decent posterior but learned a \"bad\" F (in terms of forecasting). \n- Only the posterior mean was evaluated. It is unclear how good the posterior variance is.\n- The variational posterior should be compared to the true posterior in an example where the latter is accessible analytically or via sampling.\n- Only one synthetic example is not sufficient for evaluating the method thoroughly. It is unclear how it generalizes.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Official Rebuttal to 9F96\n\nWe would like to thank Reviewer 9F96 for the positive feedback and the good suggestions. In the following, we will respond to each point separately, and refer to our general response to all reviewers whenever there is overlap with other reviewers.\n\n## Responses to Weaknesses\n\n> Since the method is extending to dynamical systems, DGMRF is a bit of \"strawman\". The method should be also compared with other filter/smoother e.g. deep KL, ensemble KL, particle smoothing and etc.\n\nWe added a spatiotemporal ST-AR baseline (see general rebuttal) to address this limitation. Unfortunately, other more complex filter/smoother approaches are more difficult to apply in the setting we are considering. For example, deep Kalman filtering with KVAE [17] requires the definition of a suitable encoder mapping from high-dimensional observations to a latent space. In settings where the pattern of missing observations varies over time, this requires some initial imputation or a smart way to inform the encoder about this pattern, and thus falls beyond the scope of this paper.\nOn the other hand, ensemble or particle filter approaches typically rely on an established dynamics model (as discussed in line 82-88). While we have such a dynamics model available for the advection-diffusion dataset, applying ensemble or particle filtering would simply converge towards the true posterior distribution, which we have access to, and therefore would, in our opinion, not form a very useful comparison.\n\n> Though the posterior was evaluated in the manuscript, the dynamical model (F) wasn't. It is very often that you got a decent posterior but learned a \"bad\" F (in terms of forecasting).\n\nThis is indeed a very interesting point. Unfortunately, we haven't had the time to evaluate the learned transition models, but we are hoping to do so during the discussion period.\n\n> Only the posterior mean was evaluated. It is unclear how good the posterior variance is.\n\nWe do evaluate the posterior variance in terms of the CRPS metric. However, as the CRPS is a combined evaluation of the mean and variance, we now also include direct evaluations of the marginal standard deviations, for the advection-diffusion dataset where the true posterior distribution is available (see Table 1 in the provided PDF).\n\n> The variational posterior should be compared to the true posterior in an example where the latter is accessible analytically or via sampling.\n\nThank you for this good suggestion. We will include this comparison in the appendix of the camera ready paper. To do this, we will perform the same evaluation as is done with the final posterior estimate (see above).  \n\n> Only one synthetic example is not sufficient for evaluating the method thoroughly. It is unclear how it generalizes.\n\nWe have added additional experiments on a real-world air quality monitoring dataset, showing that our approach generalizes well to more complex real world settings. For more details, please see the general response to all reviewers. \n\n## Responses to Questions\n\n> Were the observation parameters $H$ and $\\xi$ trainable or fixed at the true values?  \n> Were the hyperparameters such as state noise and observation noise trainable or fixed at the true values?\n\nThank you for pointing out that this crucial information is missing. In all experiments, the observation model parameters $\\mathbf{H}$ and $\\mathbf{R}$ ($\\boldsymbol{\\xi}$ is the random variable, not a parameter) are fixed. In particular, $\\mathbf{H}$ is a selection matrix defined according to the training/validation/test masks used in each experiment, and the observation noise is defined as $\\mathbf{R}=\\sigma^2\\mathbf{I}$. For the advection-diffusion experiments, $\\sigma$ matches the value used to generate the data, while for the newly added air quality dataset it is set to $\\sigma=0.01$. We will add this information to the description of the experimental setup.\n\n## Conclusion\nWe thank reviewer 9F96 again for taking the time to provide thoughtful feedback and for pointing out some missing information about the experimental setup. We hope that our responses and actions taken w.r.t. the evaluation of posterior variances and additional experiments on a real-world dataset are satisfactory. If there are any other changes you would like to see, we are more than happy to discuss them with you."}, {"review_id": "5EufkMfJ9Q", "reviewer": "Reviewer_LAkB", "summary": "The author extend work on deep Gaussian Markov random field to a spatio-temporal setting. Earlier it has been explored in spatial and graphical setting. The authors show how to retain the sparsity structure of the previous work, which is  need to be computationally feasible.\nThey also show the various types of temporal structure can be built in this framework.", "questions": "In Figure 2 to the right should w be from 6-10 in the it says 6-9,12?\nAlso how come mae is decreasing when w is increasing?\nWhy is the average value of the simulations very negative in appendix in the colorbar the value ranges from -35 to -40.\nHave you added a mean in the other models, so they can compensate for this?\nYou should also be able to get the marginal posterior distributions for your model how does they look for the simulations?", "limitations": ".", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The extension to a spatio-temporal models is of course very important as many datasets have both a spatial and temporal structure.\nAlso that the methods are computationally feasible is paramount as the complexity spatio-temporal models tend to be very large.\nEach part of the paper is clear written.", "weaknesses": "The main weakness for me is the validation of the method only a single simulated data set is rather weak.\nIf I understand it correctly you are comparing your method to either a pure time series model (AR, MA, ARMA) or a pure spatial (DGMRF?) this does not seems like a fair comparison.\nSomething simple like a AR processes with spatial noise is standard in statistics.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Official Rebuttal to LAkB\n\nWe would like to thank Reviewer LAkB for the positive assessment and good suggestions. In the following, we will respond to each point separately, and refer to our general response to all reviewers whenever there is overlap with other reviewers.\n\n## Responses to Weaknesses\n\n> The main weakness for me is the validation of the method only a single simulated data set is rather weak.\n\nWe have added additional experiments on a real-world air quality monitoring dataset, showing that our approach generalizes well to more complex real world settings. For more details, please see the general response to all reviewers. \n\n> If I understand it correctly you are comparing your method to either a pure time series model (AR, MA, ARMA) or a pure spatial (DGMRF?) this does not seems like a fair comparison. Something simple like a AR processes with spatial noise is standard in statistics.\n\nWe have added such a process as an additional baseline. For more details, please see the general response to all reviewers. \n\n## Responses to Questions\n\n> In Figure 2 to the right should w be from 6-10 in the it says 6-9,12?\n\nWell spotted! We have updated the text with the correct values w=6-10.\n\n> Also how come mae is decreasing when w is increasing?\n\nAs mentioned in line 346-347, this is an artifact of the different set of pixels used for evaluation. E.g. if the smallest mask exactly covers the \"mode\" of the advection-diffusion process (yellow pixels in Figure 2 (left)) and a model estimates all pixels to have some low \"base state\" (blue pixels), the average error between true state and model estimate will decrease as the mask size increases and starts covering more of the surrounding pixels that are close to the \"base state\".\n\n> Why is the average value of the simulations very negative in appendix in the colorbar the value ranges from -35 to -40. Have you added a mean in the other models, so they can compensate for this?\n\nThe simulated data is generated by drawing a sample from the initial state distribution and then simulating forward in time. Even though this initial distribution has zero mean, the initial sample happens to be drawn from the \"negative side\" of the distribution. We indeed made sure that the overall mean is subtracted before fitting the baseline models.\n\n> You should also be able to get the marginal posterior distributions for your model how does they look for the simulations?\n\nWe indeed have access to the true marginal posterior of the simulations. As explained in the general response to all reviewers, we have extended our evaluation to include a direct comparison between our marginal posterior estimates and these true marginal posterior distributions.\n\n## Conclusion\n\nWe thank reviewer LAkB again for taking the time to provide thoughtful feedback and for suggesting a feasible spatiotemporal baseline. We hope that we answered your questions to your satisfaction and we could resolve your concerns w.r.t. the evaluation of our method by adding experiments on a real-world dataset. If there are any other changes you would like to see, we are more than happy to discuss them with you.\n\n"}, {"review_id": "hLoxFoRLqH", "reviewer": "Reviewer_k13T", "summary": "The paper extends Gaussian Random Fields with precision-based spatiotemporal structure. ", "questions": "None", "limitations": "No issues", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Clarity: The paper is professionally written, with flawless math. The method is technical and dense, and there are many parts to it. It would have been useful to include illustrations, perhaps in the supplements.\n\nOriginality: The paper has moderate novelty. The idea of using spatial and temporal precisions to model random fields is very apt, but these ideas are also well known in general. \n\nQuality: The method is principled and well derived.", "weaknesses": "The experimental evaluation is insufficient. There is only a single simple advection example, which is effectively a toy case. There are no real world experiments. Unfortunately a paper like this really needs a real world’y experiment to show that the method has transferable performance. The paper is really good, but I’m leaning on rejection for this reason alone. \n\nThe performance is good, but incremental. I’m happy with the neural network variant performance, but there is something wrong if the advection model (that matches the true system perfectly!) can’t fit the simple dynamics, nor improve over neural network. I wonder if this is a sampling or data scarcity issue. \n\nThe results are a black box. The purpose of this work is to learn temporal causalities, and spatial couplings. Yet, neither are shown! The paper needs to illustrate what kind of structures the system has learnt, and show if they are accurate or useful, and demonstrate transferable insights.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Official Rebuttal to k13T\n\nWe would like to thank Reviewer LAkB for praising the quality and clarity of our method and for providing constructive feedback. In the following, we will respond to each point separately, and refer to our general response to all reviewers whenever there is overlap with other reviewers.\n\n## Responses to Weaknesses\n\n> The experimental evaluation is insufficient. There is only a single simple advection example, which is effectively a toy case. There are no real world experiments. Unfortunately a paper like this really needs a real world’y experiment to show that the method has transferable performance. The paper is really good, but I’m leaning on rejection for this reason alone.\n\nWe have added additional experiments on a real-world air quality monitoring dataset, showing that our approach generalizes well to more complex real world settings. For more details, please see the general response to all reviewers.\n\n> The performance is good, but incremental. I’m happy with the neural network variant performance, but there is something wrong if the advection model (that matches the true system perfectly!) can’t fit the simple dynamics, nor improve over neural network. I wonder if this is a sampling or data scarcity issue.\n\nWe are happy to report that we resolved this issue. Our updated results (see Table 1 in the provided PDF) now show a significant improvement of the ST-DGMRF models over the baselines, with the advection model with $L_{\\text{temporal}}=4$ temporal layers (matching the true transition model) clearly performing best. \nNevertheless, please note that the issue was not as severe as it seems: in the original version the table showed results for $L_{\\text{temporal}}=1$, which presents a simplification of the true transition model not accounting for longer range dependencies. \n\n> The results are a black box. The purpose of this work is to learn temporal causalities, and spatial couplings. Yet, neither are shown! The paper needs to illustrate what kind of structures the system has learnt, and show if they are accurate or useful, and demonstrate transferable insights.\n\nPlease note that we are not claiming to learn causal structures, but rather aim at exploiting prior knowledge about such structures in order to improve inferences about unobserved system states. While in general, learning the true (causal) structure of the system will definitely aid in obtaining better state estimates, in practice it may be sufficient to learn a \"good enough\" approximation. However, we agree with the reviewer that an evaluation of the learned transition and precision matrices is interesting and could provide useful insights into both the model and the system itself.\nUnfortunately, we haven't had the time to do this, but we are hoping to do so during the discussion period.\n\n## Conclusion\n\nWe thank reviewer k13T again for taking the time to provide thoughtful and constructive feedback, which has helped us to improve our experimental evaluation significantly. We hope that we have resolved your major concerns with regard to the experimental evaluation by adding additional experiments on a real-world dataset and by improving the results on the simulated dataset such that they now align well with the expected outcome. If these changes are to your satisfaction, we kindly ask you to consider revising your initial rating accordingly. We are more than happy to discuss possible solutions to any remaining issues with you.\n\n"}, {"review_id": "6K7Vwgb7YN", "reviewer": "Reviewer_Upf6", "summary": "This paper proposes a Gaussian model for graph-temporal data, with two sorts of sparseness in the precision matrix: blockwise due to an assumed Markov property in time, and within blocks due to assumed sparse graph structures. To achieve inference and parameter learning the paper proposes a two-stage approach with parameter learning done via VI and final posterior (mean) inference with CG. The variational distribution has a suitable structure that reflects the temporal sparseness of the model. \n\n\n", "questions": "With the extension propose on 267 it would seem to me that the variational distribution contains the exact posterior. Is this correct? If so, then the CG could surely be omitted if the VI was run to convergence (e.g. with an annealing learning rate). \n\nI assume the variational distribution is learned by direct gradient descent on the parameters. This would be rather inefficient if the matrices are ill conditioned. Does the natural gradient approach of [1] work here? \n\nThe motivation for the two-stage approach is that \"parameter estimation becomes computationally prohibitive in high dimensional settings\". [57] shows that doing the (truncated) CG algorithm also gives the marginal likelihood log det term with a good approximation. I feel that this is relevant work which should be discussed. \n\n[57] Scalable Log Determinants for Gaussian Process Kernel Learning, Dong, Eriksson, Nickisch, Bindel, Wilson, NeurIPS 2017\n\nTypo: 371 extend -> extent ", "limitations": "The approach is limited to Gaussian models, and in this case exact inference is available if the problem is small enough, so the method is only relevant for larger data. It is not really made clear in the experiments how large is feasible, and whether the two stage approach really works on real rather than simulated data.\n", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The presented model is a very natural and ubiquitous one, so the setting is well motivated. There are no other approaches that I am aware of that solve this particular situation. The approach is certainly valid.\n\nI found this paper easy to read and well explained. \n\nRather than solve the inference problem with a Kalman-style algorithm (as in [1]) this paper proposes a two-stage approach of using VI for parameter learning and then CG for the final posterior (mean). I initially thought this was a rather strange thing to do, but on reflection I see that it is quite a good idea in this setting. I wonder whether this has more general application. \n\nThe design of the posterior distribution is well chosen to capture the (exact?) posterior while reducing the computation complexity of the naive approach. ", "weaknesses": "The experiments section is rather weak and is comes across more as a unit test than a proper evaluation: the data is drawn from exactly the correct model class, and the primary comparison is to a model that is the same but missing the temporal dependencies. There is not comparison on real data, or on simulated data with any other approach except extremely simple AR/MA/ARMA models. The approach of [1] is applicable for the temporal part of this setting, and for a 30x30 grid it is feasible to take the 900 dimensional dense representation. This would have made a useful comparison. \n\nIn terms of novelty, the extension the spatio-temporal regime is not a huge leap, and is already arguably already in the DGMRF model class already. Most of VI is introduced in [40], so the novelty there is not that great. \n\nWhile the setting is general, to enable fast determinant the form of the spatial layer is actually quite restricted. For a general spatial graph the approach would be no better than kalman approaches.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Official Rebuttal to Upf6\n\nWe would like to thank Reviewer Upf6 for the positive assessment, the thorough feedback and interesting questions. In the following, we will respond to each point separately, and refer to our general response to all reviewers whenever there is overlap with other reviewers.\n\n## Responses to Weaknesses\n\n> The experiments section is rather weak and is comes across more as a unit test than a proper evaluation: the data is drawn from exactly the correct model class\n\nWe have added additional experiments on a real-world air quality monitoring dataset, showing that our approach generalizes well to more complex real world settings. For more details, please see the general response to all reviewers. \n\n> and the primary comparison is to a model that is the same but missing the temporal dependencies. There is not comparison on real data, or on simulated data with any other approach except extremely simple AR/MA/ARMA models. The approach of [1] is applicable for the temporal part of this setting, and for a 30x30 grid it is feasible to take the 900 dimensional dense representation. This would have made a useful comparison.\n\nWe have added an AR model with spatial noise (as suggested by reviewer LAkB) as an additional baseline, taking both spatial and temporal dependencies into account. Fore more details, please see the general response to all reviewers.\n\nWhile we agree that the approach of [1] could serve as an interesting baseline, we are not sure how feasible it is in practice. To the best of our knowledge there is no code publically available for this approach, and given the remaining time frame we do not think it is doable to implement this ourselves.\nHowever, we would happily include this baseline if the reviewer could point us to a working implementation.\n\n> While the setting is general, to enable fast determinant the form of the spatial layer is actually quite restricted. For a general spatial graph the approach would be no better than kalman approaches.\n\nIt is true that for general dense base graphs $G_{\\text{spatial}}$ and $G_{\\text{spatial}}$, there is no increase in computational efficiency. However, the idea is that in such a setting the dense graph is likely the result of a chain of more local (and thus sparse) interactions, and the dense adjacency matrix $\\mathbf{A}$ can thus be approximated by a composition of sparse matrices.\nFor further discussion and experimental validation of the spatial DGMRF approach for general graphs, we would like to point the reviewer to [40]. \n\n## Responses to Questions\n\n> With the extension propose on 267 it would seem to me that the variational distribution contains the exact posterior. Is this correct? If so, then the CG could surely be omitted if the VI was run to convergence (e.g. with an annealing learning rate).\n\nAlthough this seems to be the case at first sight, the variational distribution is defined in terms of a sparse and factorized covariance matrix, while the true posterior has a sparse precision matrix $\\mathbf{\\Omega}^+=\\mathbf{F}^T\\mathbf{S}^T\\mathbf{SF} + \\mathbf{H}^T\\mathbf{R}^{-1}\\mathbf{H}$. The corresponding covariance matrix $\\mathbf{\\Sigma}^+=(\\mathbf{\\Omega}^+)^{-1}$ will generally be dense. To facilitate fast sampling from the variational distribution, the idea is to approximate the true dense covariance matrix $\\mathbf{\\Sigma}^+$ with a sparse and factorized covariance matrix $\\mathbf{\\Lambda}$. \n\n> I assume the variational distribution is learned by direct gradient descent on the parameters. This would be rather inefficient if the matrices are ill conditioned. Does the natural gradient approach of [1] work here?\n\nThank you for this good suggestion. While we did not run into convergence issues with our experiments, it is definitely worth considering this in the future.\n\n> The motivation for the two-stage approach is that \"parameter estimation becomes computationally prohibitive in high dimensional settings\". [57] shows that doing the (truncated) CG algorithm also gives the marginal likelihood log det term with a good approximation. I feel that this is relevant work which should be discussed.\n\nThe advantage of using CG only to compute the final posterior distribution is that the associated iterations need to be run only once. In contrast, using CG to approximate log-determinants, would require running CG repeatedly until convergence during the training loop. Although CG usually converges fast, we would expect that for large systems this would still result in a major bottleneck.\n\n## Responses to Limitations\n\n> The approach is limited to Gaussian models, and in this case exact inference is available if the problem is small enough, so the method is only relevant for larger data. It is not really made clear in the experiments how large is feasible, \n\nFigure 3 in the appendix shows the general scaling behaviour of our method, supporting our theoretical analysis in the main paper. Of course, the scalability to even larger systems also depends on additional factors such as hardware constraints. To address this concern, we are planning to explore the limits or our approach during the discussion period, and report our findings as soon as possible.\n\n> and whether the two stage approach really works on real rather than simulated data.\n\nOur additional experiments on the real-world air quality dataset show that our method works well also in more challenging settings with unknown and complex dynamics (see general response to all reviewers for more details).\n\n## Conclusion\n\nWe thank reviewer Upf6 again for the time and effort spent on this review and for posing interesting and thought-provoking questions. Similar to the other reviewers, concerns were mainly raised about the experimental evaluation. We addressed this by adding experiments on a real-world dataset and by including a spatiotemporal ST-AR baseline model. If there are any other changes you would like to see, we are more than happy to discuss them with you.\n\n"}]

Author rebuttal:
# General rebuttal to all reviewers

First of all, we would like to thank all the reviewers for their thoughtful and constructive feedback. We are delighted to hear that our paper is "clear" (LAkB), "easy to read" (Upf6) and "professionally written" (k13T), and that the reviewers found our proposed method to be "well explained" (Upf6), "well motivated" (Upf6), "principled and well derived" (k13T) and "quite a good idea" (Upf6).

Concerns were raised mainly about the experimental evaluation, which was limited to a single simulated dataset and to relatively simple (purely temporal and purely spatial) baseline models.
Here, we provide a short summary of the major actions we have taken to address these concerns.

1. We agree with the reviewers that it is important to show how our method generalizes to real world settings. Therefore, we performed additional experiments on real air quality sensor data. 
2. We acknowledge the lack of a suitable spatiotemporal baseline in our evaluation. To address this limitation, we have included an AR process with unconstrained spatial noise in the transition model (as suggested by reviewer LAkB) as an additional baseline. We are happy to report that our method can improve on the performance of this spatiotemporal baseline.
3. Reviewers pointed out that while the true posterior is available for the simulated advection-diffusion dataset, we did not actually use it in the performance evaluation. This is a good point, and we have adjusted Table 2 accordingly.
4. Reviewer k13T correctly pointed out that the performance gain reported in the submitted paper is good but incremental. We are happy to report that we were able to improve the performance of all ST-DGMRF variants by a significant amount (see Table 1). An important ingredient was to replace the standard CG method by a regularized version [60], which allows for more stable posterior estimation, in particular for ST-DGMRFs with large number of temporal layers.

In the following, we provide detailed information on each of these points. Minor changes and clarifications are addressed in the responses to the individual reviewers.

## Additional experiments

We performed additional experiments on real air quality sensor data obtained from [58].
The dataset contains hourly PM2.5 measurements taken at 246 sensors distributed around Beijing. We considered a time period of T=400 hours in March 2015 and extracted relevant weather covariates from the ERA5 reanalysis [59]. The base graph is defined based on the Delaunay triangulation of sensor locations.

### Experimental setup

To define our test set, we randomly draw 10 time points $t_k$ and mask out all measurements within a spatial block (containing 50% of all sensors) for time steps $t_k, …, t_k+20$, mimicking partial network failures.

Since the transport of pollutants is a complex process involving many different temporal scales, we hypothesize that both increasing the Markov order $p$ and the number of temporal layers will improve the estimated posterior. To test this, and to validate that ST-DGMRF is indeed able to capture such higher order dependencies, we perform an ablation study where we vary $L_{\text{temporal}}$ from 1 to 4, for Markov order $p=1$ and $p=2$ respectively.

We consider two ST-DGMRF variants: one with simple "diffusion" layers, and one neural network variant taking edge features and weather covariates as inputs.

### Main results (see Table 2 and Figure 1)

- We find that ST-DGMRF estimates unobserved system states more accurately than all considered baseline models.
- As expected, increasing $p$ from 1 to 2 results in more accurate posterior estimates, for both variants.
- Similarly, increasing $L_{\text{temporal}}$ results in continuous improvement in posterior estimates, for both variants.

## Additional baseline models

The spatiotemporal AR model with unconstrained spatial noise (ST-AR) takes the form 
$x_k = \alpha\cdot x_{k-1} + \epsilon_k$, where $x_0 \sim \mathcal{N}(\mu_0, \Sigma_0)$ and $\epsilon_k \sim \mathcal{N}(0, \mathbf{Q}^{-1})$. 

We fix $\Sigma_0 = 10\cdot I$ to encode high uncertainty about the initial state $x_0$, and fit $\alpha, \mu_0$ and $\mathbf{Q}^{-1}$ to data using closed-form EM updates. We initialize the EM algorithm with $\alpha=1$, $\mu=\mathbf{0}$ and $\mathbf{Q}^{-1}=diag(\mathbf{q})$ where elements $\mathbf{q}_i \in [3, 4]$ are initialized randomly. After convergence of the EM-algorithm, the final state estimates are obtained with the Kalman smoother.

## Additional evaluations

For the advection-diffusion dataset, for which the true posterior distribution is available, we now evaluate the estimated posterior mean and std w.r.t their ground truth.
That means, we replaced the MAE and RMSE computed based on the system state (Table 2 in the submitted paper) by $MAE_{\mu}$, $RMSE_{\mu}$, $MAE_{\sigma}$ and $RMSE_{\sigma}$ (see Table 1).

## General improvements

We replaced the standard CG method with a regularized CG variant [60], resulting in more stable final posterior estimations.
This allowed us to improve the performance of ST-DGMRF variants with large numbers of temporal layers significantly (see Table 1). Note that in contrast to Table 2 in the submitted paper, showing results for $L_{\text{temporal}}=1$, we now show results for $L_{\text{temporal}}=2$ and $L_{\text{temporal}}=4$. As expected, the ST-DGMRF with advection-diffusion dynamics and $L_{\text{temporal}}=4$ performs best, as it has exactly the same form as the true transition model.

## References

[58] Y. Zheng et al. Forecasting fine-grained air quality based on big data. In Proceedings of the 21th SIGKDD conference on Knowledge Discovery and Data Mining. 2015.

[59] H. Hersbach et al. The ERA5 global reanalysis. Quarterly Journal of the Royal Meteorological Society. 2020

[60] Bai ZZ, Zhang SL. A regularized conjugate gradient method for symmetric positive definite system of linear equations. Journal of Computational Mathematics. 2002


Meta-review: The four reviewers of this paper found the work well written, the methods principled, and the approach interesting. On the other side, initially, there were multiple concerns related to experimental validation (lack thereof), the somewhat restricted setting, and the black-box nature of the results. However, following the discussion most of the concerns were addressed in a satisfactory manner. The reviewers recommend accepting this work, with the only reviewer weakly recommending rejection having a low confidence score. During the rebuttal/discussion, the authors promised to add several things to their paper (incl. additional experiments). Go through the promised changes in detail and add them to the camera-ready version of your paper.
Decision: Accept (poster)