ABM-UDE: Developing Surrogates for Epidemic Agent-Based Models via Scientific Machine Learning
Organizations: Monta Vista High School · MIT CSAIL · Sandia National Laboratories
Abstract
Agent-based epidemic models (ABMs) encode behavioral and policy heterogeneity but are too slow for nightly hospital planning. We develop county-ready surrogates that learn directly from exascale ABM trajectories using Universal Differential Equations (UDEs): mechanistic SEIR-family ODEs with a neural-parameterized contact rate (no additive residual). Our contributions are threefold: we adapt multiple shooting and an observer-based prediction-error method (PEM) to stabilize identification of neural-augmented epidemiological dynamics across intervention-driven regime shifts; we enforce positivity and mass conservation and show the learned contact-rate parameterization yields a well-posed vector field; and we quantify accuracy, calibration, and compute against ABM ensembles and UDE baselines. On a representative ExaEpi scenario, PEM-UDE reduces mean MSE by 77% relative to single-shooting UDE (3.00 vs. 13.14) and by 20% relative to MS-UDE (3.75). Reliability improves in parallel: empirical coverage of ABM -% and -% bands rises from 0.68/0.43 (UDE) and 0.79/0.55 (MS-UDE) to 0.86/0.61 with PEM-UDE and 0.94/0.69 with MS+PEM-UDE, indicating calibrated uncertainty rather than overconfident fits. Inference runs in seconds on commodity CPUs (20-35 s per 90-day forecast), enabling nightly ''what-if'' sweeps on a laptop. Relative to a 100 CPU-hour ABM reference run, this yields lower wall-clock per scenario. This closes the realism-cadence gap, supports threshold-aware decision-making (e.g., maintaining ICU occupancy %), preserves mechanistic interpretability, and enables calibrated, risk-aware scenario planning on standard institutional hardware. Beyond epidemics, the ABMUDE recipe provides a portable path to distill agent-based simulators into fast, trustworthy surrogates for other scientific domains.
Figures & tables
| Method | Final MSE (Dataset 1) |
|---|---|
| UDE (baseline) | 16.85 |
| MS-UDE | 3.76 |
| PEM-UDE | 3.01 |
| MS+PEM-UDE | 2.93 |
| Method | Latency |
|---|---|
| UDE (baseline) | 21 s ( 24 s ) |
| MS-UDE (inference) | 23 s ( 26 s ) |
| PEM-UDE ( ) | 31 s ( 35 s ) |
| MS+PEM-UDE ( ) | 32 s ( 35 s ) |
| ExaEpi ABM (reference) | 100 h = 360,000 s |
| Method | 10–90 % Coverage | 25–75 % Coverage |
|---|---|---|
| UDE (baseline) | 0.68 | 0.43 |
| MS-UDE | 0.79 | 0.55 |
| PEM-UDE | 0.86 | 0.61 |
| MS+PEM-UDE (ensemble) | 0.94 | 0.69 |
| Compartment | MS+PEM vs. PEM | MS+PEM vs. MS | MS+PEM vs. UDE |
|---|---|---|---|
| Exposed ( ) | |||
| Presymptomatic ( ) | |||
| Symptomatic ( ) | |||
| Asymptomatic ( ) | |||
| Never-infected / Susceptible-like pool ( ) | |||
| Deaths ( ) |