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

arXiv: https://arxiv.org/abs/2602.21588 · HTML: https://mlsanity.com/paper/2602.21588/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 7 of its 46 references that are in the map.
- First posted on arXiv 25 Feb 2026.

## Builds on (most cited first, 7 of 7)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Universal Differential Equations for Scientific Machine Learning](https://mlsanity.com/paper/2001.04385/) (2020)
- [Towards the Systematic Reporting of the Energy and Carbon Footprints of Machine Learning](https://mlsanity.com/paper/2002.05651/) (2020)
- [Projection onto the probability simplex: An efficient algorithm with a simple proof, and an application](https://mlsanity.com/paper/1309.1541/) (2013)
- [Bayesian calibration of stochastic agent based model via random forest](https://mlsanity.com/paper/2406.19524/) (2024)
- [Pandemics In Silico: Scaling an Agent-Based Simulation on Realistic Social Contact Networks](https://mlsanity.com/paper/2401.08124/) (2024)

## Abstract (arXiv)

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 kappaφ(u,t) (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 10-90% and 25-75% 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 ~104× lower wall-clock per scenario. This closes the realism-cadence gap, supports threshold-aware decision-making (e.g., maintaining ICU occupancy <75%), preserves mechanistic interpretability, and enables calibrated, risk-aware scenario planning on standard institutional hardware. Beyond epidemics, the ABM→UDE recipe provides a portable path to distill agent-based simulators into fast, trustworthy surrogates for other scientific domains.