Learning Calibrated and Transferable Node-Level Infection Risk in Stochastic Epidemics on Complex Networks

Technical Summary: Learning Calibrated and Transferable Node-Level Infection Risk in Stochastic Epidemics on Complex Networks

Problem Statement
The paper addresses the challenge of predicting future infection risk for susceptible nodes in stochastic epidemics on complex networks. Traditional compartmental models often assume homogeneous mixing, while network epidemiology captures heterogeneity but struggles to jointly account for epidemic state, network structure, and stochastic uncertainty. The authors pose a specific probabilistic question: given a graph $G$, an observed epidemic state at time $t_0$, and epidemiological parameters, can the probability that a currently susceptible node $i$ becomes infected be learned as a calibrated graph function? This quantity, $p_i(t_0)$, is inherently stochastic, as independent realizations from the same state may infect different subsets of nodes. The study further investigates whether such learned risk functions can transfer to unseen network structures (topology-level out-of-distribution generalization) and whether they can support effective targeted interventions.

Methodology
The study employs a workflow combining stochastic simulation, graph learning, and mechanistic auditing:

1. Data Generation: Stochastic SIR epidemics were simulated on four network families: Erdős–Rényi (ER), Barabási–Albert (BA), Watts–Strogatz (WS), and Stochastic Block Model (SBM). Each graph contained $N=100$ nodes with a mean degree of $\approx 6$.
2. Target Definition: For each observed state (at $t_0=5$), the authors generated "soft" target probabilities using repeated Monte-Carlo continuations ($M=100$ for training, $M=500$ for testing). The target for a susceptible node is the fraction of continuations in which it eventually becomes infected.
3. Model Architecture: A parameter-conditioned Graph Neural Network (GNN) was trained to estimate these probabilities. The model uses three message-passing layers with sum aggregation (chosen to preserve the count of infectious neighbors, a key factor in the infection probability equation). Inputs include the observed S/I/R states, graph structure, and epidemiological parameters ($\beta, \gamma, \lambda_1, \rho$).
4. Evaluation Strategy:
   • Graph-Disjoint Split: Testing on unseen graph realizations from the same topology families.
   • Leave-One-Topology-Out (LOTO): Testing on an entire topology family absent from training to evaluate structural distribution shift.
   • Baselines: Performance was compared against ridge regression, gradient boosting (on handcrafted features), and untrained centrality/exposure metrics.
5. Explainability & Intervention: Integrated Gradients and edge occlusion were used to audit the model's reliance on epidemiological mechanisms. Risk-guided interventions (vaccinating top-ranked nodes) were simulated and compared against a Monte-Carlo oracle and other heuristics.

Key Results

• Calibration and Accuracy: On 13,097 susceptible nodes from unseen graph realizations, the GNN achieved a Brier score of $7.29 \times 10^{-4}$ and an Expected Calibration Error (ECE) of 0.003. It outperformed gradient boosting and ridge regression in absolute probability estimation, though boosting remained competitive for ranking nodes within a single epidemic episode.
• Structural Transfer (OOD): The model generalized well to unseen ER, WS, and SBM networks. However, performance degraded significantly when tested on unseen BA (scale-free) networks, which exhibited higher degree heterogeneity (Coefficient of Variation = 0.819 vs. <0.40 for others) and larger maximum degrees. This suggests a boundary of transferability defined by structural heterogeneity.
• Mechanistic Audit: Explainability analysis confirmed the model learned epidemiologically meaningful logic. Integrated Gradients attributed 76.8% of the prediction importance to the infectious state compartment. Edge occlusion showed that removing edges to infectious neighbors reduced predicted risk approximately 6.7 times more than removing edges to susceptible neighbors.
• Intervention Efficacy: Targeting the top 5% of susceptible nodes based on GNN risk reduced the final attack rate by 0.0481, closely matching the Monte-Carlo oracle (0.0480) and outperforming degree-based (0.0357) and random (0.0160) targeting. The advantage was most pronounced at low intervention budgets (5–10%).

Significance and Claims
The paper claims to reframe epidemic graph learning from trajectory reproduction to the prediction of stochastic individual risk. Its primary contributions are:

1. Probabilistic Formulation: Establishing a framework where the target is a calibrated conditional probability derived from Monte-Carlo ensembles, rather than a deterministic state or trajectory.
2. Rigorous OOD Evaluation: Demonstrating that while GNNs generalize well to unseen instances of familiar topologies, they face specific failure modes when transferring to topologies with extreme structural heterogeneity (e.g., scale-free networks), emphasizing the need for explicit structural out-of-distribution testing.
3. Operational Utility: Showing that a trained GNN can "amortize" the computational cost of repeated stochastic simulations, providing near-oracle intervention decisions with a single forward pass.
4. Mechanistic Validity: Using explainability not to generate intervention policies, but to audit the predictor, confirming that the model relies on infectious neighbors and state information consistent with SIR dynamics.

The authors maintain a modest stance, noting that the study uses synthetic static networks and assumes known parameters. They identify the need for future work on temporal/weighted networks, parameter uncertainty, and empirical validation on real-world contact data. The work positions GNNs in epidemic modeling as tools for estimating stochastic individual risk and characterizing the boundaries of structural generalization.