🤖 machine learning

ABM-UDE: Developing Surrogates for Epidemic Agent-Based Models via Scientific Machine Learning

This paper introduces ABM-UDE, a scientific machine learning framework that utilizes Universal Differential Equations with multiple shooting and observer-based prediction-error methods to distill slow, exascale agent-based epidemic models into fast, calibrated, and interpretable surrogates capable of enabling real-time, risk-aware hospital planning on standard hardware.

Original authors: Sharv Murgai, Utkarsh Utkarsh, Kyle C. Nguyen, Alan Edelman, Erin C. S. Acquesta, Christopher Vincent Rackauckas

Published 2026-02-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Sharv Murgai, Utkarsh Utkarsh, Kyle C. Nguyen, Alan Edelman, Erin C. S. Acquesta, Christopher Vincent Rackauckas

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are the mayor of a city, and a sudden storm (an epidemic) is approaching. You need to know: Will the shelters overflow? Do we need to open more hospitals? When should we close the schools?

To answer this, you have two tools:

The Super-Computer Simulation (ABM): This is like a massive, hyper-realistic video game where every single citizen (millions of them) is an individual character. They have jobs, families, habits, and they interact randomly. It's incredibly accurate, but it takes 100 hours of super-computer time to run just one 3-month scenario. By the time the answer is ready, the storm has already passed.
The Simple Spreadsheet (ODE): This is a quick, rough estimate. It treats everyone as an average person. It runs in seconds, but it's often wrong because it misses the messy details of real human behavior.

The Problem: We need the accuracy of the Super-Computer Simulation but the speed of the Spreadsheet. We need to run hundreds of "What if?" scenarios overnight to make life-saving decisions, but we can't wait 100 hours for each one.

The Solution: The "Smart Proxy" (ABM-UDE)
The authors of this paper built a Smart Proxy. Think of it as a "ghost" version of the massive simulation that learns from the big one but runs on a regular laptop.

Here is how they did it, using some creative analogies:

1. The "Hybrid Engine" (Universal Differential Equations)

Imagine a car engine. The pistons, gears, and fuel lines are the mechanical rules (the known science of how viruses spread). These are fixed and reliable.
However, the driver's behavior (how fast they press the gas, when they brake) is unpredictable. In the old days, scientists tried to guess the driver's behavior with a simple formula, which often failed.

The authors replaced the "driver's behavior" part with a Neural Network (a tiny AI brain).

The Engine: The known laws of physics and biology (e.g., if you get sick, you can infect others).
The AI Brain: A smart component that learns exactly how the virus behaves in this specific city by watching the Super-Computer Simulation.
The Result: A model that respects the laws of physics but adapts to the messy reality of human behavior. It's like a car that knows the rules of the road but has a co-pilot who learned from watching millions of other drivers.

2. The "Stabilization" Tricks (Multiple Shooting & PEM)

Training this AI brain is tricky. If you try to teach it by showing it a 90-day movie all at once, it gets confused. It might learn the first week perfectly but drift off course by week 10, like a GPS that slowly loses signal.

The authors used two clever tricks to keep the AI on track:

Multiple Shooting (The "Checkpoints"): Instead of watching the whole 90-day movie at once, they cut it into small 10-day clips. They train the AI on the first clip, check if it's right, then reset and train on the next clip. It's like learning a long dance routine by mastering one move at a time, checking your form, and then moving to the next move. This prevents the AI from getting "lost" in the long run.
Prediction Error Method (The "Reality Check"): Imagine the AI is trying to predict the weather. Every day, it makes a guess, and then you show it the actual weather. If it was wrong, you gently nudge its internal settings to correct the mistake for tomorrow. This "nudge" (called an observer) keeps the AI anchored to reality, ensuring it doesn't drift away from the truth.

3. The Results: From "Supercomputer" to "Laptop"

The results were dramatic:

