Title: Deep Gaussian Markov Random Fields for Graph-Structured Dynamical Systems
Venue: NeurIPS 2023
Accepted: yes
Keywords: probabilistic inference

graphical models

spatiotemporal dynamical systems

state-space models
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

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.

1 Introduction

Consider the problem of monitoring air pollution, the leading environmental risk factor for mortality worldwide [11, 42]. In the past years, sensor networks have been installed across major metropolitan areas, measuring the concentration of pollutants across time and space [27, 38]. These measurements are, however, prone to noise or might be missing completely due to hardware failures or limited sensor coverage. To identify sources of pollution or travel routes with limited exposure, it is 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, or transportation systems, where animal movements, the spread of diseases, or traffic load need to be estimated from imperfect data to facilitate scientific discovery and decision-making. From a probabilistic inference perspective, all these problems amount to estimating the posterior distribution over the latent states of a spatiotemporal system given partial and noisy multivariate time series data.

To make inference in these systems feasible, it is common to assume a state-space model, where the latent states evolve according to a 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 matrixmatrix multiplications and inversions. Thus, as 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

optimization loop [21, 40, 9]. If, additionally, the access to historical training data is limited, it is essential to 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 in a Dynamic Bayesian network (DBN) [12, 39, 56], an interpretable and data-efficient graphical model for multivariate time series data. If the joint space-time 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 [45, 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.

Combining the favorable statistical properties of classical graphical models with the flexibility and scalability of deep learning offers promising solutions in this regard [57]. Deep Gaussian Markov random fields (DGMRF), for example, integrate principled inference in Gaussian Markov random fields (GMRF) with convolutional [50] or graph layers [43] to define a new flexible family of GMRFs that can be learned efficiently by maximizing a variational lower bound, while facilitating principled Bayesian inference that remains computationally feasible even for high-dimensional complex models with dense dependency structures. While in principle DGMRFs can be applied to spatiotemporal systems by treating each time step independently, they are by design limited to spatial structures.

In this paper we extend the DGMRF framework to spatiotemporal systems with graph-structured states, transitions and noise terms. To this end, we formulate a 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 structure 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.

1.1 Related work

The simplest approach to scaling Kalman filtering and smoothing to high-dimensional state variables is 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 [17, 18, 26, 31], instead approximate state distributions with Monte Carlo estimates, balancing computational complexity and accuracy. This is widely used in the geosciences [53, 55] 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 is only a poor approximation and complex error covariances need to be estimated [15, 52]. Alternatively, expectation maximization

(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. Alternatively, variational inference allows for fast inference in (deep) state space models [20, 3, 28, 34, 54] by approximating the true posterior with a 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-making.

Gaussian processes (GP) are another widely used model class for spatiotemporal inference [35]. In contrast to 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, 13]. However, scaling GPs to 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, 10, 1, 58, 23] which, as discussed before, remains computationally prohibitive for large state spaces. Extensions to graph-structured domains [8, 41] rely on the relation between GPs and stochastic differential equations (SDE) [36, 51], 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 is limited or even unavailable. Our contributions can be summarized as follows:

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.

We show that the multi-layer space-time formulation facilitates efficient learning and inference, scaling linearly w.r.t. the number of layers, time series length, and state dimensionality.

In experiments on synthetic and real world data, we demonstrate that our approach provides accurate state and uncertainty estimates, comparing favorably to other scalable approaches relying on ensembles or simplifications of the dependency structure.

2 Preliminaries

2.1 Problem formulation

Consider a single sequence of high-dimensional measurements y 0: K = ( y 0 , . . . , y K ) taken at consecutive time points t 0 , . . . , t K . Each y k ∈ R M k represents partial and noisy observations (e.g. from spatially distributed sensors) generated according to a linear Gaussian noise model

from the latent state x k ∈ R N of an underlying dynamical process of dimensionality N ≥ M k (e.g. air pollution along N roads within an urban area). It is common to assume either R k = diag ( σ 2 k, 1 , . . . , σ 2 k,M k ) , representing independent sensors with varying noise levels, or R k = σ 2 I M k , representing a fixed noise level across all sensors and time points [9]. Further, the observation matrices H k ∈ R M k × N are typically very sparse with O ( M k ) number of non-zero entries.

Assuming that all H k are known, we aim at reconstructing the latent system states x 0: K = ( x 0 , . . . , x K ) ∈ R ( K +1) × N from data y 0: K . Since observations are partial and noisy, x 0: K cannot be recovered without ambiguity. It is therefore essential to treat the problem probabilistically and incorporate available domain knowledge into the prior distribution p ( x 0: K | θ ) . The task then amounts to (i) specifying a suitable prior that captures the spatiotemporal structure of x 0: T , (ii) estimating unknown parameters, and (iii) computing the posterior distribution p ( x 0: K | y 0: K , ˆ θ ) .

Notation We use x = vec ( x 0: K ) ∈ R ( K +1) N and y = vec ( y 0: K ) ∈ R M with M = ∑ K k =0 M k to denote vectorized states and measurements. Similarly, we use H = diag ( H 0 , . . . , H K ) ∈ R M × ( K +1) N and R = diag ( R 0 , . . . , R K ) ∈ R M × M to 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 we build in Section 3 to formulate a flexible multi-layer spatiotemporal prior.

2.2 Linear dynamical system

A discrete-time linear dynamical system (LDS) is a state-space model that defines the prior over states x 0: K in terms of linear Gaussian transition models

and an initial distribution x 0 ∼ N ( µ 0 , Q -1 0 ) . The classical Kalman filter [29] and its variants [47, 37] use recursive algorithms that exploit the temporal independence structure of the prior and its conjugacy to the observation model (see Eq. (1)) to compute marginal posterior distributions p ( x k | y 0: k ) in O ( KN 3 ) . The same formulas can be used within the EM-algorithm to learn unknown parameters µ 0 , Q 0 , F k , c k , Q -1 k in O ( JKN 3 ) 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 Markov random fields

A multivariate Gaussian N ( µ , Ω -1 ) forms a GMRF [48] with respect to an undirected graph G = ( V , E ) if the edges E correspond to the non-zero entries of the precision matrix Ω , i.e. Ω ij = 0 ⇔ ( i, j ) ∈ E . Given a linear Gaussian observation model y ∼ N ( Hx , R ) , the posterior can be written in closed form p ( x | y , ˆ θ ) = N ( µ + , ( Ω + ) -1 ) with precision matrix Ω + = Ω + H T R -1 H and mean µ + = ( Ω + ) -1 ( Ω µ + H T R -1 y ) .

̸

However, naively computing µ + by matrix inversion requires O ( K 3 N 3 ) computations and O ( K 2 N 2 ) memory, which is infeasible for high-dimensional systems. Instead, the conjugate gradient method [6] can be used to iteratively solve the sparse linear system Ω + µ + = Ω µ + H T R -1 y for µ + . The same approach allows to generate samples ˆ x ∼ N ( µ + , ( Ω + ) -1 ) and obtain Monte Carlo estimates of marginal variances [44]. In practice, however, it remains difficult to design suitable 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

