STOAT: Spatial-Temporal Probabilistic Causal Inference Network

TL;DR

STOAT addresses the challenge of forecasting spatial-temporal causal time series with calibrated uncertainty by integrating a spatially informed causal adjustment mechanism into a deep probabilistic forecasting framework. It extends Difference-in-Differences through a learnable spatial relation matrix to obtain causally adjusted representations that capture both treatment effects and spatial spillovers, then feeds these into a neural encoder-decoder to produce distributional forecasts across regions. Empirical evaluation on multi-country COVID-19 data shows STOAT variants achieving superior CRPS and Energy Score compared with leading probabilistic forecasters, with clear gains during periods of strong interventions and spatial dependence. Ablation studies confirm the complementary value of spatial causal inference and covariate integration for robust probabilistic forecasting in STC-TS. The framework offers a principled approach to epidemic management and other spatial-temporal tasks where interventions and spatial connectivity drive uncertainty, while pointing to future work on dynamic spatial dependencies and mixture models for greater flexibility.

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-$t$, 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.

Figures (6)

• 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.
• 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.
• 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-t, Gaussian, and Laplace distributions, respectively.
• 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-$t$, Gaussian, and Laplace) are compared against baseline models (DeepVAR, DeepAR, DSSM, DeepFactor, MTSNet).
• 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.