Speed: The old Super-Computer simulation took 100 hours to run one scenario. The new Smart Proxy runs on a standard laptop in about 30 seconds. That is 10,000 times faster.
Accuracy: It didn't just get faster; it got smarter. By using the "Checkpoints" and "Reality Checks," the new model was 77% more accurate than the standard AI approach.
Reliability: Most importantly, the model knows when it's unsure. It doesn't just give a single number; it gives a range (e.g., "The ICU will be 60% full, give or take 5%"). This helps doctors and mayors plan for the worst-case scenario without panicking.

Why This Matters

Before this, local health officials had to choose between slow, accurate data (useless for daily planning) or fast, inaccurate guesses (dangerous for decision-making).

This new method allows a county health officer to sit at their desk, type in a question like "What happens if we open schools next Monday?" and get a highly accurate, risk-aware answer in seconds. They can run hundreds of "What if?" scenarios overnight to find the perfect balance of safety and normalcy, potentially saving thousands of lives by acting before the hospitals get overwhelmed.

In short: They took a massive, slow, perfect simulation, distilled its "wisdom" into a tiny, fast, smart AI, and taught it how to stay on track so it can help us make life-saving decisions in real-time.

1. Problem Statement

The Gap Between Realism and Operational Speed:
Public health decision-making (e.g., hospital staffing, ICU capacity management) requires high-fidelity forecasts that account for human behavior, spatial heterogeneity, and policy interventions. Agent-Based Models (ABMs) are the "gold standard" for this realism but are computationally prohibitive.

Computational Cost: A single 70–90 day ABM simulation (e.g., CityCOVID) can take ~100 CPU-hours on standard hardware or require thousands of cores on supercomputers.
Operational Need: Health officials need to run hundreds of "what-if" scenarios nightly (e.g., "What if mobility increases by 10%?") to make timely decisions. ABMs cannot provide this cadence.
Limitation of Current Surrogates: Traditional compartmental ODE models (SEIR) are fast but rely on homogeneous mixing assumptions, failing to capture local surges or complex behavioral shifts. Existing hybrid surrogates often lack stability during intervention shocks or require High-Performance Computing (HPC) for inference.

Goal: Develop a "county-ready" surrogate that learns directly from exascale ABM trajectories, offering ABM-level accuracy with ODE-level speed on commodity hardware.

2. Methodology: The ABM→UDE Pipeline

The authors propose a framework using Universal Differential Equations (UDEs) to bridge the gap. Instead of adding a free-form neural residual to the ODE (which breaks physical constraints), they embed a neural network inside the mechanistic equations to learn the contact rate.

A. Model Formulation

Mechanistic Core: A SEIR-family ODE system ( $S, E, I_{\sim s}, I_s, I_a, R, D$ ) governing disease progression.
Neural Parameterization: The time-dependent contact rate $\kappa(t)$ $κ (t)$ is replaced by a bounded neural map $\kappa_\phi(u, t)$ $κ_{ϕ} (u, t)$ :
$\dot{u} = f(u, t; \theta, \kappa_\phi(u, t))$
$\kappa_\phi(u, t) = \kappa_{max} \sigma(N_\phi(\tilde{u}(t)))$
- Constraints: The architecture enforces positivity (no negative populations) and mass conservation (total population $N$ is constant) by construction. Only the contact rate is learned; the flow structure remains mechanistic.

B. Training Strategies for Stability

Training neural-augmented ODEs over long horizons (89 days) with abrupt intervention changes (e.g., lockdowns) is unstable. The paper evaluates and combines three stabilization techniques:

Vanilla UDE (Single Shooting): Standard backpropagation through the entire trajectory. Prone to error accumulation and drift, especially around intervention change-points.
Multiple Shooting (MS): The time horizon is partitioned into $M$ windows. Each window solves a short Initial Value Problem (IVP). Continuity between windows is enforced via a soft penalty. This reduces the effective backpropagation horizon, stabilizing training against regime shifts.
Prediction-Error Method (PEM): An observer-based approach. A correction term $K(y(t) - Hu)$ is added to the ODE, where $y(t)$ is the observed data (interpolated ABM trajectory) and $K$ is a learnable gain matrix. This continuously re-anchors the model state to the data, preventing drift.
MS+PEM-UDE (Hybrid): The proposed best method. It applies the observer correction within each multiple-shooting window. MS handles the global regime shifts, while PEM handles local error accumulation.

