Graph Neural Network Surrogates to leverage Mechanistic Expert Knowledge towards Reliable and Immediate Pandemic Response

TL;DR

The paper addresses the need for rapid, reliable forecasts in time-critical pandemic settings by pairing a validated spatially resolved mechanistic metapopulation model with a graph neural network (GNN) surrogate. It trains a seven-layer ARMAConv GNN on a 400-node county graph to reproduce demographically and spatially resolved trajectories under outbreak and persistent-threat regimes, achieving $MAPE$ in the range of $10\%-27\%$ over forecast horizons of $30$–$90$ days. The surrogate delivers near-constant runtime as horizon increases and accelerates inference by up to $2.87\times 10^{4}$ fold compared with the mechanistic model, enabling interactive, web-based scenario exploration. This work demonstrates that GNN surrogates can translate complex metapopulation dynamics into fast, interpretable decision-support tools with explicit ties to expert models, and outlines future directions such as incorporating edge weights and vaccination dynamics for broader applicability.

Abstract

During the COVID-19 crisis, mechanistic models have guided evidence-based decision making. However, time-critical decisions in a dynamical environment limit the time available to gather supporting evidence. We address this bottleneck by developing a graph neural network (GNN) surrogate of a spatially and demographically resolved mechanistic metapopulation simulator. This combined approach advances classical machine learning approaches which are often black box. Our design of experiments spans outbreak and persistent-threat regimes, up to three contact change points, and age-structured contact matrices on a 400-node spatial graph. We benchmark multiple GNN layers and identify an ARMAConv-based architecture that offers a strong accuracy-runtime trade-off. Across 30-90 day horizons and up to three contact change points, the surrogate attains 10-27 % mean absolute percentage error (MAPE) while delivering (near) constant runtime with respect to the forecast horizon. Our approach accelerates evaluation by up to 28,670 times compared with the mechanistic model, allowing responsive decision support in time-critical scenarios and straightforward web integration. These results show how GNN surrogates can translate complex metapopulation models into immediate, reliable tools for pandemic response.

Figures (8)

• Figure 1: Disease state transition model. The disease state transition model comprises eight infection states with possibility for recovery or aggravation after each disease state following exposure state (except for death and recovery itself).
• Figure 2: Initial conditions over five input days and trajectories over the prediction horizon. A) Initial conditions for outbreak range scenarios. B) Comparison between initial conditions for outbreak range and persistent-threat range scenarios. C) The two hundred first trajectories for the outbreak scenarios. D) The two hundred first trajectories for the persistent-threat scenarios.
• Figure 3: Non-spatial vs. spatial input encoding. A) Non-spatial 2D input (5 $\times$ 162): five input days with 48 (age groups $\times$ compartments) values plus up to three $6\times6$ contact matrices, their change days, and reduction factors. B) Spatial GNN input: for each node (400 counties) the same features as in Panel A are broadcast and paired with the adjacency matrix.
• Figure 4: MAPE results of the grid search for the simple neural networks and simple input data structure. A) Validation MAPE results (in log-scale) for a grid search for simple models with six age groups and no contact change on a horizon over 30 days. B) Validation MAPE results (in log-scale) for the grid search for different optimizers and activation functions. C) Test MAPE results (in original scale) for the final LSTM model on different prediction horizons (30, 60, and 90 days) and for different numbers of contact changes (0, 1, 2, 3 changes in 30 days).
• Figure 5: Spatially distributed disease dynamics and grid search of different GNN infrastructures. A) Four random initializations of heterogeneous disease dynamics over the 400 counties of Germany. B) Validation MAPEs for the four different GNN layers. C) Training times corresponding to Panel B). D) Extended grid search for ARMAConv model. E) Training times corresponding to Panel D). F) Further hyperparameter tuning for ARMAConv model with seven layers and 512 channels per layer.