CausalGNN: Causal-Based Graph Neural Networks for Spatio-Temporal Epidemic Forecasting

Abstract

Infectious disease forecasting has been a key focus in the recent past owing to the COVID-19 pandemic and has proved to be an important tool in controlling the pandemic. With the advent of reliable spatiotemporal data, graph neural network models have been able to successfully model the inter-relation between the cross-region signals to produce quality forecasts, but like most deep-learning models they do not explicitly incorporate the underlying causal mechanisms. In this work, we employ a causal mechanistic model to guide the learning of the graph embeddings and propose a novel learning framework -- Causal-based Graph Neural Network (CausalGNN) that learns spatiotemporal embedding in a latent space where graph input features and epidemiological context are combined via a mutually learning mechanism using graph-based non-linear transformations. We design an attention-based dynamic GNN module to capture spatial and temporal disease dynamics. A causal module is added to the framework to provide epidemiological context for node embedding via ordinary differential equations. Extensive experiments on forecasting daily new cases of COVID-19 at global, US state, and US county levels show that the proposed method outperforms a broad range of baselines. The learned model which incorporates epidemiological context organizes the embedding in an efficient way by keeping the parameter size small leading to robust and accurate forecasting performance across various datasets.