C. Data and Setup

Source: ExaEpi, an exascale ABM simulator.
Dataset: One representative scenario (Dataset 1) spanning 89 days with a lockdown intervention at day 49.
Ensemble: To quantify uncertainty, 100 independent models were trained for each method using different random seeds.
Hardware: Training and inference performed on a commodity Apple Macbook Air (M4 CPU).

3. Key Contributions

Novel UDE Architecture: A mechanistic SEIR model where the contact rate is a neural parameter, preserving biological interpretability and physical constraints (mass conservation) without additive residuals.
Stabilized Training Framework: The adaptation of Multiple Shooting and Observer-based PEM specifically for neural-augmented epidemiological models. This addresses the instability caused by intervention shocks and long-horizon drift.
Operational Feasibility: Demonstrating that high-fidelity surrogates can be trained and run on standard institutional hardware (laptops) rather than supercomputers.
Rigorous Evaluation: A comprehensive benchmark comparing trajectory accuracy, uncertainty calibration (coverage), and computational efficiency against ABM ground truth.

4. Results

The study evaluated four variants: Vanilla UDE, MS-UDE, PEM-UDE, and MS+PEM-UDE.

A. Accuracy (Mean Squared Error)

Vanilla UDE: High error (MSE $\approx$ 16.85) due to drift.
MS-UDE: Significant improvement (MSE $\approx$ 3.76).
PEM-UDE: Strong performance (MSE $\approx$ 3.01).
MS+PEM-UDE: Best performance with the lowest MSE (2.93).
- Improvement: 77% reduction in MSE compared to Vanilla UDE; 20% reduction compared to MS-UDE.
- The hybrid method most effectively captured the dynamics post-intervention (day 49).

B. Uncertainty Calibration

Reliability is measured by how often the surrogate's prediction falls within the ABM's stochastic quantile bands (10–90% and 25–75%).

Vanilla UDE: Poor calibration (0.68 / 0.43 coverage).
MS+PEM-UDE: Excellent calibration (0.94 / 0.69 coverage).
- This indicates the surrogate provides calibrated uncertainty rather than overconfident fits, crucial for risk-aware decision-making.

C. Computational Efficiency

Inference Speed: The trained UDE surrogate predicts a ~90-day horizon in 20–35 seconds on a single CPU thread.
Speedup: Compared to the ABM reference run (~100 CPU-hours), the surrogate offers a $\sim 10^4$ (4 orders of magnitude) reduction in wall-clock time.
Throughput: Enables nightly "what-if" sweeps of hundreds of scenarios on a laptop.

5. Significance and Impact

Closing the Realism-Cadence Gap: The framework allows local health systems to utilize ABM-level realism for daily decision-making without access to supercomputers.
Risk-Aware Planning: By providing calibrated uncertainty, the model helps planners identify policies that keep ICU occupancy below critical thresholds (e.g., 75%), directly addressing the mortality spikes associated with capacity overload.
Generalizability: The "ABM→UDE" recipe is presented as a portable path for distilling complex agent-based simulators into fast, trustworthy surrogates for other scientific domains beyond epidemics.
Interpretability: Unlike "black box" deep learning models, the UDE retains mechanistic interpretability, allowing stakeholders to understand why a forecast was made (via the learned contact rate $\kappa_\phi$ ).

Conclusion: The paper demonstrates that MS+PEM-UDE is a robust, scalable, and operationally viable solution for epidemic forecasting, successfully compressing exascale ABM complexity into a lightweight, physics-informed surrogate suitable for real-time public health strategy.