Epidemiology-informed Graph Neural Network for Heterogeneity-aware Epidemic Forecasting
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.