category: literaturenote citekey: yangstoatspatialtemporalprobabilistic2025 title: "STOAT: Spatial-Temporal Probabilistic Causal Inference Network" authors: "Yang, Yang; Yin, Du; Xue, Hao; Salim, Flora" year: 2025 date: 2025-00-00 2025 url: "https://arxiv.org/abs/2506.09544v2" zotero_key: MDKR9Z7F zotero_storage: X4F3D9NB collections: transpordi_planeerimine folder: 001_artiklid firstAuthor: "Yang, Yang"

status: converted

STOAT: Spatial-Temporal Probabilistic Causal Inference Network

Yang Yang University of New South Wales Sydney, NSW, Australia yang.yang26@unsw.edu.au

Hao Xue University of New South Wales Sydney, NSW, Australia hao.xue1@unsw.edu.au

Abstract

Spatial-temporal causal time series (STC-TS) involve region-specific temporal observations driven by causally relevant covariates and interconnected across geographic or network-based spaces. Existing methods often model spatial and temporal dynamics independently and overlook causality-driven probabilistic forecasting, limiting their predictive power. To address this, we propose STOAT (Spatial-Temporal Probabilistic Causal Inference Network), a novel framework for probabilistic forecasting in STC-TS. The proposed method extends a causal inference approach by incorporating a spatial relation matrix that encodes interregional dependencies (e.g. proximity or connectivity), enabling spatially informed causal effect estimation. The resulting latent series are processed by deep probabilistic models to estimate the parameters of the distributions, enabling calibrated uncertainty modeling. We further explore multiple output distributions (e.g., Gaussian, Student's-, Laplace) to capture region-specific variability. Experiments on COVID-19 data across six countries demonstrate that STOAT outperforms state-of-the-art probabilistic forecasting models (DeepAR, DeepVAR, Deep State Space Model, etc.) in key metrics, particularly in regions with strong spatial dependencies. By bridging causal inference and geospatial probabilistic forecasting, STOAT offers a generalizable framework for complex spatial-temporal tasks, such as epidemic management.

Keywords

Spatial-temporal modeling, Causal inference, Probabilistic forecasting, Uncertainty Analysis

ACM Reference Format:

Yang Yang, Du Yin, Hao Xue, and Flora Salim. 2025. STOAT: Spatial-Temporal Probabilistic Causal Inference Network. In Proceedings of Proceedings of the 33rd ACM SIGSPATIAL International Conference on Advances in Geographic Information Systems (SIGSPATIAL '25). ACM, New York, NY, USA, 12 pages.https://doi.org/XXXXXXX.XXXXXXX

Permission to make digital or hard copies of all or part of this work for personal or classroom use is granted without fee provided that copies are not made or distributed for profit or commercial advantage and that copies bear this notice and the full citation on the first page. Copyrights for components of this work owned by others than the author(s) must be honored. Abstracting with credit is permitted. To copy otherwise, or republish, to post on servers or to redistribute to lists, requires prior specific permission and/or a fee. Request permissions from permissions@acm.org.

SIGSPATIAL '25, Minneapolis, MN, USA

https://doi.org/XXXXXXX.XXXXXXX

Du Yin University of New South Wales Sydney, NSW, Australia du.yin@unsw.edu.au

Flora Salim University of New South Wales Sydney, NSW, Australia flora.salim@unsw.edu.au

1 Introduction

Spatial-temporal causal time series (STC-TS) consist of regional time series that are influenced by causally relevant covariates and interconnected through spatial dependencies, as shown in Figure 1. A representative case is epidemic forecasting across countries, where region-level outcomes such as case counts or hospitalizations are influenced by external factors including government interventions, vaccination rates, mobility patterns, population density, and virus evolution [8, 9, 12, 25]. These covariates are heterogeneous, time-varying, and exhibit complex causal impacts on the temporal trajectory of the target variables. For epidemic forecasting, the target variables such as confirmed or hospitalized cases form regionspecific time series, while external covariates evolve alongside, often asynchronously and irregularly. The complex causal interactions between such covariates (for instance, the joint effect of virus mutations and policy responses) remain difficult to model explicitly, particularly when considering their varying impacts across different regions [21]. Furthermore, regional outcomes are not isolated phenomena. Spatial dependencies arise naturally as outbreak patterns often exhibit spillover effects due to cross-border mobility, trade connections, and latent transmission pathways [32, 40]. Understanding spatio-temporal dynamics therefore requires modeling both the causal effects of covariates within each region and the spatial influences propagating between regions. These interconnected challenges in STC-TS settings give rise to three fundamental modeling requirements: cross-region causal inference, spatial-temporal causal representation learning, and causally informed probabilistic forecasting.

First Challenge: Cross-region causal inference. Factors and events across different regions (such as policy changes, mobility shifts, and vaccination progress) often have interdependent spillover effects on the outcome variables through spatial [24]. However, existing deep probabilistic forecasting models do not explicitly model how these cross-regional interactions and spatial dependencies jointly affect target dynamics across multiple regions. For example, DeepAR [34] and Deep State Space Models [28] treat covariates as auxiliary inputs without capturing spatial causal relationships between regions, lacking interpretability in understanding how interventions propagate across spatial networks. Classical Difference-in-Differences (DiD) [22, 29, 31] enables causal inference but is typically designed for simple treatment-control

setups without considering spatial spillover effects or the complex interdependencies inherent in spatial networks. In contrast to these approaches, our framework explicitly models spatial causal mechanisms through a spatial lag term that quantifies how neighboring regions' outcomes causally influence each target region, while simultaneously incorporating treatment effects within a unified causal inference structure. Bridging this gap by developing spatially-aware causal inference frameworks that can both capture cross-regional spillover effects and provide interpretable causal estimates for complex spatio-temporal time series represents a critical yet largely unaddressed challenge in the intersection of causal inference and deep learning.

Second Challenge: Spatial-Temporal Causal Representation Learning. Classical approaches in epidemiology, including compartmental models (e.g., SIR [11]) and statistical methods such as Difference-in-Differences [26, 38], Instrumental Variables (IV) [7], and Propensity Score Matching (PSM) [19], lack the capacity to capture high-dimensional, non-linear, and dynamic causal dependencies. These methods typically assume static confounding structures and fail to model complex spatio-temporal interactions in a data-driven manner. Meanwhile, existing deep temporal models often treat covariates as exogenous inputs without accounting for structural confounding or causal entanglement. Spatial dependencies, if considered at all, are rarely modeled through a causal lens. Consequently, both classical and modern approaches remain limited in their ability to infer structured causal dynamics over time and across regions. To address this, we argue that accurate forecasting in spatial-temporal causal time series (STC-TS) requires latent representations that encode both temporal dependencies and spatially heterogeneous causal effects.

Third Challenge: Causally Informed Probabilistic Forecasting. Causally informed probabilistic forecasting is critical in high-stakes scenarios like epidemic control, where uncertainty quantification is as important as point predictions [41]. However, existing methods often rely on deterministic predictions and fail to adaptively model the distributional dynamics induced by timevarying covariates or interventions, overlooking the causal relationships between covariates and targets [13]. As a result, predictive intervals are miscalibrated and insensitive to policy changes, regional disparities, and intervention timing. Furthermore, these models typically do not distinguish between the different sources of uncertainty, such as aleatoric, epistemic, and structural uncertainty, which limits both interpretability and the robustness of downstream decision-making.

To address the unique challenges of spatial-temporal causal time series (STC-TS) forecasting, we introduce STOAT, a unified framework that integrates causal inference, spatial-temporal representation learning, and probabilistic forecasting. STOAT is built on the premise that reliable forecasting under covariate interventions and spatial spillovers requires more than standard temporal modeling; it requires structured causal reasoning and calibrated uncertainty estimation. The key design principle behind STOAT is to first disentangle causal effects from spatio-temporal confounding through a dedicated Spatial-Temporal Causal Inference Module, which extends the classical Difference-in-Differences (DiD) [22, 29, 31] using a learnable spatial relation matrix. This module yields causally

adjusted representations that reflect both treatment effects and spatial spillovers across regions. These representations are then fed into a Deep Probabilistic Forecasting Module, which captures temporal dependencies using neural sequence models and projects the latent states into structured output distributions. This enables STOAT to produce calibrated probabilistic forecasts that account for both intervention dynamics and spatial correlations. Figure 2 provides an overview of the STOAT architecture, which consists of two interconnected modules: (1) a spatial-temporal causal inference pathway that generates covariate-adjusted representations, and (2) a probabilistic forecasting pathway that learns temporal dynamics and generates distributional predictions. The primary contributions of this work are summarized as follows:

Cross-region causal inference. We develop a spatiallyinformed causal inference mechanism that extends classical DiD using a spatial relation matrix, enabling STOAT to estimate region-specific treatment effects while accounting for spatial confounding and spillover influences. (Addressing Challenge 1)
Causal representation learning for STC-TS. By integrating causal adjustment into a neural encoder-decoder architecture, STOAT jointly models intra-regional temporal dynamics and inter-regional causal dependencies. This allows the model to learn high-dimensional spatial-temporal representations that go beyond static causal estimation. (Addressing Challenge 2)
Causally informed probabilistic forecasting. Building on causally adjusted representations, STOAT introduces structured probabilistic projection heads with support for Gaussian, Student's-, and Laplace distributions. This enables calibrated, interpretable, and region-aware uncertainty estimation that reflects intervention effects and covariatedriven dynamics. (Addressing Challenge 3)

2 Related Work

Here, we review the related literature on causal inference for time series, deep probabilistic forecasting, and sequence modeling in the context of epidemic prediction. We highlight how existing approaches fall short in modeling spatial-temporal causal dependencies and probabilistic uncertainty jointly.

Causal Inference in Time Series. Causal inference has been extensively studied in econometrics and policy analysis, especially through the framework of Difference-in-Differences (DiD) [1, 6, 23, 30]. DiD has been widely applied to time series data to estimate treatment effects over time by comparing treated and control groups. Recent extensions include generalized DiD methods for staggered treatment adoption [37] and synthetic control-based approaches [2]. However, these methods are typically restricted to one-dimensional or pairwise group comparisons, and do not account for spatial interactions across regions. They also assume static treatment effects and rarely integrate with representation learning. While some recent work incorporates DiD with neural networks [35, 43], they lack structured spatial modeling and do not support multi-region causal adjustment. Our work addresses

Figure 1: Illustration of the Spatial-temporal causal time series (STC-TS) setting. The bottom-left pink block (causal dependencies) depicts intra-regional causal inference, where time-varying covariates (e.g., policy interventions) influence observed outcomes within each region. Black-and-white nodes represent causally relevant covariates at specific time steps. The top gray block (temporal dependencies) shows the evolution of observed variables across three regions, with colored points (red, yellow, green) denoting distinct temporal trends. The bottom-right green block (spatial dependencies) illustrates inter-regional relationships, highlighting spatially structured influences and spillover effects. Together, the figure captures the intertwined spatial, temporal, and causal structures that define the STC-TS framework.

this gap by integrating spatial causal inference into a deep temporal forecasting framework using a structured adjacency matrix for regional spillover modeling.

Deep Probabilistic Time-Series Forecasting. Several deep learning models have been proposed for probabilistic time series forecasting by combining recurrent neural networks with statistical likelihood estimation. DeepAR [34] and MTSNet [42] model each time series independently under a univariate distribution, using shared RNN parameters to learn sequential dependencies. Deep State Space Models (DSSM) [28] incorporate a linear state-space formulation with recurrent updates, but similarly assume independent marginal distributions for each series. DeepVAR [33] enhances this with Gaussian copula processes to capture dependencies between target variables but does not model structured causal relationships across space or time. DeepFactor [39] decomposes global and local components using latent variables, yet still treats covariates as exogenous inputs without causal interpretation. None of these models incorporate causal inference mechanisms or address confounding in the presence of spatially distributed covariates. In contrast, STOAT explicitly adjusts for covariate-driven biases through spatial causal representations and generates calibrated predictive distributions informed by intervention dynamics.

Sequence Modeling for Epidemic Forecasting. Sequenceto-sequence (seq2seq) models and encoder-decoder architectures have been widely applied in epidemic forecasting [3, 27, 36]. These models, including GRU-based and transformer variants, are primarily trained for point prediction, optimizing loss functions such

as MSE or MAE. While effective in short-term forecasting, they do not quantify predictive uncertainty or incorporate causal reasoning. Classical epidemiological models, such as SIR and SEIR variants [9, 10, 17, 20], provide mechanistic insights into disease transmission, but they rely on fixed parametric assumptions and are difficult to scale or integrate with heterogeneous covariates. Moreover, these models often do not support flexible uncertainty modeling or data-driven representation learning. Our work bridges this gap by combining causal inference with probabilistic deep learning. STOAT allows for region-aware, covariate-adjusted forecasting with interpretable uncertainty, offering a unified approach to modeling spatial-temporal dynamics in complex real-world settings. Compared to existing approaches, STOAT differs in three key aspects:

(1) Spatial Causal Adjustment: STOAT explicitly incorporates causal inference into deep forecasting by extending Difference-in-Differences (DiD) with a spatial relation matrix, enabling causal adjustment across interconnected regions—something not supported by classical DiD or deep forecasting models.
(2) Causal Representation Learning: Unlike existing deep probabilistic models (e.g., DeepAR, DeepVAR), which treat covariates as exogenous and non-causal, STOAT constructs causally adjusted latent representations that reflect both spatial spillovers and intervention effects over time.
(3) Causally Informed Probabilistic Forecasting: While prior deep forecasting models often estimate point values or assume fixed likelihoods, STOAT performs distributional prediction by estimating the parameters of output distributions (e.g., Gaussian, Student's-, Laplace), conditioned on causally disentangled temporal dynamics.

3 STOAT: Spatial-Temporal Probabilistic Causal Inference Network

3.1 Overview of STOAT

To solve the STC-TS probabilistic forecasting problem, STOAT employs two key modules as shown in Figure 2. The Spatial-Temporal Causal Inference Module first performs causallyinformed representation learning by extending the Difference-in-Differences framework with spatial relation matrices, enabling region-aware causal adjustment that disentangles treatment effects from confounding factors across interconnected spatial units. The causally adjusted latent representations are then fed into the Deep Probabilistic Projection Module, which employs neural encoder-decoder architectures to capture complex temporal dependencies and generates distributional forecasts by estimating parameters of flexible output distributions (Gaussian, Student's-, Laplace), thereby providing calibrated uncertainty quantification that reflects both spatial spillover effects and intervention-driven dynamics in the forecasting process.

Spatial-temporal causal time series (STC-TS) consisting of target variables time-series y and multiple causally relevant covariates time-series C exhibit both inter-regional spatial dependencies and intra-regional temporal couplings, as shown in Figure 1. For the considered epidemic forecasting data, we assume y represents a target vector composed of the values of the multivariate time series at time step , e.g., confirmed cases from regions/countries. C represents the causally relevant covariates corresponding to the target multivariate time series, e.g., policy stringency index [15, 16], vaccination rates, ICU patient and reproduction number [5] from regions/countries, where these covariates serve as causal factors (treatments) that influence the target outcomes. STOAT characterizes the STC-TS probabilistic forecasting problem in Eq. (1),

$$P(\mathbf{y}_{t_0+1:t0+m} \mid \mathbf{y}{1:t0}, \mathbf{C}{1:t_0+m}, \mathbf{S}, \mathbf{\Theta}),$$ (1)

where Θ denotes the parameters of STOAT, S is the spatial relation matrix capturing inter-regional spillover effects and confounding relationships, the historical time series steps are {1, 2, . . . , 0}, and the prediction time length is ∈ N. STOAT models the conditional joint multivariate probability distributions of the future multivariate time series y0+1:0+, given the historical multivariate time series y1:⁰ , their causally relevant covariates C1:0+, and the spatial relation matrix S, which is incorporated within the model architecture to account for inter-regional dependencies during the causal adjustment and forecasting processes.

3.2 Spatial-Temporal Causal Inference Module

The Spatial-Temporal Causal Inference Module performs causallyinformed representation learning by extending the Difference-in-Differences [23] framework with spatial relation matrices, enabling region-aware causal adjustment representation learning.

3.2.1 Spatial Relation Matrix Construction. Given the geographic coordinates (lat , lon) for each region , we construct a spatial relation matrix S ∈ R ^× to capture the spatial dependencies between regions. We first compute the geodesic distance , between each pair of regions and , then define the spatial relation matrix as:

$$\mathbf{S} = \begin{bmatrix} 0 & S{1,2} & \cdots & S{1,N} \ S{2,1} & 0 & \cdots & S{2,N} \ \vdots & \vdots & \ddots & \vdots \ S{N,1} & S{N,2} & \cdots & 0 \end{bmatrix}, \quad S{i,j} = \begin{cases} \frac{1}{d{i,j}^{H}} & \text{if } i \neq j \ 0 & \text{if } i = j \end{cases}$$ (2)

The parameter > 0 controls the spatial decay rate, with higher values of leading to more localized spatial effects. The rownormalization ensures that ^Í =1 , = 1 for all , making S a row-stochastic matrix that preserves the interpretation of spatial weights as proportional influences.

While we focus on distance-based relationships in this formulation, the spatial relation matrix S can be constructed using alternative proximity measures such as economic connectivity, transportation networks, administrative boundaries, or hybrid approaches that combine multiple spatial relationship indicators.

3.2.2 Spatial Causal Inference. The spatial causal inference mechanism in STOAT extends the classical Difference-in-Differences (DiD) framework by incorporating spatial dependencies through the spatial relation matrix S constructed in Section 3.2.1. This spatiallyinformed approach enables region-aware causal adjustment that accounts for both direct treatment effects and spatial spillover influences across interconnected regions.

Spatially-Enhanced DiD Formulation. To capture the complex interplay between spatial dependencies and causal effects in

epidemic dynamics, we formulate the core spatial causal inference model as:

$$y{i,t} = \rho \sum{j=1}^{N} S{i,j} y{j,t} + \beta_0 + \beta_1 T_i + \beta_2 \text{Post}_t + \delta(T_i \cdot \text{Post}t) + \mathbf{y}^{\top} \mathbf{c}{i,t} + \epsilon_{i,t}$$ (3)

where , represents the observed outcome (e.g., COVID-19 cases) for region at time . The spatial lag term Í =1 ,, captures the weighted average of neighboring regions' outcomes based on the spatial relation matrix, with coefficient controlling the strength of spatial spillover effects. The treatment indicator is a binary variable denoting whether region has adopted the intervention, while Post indicates the post-treatment period. The parameter represents the treatment effect, capturing the causal impact of interventions on target outcomes.

3.2.3 Multi-dimensional Covariates Construction. Building upon the spatially-enhanced DiD framework in Eq. (3), we incorporate epidemiological covariates through a structured multi-dimensional vector. The covariate vector c, encompasses four key epidemiological dimensions:

$$\mathbf{c}{i,t} = \begin{bmatrix} R{i,t} \ M{i,t} \ V{i,t} \ I_{i,t} \end{bmatrix}, \quad \boldsymbol{\gamma} = \begin{bmatrix} \gamma_1 \ \gamma_2 \ \gamma_3 \ \gamma_4 \end{bmatrix}$$ (4)

where , denotes the Reproduction Number capturing transmission dynamics, , represents the Mitigation Stringency Index measuring policy intervention intensity, , indicates the Vaccination Share representing immunization coverage, and , represents the ICU Patient Number indicating healthcare system strain. The corresponding parameter vector captures the marginal effects of each covariate dimension.

The covariate contribution in Eq. (3) is explicitly modeled as:

$$\gamma^{\top} \mathbf{c}_{i,t} = \gamma1 R{i,t} + \gamma2 M{i,t} + \gamma3 V{i,t} + \gamma4 I{i,t}$$ (5)

This formulation allows for the simultaneous estimation of individual covariate effects while maintaining the spatial-temporal structure inherent in the epidemic forecasting problem.

3.2.4 Parameter Estimation and Identification. The estimation of parameters in Eq. (3) follows a two-stage approach. First, we estimate the spatial autoregressive coefficient using instrumental variable techniques to address endogeneity in the spatial lag term. Subsequently, we apply ordinary least squares (OLS) to estimate the remaining parameters {0, 1, 2, ,}, given the estimated ˆ. The identification of relies on the parallel trends assumption, which requires that treated and control regions would have followed similar trajectories in the absence of treatment.
3.2.5 Causal Adjustment and Target Reconstruction. Following the parameter estimation from Eq. (3), STOAT performs causal adjustment by constructing counterfactual outcomes that isolate the underlying temporal dynamics from intervention effects. The causally adjusted target variable is obtained as:

$$\tilde{y}{i,t} = y{i,t} - \hat{\delta} \cdot (T_i \cdot \text{Post}_t)$$ (6)

where ˆ represents the estimated treatment effect from Eq. (3). This adjustment yields causally corrected target variables ˜, that

Figure 2: Overview of the STOAT architecture, which consists of two interconnected modules: a spatial-temporal causal inference pathway that generates covariate-adjusted representations, and a probabilistic forecasting pathway that learns temporal dynamics and generates distributional predictions. In what follows, we detail each component of the STOAT framework.

represent the counterfactual outcomes in the absence of treatment effects, thereby isolating the underlying epidemiological dynamics from policy-driven changes.

To construct the input tensor for the subsequent deep probabilistic forecasting module, we combine the causally adjusted targets with the spatial lag effects:

$$\mathbf{z}{i,t} = \tilde{y}{i,t} + \hat{\rho} \sum{j=1}^{N} S{i,j} \tilde{y}_{j,t}$$ (7)

The reconstructed variable ˜, and spatially adjusted input z, explicitly incorporate information from the spatial relation matrix S through spatial lag effects and the causally relevant covariates c through causal adjustments, effectively capturing the influence of spatial dependencies and covariate-driven interventions.

Thus, for subsequent deep probabilistic forecasting, we leverage z, directly, encapsulating the information of spatial dependencies and covariate-driven causal effects. Therefore, the conditional probability formulation can be simplified from its original form:

$$P(\mathbf{y}_{t_0+1:t0+m} \mid \mathbf{y}{1:t0}, \mathbf{C}{1:t_0+m}, \mathbf{S}, \boldsymbol{\Theta})$$ (8)

to the simplified form:

$$P\left(\mathbf{y}_{t_0+1:t0+m} \mid \mathbf{y}{1:t0}, \mathbf{z}{i,1:t_0}, \mathbf{\Theta}\right)$$ (9)

where represents the parameters of the deep probabilistic neural network shared across all regions.

3.2.6 Interpretability and Causal Parameters. The spatial causal inference mechanism provides three key interpretable parameters that offer insights into the underlying epidemic dynamics:

The treatment effect parameter quantifies the average causal impact of interventions on target outcomes across regions. A significant negative (e.g., < −0.1) indicates that policy interventions

effectively reduce case counts, while ≈ 0 suggests limited intervention effectiveness. The covariate effect vector from Eq. (5) captures the marginal contribution of each epidemiological factor. For instance, ¹ > 0 confirms that higher reproduction numbers drive increased transmission, while ² < 0 validates that stricter policies reduce case growth rates. The magnitude of |³ | reflects the relative importance of vaccination coverage in epidemic control. The spatial spillover parameter measures inter-regional dependency strength. Values of > 0.3 indicate strong positive spatial correlation, suggesting that neighboring regions' epidemic patterns significantly influence local outcomes through cross-border mobility or policy coordination effects.

3.2.7 Integration with Deep Probabilistic Forecasting. The causally adjusted spatial representations z, from Eq. (7) serve as inputs to the deep probabilistic projection module described in Section 3.3. This integration ensures that subsequent probabilistic forecasting operates on causally disentangled temporal patterns while preserving essential spatial dependencies. The framework enables STOAT to generate predictions that reflect both underlying epidemiological dynamics and expected intervention effects, providing a principled foundation for uncertainty quantification in policy-sensitive forecasting scenarios.

3.3 Deep Probabilistic Forecasting Module

The deep probabilistic forecasting module in STOAT leverages the causally adjusted spatial-temporal representations z, derived from the spatial causal inference module (Eq. (7)). These representations encapsulate both intra-regional temporal dynamics and inter-regional spatial dependencies, effectively integrating the effects of the causally relevant covariates c, and the spatial relation

matrix S. By using z,, the forecasting problem is simplified to model the conditional joint probability distribution of future multivariate time series y0+1:0+, given the historical target series y1:⁰ and the causally adjusted representations z,1:⁰ . The probabilistic forecasting problem is formulated as:

$$P\left(\mathbf{y}_{t_0+1:t0+m} \mid \mathbf{y}{1:t0}, \mathbf{z}{i,1:t_0}, \mathbf{\Theta}\right),$$ (10)

where Θ denotes the learnable parameters of STOAT's neural architecture, the historical time steps span {1, 2, . . . , 0}, and the prediction horizon is ∈ N. The use of z,1:⁰ as a conditioning variable simplifies the model by embedding the spatial and causal effects into a unified representation, enabling robust forecasting that accounts for causal interventions and spillover effects.

Assuming the time series points are conditionally independent given the latent representations, the joint distribution in Eq. (10) can be factorized as:

$$P\left(\mathbf{y}{t{0}+1:t{0}+m} \mid \mathbf{y}{1:t{0}}, \mathbf{z}{i,1:t{0}}, \mathbf{\Theta}\right) = \prod{t=t{0}}^{t{0}+m-1} P\left(\mathbf{y}{t+1} \mid \mathbf{y}{1:t}, \mathbf{z}_{i,1:t}, \mathbf{\Theta}\right).$$ (11)

To model this conditional distribution, STOAT employs a neural encoder-decoder architecture that processes the causally adjusted representations z,. Specifically, the latent representation at time for region , denoted h,, is generated by a neural transformation function ℎ(·) that captures temporal dependencies:

$$\mathbf{h}{i,t} = h\left(\mathbf{z}{i,1:t}, \mathbf{y}{1:t}, \mathbf{h}{i,t-1}, \mathbf{\Theta}_h\right),$$ (12)

where z,1: is the sequence of causally adjusted representations for region up to time , y1: includes the historical target series, h,−1 represents the hidden state from the previous time step, and Θℎ denotes the parameters of the neural transformation. The function ℎ(·) is implemented using a stacked recurrent neural network (e.g., LSTM or GRU) to capture long-term temporal dependencies, leveraging the embedded spatial and causal information in z,1:

The conditional probability in Eq. (11) is modeled as a product of likelihood functions, where the parameters of the output distribution are determined by a projection layer (·). For each region and time + 1, the likelihood is defined as:

$$\prod_{l=t_0}^{t0+m-1} P\left(\mathbf{y}{t+1} \mid \mathbf{y}{1:t}, \mathbf{z}{i,1:t}, \mathbf{\Theta}\right) = \prod_{l=t_0}^{t0+m-1} \ell\left(y{i,t+1} \mid k\left(\mathbf{h}_{i,t}, \mathbf{\Theta}_k\right)\right),$$ (13)

where ℓ(·) is the likelihood function, and (·) maps the latent representation h, to the parameters of the output distribution, with learnable parameters Θ . The global parameter set Θ comprises Θℎ and Θ

3.3.1 Probabilistic Projection. To capture the diverse statistical properties of STC-TS data (e.g., epidemic case counts), STOAT supports multiple output distributions, including Gaussian, Laplace, and Student's-, to model region-specific variability and uncertainty. The projection layer (·) maps the latent representation h, to the parameters of these distributions. The likelihood functions and corresponding projection mappings are defined as follows:

Laplace distribution likelihood function:

$$\ellL(y{i,t} \mid \mu, \sigma) = \frac{1}{2\sigma} \exp\left(-\frac{|y_{i,t} - \mu|}{\sigma}\right),$$ (14)

Gaussian distribution likelihood function:

$$\ellG(y{i,t} \mid \mu, \sigma) = \frac{1}{\sqrt{2\pi\sigma^2}} \exp\left(-\frac{(y_{i,t} - \mu)^2}{2\sigma^2}\right),\tag{15}$$

Student's- distribution likelihood function:

$$\ellS(y{i,t} \mid \nu, \mu, \sigma) = \frac{\Gamma\left(\frac{\nu+1}{2}\right)}{\Gamma\left(\frac{\nu}{2}\right)\sqrt{\pi\nu\sigma}} \left(1 + \frac{1}{\nu}\left(\frac{y_{i,t} - \mu}{\sigma}\right)^2\right)^{-\frac{\nu+1}{2}},$$ (16)

3.3.2 Likelihood Loss. To optimize the parameters Θ = {Θℎ, Θ }, STOAT employs a negative log-likelihood loss function, defined as:

$$\mathcal{L} = -\sum{i=1}^{N} \sum{t=t_0}^{t0+m-1} \log \ell \left( y{i,t+1} \mid k \left( \mathbf{h}_{i,t}, \Theta_k \right) \right). \tag{17}$$

This loss is minimized using stochastic gradient descent, enabling the model to learn parameters that maximize the likelihood of the observed data under the specified distribution.

During training, both the conditioning range {1, 2, . . . , 0} and the prediction range {⁰ +1, . . . , ⁰ +} are available. For prediction, only the historical data , ( ≤ 0) and causally adjusted representations z,1:⁰ are used as inputs. The forecasted values ˜, ( > 0) are sampled from the estimated distribution:

$$\tilde{\mathbf{y}}_{i,t_0+1:t0+m} \sim P\left(\mathbf{y}{i,t_0+1:t0+m} \mid \mathbf{y}{i,1:t0}, \mathbf{z}{i,1:t_0}, \boldsymbol{\Theta}\right).$$ (18)

These samples are generated via ancestral sampling, where each ˜,+¹ is drawn from the distribution parameterized by (h,, Θ ) and used as input for the next time step until = 0+. This process enables STOAT to produce probabilistic forecasts with calibrated uncertainty, leveraging the causally adjusted representations z, to reflect both spatial dependencies and intervention effects.

The overall algorithm is shown in Appendix A.1

4 Experimental Evaluation

To assess the performance of the proposed STOAT framework on spatial-temporal causal time series (STC-TS) data, we conduct experiments using a comprehensive COVID-19 dataset. We compare STOAT against state-of-the-art probabilistic time-series forecasting methods, focusing on its ability to model inter-regional spatial dependencies and causal effects.

4.1 Experimental Setup

4.1.1 COVID-19 Dataset with Spatial and Causal Factors. The experiments utilize a COVID-19 dataset comprising target time series and causally relevant factors from six countries across different continents: UK, US, Canada, France, Spain, and Italy. These countries were selected due to their diverse epidemic trajectories, varied nonpharmaceutical intervention (NPI) policies, and geographic distribution across multiple continents, enabling evaluation of STOAT's generalization across different spatial and temporal contexts.
- Target time series: The dataset includes daily confirmed COVID-19 cases, sourced from the Johns Hopkins University COVID-19 data repository. Data spans from June 1, 2020, to June 30, 2022, covering a period of significant epidemic variability.
Causal Factors: The causal factors include ICU admissions, vaccination booster rates, Stringency Index [16], and daily reproduction number, representing virus mutation. These factors constitute the covariate matrix C. Omicron proportion data is sourced from the GISAID initiative1 , vaccination data from Our World in Data2 , Stringency Index from the Oxford Covid-19 Government Response Tracker, and reproduction rates from a public repository [1]3 .
Spatial dependencies: Inter-regional couplings are modeled using the spatial relation matrix S, constructed from geodesic distances between countries' geographic coordinates (latitude and longitude).

The dataset from June 1, 2020 is used for training, with the goal of predicting daily confirmed cases and ICU admissions over forecast horizons of 5, 7, and 10 days. Due to the unique combination of spatial dependencies and multiple causal factors in this dataset, no other comparable time-series datasets are currently available for evaluation.

4.1.2 Baseline Models. We compare STOAT with leading probabilistic time-series forecasting models implemented in the GluonTS toolkit [4]: DeepAR, Deep State Space Model (DSSM), DeepFactor (DF), MTSNet [42], and DeepVAR. Non-probabilistic methods (e.g., ARIMA, LSTM, etc.) are excluded, as they do not produce probabilistic forecasts suitable for comparison. As STOAT focuses on probabilistic forecasting to evaluate distribution accuracy, we do not compare with point-based spatio-temporal forecasting models.
- DeepAR [34]: Utilizes a two-layer RNN to extract temporal features, mapping hidden states to distribution parameters for probabilistic predictions.
- DSSM [28]: Combines a linear state-space model with RNNs, incorporating Gaussian noise for tractable likelihood-based training.
- DF [39]: Employs RNNs for global time-series representations, using Kalman filtering to generate Gaussian-based probabilistic forecasts.
- MTSNet [42]: Leverages global learning framework with RNNs to model multiple time series, outputting parameters for probabilistic distributions.
- DeepVAR [33]: Extends vector autoregression with deep learning, using RNNs to model multivariate dependencies with Copula function and produce probabilistic forecasts.

To evaluate the contributions of STOAT's spatial causal inference and causal factor integration, we test the following variants:

STOAT-Laplace: Uses Laplace distribution with full spatial causal adjustment (via S) and causal factors C.
STOAT-Laplace-NoSpatial: Uses Laplace distribution without spatial causal adjustment (i.e., bypassing S).
STOAT-Laplace-NoFactors: Uses Laplace distribution without causal factors C.
STOAT-Gaussian: Uses Gaussian distribution with full spatial causal adjustment and causal factors.

Figure 3: Visualization of 10-day probabilistic forecasts for COVID-19 case data during peak periods in Canada (left) and Italy (right) using the STOAT model. The yellow, green, and red lines represent the predictive distributions based on Student's-, Gaussian, and Laplace distributions, respectively.

STOAT-Gaussian-NoSpatial: Uses Gaussian distribution without spatial causal adjustment.
STOAT-Gaussian-NoFactors: Uses Gaussian distribution without causal factors.
STOAT-StudentT: Uses Student's- distribution with full spatial causal adjustment and causal factors.
STOAT-StudentT-NoSpatial: Uses Student's- distribution without spatial causal adjustment.
STOAT-StudentT-NoFactors: Uses Student's- distribution without causal factors.
4.1.3 Evaluation Metrics. We evaluate STOAT and baseline models using metrics designed to assess probabilistic forecasting performance, focusing on the accuracy of predicted distributions. Specifically, we employ the Continuous Ranked Probability Score (CRPS), the Weighted Quantile Loss (WQL), and the Energy Score to quantify how well each model's predictive distribution aligns with observed outcomes and the details are shown in Appendix A.2.
4.1.4 Settings. All experiments are conducted on same GPU to ensure fair comparison. We perform forecasting experiments with prediction horizons of 3, 5, and 7 days, using five different random seeds to ensure robustness. The context window length ⁰ and prediction horizon are set in a 5:1 ratio (e.g., 15 days for training and 3 days for prediction, 25 days for training and 5 days for prediction, or 35 days for training and 7 days for prediction). Hyperparameters are optimized via a back-test strategy before forecasting. For each prediction task, 100 samples are drawn from the probabilistic decoder to generate forecasts. These settings are consistently applied to STOAT and all baseline models. An ablation study is performed to evaluate the contributions of spatial causal inference (via S and z) and causal factors (C) to forecasting accuracy.

¹Available at covariants.org.

²Available at github.com/owid/covid-19-data.

³Available at github.com/crondonm/TrackingR.

Time Period	Metrics	DeepVAR	DeepAR	DSSM	DeepFactor	MTSNet	STOAT StudentT	STOAT Gaussian	STOAT Laplace
5 Days	Quantile Loss [0.1] Quantile Loss [0.5]	0.1894 0.3742	0.1723 0.3421	0.2015 0.3927	0.1976 0.3814	0.1845 0.3598	0.1543 0.3476	0.1421 0.3378	0.1478 0.3312
	Quantile Loss [0.9]	0.2653	0.2512	0.2789	0.2478	0.2634	0.2356	0.2412	0.2389
	Coverage [0.1]	0.2000	0.3000	0.4000	0.1000	0.2000	0.1000	0.1000	0.1000
	Coverage [0.5]	0.4000	0.6000	0.3000	0.7000	0.5000	0.4000	0.5000	0.5000
	Coverage [0.9]	0.8000	0.9000	0.6000	0.8000	0.7000	0.9000	0.8000	0.9000
	CRPS	0.3298	0.2954	0.3512	0.3176	0.3289	0.2897	0.2745	0.2798
	Energy Score	1.2789	1.1432	1.4123	1.3214	1.2987	1.1012	1.0789	1.0876
7 Days	Quantile Loss [0.1]	0.2187	0.2056	0.2312	0.2234	0.2112	0.1898	0.1789	0.1845
	Quantile Loss [0.5]	0.5245	0.4876	0.5467	0.5321	0.4987	0.4723	0.4612	0.4534
	Quantile Loss [0.9]	0.2812	0.2678	0.2934	0.2712	0.2545	0.2612	0.2589	0.2567
	Coverage [0.1]	0.3000	0.2000	0.4000	0.1000	0.2000	0.1000	0.1000	0.1000
	Coverage [0.5]	0.4000	0.6000	0.3000	0.7000	0.5000	0.4000	0.5000	0.5000
	Coverage [0.9]	0.7000	0.9000	0.6000	0.8000	0.8000	0.9000	0.8000	0.9000
	CRPS	0.3623	0.3478	0.3845	0.3712	0.3543	0.3176	0.3056	0.3123
	Energy Score	1.4678	1.3765	1.5345	1.4567	1.4321	1.2234	1.1876	1.1987
10 Days	Quantile Loss [0.1]	0.2987	0.2765	0.3123	0.2876	0.2934	0.2567	0.2456	0.2512
	Quantile Loss [0.5]	0.7234	0.6987	0.7456	0.7321	0.7123	0.6654	0.6456	0.6234
	Quantile Loss [0.9]	0.3789	0.3567	0.3923	0.3678	0.3545	0.3321	0.3289	0.3212
	Coverage [0.1]	0.3000	0.2000	0.4000	0.0000	0.2000	0.1000	0.1000	0.1000
	Coverage [0.5]	0.3000	0.7000	0.2000	0.6000	0.5000	0.4000	0.4000	0.4000
	Coverage [0.9]	0.6000	0.7000	0.5000	0.7000	0.8000	0.9000	0.8000	0.9000
	CRPS	0.4876	0.4654	0.5123	0.4987	0.4765	0.4321	0.4123	0.4234
	Energy Score	1.8765	1.7898	1.9456	1.8765	1.8321	1.5432	1.4987	1.5123

Table 1: Comparison of probabilistic forecasting performance across models for 5-day, 7-day, and 10-day horizons on COVID-19 case data, averaged over six countries. Metrics include Quantile Loss at 0.1, 0.5, and 0.9 quantiles, Coverage at corresponding quantiles, Continuous Ranked Probability Score (CRPS), and Energy Score. The proposed STOAT variants (Student's-, Gaussian, Laplace) are compared against baseline models (DeepVAR, DeepAR, DSSM, DeepFactor, MTSNet). Bold values indicate the best performance for each metric and time period.

4.2 Comparison Results

This study evaluates probabilistic forecasting performance during the peak phase of the COVID-19 pandemic, focusing on scenarios with significant external causal interventions, using case data from six countries: Canada, France, Italy, UK, US, and Spain. We compare the proposed STOAT model variants (STOAT-StudentT, STOAT-Gaussian, STOAT-Laplace) against baseline models, including DeepVAR, DeepAR, DSSM, DeepFactor, and MTSNet, over 5 day, 7-day, and 10-day horizons, with results averaged across these countries. The experimental outcomes, presented in Table1, assess performance using Quantile Loss (QL) at 0.1, 0.5, and 0.9 quantiles, Coverage at corresponding quantiles, Continuous Ranked Probability Score (CRPS), and Energy Score. These metrics evaluate distinct aspects of probabilistic forecasts: QL measures accuracy at specific quantiles (lower is better), Coverage assesses calibration by the proportion of true observations within prediction intervals (ideal when matching nominal levels, e.g., 0.9 for 90% intervals), CRPS quantifies overall distributional accuracy (lower is better), and Energy Score evaluates multivariate forecast quality (lower is better). The results demonstrate STOAT's superior performance during peak pandemic periods with strong causal interventions, driven by its spatial causal inference and causal factor integration, though its advantages are less pronounced in scenarios with weak causal signals.

The STOAT variants consistently outperform baselines in key metrics, particularly during peak periods. STOAT-Gaussian achieves the lowest CRPS (0.2745, 0.3056, 0.4123 for 5, 7, and 10 days, respectively) and Energy Score (1.0789, 1.1876, 1.4987), indicating

superior probabilistic accuracy and multivariate forecast quality. It also excels in QL 0.1, reflecting precise predictions in the lower tail. STOAT-Laplace leads in QL [0.5] and [0.9] for the 10-day horizon (0.6234 and 0.3212), leveraging its robustness to outliers, while STOAT-StudentT achieves the best QL [0.9] for 5 days (0.2356) and matches the optimal Coverage 0.9 across multiple horizons, benefiting from its heavy-tailed distribution to capture extreme case surges. For Coverage, STOAT variants are well-calibrated, achieving ideal values (e.g., 0.1000 for Coverage [0.1], 0.9000 for Coverage [0.9]), unlike DSSM, which shows under-coverage (e.g., 0.6000 for Coverage [0.9] at 5 days), likely due to its Gaussian assumptions failing to model extreme variations driven by interventions like lockdowns.

However, STOAT's performance is not uniformly superior. DeepAR achieves the lowest QL [0.5] for 5 days (0.3421), and MTSNet matches the optimal Coverage 0.5 for 5 and 7 days, indicating competitive performance in central quantiles. These results highlight the sensitivity of probabilistic forecasting to distributional assumptions and data characteristics across the six countries during peak phases. STOAT's causal modules (Z, C) enhance performance by capturing covariates such as mobility restrictions and public health interventions, which are critical during peak periods in countries like the US and Italy. In contrast, baselines like DSSM and DeepFactor, with simpler RNN or state-space architectures, struggle with spatial dependencies, leading to higher QL (e.g., 0.2015 and 0.1976 for QL [0.1] at 5 days) and CRPS (0.3512 and 0.3176). In scenarios with weak causal signals, such as non-peak periods in stable regions like Canada, STOAT variants perform

comparably to DeepAR and MTSNet, suggesting that the causal modules' benefits are context-dependent.

The flexibility of STOAT's architecture, incorporating Student's- , Gaussian, and Laplace distributions, allows it to adapt to heterogeneous time-series patterns across the six countries. For instance, STOAT-Gaussian's low CRPS reflects its balance of calibration and sharpness, while STOAT-Laplace's QL performance at higher quantiles suits extreme values. Figure3 illustrates STOAT-Laplace's 10 day forecast for Canada and Italy, showing median predictions and 50% and 90% confidence intervals, effectively capturing peak case trends influenced by interventions. However, challenges remain in modeling small, highly variable datasets with cross-country inconsistencies, where no single distribution consistently outperforms across all metrics and horizons. Overall, STOAT's causal integration enhances its robustness for peak pandemic forecasting under significant external interventions, positioning it as a competitive approach in probabilistic time-series modeling.

Figure 4: Comparison of probabilistic forecasting models across three forecast horizons (5-day, 7-day, and 10-day) on COVID-19 case data averaged over six countries. The performance is evaluated using the main metrics including CRPS and Energy Score. STOAT variants (Student's-, Gaussian, and Laplace) are compared against baseline models (DeepVAR, DeepAR, DSSM, DeepFactor, MTSNet).

4.3 Ablation Study

To evaluate the contributions of spatial causal inference and causal factor integration within STOAT, we perform ablation experiments on COVID-19 case data, averaged over peak periods in six countries. Across three probabilistic output distributions (Laplace, Gaussian, Student's- distribution), we compare the full model against ablated variants that remove either the spatial relation matrix or the set of causal covariates.

As shown in Table 2, the complete STOAT model consistently achieves the best overall performance in terms of CRPS and Energy Score across all forecast horizons. For instance, STOAT-Gaussian achieves a CRPS of 0.2745 and an Energy Score of 1.0789 on the 5-day horizon, outperforming both -NoSpatial and -NoFactors variants. However, we observe that in some cases, ablated models yield lower errors on individual metrics. For example, STOAT-StudentT-NoSpatial attains a lower Quantile Loss 0.1 than the full Student's- variant (0.1543), and STOAT-Gaussian-NoFactors achieves a lower Quantile Loss 0.1 compared to the full Gaussian model (0.1421) on the same horizon. This phenomenon is most pronounced on short-term forecasts, where reduced model

Figure 5: Performance Comparison of STOAT-Gaussian Model on 5-day, 7-day, and 10-day Forecast Horizons.The radar plot demonstrates that the STOAT-Gaussian model performs better on short-term forecasts, with lower Quantile Loss, CRPS, and Energy Score for the 5-day horizon compared to longer horizons, reflecting the expected increase in forecasting uncertainty over time.

Figure 6: Absolute performance of the STOAT model and its ablation variants across three distribution choices (Student's- , Gaussian, Laplace) and three forecast horizons (5, 7, 10 days). For each distribution, the full model ("Full") is compared with its "NoSpatial" and "NoFactors" ablation variants, reporting both CRPS (bars, left axis) and Energy Score (lines, right axis). The results show that removing either spatial causal adjustment or causal factor integration generally leads to degraded predictive performance, confirming the contribution of both components to overall model accuracy.

complexity or incidental alignment with test-set fluctuations can occasionally result in marginally better pointwise quantile performance.

Nevertheless, as the prediction horizon increases, the benefits of joint spatial and causal modeling become more apparent. For example, on the 10-day forecast, the full STOAT-Gaussian model yields a CRPS of 0.4123, compared to 0.4204 and 0.4229 for the

Time Period	Metrics	STOAT StudentT	STOAT StudentT NoSpatial	STOAT StudentT NoFactors	STOAT Gaussian	STOAT Gaussian NoSpatial	STOAT Gaussian NoFactors	STOAT Laplace	STOAT Laplace NoSpatial	STOAT Laplace NoFactors
5 Days	Quantile Loss [0.1]	0.1543	0.1512	0.1585	0.1421	0.1439	0.1400	0.1478	0.1490	0.1485
	Quantile Loss [0.5]	0.3476	0.3487	0.3325	0.3378	0.3310	0.3383	0.3312	0.3325	0.3334
	Quantile Loss [0.9]	0.2356	0.2370	0.2374	0.2412	0.2428	0.2430	0.2389	0.2340	0.2416
	Coverage [0.1]	0.1000	0.1000	0.1000	0.1000	0.1000	0.1000	0.1000	0.1000	0.1000
	Coverage [0.5]	0.4000	0.4000	0.4000	0.5000	0.5000	0.5000	0.5000	0.5000	0.5000
	Coverage [0.9]	0.9000	0.9000	0.8000	0.8000	0.8000	0.7000	0.9100	0.9000	0.8000
	CRPS	0.2897	0.2920	0.2918	0.2745	0.2770	0.2730	0.2798	0.2807	0.2812
	Energy Score	1.1012	1.1035	1.0970	1.0789	1.0801	1.0812	1.0876	1.0871	1.0878
7 Days	Quantile Loss [0.1]	0.1898	0.1932	0.1880	0.1789	0.1812	0.1831	0.1845	0.1825	0.1886
	Quantile Loss [0.5]	0.4723	0.4690	0.4758	0.4612	0.4628	0.4609	0.4534	0.4557	0.4565
	Quantile Loss [0.9]	0.2612	0.2630	0.2638	0.2589	0.2599	0.2602	0.2567	0.2563	0.2585
	Coverage [0.1]	0.1000	0.1000	0.1000	0.1000	0.1000	0.1200	0.1000	0.1000	0.1000
	Coverage [0.5]	0.4000	0.4000	0.4000	0.5000	0.5000	0.4000	0.5000	0.4000	0.4000
	Coverage [0.9]	0.9000	0.8000	0.8000	0.8000	0.8000	0.7000	0.9000	0.9100	0.8000
	CRPS	0.3176	0.3210	0.3225	0.3056	0.3079	0.3087	0.3123	0.3139	0.3146
	Energy Score	1.2234	1.2298	1.2112	1.1876	1.1922	1.1935	1.1987	1.1962	1.2029
10 Days	Quantile Loss [0.1]	0.2567	0.2685	0.2708	0.2456	0.2517	0.2532	0.2512	0.2605	0.2450
	Quantile Loss [0.5]	0.6654	0.6752	0.6385	0.6456	0.6564	0.6589	0.6234	0.6388	0.6410
	Quantile Loss [0.9]	0.3321	0.3290	0.3430	0.3289	0.3348	0.3364	0.3212	0.3281	0.3297
	Coverage [0.1]	0.1000	0.1000	0.1000	0.1000	0.1000	0.1000	0.1000	0.1000	0.1000
	Coverage [0.5]	0.4000	0.4100	0.3000	0.4000	0.4000	0.3000	0.4000	0.3000	0.3000
	Coverage [0.9]	0.9000	0.7000	0.7000	0.8000	0.7000	0.7000	0.9000	0.8000	0.8000
	CRPS	0.4321	0.4419	0.4452	0.4123	0.4204	0.4229	0.4234	0.4341	0.4365
	Energy Score	1.5432	1.5603	1.5632	1.4987	1.5199	1.5225	1.5123	1.5301	1.5336

Table 2: Ablation study of STOAT model variants: performance comparison of STOAT-StudentT, STOAT-Gaussian, and STOAT-Laplace distributions, including -NoSpatial and -NoFactors ablations, across 5-day, 7-day, and 10-day horizons on COVID-19 case data. Some individual ablation branches occasionally outperform the full model on specific metrics, yet the full models consistently demonstrate optimal overall performance (CRPS, Energy Score). Metrics in bold indicate the best performance among the STOAT variants.

-NoSpatial and -NoFactors ablations, respectively, with a similar trend for the Laplace and Student's- settings. These results highlight that while ablated models may occasionally surpass the full model on isolated quantile losses, the overall probabilistic calibration and sharpness—as captured by proper scoring rules—remain superior for the complete spatial-causal STOAT. This underscores the importance of incorporating both spatial dependencies and causally relevant covariates for robust spatial-temporal probabilistic forecasting.

Figure 6 illustrates the Energy Scores for STOAT, STOAT-NoSpatial, STOAT-NoFactors across 3, 5, and 7-day predictions. STOAT achieves the lowest Energy Scores (e.g., 0.294 for 3-day, 0.321 for 7-day), outperforming both its variants and baselines. All models show better performance on shorter horizons, reflecting the challenges of long-term forecasting for volatile STC-TS data. Notably, STOAT-NoSpatial outperforms STOAT-NoFactors by 2.5% in CRPS for 7 day predictions, suggesting that spatial causal inference (via S ) contributes more to long-term accuracy than causal factors alone. These findings highlight the synergistic effect of STOAT's modules in robust probabilistic forecasting.

4.4 Conclusion

Spatial-temporal causal time series (STC-TS) are prevalent in domains like epidemiology, where region-specific observations are driven by causally relevant covariates and interconnected through spatial dependencies. Existing deep probabilistic methods often treat covariates as exogenous inputs, neglecting spatial causal relationships and producing miscalibrated forecasts. We propose

STOAT, a novel framework that addresses these limitations through three key contributions: (1) spatial causal adjustment via extended Difference-in-Differences with learnable spatial relation matrices, (2) causal representation learning that constructs causally adjusted latent representations reflecting spatial spillovers and intervention effects, and (3) causally informed probabilistic forecasting with flexible output distributions conditioned on causally disentangled dynamics. Experiments on COVID-19 data from six countries demonstrate STOAT's superior performance, with STOAT-Gaussian achieving the lowest CRPS and Energy Score across forecast horizons, outperforming state-of-the-art deep probabilistic forecasting baselines including DeepAR, DSSM, DeepFactor, MTSNet, and DeepVAR. Ablation studies confirm the synergistic effects of spatial causal inference and causal factors. However, STOAT assumes static spatial correlations and shows diminished benefits in scenarios with weak causal signals. Future work will focus on dynamic spatial modeling, mixture models for heterogeneous dependencies, and real-time causal intervention integration for broader STC-TS applications.

5 Acknowledgments

This research is supported by the Australian Commonwealth Scientific and Industrial Research Organisation (CSIRO) and the the United States National Science Foundation under Grants No. 2302968, No. 2302969, and No. 2302970, titled "Collaborative Research: NSF-CSIRO: HCC: Small: Understanding Bias in AI Models for the Prediction of Infectious Disease Spread".

References

[1] Alberto Abadie. 2005. Semiparametric difference-in-differences estimators. The review of economic studies 72, 1 (2005), 1–19.
[2] Alberto Abadie, Alexis Diamond, and Jens Hainmueller. 2010. Synthetic control methods for comparative case studies: Estimating the effect of California's tobacco control program. Journal of the American statistical Association 105, 490 (2010), 493–505.
[3] Madini O Alassafi, Mutasem Jarrah, and Reem Alotaibi. 2022. Time series predicting of COVID-19 based on deep learning. Neurocomputing 468 (2022), 335–344.
[4] Alexander Alexandrov, Konstantinos Benidis, Michael Bohlke-Schneider, Valentin Flunkert, Jan Gasthaus, Tim Januschowski, Danielle C Maddix, Syama Rangapuram, David Salinas, Jasper Schulz, et al. 2020. Gluonts: Probabilistic and neural time series modeling in python. Journal of Machine Learning Research 21, 116 (2020), 1–6.
[5] Yousef Alimohamadi, Maryam Taghdir, and Mojtaba Sepandi. 2020. Estimate of the basic reproduction number for COVID-19: a systematic review and metaanalysis. Journal of Preventive Medicine and Public Health 53, 3 (2020), 151.
[6] Susan Athey and Guido W Imbens. 2017. The state of applied econometrics: Causality and policy evaluation. Journal of Economic perspectives 31, 2 (2017), 3–32.
[7] Michael Baiocchi, Jing Cheng, and Dylan S Small. 2014. Instrumental variable methods for causal inference. Statistics in medicine 33, 13 (2014), 2297–2340.
[8] Longbing Cao. 2022. AI in combating the COVID-19 pandemic. IEEE Intelligent Systems 37, 2 (2022), 3–13.
[9] Longbing Cao and Qing Liu. 2022. COVID-19 Modeling: A Review. medRxiv (2022), 1–105. doi:10.1101/2022.08.22.22279022
[10] José M Carcione, Juan E Santos, Claudio Bagaini, and Jing Ba. 2020. A simulation of a COVID-19 epidemic based on a deterministic SEIR model. Frontiers in public health 8 (2020), 230.
[11] Ian Cooper, Argha Mondal, and Chris G Antonopoulos. 2020. A SIR model assumption for the spread of COVID-19 in different communities. Chaos, Solitons & Fractals 139 (2020), 110057.
[12] Jing Du, Haley Stone, Yang Yang, Ashna Desai, Hao Xue, Andreas Züfle, Chandini Raina MacIntyre, and Flora D Salim. 2025. BLUE: Bi-layer Heterogeneous Graph Fusion Network for Avian Influenza Forecasting. arXiv preprint arXiv:2505.22692 (2025).
[13] Tilmann Gneiting and Matthias Katzfuss. 2014. Probabilistic forecasting. Annual Review of Statistics and Its Application 1, 1 (2014), 125–151.
[14] Tilmann Gneiting, Adrian E Raftery, Anton H Westveld III, and Tom Goldman. 2005. Calibrated probabilistic forecasting using ensemble model output statistics and minimum CRPS estimation. Monthly weather review 133, 5 (2005), 1098–1118.
[15] Daniel Gros, Alexandre Ounnas, and Timothy Yu-Cheong Yeung. 2021. A new COVID policy stringency index for Europe. Covid Economics 115 (2021).
[16] Thomas Hale, Noam Angrist, Rafael Goldszmidt, Beatriz Kira, Anna Petherick, Toby Phillips, Samuel Webster, Emily Cameron-Blake, Laura Hallas, Saptarshi Majumdar, et al. 2021. A global panel database of pandemic policies (Oxford COVID-19 Government Response Tracker). Nature human behaviour 5, 4 (2021), 529–538.
[17] Shaobo He, Yuexi Peng, and Kehui Sun. 2020. SEIR modeling of the COVID-19 and its dynamics. Nonlinear dynamics 101 (2020), 1667–1680.
[18] Alexander Jordan, Fabian Krüger, and Sebastian Lerch. 2019. Evaluating probabilistic forecasts with scoringRules. Journal of Statistical Software 90 (2019), 1–37.
[19] Liam T Kane, Taolin Fang, Matthew S Galetta, Dhruv KC Goyal, Kristen J Nicholson, Christopher K Kepler, Alexander R Vaccaro, and Gregory D Schroeder. 2020. Propensity score matching: a statistical method. Clinical spine surgery 33, 3 (2020), 120–122.
[20] Leonardo López and Xavier Rodo. 2021. A modified SEIR model to predict the COVID-19 outbreak in Spain and Italy: simulating control scenarios and multi-scale epidemics. Results in physics 21 (2021), 103746.
[21] Atalanti Mastakouri and Bernhard Schölkopf. 2020. Causal analysis of Covid-19 spread in Germany. Advances in Neural Information Processing Systems 33 (2020), 3153–3163.
[22] David Metcalfe, Cheryl K Zogg, Andrew Judge, Daniel C Perry, B Gabbe, K Willett, and ML Costa. 2019. Pay for performance and hip fracture outcomes an interrupted time series and difference-in-differences analysis in England and Scotland. The bone & joint journal 101, 8 (2019), 1015–1023.
[23] Raha Moraffah, Paras Sheth, Mansooreh Karami, Anchit Bhattacharya, Qianru Wang, Anique Tahir, Adrienne Raglin, and Huan Liu. 2021. Causal inference for time series analysis: Problems, methods and evaluation. Knowledge and Information Systems 63 (2021), 3041–3085.
[24] Elias Naumann, Katja Möhring, Maximiliane Reifenscheid, Alexander Wenz, Tobias Rettig, Roni Lehrer, Ulrich Krieger, Sebastian Juhl, Sabine Friedel, Marina Fikel, et al. 2020. COVID-19 policies in Germany and their social, political, and psychological consequences. European Policy Analysis 6, 2 (2020), 191–202.
[25] Sam Nolan, Shakila Khan Rumi, Christoph Anderson, Klaus David, and Flora D Salim. 2020. Exploring the impact of COVID-19 lockdown on social roles and
emotions while working from home. arXiv preprint arXiv:2007.12353 (2020).
[26] Henrik Ohlsson and Kenneth S Kendler. 2020. Applying causal inference methods in psychiatric epidemiology: A review. JAMA psychiatry 77, 6 (2020), 637–644.
[27] Refat Khan Pathan, Munmun Biswas, and Mayeen Uddin Khandaker. 2020. Time series prediction of COVID-19 by mutation rate analysis using recurrent neural network-based LSTM model. Chaos, Solitons & Fractals 138 (2020), 110018.
[28] Syama Sundar Rangapuram, Matthias W Seeger, Jan Gasthaus, Lorenzo Stella, Yuyang Wang, and Tim Januschowski. 2018. Deep state space models for time series forecasting. Advances in neural information processing systems 31 (2018).
[29] Sarah Rothbard, James C Etheridge, and Eleanor J Murray. 2024. A tutorial on applying the difference-in-differences method to health data. Current Epidemiology Reports 11, 2 (2024), 85–95.
[30] Jakob Runge, Andreas Gerhardus, Gherardo Varando, Veronika Eyring, and Gustau Camps-Valls. 2023. Causal inference for time series. Nature Reviews Earth & Environment 4, 7 (2023), 487–505.
[31] Andrew M Ryan, Evangelos Kontopantelis, Ariel Linden, and James F Burgess Jr. 2019. Now trending: Coping with non-parallel trends in difference-in-differences analysis. Statistical methods in medical research 28, 12 (2019), 3697–3711.
[32] Flora D. Salim. 2021. Learning Spatio-Temporal Behavioural Representations for Urban Activity Forecasting. In Companion Proceedings of The Web Conference 2021. Association for Computing Machinery, New York, NY, USA, 347–348. doi:10. 1145/3442442.3451892
[33] David Salinas, Michael Bohlke-Schneider, Laurent Callot, Roberto Medico, and Jan Gasthaus. 2019. High-dimensional multivariate forecasting with low-rank gaussian copula processes. Advances in neural information processing systems 32 (2019).
[34] David Salinas, Valentin Flunkert, Jan Gasthaus, and Tim Januschowski. 2020. DeepAR: Probabilistic forecasting with autoregressive recurrent networks. International journal of forecasting 36, 3 (2020), 1181–1191.
[35] Patrick Schwab, Lorenz Linhardt, Stefan Bauer, Joachim M Buhmann, and Walter Karlen. 2020. Learning counterfactual representations for estimating individual dose-response curves. In Proceedings of the AAAI Conference on Artificial Intelligence, Vol. 34. 5612–5619.
[36] Farah Shahid, Aneela Zameer, and Muhammad Muneeb. 2020. Predictions for COVID-19 with deep learning models of LSTM, GRU and Bi-LSTM. Chaos, Solitons & Fractals 140 (2020), 110212.
[37] Liyang Sun and Sarah Abraham. 2021. Estimating dynamic treatment effects in event studies with heterogeneous treatment effects. Journal of econometrics 225, 2 (2021), 175–199.
[38] Eric J Tchetgen Tchetgen, Chan Park, and David B Richardson. 2024. Universal difference-in-differences for causal inference in epidemiology. Epidemiology 35, 1 (2024), 16–22.
[39] Yuyang Wang, Alex Smola, Danielle Maddix, Jan Gasthaus, Dean Foster, and Tim Januschowski. 2019. Deep factors for forecasting. In International conference on machine learning. PMLR, Proceedings of Machine Learning Research (PMLR), Long Beach, California, USA, 6607–6617.
[40] Zhaonan Wang, Renhe Jiang, Hao Xue, Flora D Salim, Xuan Song, Ryosuke Shibasaki, Wei Hu, and Shaowen Wang. 2024. Learning spatio-temporal dynamics on mobility networks for adaptation to open-world events. Artificial Intelligence 335 (2024), 104120.
[41] Yuying Xie, Chaoshun Li, Mengying Li, Fangjie Liu, and Meruyert Taukenova. 2023. An overview of deterministic and probabilistic forecasting methods of wind energy. Iscience 26, 1 (2023).
[42] Yang Yang and Longbing Cao. 2023. MTSNet: Deep probabilistic crossmultivariate time series modeling with external factors for COVID-19. In 2023 International Joint Conference on Neural Networks (IJCNN). IEEE, IEEE, 1–10.
[43] Meng Zhou, Sixian Huang, Wei Tu, and Donggen Wang. 2024. Machine learningbased causal inference for evaluating intervention in travel behaviour research: A difference-in-differences framework. Travel Behaviour and Society 37 (2024), 100852.

A Appendix

A.1 Algorithm

Algorithm 1: STOAT Framework: Spatial-Temporal Causal Inference and Probabilistic Forecasting

Input: Target time series y ∈ R × , Covariates c ∈ R ^× ^× , Intervention indicator T ∈ {0, 1} × , Post-intervention indicator Post ∈ {0, 1} , Spatial coordinates {(lat , lon)} =1 , Lookback window size

Output: Probabilistic forecast parameters , Causal effect , Spatial effect , Covariate effects

Step 1: Constructing spatial relation matrix

for ← 1 to do for ← 1 to do if ≠ then , ← GeodesicDistance(lat , lon , lat , lon) , ← ¹ , else , ← 0 Normalize: · ^← ^Í · ,

Step 2: Computing spatial lag

for ← 1 to do

$$\begin{vmatrix} \mathbf{for} \ t \leftarrow 1 \mathbf{to} \ N \mathbf{do} \ \mathbf{for} \ t \leftarrow 1 \mathbf{to} \ T \mathbf{do} \ \mathbf{Lag}{i,t} \leftarrow \sum{j=1}^{N} S{i,j} y{j,t} \end{vmatrix}$$

Step 3: Estimating causal and spatial effects

Fit spatial DiD model (Eq. (3)):

Estimate ˆ, ˆ, using instrumental variables and OLS ;

// Obtain causal effect ˆ, spatial effect ˆ,

covariate effects Step 4: Causal adjustment

for ← 1 to do

$$\begin{aligned} & \textbf{for } t \leftarrow 1 \textbf{ to } T \textbf{ do} \ & & \tilde{y}{i,t} \leftarrow y{i,t} - \hat{\delta} \cdot (T_i \cdot \text{Post}t) \ & & z{i,t} \leftarrow \tilde{y}{i,t} + \hat{\rho} \sum{j=1}^{N} S{i,j} \tilde{y}{j,t} \; ; \ & & \text{adjusted representation (Eq. (7))} \end{aligned}$$

Step 5: Preparing inputs for probabilistic forecasting for ← 1 to do

for ← to − 1 do Prepare inputs z,−: , y,−: ; // Lookback window for causally adjusted and target series Set target ,+¹ ← ˜,+¹ ; // Target for next time step

Step 6: Training deep probabilistic model Step 7: Generating probabilistic forecasts return predicted distribution P(.)

A.2 Metrics

• Continuous Ranked Probability Score (CRPS) [14]: Quantifies the discrepancy between predicted and empirical cumulative distribution functions:

CRPS $$(P, x{i,t}) = \int{-\infty}^{\infty} (P(y) - 1{y \ge x_{i,t}})^2 dy,$$ (19)

where is the predicted cumulative distribution, and 1{ ≥ , } is the indicator function for observed values ,. CRPS evaluates the overall quality of probabilistic forecasts by measuring how well the predicted distribution matches the observed values. Lower CRPS values indicate better performance, with perfect forecasts achieving CRPS = 0.

• Weighted Quantile Loss (WQL): Measures accuracy at specific quantiles of the predicted distribution:

$$\begin{aligned} \text{WQL}[\tau] &= 2 \frac{\sum{i,t} Q{i,t}^{(\tau)}}{\sum{i,t} |x{i,t}|}, \ Q{i,t}^{(\tau)} &= \begin{cases} (1-\tau)|q{i,t}^{(\tau)} - x{i,t}| & \text{if } q{i,t}^{(\tau)} > x{i,t}, \ \tau|q{i,t}^{(\tau)} - x_{i,t}| & \text{otherwise,} \end{cases} \end{aligned} \tag{20}$$

where ∈ {0.1, 0.5, 0.9} represents quantiles, and ( ) , is the predicted value at quantile . WQL assesses the precision of quantile predictions, with asymmetric penalties that reflect the importance of over- and under-estimation. Lower WQL values indicate better quantile forecast accuracy.

• Coverage: Measures the proportion of observations that fall within predicted confidence intervals:

Coverage $$[\alpha] = \frac{1}{N} \sum{i,t} 1 \left{ q{i,t}^{(\alpha/2)} \le x{i,t} \le q{i,t}^{(1-\alpha/2)} \right},$$ (21)

where is the total number of observations, is the significance level (e.g., = 0.1 for 90% intervals), and (· ) , are the predicted quantiles. The best performance occurs when Coverage matches the nominal level (i.e., Coverage[0.1] ≈ 0.9 and Coverage[0.5] ≈ 0.5), indicating well-calibrated prediction intervals.

• Energy Score [18]: Evaluates multivariate probabilistic forecast accuracy, sensitive to inter-variable correlations:

$$\begin{split} \text{Energy}({\tilde{x}_t, xt}) &= \mathbb{E}{\tilde{x}_t} |\tilde{x}_t - xt|{Fro}^{\beta} \ &- \frac{1}{2} \mathbb{E}_{\tilde{x}_t, \tilde{x}_t'} |\tilde{x}_t - \tilde{x}t'|{Fro}^{\beta}, \end{split} \tag{22}$$

where ˜ and ˜ ′ are independent samples from the predicted distribution, is the observed value, = 1, and ∥ · ∥ denotes the Frobenius norm. Energy Score captures both the accuracy of multivariate forecasts and the quality of predicted dependence structures between variables. Lower Energy Score values indicate better multivariate forecast performance, with the metric being particularly valuable for high-dimensional time series.

yangstoatspatialtemporalprobabilistic2025.md 81 KB История Исходник