where g θ = g ( L ) θ ◦· · · ◦ g (1) θ is a composition of L simple linear layers with convenient computational properties. This implicitly defines a GMRF x ∼ N ( µ , Ω -1 ) with µ = -G -1 θ b θ and Ω = G T θ 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.

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 h ( l ) = G ( l ) θ h ( l -1) + b ( l ) θ is defined on a sparse base graph ¯ G with adjacency matrix A and degree matrix D such that

with parameters θ l = ( α l , β l , γ l , b l ) . This construction allows for fast log-determinant computations during learning, and GPU-accelerated computation of g θ ( x ) using existing software for graph neural networks. Note that for L layers defined on ¯ G , the sparsity pattern of Ω corresponds to the sparsity pattern of ( A + D ) 2 L . And thus, the resulting model defines a GMRF w.r.t the 2 L -hop graph of ¯ G .

3 Spatiotemporal DGMRFs

To extend the DGMRF framework to spatiotemporal systems, we first formulate a 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 a flexible model architecture that facilitates efficient learning and principled Bayesian inference in high-dimensional dynamical systems, scaling favorably compared to Kalman filter based methods.

3.1 Graph-structured dynamical prior

Consider a 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 i of the system state x k can be associated with a node i ∈ V in a multigraph G = ( V , E spatial , E temporal ) , where the set of undirected spatial edges E spatial ⊆ V × V defines the sparsity pattern of noise precision (inverse covariance) matrices

̸

̸

and the set of directed temporal edges E temporal ⊆ V × V defines the sparsity pattern of the state transition matrices

̸

This defines a DBN over x 0: K encoding conditional independencies of the form

where n temporal ( i ) denotes the set of neighbors j of i for which ( j, i ) ∈ E temporal. Intuitively, E temporal represent causal effects over time, while E spatial 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 a multivariate Gaussian prior N ( µ , Ω -1 ) on x with sparse precision matrix Ω ∈ R ( K +1) N × ( K +1) N . Importantly, Ω can be shown to factorize as F T QF with block diagonal matrix Q and unit lower block bidiagonal matrix F defined as

where empty positions represent zero-blocks. Further, while the mean µ ∈ R ( K +1) N needs to be computed iteratively as µ k = F k µ k -1 + c k , the information vector η = Ω µ can be expressed compactly as

with c = [ µ 0 , c 1 , . . . , c K ] ∈ R ( K +1) N . See Appendix A.1 for detailed derivations.

3.1.2 DGMRF formulation

We now reformulate x ∼ N ( µ , Ω -1 ) as a DGMRF. Importantly, in contrast to [50, 43], the graphstructured dynamical prior allows us to impose additional structure, reflecting the spatiotemporal nature of the underlying system. In particular, since Q 0 , Q 1 , . . . , Q K are symmetric positive definite, we can factorize Q = S T S with symmetric block-diagonal matrix S , and thus Ω = F T S T SF . Together with Eq. (9), this results in a GMRF defined by

Finally, we separate g θ into a temporal map f θ : R ( K +1) N → R ( K +1) N and a spatial map s θ : R ( K +1) N → R ( K +1) N defined as

Note that this results in an overall bias b θ = Sb f + b s and thus c = -( b f + S -1 b s ) , which allows for modelling long-range spatial dependencies in µ 0 and c 1 , . . . , c K . The combined transformation z = ( s θ ◦ f θ )( x ) essentially defines a standard two-layer DGMRF, which describes a graph-structured dynamical system if the parameters θ = ( b s , b f , F 1 , . . . , F K , Q 0 , . . . , Q K ) are subject to sparsity constraints (5) and (6).

3.2 Parameterization

Learning the unknown parameters θ directly from a single sequence y 0: K will result in a highly overparameterized model that is unsuitable for the problem of interest. In addition, careful parameterization of g θ can greatly reduce the computational complexity of the transformation and the associated log-determinant computations required during learning [50, 43]. To achieve a good trade-off between data-efficiency, expressivity and scalability, we define both s θ and f θ in terms of simple layers with few parameters and convenient computational properties.

3.2.1 Spatial layer(s)

To enable fast log-determinant computations for arbitrary graph structures, we follow [43] and parameterize s θ in terms of K + 1 independent DGMRFs z k = S k h k + ( b s ) k , each defining a GMRF w.r.t G spatial = ( V , E spatial ) . Note that due to the multi-layer construction (see Section 2.4), G spatial is implicitly defined by the base graph ¯ G spatial and the number of layers L , with special case G spatial = ¯ G spatial if L spatial = 1 .

3.2.2 Temporal layer(s)

Compared to s θ , the definition of f θ is much less constrained as it does not affect the log-determinant computation (see Section 3.3.1), and should hence incorporate as much domain knowledge into transition matrices F k as possible. This could be, for example, based on conservation laws or knowledge about relevant covariates. To increase the flexibility in cases where the dynamics are unknown or involve long-range dependencies, each F k can again be decomposed into simpler layers F ( L temporal ) k · · · F (1) k , each defined according to a temporal base graph ¯ G temporal. Table 1 provides several example layers to give an idea of what is possible.

Note that any function can be used to define the entries of F ( l ) k , including neural networks taking available covariates or node/edge features as inputs [46]. The associated parameters can simply be included in θ . Similarly, non-linear dynamics can be approximated either through linearization akin to the Extended Kalman Filter [37], or through neural linearization [19, 7]. We, however, recommend sharing parameters where appropriate to avoid overparameterization.

Higher order Markov processes It is straight forward to extend f θ to describe a p -th order Markov process

where x -τ = 0 for τ = 1 , . . . , p . This can be done by introducing edges E temporal (2) , . . . , E temporal ( p ) and adding the corresponding higher-order transition matrices ( F τ,τ , . . . , F K,τ ) to the τ -th lower block diagonal of F (see Appendix A.2 for derivations).

3.3 Learning and inference

Although the marginal likelihood p ( y | θ ) is available in closed form (see Section 2.3), maximum likelihood parameter estimation becomes computationally prohibitive in high dimensional settings

[50]. Instead, a variational approximation is used to learn parameters ˆ θ for large-scale DGMRFs. Then, the conjugate gradient method allows to efficiently compute the exact posterior mean and to draw samples from p ( x | y , ˆ θ ) .

3.3.1 Scalable parameter estimation

Given a variational distribution q ϕ ( x ) = N ( ν ϕ , Λ ϕ ) , the parameters { θ , ϕ } are optimized jointly by maximizing the Evidence Lower Bound (ELBO)

using stochastic gradient descent, where the expectation is replaced by a Monte-Carlo estimate based on samples ˆ x ∼ q ϕ .

Variational distribution To facilitate efficient log-determinant computations and sampling via the reparameterization trick [33], we follow [43] and define q ϕ as an affine transformation x = P ϕ z + ν ϕ with z ∼ N ( 0 , I ) , resulting in Λ ϕ = P ϕ ( P ϕ ) T . Then, P ϕ can be defined as a block-diagonal matrix diag ( P 0 , . . . , P K ) with P k = diag ( ρ k ) ˜ S k diag ( ψ k ) , where ρ k , ψ k ∈ R N + are variational parameters and ˜ S k is based on spatial DGMRF layers as discussed in Section 3.2.1. To relax the temporal independence assumptions in Λ ϕ , we propose an extension P ϕ = diag ( P 0 , · · · , P K ) ˜ F , where ˜ F has a similar structure to F (see Eq. (8)). Again, the design of ˜ F k is highly flexible as it does not enter the log-determinant computation.

