🤖 AI Summary
Existing spatiotemporal graph neural networks (STGNNs) for epidemic forecasting assume that regions with similar historical features exhibit consistent future infection trends, overlooking mechanism heterogeneity induced by unobserved factors—such as healthcare capacity, viral variants, and human mobility—across geographic and temporal dimensions. To address this, we propose Mechanism-Aware STGNN (HeatGNN), the first framework to embed SEIR-inspired mechanistic models into graph neural networks, enabling learnable and interpretable, location-specific propagation modeling via a time-varying mechanism affinity graph. Our method integrates dynamic graph learning with heterogeneous spatiotemporal graph convolution, ensuring linear scalability. Extensive experiments on three benchmark datasets demonstrate statistically significant improvements over state-of-the-art baselines, validating HeatGNN’s capability to effectively capture mechanism heterogeneity and its feasibility for large-scale deployment.
📝 Abstract
Among various spatio-temporal prediction tasks, epidemic forecasting plays a critical role in public health management. Recent studies have demonstrated the strong potential of spatio-temporal graph neural networks (STGNNs) in extracting heterogeneous spatio-temporal patterns for epidemic forecasting. However, most of these methods bear an over-simplified assumption that two locations (e.g., cities) with similar observed features in previous time steps will develop similar infection numbers in the future. In fact, for any epidemic disease, there exists strong heterogeneity of its intrinsic evolution mechanisms across geolocation and time, which can eventually lead to diverged infection numbers in two ``similar'' locations. However, such mechanistic heterogeneity is non-trivial to be captured due to the existence of numerous influencing factors like medical resource accessibility, virus mutations, mobility patterns, etc., most of which are spatio-temporal yet unreachable or even unobservable. To address this challenge, we propose a Heterogeneous Epidemic-Aware Transmission Graph Neural Network (HeatGNN), a novel epidemic forecasting framework. By binding the epidemiology mechanistic model into a GNN, HeatGNN learns epidemiology-informed location embeddings of different locations that reflect their own transmission mechanisms over time. With the time-varying mechanistic affinity graphs computed with the epidemiology-informed location embeddings, a heterogeneous transmission graph network is designed to encode the mechanistic heterogeneity among locations, providing additional predictive signals to facilitate accurate forecasting. Experiments on three benchmark datasets have revealed that HeatGNN outperforms various strong baselines. Moreover, our efficiency analysis verifies the real-world practicality of HeatGNN on datasets of different sizes.