Log-determinant computations [50, 43] proposed specific lattice and graph layers for which the associated log-determinants log | det ( G θ ) | are computationally scalable. Conveniently, the spatiotemporal case does not require a new type of layer. Instead, we can build directly on top of existing spatial layers. Specifically, using Eq. (8)&(10), the log-determinant log | det ( G θ ) | simplifies to

where ( i ) follows from S being block-diagonal and ( ii ) follows from F being unit-lower triangular with det( F ) = 1 . Note that each S k is defined by a spatial DGMRF of dimension N (see Section 3.2.1), for which efficient log-determinant methods have been developed. The same arguments hold for the variational distribution proposed above. Finally, due to the diagonality assumptions on R t (see Section 2.1), 1 2 log | det ( R k ) | simplifies to ∑ M k i =1 log( σ i ) .

Computational complexity During training, log | det ( G θ ) | and log | det ( Λ ϕ ) | can be computed in O ( KNL spatial ) using Eq. (16) together with the methods proposed in [43]. Note that no complexity is added when introducing conditional dependencies between time steps. The necessary preprocessing steps, i.e. computing eigenvalues or traces, only need to be performed for the spatial base graph and are thus independent of the number of transitions K . Finally, assuming an average of d spatial and d temporal edges per node in ¯ E spatial and ¯ E temporal, the transformations s θ and f θ scale linearly with the number of nodes N . In particular, g θ ( x ) = ( s θ ◦ f θ )( x ) can be computed in O ( KNd spatial L spatial + KNd temporal L temporal ) , which dominates the computational complexity of the training loop. In addition, we can leverage massively parallel GPU computations to speed up this process even further.

3.3.2 Exact inference with conjugate gradients

As discussed in Section 2.3, the conjugate gradient method can be used to iteratively compute the posterior mean and obtain Monte Carlo estimates of marginal variances. Importantly, each iteration

is dominated by a single matrix-vector multiplication of the form

This amounts to a series of sparse matrix-vector multiplications with total computational complexity O ( KNd spatial L spatial + KNd temporal L temporal ) , which is again linear in the number of nodes, time steps, and layers, and can be implemented in parallel on a GPU.

4 Experiments

We implemented ST-DGMRF in Pytorch and Pytorch Geometric, and conducted experiments on a consumer-grade GPU, leveraging parallel computations in both spatial and temporal layers. In all experiments, we optimize parameters for 10 000 iterations using Adam [32] with learning rate 0 . 01 , and draw 100 posterior samples to estimate marginal variances. Unless specified otherwise, we use L spatial = 2 , L temporal = 4 and p = 1 , and define the variational distribution based on one spatial layer and one temporal diffusion layer. Additional details on our experiments are provided in Appendix B.

4.1 Advection-diffusion process

We start with a synthetic dataset for which we have access to both the ground truth posterior distribution and transition matrix. The dataset consists of K = 20 system states that are sampled from an ST-DGMRF with time-invariant transition matrix F adv-diff. This matrix is defined according to the third-order Taylor approximation of an advection-diffusion process with constant velocity and diffusion parameters, discretized on a 30 × 30 lattice with periodic boundary conditions. From the sampled state trajectory, we generate observations with varying amount of missing data by removing pixels within a square mask of width w for 10 consecutive time steps and adding noise with σ = 0 . 01 . For all experiments, we use the masked pixels as test set, and 10% of the observed pixels as validation set for hyperparameter tuning.

Two ST-DGMRF variants are considered: (1) using advection-diffusion matrices F ( l ) k 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 ij pointing from pixel i to j as input. In both cases, ¯ G temporal and ¯ G spatial are defined as the 4-nearest neighbor graph. Further, as the underlying dynamics are timeinvariant, we enforce F ( l ) k = F ( l ) k +1 and S ( l ) k = S ( l ) k +1 ∀ k ≥ 1 .

Baselines We compare our approach to a range of baselines accounting for varying degrees of spatial and/or temporal dependencies. In particular, we consider the original DGMRF applied to all time steps independently, an ARMA state-space model assuming spatial independence among the state variables, and a spatiotemporal AR state-space model (ST-AR) with spatially correlated error terms ϵ k for which an unconstrained covariance matrix Q -1 is estimated using the EM algorithm. For both ARMA and ST-AR, we use the standard Kalman smoother [47] to obtain posterior estimates. Additionally, we consider two Ensemble Kalman Smoother (EnKS) variants, one with an advection-diffusion transition model matching the

Table 2: Performance on the advectiondiffusion data with w = 9 . We report the mean over 5 runs with different random seeds.

true data-generating process, and one using state augmentation to estimate the velocity and diffusion parameters jointly with the system states. Note that, in contrast to the ST-DGMRF approach, we consider initial and transition noise parameters to be fixed in order to avoid divergence of the EnKS.

Performance evaluation We evaluate the estimated posterior mean and marginal standard deviations in terms of the root-mean-square-error (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

accurate estimates than the purely spatial DGMRF, the purely temporal ARMA model, and the STAR model with highly simplified transitions and unstructured noise terms. As expected, the two EnKS variants with fixed noise parameters provide the most accurate uncertainty estimates. However, the CRPS scores indicate that overall our ST-DGMRF approach results in better calibrated posterior distributions. More detailed results are reported in Appendix C.

systemstate

data

1.75

-34

data

1.50

system state

ST-AR

ST-AR

node (pixel) state

-35

DGMRF

ST-DGMRF

1.25

x

→X=

EnKS (true dynamics)

EnKS(estimated dynamics)

1.00

--

DGMRF

-

ST-DGMRF(advection-diffusion)

-36

RMSE

0.75

ST-DGMRF (neuralnetwork)

DGMRF

ST-DGMRF

0.50

37

0.25

-38

0.00

0.25

0

2

4

6

8

10

12

14

16

18

6

7

8

6

10

11

12

time

maskwidthw

Increasing mask size To evaluate the robustness to missing data, we analyze how posterior estimates change with increasing mask size. Figure 2 (right) shows the effect on RMSE µ when varying w from 6 to 12. Clearly, the purely spatial DGMRF suffers the most from expanding the unobserved region. In contrast, the ST-DGMRF variants continue to provide accurate state estimates that are on par with the EnKS using ground truth dynamics. Note that the decreasing errors for STAR can be attributed to the changing set of pixels used for evaluation.

truestencil

Ltemporal = 1

Ltemporal = 2

Ltemporal = 4

1.00

0.95

ears

=

0.90

0.85

P

ST-DGMRF transition

0.80

advection-diffusion

neural network

0.75

-0.25

0.00

0.25

0.00

0.05

0.10

0.15

0.20

0.25

1

2

Ltemporal

3

4

5

6

transitionweight

absoluteerror

Evaluation of learned transition models Finally, we assess how well the temporal ST-DGMRF layers can approximate the true data-generating dynamics, given varying levels of complexity in the transition model. For this purpose, we extract the stencil, i.e. the weights assigned to nearby pixels, from the ground truth and learned transition matrices and compare them in terms of absolute error and Pearson correlation. We find that the learned transition weights converge rapidly towards the true weights as the number of temporal layers increases (see Figure 3). For small L temporal, the neural network based layers show a stronger agreement with the true dynamics, suggesting that some flexibility in the temporal layers helps to compensate for simplifications in the transition structure.

4.2 Air quality data

To test our method on a real world system exhibiting more complex dynamics and graph structures, we conduct experiments on an air quality dataset obtained from [59]. The dataset contains hourly PM2.5 measurements from 246 sensors distributed around Beijing, China, covering a time period of K = 400 hours. To ensure that predicted PM2.5 concentrations are non-negative, we model both system states and observations in log-space. The spatial and temporal base graphs are defined

based on the Delaunay triangulation, with edge weights proportional to the inverse distance between sensors. To mimic a realistic scenario of local network failures, we randomly choose 10 time points t k and mask out all data points within a spatial block containing 50% of all sensors, for time steps t k , . . . , t k +20 (see Figure 4). As before, we use the masked data for model evaluation, and use 10% of the remaining data as validation set. The observation noise is assumed to be uniform with σ = 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 ( l ) k are then formed in the same way as in Section 4.1. In addition, we consider a simplified diffusion transition model (see Table 1) which cannot capture any directional transport processes.

Results Next to the baselines introduced in Section 4.1, we consider a multi-layer perceptron (MLP) mapping local weather features to the corresponding log-transformed PM2.5 concentration. Following [50], we also include weather features into the spatial DGMRF model by adding linear effects to the measurement model. Table 3 reports the resulting CRPS and the RMSE with respect to the masked out data. Clearly, the neural network based STDGMRF, accounting for time-varying directional transport processes, provides the most accurate state estimates. However, even with highly simplified diffusion transitions our approach provides better estimates than the considered baselines. We attribute this to the expressive DGMRF noise terms which can capture complex error structures

resulting from inaccuracies in the transition model. Increasing the Markov order from p = 1 to 2 clearly improves the resulting posterior estimates for both ST-DGMRF variants, reflecting the complexity of the modeled process. A similar trend is observed as we increase L temporal (see Appendix C).

inputdata

testdata

ST-DGMRFmean

ST-DGMRF sample

2

0.6

ST-DGMRF std

log(PM2.5)

0.4

0

0.2

0.0

50

100

150

200

250

300

350

400

time stepk

5 Conclusion

We have presented ST-DGMRF, an extension to Deep Gaussian Markov Random Fields for inference in spatiotemporal dynamical systems with partial and noisy observations, (partially) unknown dynamics, and limited historical data. Our reformulation of graph-structured state-space models as multi-layer space-time GMRFs enables computationally efficient learning and inference even in high-dimensional settings, scaling linear w.r.t. the number of both time steps and state variables. Empirically, ST-DGMRF provides more accurate posterior estimates than other scalable approaches relying on simplifications of the dependency structure. While our approach relies on linear Gaussian assumptions, which can be restrictive for real systems, we find that expressive DGMRF noise terms can compensate (to a certain extend) for approximations in the transition model. In the future, non-linearities between the layers could be explored as discussed in [50]. Further, we see potential in defining more flexible time-varying transition matrices based on a hierarchy of latent variables that again follow a ST-DGMRF.

References

V. Adam, S. Eleftheriadis, A. Artemev, N. Durrande, and J. Hensman. Doubly sparse variational gaussian processes. In International Conference on Artificial Intelligence and Statistics (AISTATS) . PMLR, 2020.

J. L. Anderson. An ensemble adjustment Kalman filter for data assimilation. Monthly Weather Review , 129(12):2884-2903, 2001.

E. Archer, I. M. Park, L. Buesing, J. Cunningham, and L. Paninski. Black box variational inference for state space models. arXiv preprint arXiv:1511.07367 , 2015.

Z. Z. Bai and S. L. Zhang. A regularized conjugate gradient method for symmetric positive definite system of linear equations. Journal of Computational Mathematics , pages 437-448, 2002.

D. Barber, A. T. Cemgil, and S. Chiappa. Bayesian time series models . Cambridge University Press, 2011.

R. Barrett, M. Berry, T. F. Chan, J. Demmel, J. Donato, J. Dongarra, V. Eijkhout, R. Pozo, C. Romine, and H. Van der Vorst. Templates for the solution of linear systems: building blocks for iterative methods . SIAM, 1994.

P. Becker, H. Pandya, G. Gebhardt, C. Zhao, C. J. Taylor, and G. Neumann. Recurrent Kalman networks: Factorized inference in high-dimensional deep feature spaces. In International Conference on Machine Learning (ICML) . PMLR, 2019.

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.

J. Brajard, A. Carrassi, M. Bocquet, and L. Bertino. Combining data assimilation and machine learning to emulate a dynamical model from sparse and noisy observations: A case study with the Lorenz 96 model. Journal of Computational Science , 44:101171, 2020.

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 on Machine Learning for Signal Processing (MLSP) , volume 30. IEEE, 2020.

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-year trends of the global burden of disease attributable to ambient air pollution: an analysis of data from the Global Burden of Diseases Study 2015. The Lancet , 389(10082):1907-1918, 2017.

P. Dagum, A. Galper, and E. Horvitz. Dynamic network models for forecasting. In Uncertainty in Artificial Intelligence (UAI) . Elsevier, 1992.

Z. Dai, A. Damianou, J. González, and N. Lawrence. Variational auto-encoded deep Gaussian processes. arXiv preprint arXiv:1511.06455 , 2015.

A. Damianou and N. D. Lawrence. Deep Gaussian processes. In International Conference on Artificial Intelligence and Statistics (AISTATS) . PMLR, 2013.

T. DelSole and X. Yang. State and parameter estimation in stochastic dynamical models. Physica D: Nonlinear Phenomena , 239(18):1781-1788, 2010.

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.

G. Evensen. Sequential data assimilation with a nonlinear quasi-geostrophic model using Monte Carlo methods to forecast error statistics. Journal of Geophysical Research: Oceans , 99(C5):10143-10162, 1994.

G. Evensen and P. J. Van Leeuwen. An ensemble Kalman smoother for nonlinear dynamics. Monthly Weather Review , 128(6):1852-1867, 2000.

M. Fraccaro, S. Kamronn, U. Paquet, and O. Winther. A disentangled recognition and nonlinear dynamics model for unsupervised learning. In Advances in Neural Information Processing Systems (NeurIPS) , 2017.

Z. Ghahramani and G. E. Hinton. Variational learning for switching state-space models. Neural Computation , 12(4):831-864, 2000.

Z. Ghahramani and S. Roweis. Learning nonlinear dynamical systems using an EM algorithm. In Advances in Neural Information Processing Systems (NeurIPS) , 1998.

T. Gneiting and A. E. Raftery. Strictly proper scoring rules, prediction, and estimation. Journal of the American statistical Association , 102(477):359-378, 2007.

O. Hamelijnck, W. Wilkinson, N. Loppi, A. Solin, and T. Damoulas. Spatio-temporal variational Gaussian processes. In Advances in Neural Information Processing Systems (NeurIPS) , 2021.

J. Hartikainen, J. Riihimäki, and S. Särkkä. Sparse spatio-temporal Gaussian processes with general likelihoods. In International Conference on Artificial Neural Networks (ICANN) , 2011.

H. Hersbach, B. Bell, P. Berrisford, S. Hirahara, A. Horányi, J. Muñoz-Sabater, J. Nicolas, C. Peubey, R. Radu, D. Schepers, et al. The ERA5 global reanalysis. Quarterly Journal of the Royal Meteorological Society , 146(730):1999-2049, 2020.

P. L. Houtekamer and F. Zhang. Review of the ensemble Kalman filter for atmospheric data assimilation. Monthly Weather Review , 144(12):4489-4532, 2016.

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.

M. J. Johnson, D. K. Duvenaud, A. 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 (NeurIPS) , 2016.

R. E. Kalman. A new approach to linear filtering and prediction problems. 1960.

M. Katebi and M. Johnson. Predictive control design for large-scale systems. Automatica , 33(3):421-425, 1997.

M. Katzfuss and N. Cressie. Spatio-temporal smoothing and EM estimation for massive remote-sensing data sets. Journal of Time Series Analysis , 32(4):430-446, 2011.

D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

D. P. Kingma and M. Welling. Auto-encoding variational bayes. In International Conference on Learning Representations (ICLR) , 2014.

R. Krishnan, U. Shalit, and D. Sontag. Structured inference networks for nonlinear state space models. In AAAI Conference on Artificial Intelligence , 2017.

P. C. Kyriakidis and A. G. Journel. Geostatistical space-time models: a review. Mathematical geology , 31:651-684, 1999.

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-498, 2011.

L. Ljung. Asymptotic behavior of the extended Kalman filter as a parameter estimator for linear systems. IEEE Transactions on Automatic Control , 24(1):36-50, 1979.

M. B. Marinov, I. Topalov, E. Gieva, and G. Nikolov. Air quality monitoring in urban environments. In 39th International Spring Seminar on Electronics Technology (ISSE) . IEEE, 2016.

K. P. Murphy. Dynamic Bayesian networks: representation, inference and learning . University of California, Berkeley, 2002.

D. Nguyen, S. Ouala, L. Drumetz, and R. Fablet. EM-like learning chaotic dynamics from noisy and partial observations. arXiv preprint arXiv:1903.10335 , 2019.

A. V. Nikitin, S. John, A. Solin, and S. Kaski. Non-separable spatio-temporal graph kernels via SPDEs. In International Conference on Artificial Intelligence and Statistics (AISTATS) . PMLR, 2022.

S. Ohlwein, R. Kappeler, M. Kutlar Joss, N. Künzli, and B. Hoffmann. Health effects of ultrafine particles: a systematic literature review update of epidemiological evidence. International Journal of Public Health , 64:547-559, 2019.

J. Oskarsson, P. Sidén, and F. Lindsten. Scalable deep Gaussian Markov random fields for general graphs. In International Conference on Machine Learning (ICML) . PMLR, 2022.

G. Papandreou and A. L. Yuille. Gaussian sampling by local perturbations. In Advances in Neural Information Processing Systems (NeurIPS) , 2010.

J. Pearl. Probabilistic reasoning in intelligent systems: networks of plausible inference . Morgan Kaufmann, 1988.

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.

H. E. Rauch, F. Tung, and C. T. Striebel. Maximum likelihood estimates of linear dynamic systems. AIAA journal , 3(8):1445-1450, 1965.

H. Rue and L. Held. Gaussian Markov random fields: theory and applications . CRC press, 2005.

S. Sarkka, A. Solin, and J. Hartikainen. Spatiotemporal learning via infinite-dimensional Bayesian filtering and smoothing: A look at Gaussian process regression through Kalman filtering. IEEE Signal Processing Magazine , 30(4):51-61, 2013.

P. Sidén and F. Lindsten. Deep Gaussian Markov random fields. In International Conference on Machine Learning (ICML) . PMLR, 2020.

A. Solin and S. Särkkä. Explicit link between periodic covariance functions and state space models. In International Conference on Artificial Intelligence and Statistics (AISTATS) . PMLR, 2014.

J. R. Stroud and T. Bengtsson. Sequential state and variance estimation within the ensemble Kalman filter. Monthly Weather Review , 135(9):3194-3208, 2007.

I. Szunyogh, E. J. Kostelich, G. Gyarmati, E. Kalnay, B. R. Hunt, E. Ott, E. Satterfield, and J. A. Yorke. A local ensemble transform Kalman filter data assimilation system for the NCEP global model. Tellus A: Dynamic Meteorology and Oceanography , 60(1):113-130, 2008.

H. Wang, A. Bhattacharya, D. Pati, and Y. Yang. Structured variational inference in Bayesian state-space models. In International Conference on Artificial Intelligence and Statistics (AISTATS) . PMLR, 2022.

J. S. Whitaker, T. M. Hamill, X. Wei, Y. Song, and Z. Toth. Ensemble data assimilation with the NCEP global forecast system. Monthly Weather Review , 136(2):463-482, 2008.

C. K. Wikle and M. B. Hooten. A general science-based framework for dynamical spatio-temporal models. TEST , 19:417-451, 2010.

C. K. Wikle and A. Zammit-Mangion. Statistical deep learning for spatial and spatiotemporal data. Annual Review of Statistics and Its Application , 10:247-270, 2023.

W. Wilkinson, A. Solin, and V. Adam. Sparse algorithms for Markovian Gaussian processes. In International Conference on Artificial Intelligence and Statistics (AISTATS) . PMLR, 2021.

Y. Zheng, X. Yi, M. Li, R. Li, Z. Shan, E. Chang, and T. Li. Forecasting fine-grained air quality based on big data. In SIGKDD Conference on Knowledge Discovery and Data Mining , 2015.

Appendices

Appendix A provides derivations supporting Section 3 in the main paper. In Appendix B, we explain our experimental setup, including dataset preparation and model implementation, in more detail. Finally, Appendix C provides additional results supporting our claims regarding the scalability of our method, together with additional results from the experiments presented in Section 4.

A ST-DGMRF derivations

In this section we provide detailed derivations of the ST-DGMRF joint distribution, for both firstorder transition models (Section A.1) and higher-order transition models (Section A.2).

A.1 Joint distribution

The LDS (see Section 2.2 and 3.1 in the main paper) defines a joint distribution over system states x 0: K that factorizes as

with x k , µ 0 , c k ∈ R N and F k , Q k ∈ R N × N . As a product of Gaussian distributions, p ( x 0: K ) can be written as a joint Gaussian N ( µ , Ω -1 ) with mean µ ∈ R ( K +1) N and precision (inverse covariance) matrix Ω ∈ R ( K +1) N × ( K +1) N . Here, we derive expressions for µ and Ω in terms of µ 0 , c k , F k , Q k .

First, note that Eq. (18) can be written as a set of linear equations

Moving all x k -terms to the left-hand side, we can rewrite this as a matrix-vector multiplication

with block-matrix F ∈ R ( K +1) N × ( K +1) N and vectorized x = vec ( x 0 , . . . , x K ) ∈ R ( K +1) N , c = vec ( µ 0 , c 1 , . . . , c K ) ∈ R ( K +1) N and ϵ = vec ( ϵ 0 , . . . , ϵ K ) ∈ R ( K +1) N . Empty positions in F represent zero-blocks.

Now, we can express x as an affine transformation of ϵ

where F -1 exists because det( F ) = 1 . Since ϵ is distributed as ϵ ∼ N ( 0 , Q -1 ) with Q = diag ( Q 0 , Q 1 , . . . , Q K ) , and c is deterministic, we can use the affine property of Gaussian distributions to obtain the joint distribution

Thus, the joint precision matrix Ω factorizes as

and has a block-tridiagonal structure

Note that for matrix-vector multiplications of the form Ωx , 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 -1 c without expensive matrix inversion, the components µ k have to be computed iteratively as

In contrast, the information vector η = Ω µ can be expressed compactly as

which can be computed efficiently using sparse and parallel matrix-vector multiplications on a GPU. We make use of this property in the DGMRF formulation and in the conjugate gradient method.

A.2 Extension to higher-order Markov processes

We can easily adjust the joint distribution to accommodate higher-order processes with dependencies on multiple past time steps.

For a p -th order Markov process, the dynamics are defined by equations

x K = F K, 1 x K -1 + F K, 2 x K -2 + · · · + F K,p x K -p + c K + ϵ K ϵ K ∼ N ( 0 , Q -1 K ) . Following the same steps as before, this results in a linear system

This means that the matrix F is extended by adding the higher-order transition matrices ( F τ,τ , . . . , F K,τ ) to the τ -th lower block diagonal of F for all τ = 1 , . . . , p . The expressions for Ω , µ and η remain the same (using the extended F ), resulting in a block p -diagonal precision matrix.

B Experimental details

B.1 Advection-diffusion process

The advection-diffusion dataset is a random sample from a ST-DGMRF for which the transition matrices are defined according to an advection-diffusion process

with constant diffusion coefficient D and velocity vector v = [ u, v ] T . The process is discretized on a 30 × 30 lattice with grid cell size ∆ x = ∆ y = 1 and periodic boundary conditions. The spatial discretization results in a system of ordinary differential equations

where ρ ( t ) is a vector containing the system states of all grid cells. Using a finite difference discretization, matrix M is defined as

where d ( i, j ) denotes the distance between cell i and j , and n ij denotes the unit vector pointing from lattice cell i to its neighbor j . For example, for cell i = ( s x , s y ) and cell j = ( s x , s y -1) it is n ij = [0 , -1] T , and thus n T ij v = -v .

Eq. 28 is converted into a discrete-time dynamical system by approximating

using a third-order Taylor series expansion. For simplicity, we use time resolution ∆ t = 1 resulting in

B.1.1 Process simulation

We sample the initial state ρ 0 from a GMRF with µ 0 = 0 and precision matrix Q 0 = S T 0 S 0 with S 0 = ( D -A ) , where A is the adjacency matrix of the 4-nearest neighbour graph G lattice, and D = 4 · I is the corresponding degree matrix. This corresponds to a 1-layer DGMRF with parameters α = 1 , β = -1 , γ = 1 and b = 0 .

Starting from ρ 0 , we iteratively sample the next system state according to

with time-invariant transition matrix F k = ( F adv-diff ) 4 , where F adv-diff is defined according to Eq. (31). This effectively aggregates four simulation steps into one, i.e. ∆ t k = ( t k +1 -t k ) = 4 , resulting in larger differences between consecutive system states. For the noise terms ϵ k , we use a time-invariant precision matrix Q k = S T k S k where S k = (10 · I -A ) . This corresponds to a 1-layer DGMRF with parameters α = 10 4 , β = -1 , γ = 1 and b = 0 .

Wesimulate for K = 20 time steps, using D = 0 . 01 and v = [ -0 . 3 , 0 . 3] T , and generate observations by masking out grid cells within a square of width w ∈ { 6 , . . . , 12 } for 10 consecutive time steps and applying white noise with standard deviation σ = 0 . 01 .

k = 1

k=2

k = 3

k= 4

k = 5

k=6

k = 7

k=9

k = 10

k = 11

k = 12

k = 13

k = 14

k = 15

k =16

k = 17

k = 18

k = 19

epep

s

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 G lattice, and temporal bias terms are, similar to spatial bias terms, defined as b ( l ) f = b ( l ) f 1 .

Variant 1 If prior knowledge is available in the form of a parameterized transition model, the STDGMRFtransition matrices can be parameterized accordingly. Here, we consider temporal layers F ( l ) k that are a simplified first-order approximation to the true transition matrix F adv-diff used to generate the data (see Eq. (31)), i.e. F ( l ) k = I + M ( l ) , with time-invariant learnable diffusion coefficients D ( l ) and velocity vectors v ( l ) . To ensure that the diffusion coefficient is non-negative, we model it as D ( l ) = ( d ( l ) ) 2 . This leaves us with four learnable parameters d ( l ) , u ( l ) , v ( l ) and b ( l ) f per temporal layer.

Variant 2 If only partial knowledge about the underlying dynamics is available, the unknown parts can, for example, be replaced by a small neural network. Here, we consider temporal layers of the form F ( l ) k = I + M ( l ) with

where we define ϕ ( l ) ij, 1 , ϕ ( l ) ij, 2 = f ( l ) MLP ( n ij ) where f ( l ) MLP : R 2 → R 2 is a multilayer perceptron (MLP) with one hidden layer of width 16 with ReLU non-linearity, and Tanh output non-linearity. Again, we define the transition model to be time-invariant and share MLP parameters across time and space. This amounts to 83 learnable parameters per temporal layer.

Log-determinant computations In our experiments, the spatial base graph G lattice is small enough to pre-compute eigenvalues exactly and use the eigenvalue method for log-determinant computations proposed in [43].

Variational distribution For the variational distribution, we also consider two variants, one without temporal dependencies (equivalent to the DGMRF baseline) and one with a single temporal layer with time-invariant diffusion transition matrices ˜ F k = λ I + ω ( 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 observation noise level of σ = 0 . 01 . The observation matrices H k are defined as selection matrices matching the training masks during the learning phase and the training plus validation masks during the testing phase.

B.2 Air quality data

The air quality dataset is based on hourly PM2.5 measurements obtained from [59]. We consider 246 sensors within the metropolitan area of Beijing, China, for which we extracted time series of K = 400 hours between 13 March 2015 at 12pm and 30 March 2015 at 3am. Relevant weather covariates (surface temperature, as well as u and v wind components at 10 meters above ground level) were extracted from the ERA5 reanalysis dataset [25].

B.2.1 Data preprocessing

We define both the spatial and the temporal base graph based on the Delaunay triangulation of sensor locations, G Delaunay, where we disregard edges between sensors that are more than 160 kilometers apart. Edge weights are defined as the inverse distance between sensors, normalized to range between 0 and 1. The raw PM2.5 measurements are log-transformed and standardized to zero mean and unit variance. Finally, we remove clear outliers where the transformed values jump up and down by more than a threshold of δ = 2 . 0 within three consecutive time steps. The ERA5 covariates are normalized to range between -1 and 1.

To mimic a realistic setting of repeatedly occurring partial network failures, we define our test set by masking out all measurements within a predefined spatial block (containing 50% of all sensors) within 10 randomly placed windows of 20 time steps (see Figure 6). Note that these windows may overlap, resulting in fewer periods of missing data with variable length. The masked out measurements are used for the final model evaluation.

log(PM2.5)

2

0

2

log(PM2.5)

2

0

0

20

40

timestepk

B.2.2 ST-DGMRF parameterization

As with the advection-diffusion dataset, we consider two ST-DGMRF variants with different types of temporal layers. In both cases, spatial and temporal layers are defined based on the Delaunay triangulation described in Section B.2.1, and temporal bias terms b f are defined in terms of a neural network mapping local weather covariates to temporally and spatially varying biases. We use a simple MLP with one hidden layer of width 16 with ReLU activations and no output non-linearity. The MLP parameters are shared over both space and time.

Variant 1 This variant accounts for directional transport processes, adopting a transition model similar to the neural network model used in the advection-diffusion experiments. In particular, we consider temporal layers of the form F ( l ) k = I + M ( l ) k with

where we define ϕ ( l ) k,ij , ψ ( l ) k,ij = f ( l ) MLP ( n ij , w ij , ( u k ) i ) where f ( l ) MLP : R 6 → R 2 is a MLP with one hidden layer of width 16 with ReLU activations, and Tanh output non-linearity. w ij are the edge weights of the base graph (see Section B.2.1), and ( u k ) i is the vector of weather covariates for node i at time k . Since we use these time-dependent covariates as input to the MLP, the resulting transition model is not time-invariant anymore. However, the parameters of the MLP remain shared across time and space. As before, diffusion parameter d ( l ) is assumed to be spatially and temporally invariant.

Variant 2 The second variant uses highly simplified diffusion temporal layers of the form F ( l ) k = λ ( l ) I + ω ( l ) ( A -D ) with spatially and temporally invariant parameters λ ( l ) and ω ( l ) .

Log-determinant computations Again, the spatial base graph G Delaunay is small enough to precompute eigenvalues exactly and use the eigenvalue method for log-determinant computations [43].

Variational distribution As with the advection-diffusion dataset, we consider two variants for the variational distribution, one without temporal dependencies and one with a single temporal diffusion layer.

Observation model All ST-DGMRF variants assume a temporally and spatially invariant observation noise level of σ = 0 . 01 . The observation matrices H k are defined as selection matrices matching the training masks during the learning phase and the training plus validation masks during the testing phase.

B.3 Baseline models

B.3.1 DGMRF

We apply the DGMRF for general graphs introduced by [43] to each time frame of the time series, not accounting for temporal dependencies. The DGMRF parameters are not shared across time, allowing for dynamically changing spatial covariance patterns. We use one spatial layer in the variational distribution, as proposed in [43], and run a hyperparameter search over L spatial ∈ { 1 , 2 , 3 } with L spatial = 2 performing best.

Including covariates In our experiments on the air quality dataset, for which we have access to relevant covariates, we follow [50] and add linear effects to the measurement model. Note that the vector of coefficients if shared across both space and time.

B.3.2 ARMA

We implemented ARMA( p, q ) models with p = 1 and q = 1 for the advection-diffusion data, and with p = 2 and q = 2 for the air quality data, using the Python statsmodels package. For each node in the test set, maximum likelihood parameter estimation is performed based on the observed time points. Given the estimated model coefficients, we obtain posterior mean and variance estimates using the standard Kalman smoother [47]. As the maximum likelihood estimates are deterministic, we do not provide standard deviations of the evaluation metrics for these models.

B.3.3 ST-AR

The spatiotemporal autoregressive (ST-AR) model takes the form x k = α · x k -1 + ϵ k , with initial state x 0 ∼ N ( µ 0 , Σ 0 ) and unconstrained spatial error terms ϵ k ∼ N ( 0 , Q -1 ) . We fix Σ 0 = 10 · I to encode high uncertainty about the initial state x 0 , and fit α, µ 0 and Q -1 to the data using closedform EM updates. The EM algorithm is initialized with α = 1 , µ 0 = 0 and Q -1 = diag ( q ) where elements q i are drawn randomly from the interval [5 , 6] . After convergence of the EM-algorithm, the final state estimates are obtained with the Kalman smoother [47].

B.3.4 EnKS

We consider an Ensemble Kalman Smoother (EnKS) variant for which the transition model matches the true data-generating process of the advection-diffusion dataset, as well as an EnKS variant for which we use a state augmentation approach to estimate unknown parameters v and d = √ D jointly with the system states. For both variants, we use 10 4 ensemble members (the maximum feasible on our machine). We fix the initial state distribution to x 0 ∼ N ( 0 , 10 · I ) , and sample transition noise terms as ϵ k ∼ N ( 0 , 0 . 1 · I ) .

For the state augmentation approach, we define the initial distribution over velocities v as N ( µ v , 0 . 1 · I ) , where µ v is randomly drawn from [ -1 , 1] for each repeated run of the EnKS. Similarly, the initial distribution for diffusion parameter d is defined as N ( µ d , 0 . 01) where µ d is randomly drawn from [0 , 0 . 2] for each repeated run of the EnKS. Finally, the transition noise terms for parameters v and d are sampled from N ( 0 , 0 . 01 · I ) .

B.3.5 MLP

For the air quality dataset, the MLP baseline maps local weather covariates ( u k ) i ∈ R 3 to logtransformed PM2.5 measurements. We use one hidden layer of width 16 with ReLU activations and no output non-linearity. The MLP parameters are shared over both space and time.

B.4 Regularized Conjugate Gradients

We use a regularized variant of the conjugate gradient (CG) method [4] to avoid slow convergence in the case of ill-conditioned matrices. Instead of directly solving a potentially ill-conditioned linear system Ax = b , the idea is to iteratively solve a sequence of regularized (i.e. well conditioned) linear systems

At each iteration, the solution from the previous iteration x ( i ) is used to obtain the next solution x ( i +1) . Eventually, this sequence will converge towards the true solution x ∗ of the original system Ax = b .

We start with ν = 10 and decrease it every 10 iterations by factor 10 . In each iteration, the standard CG method is employed to iteratively solve the regularized linear system until the residual norm drops below a threshold of 10 -7 or a maximum of 200 inner CG iterations is reached. This inner loop is repeated until the norm of the residuals

drops below a threshold of 10 -7 or a maximum of 100 outer iterations is reached. The initial guess x (0) is given by the mean of the variational distribution q ϕ ( x ) .

C Additional results

In this section, we present additional results regarding the scalability of our approach (Section C.1), and provide more detailed results for the experiments in Section 4 of the main paper (Section C.2 and C.3). Finally, in Section C.4 we provide estimates of the total computation time required for our experiments.

C.1 Scalability

To empirically demonstrate the scalability of our method, we generate additional advection-diffusion datasets with varying lattice size and compare the runtime of our ST-DGMRF approach to a naive Kalman smoother (KS) [47] approach. To this end, we consider a model with advection-diffusion transition matrix using L temporal = 2 temporal and L spatial = 2 spatial layers. To avoid additional matrix inversions in the KS approach, we set the spatial and temporal bias terms b s , b f to zero, resulting in µ 0 = 0 and c k = 0 . We train the model for 1 000 iterations and measure the average wall clock time per iteration. In addition, we measure the wall clock time needed to perform inference with the trained model.

ST-DGMRF For the ST-DGMRF approach, we proceed as before using a variational distribution with one temporal diffusion layer during training. For better comparability, we employed the standard (non-regularized) CG method (with a tolerance of 10 -7 ) to compute the posterior mean and marginal variances (based on 100 CG samples) and provide measurements of the average time per CG iteration instead of the total time needed to perform inference. Multiplying this with the average number of CG iterations needed until convergence results in an estimate of the average total time for inference.

Kalman smoother For the KS approach, instead of approximating the true posterior with a variational distribution and estimating the ELBO based on Monte-Carlo samples, we use the KS to obtain exact marginal posterior estimates, which are used to compute expectations in closed form. The marginal covariance and transition matrices required for the KS equations are extracted from the ST-DGMRF model in every iteration. The associated parameters are then optimized via a form of Generalized EM-algorithm [16], where in each iteration a single gradient ascent step is taken.

Results Figure 7 shows how ST-DGMRF and KS training and inference scale as the number of nodes N in the system increases. Clearly, the time per training iteration increases super-linearly when the KS is used to obtain exact marginal posterior distributions, while the variational ST-DGMRF training time increases only marginally and remains below the fastest KS iteration for all tested N . In addition, KS memory requirements (due to storing K dense N × 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 O ( 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

1.75

ST-DGMRF training

2.5

ST-DGMRFinference

200

KS training

KSinference

1.50

2.0

150

ations

1.25

1.00

1.5

iterat

E

e

time

e

100

ü

0.75

1.0

0.50

#

50

0.5

0.25

0.00

0.0

0

0

250

500

750

1000

1250

1500

0

250

500

750

1000

1250

1500

numberofnodesN

number of nodesN

total computation time for CG inference also scales linearly in N . In contrast, KS inference again scales super-linearly. This means that even if the KS approach would remain feasible in terms of memory requirements, its computation time will quickly approach, and eventually exceed, the time needed for CG inference.

C.2 Advection-diffusion experiments

Table 4 summarizes all results for the advection-diffusion dataset with mask size w = 9 , including standard deviations for all metrics based on 5 runs with different random seeds. As discussed in the main paper, the ST-DGMRF variants provide more accurate posterior estimates than the baselines relying on simplified spatiotemporal dependency structures.

Ablation results Table 4 contains additional results for the ST-DGMRF variants using different settings for the variational distribution (see Section B.1.2) For this dataset, we do not find a significant effect of introducing temporal dependencies in the variational distribution. Further, Figure 8 shows additional results for the ST-DGMRF variants when varying the number of temporal layers L temporal . For all metrics, the performance improves significantly as we start adding temporal layers and stabilizes around L temporal = 3 . Note that around the same point, both ST-DGMRF variants converge towards the EnKS using the true data-generating dynamics, in terms of the 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.e. L spatial) of the noise terms can further reduce this gap.

EnKS(estimated dynamics)

ST-DGMRF (neuralnetwork)

EnKS (true dynamics)

ST-DGMRF(advection-diffusion)

1.50

0.8

1.25

0.8

1.00

0.6

0.6

RM

S

0.75

RMS

MSEa

0.4

CRPS

0.4

0.50

0.25

0.2

0.2

0.00

0.0

0.0

0

1

2

3

4

5

6

0

1

2

3

4

5

6

0

1

2

3

4

5

6

Ltemporal

Ltemporal

Ltemporal

C.3 Air quality experiments

Table 5 summarizes all results for the air quality dataset. It contains additional results for the STDGMRF variants using different settings for the variational distribution (see Section B.2.2), and provides standard deviations for all metrics based on 5 runs with different random seeds.

Ablation results We find that, in contrast to our experiments on the advection-diffusion data, accounting for temporal dependencies in the variational distribution is clearly beneficial in the real world setting. Especially for the ST-DGMRF with neural network based transitions, adding the temporal diffusion layer results in significantly improved posterior estimates, and at the same time reduces the variability across different runs. Further, we find that at least two temporal layers are needed to achieve good posterior estimates that improve on the baselines (see Figure 9).

Example model outputs Figure 10 shows state estimates and associated uncertainties together with sensor measurements for two example sensors within the masked out area of the network, for ST-DGMRF, DGMRF and ARMA respectively. For all three models, state estimates are obtained by conditioning on the input data points (used for training), resulting in low errors and uncertainties for observed time points and higher errors and uncertainties for masked out time points. Moreover, for both ST-DGMRF and DGMRF, higher uncertainties coincide with larger errors and larger fluctuations in the measurements (top), while more accurate state estimates come with smaller uncertainties (bottom). Finally, Figure 11 visualizes how spatial and temporal ST-DGMRF layers transform samples from the estimated posterior over system states into (approximately) independent Gaussian noise, as derived in Section 3.1.2 and visualized in Figure 1. Clearly, temporal layers remove daily patterns and overall trends, while spatial layers remove dependencies between close-by sensors and increase temporal fluctuations.

0.425

0.75

ST-DGMRF transition

0.400

一

neuralnetwork

0.70

diffusion

0.375

RMS

E

15

0.65

S

P

0.350

0.325

0.60

0.300

0.55

0

1

2

3

4

0

1

2

3

4

Ltemporal

Ltemporal

input data

testdata

ST-DGMRF(neuralnetwork)

DGMRF

ARMA

4

log(PM2.5)

2

0

2

4

log(PM2.5)

2

0

2

0

20

40

60

timestepk

C.4 Total compute

Most computations were performed on a Nvidia Titan X GPU. On top of the final experiments, we performed hyperparameter sweeps and additional test runs. Here, we provide estimates of the total compute time grouped by experiment:

Advection-diffusion dataset:

Performance comparison & ablations: ca. 50 GPU hours

Varying mask size: ca. 30 GPU hours per w , resulting in ca. 210 GPU hours in total

Scalability: ca. 15 GPU hours

Air quality dataset:

Performance comparison & ablations: ca. 100 GPU hours

4

×

0

-2

2.0

0

50

100

150

200

250

300

350

400

1.5

1.0

2

0.5

×

0

0.0

-0.5

4

-1.0

0

50

100

150

200

250

300

350

400

-1.5

4

-2.0

(x)(

2

0

。

5

-2

0

50

100

150

200

250

300

350

400

time stepk

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nWith 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\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

comments: SUMMARY:
The paper extends Gaussian Random Fields with precision-based spatiotemporal structure.

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.

Originality: 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. 

Quality: 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. 

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. 

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.

QUESTIONS:
None

LIMITATIONS:
No issues
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-04

comments: 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.
They also show the various types of temporal structure can be built in this framework.

STRENGTHS:
The extension to a spatio-temporal models is of course very important as many datasets have both a spatial and temporal structure.
Also that the methods are computationally feasible is paramount as the complexity spatio-temporal models tend to be very large.
Each 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.
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.

QUESTIONS:
In Figure 2 to the right should w be from 6-10 in the it says 6-9,12?
Also how come mae is decreasing when w is increasing?
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?
You should also be able to get the marginal posterior distributions for your model how does they look for the simulations?

LIMITATIONS:
.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

comments: 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.

STRENGTHS:
- The method extends DGMRF to account for temporal dependency in data.
- It has been shown that the method has low computational complexity and good scalability.
- 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.
- 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). 
- Only the posterior mean was evaluated. It is unclear how good the posterior variance is.
- The variational posterior should be compared to the true posterior in an example where the latter is accessible analytically or via sampling.
- Only one synthetic example is not sufficient for evaluating the method thoroughly. It is unclear how it generalizes.

QUESTIONS:
- Were the observation parameters $H$ and $\xi$ trainable or fixed at the true values?
- 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-24