Machine Learning Sandwich for Agent-Based Models
Automatic calibration and post-processing adjustment for ABMs
Motivation: Agent-Based Models in Science and Policy
Agent-Based Models [ABMs] are an important tool for scientific research and policy-making:
- For scenario modeling: “What if we do X?”
- For forecasting: “How many cases will we have in the next week?”
Calibration can be time-consuming: grid search, Markov models, etc.
There’s no consensus on the best calibration approach.
ABMs (and models in general) are limited in their predictive power.
Schematic of compartmental and ABMs used by the US Centers for Disease and Control Prevention during a Measles Outbreak Response CDC (2024).
Machine Learning Sandwich for ABMs
Using ML for Calibration
Machine Learning is Broken
- After all the data pouring, attention to causal inference and mechanistic models is coming back (Baker et al. 2018; Pearl 2019)
- The case of Google Flu Trends: Paper reported a 0.97 correlation (Ginsberg et al. 2009), but predictions overshoot by 100% (Kandula and Shaman 2019; Lazer et al. 2014).
How to MechML: Correcting prediction errors
Train a model that predicts ABM forecast errors. The model takes the following form:
\[ Loss(\boldsymbol{\omega}): \left\{\mathcal{M}(\boldsymbol{\theta}), \boldsymbol{x}\right\} \to \widehat{\boldsymbol{\varepsilon}} \equiv \left\lVert \widehat{\boldsymbol{\varepsilon}} - f\left(\boldsymbol{\omega}, \mathcal{M}(\boldsymbol{\theta}), \boldsymbol{x}\right)\right\rVert{}_p, \]
where \(\widehat{\boldsymbol{\varepsilon}}\equiv \left(\mathcal{M}(\boldsymbol{\theta}) - y_{obs}\right)\) ABM forecast error, \(\boldsymbol{y}_{obs}\) is the observed data, \(f(\cdot)\) is a non-linear function, \(\boldsymbol{x}\) are additional features for the model, and \(\boldsymbol{\omega}\) is an array of weights associated with the ML o predict \(\boldsymbol{\varepsilon}\).
How to MechML: Predictions as feature
Use the ABM predictions as a feature in the ML model. The model takes the following form:
\[ Loss(\boldsymbol{\omega}) \equiv \left\lVert \boldsymbol{y}_{obs} - f\left(\boldsymbol{\omega}, \mathcal{M}(\boldsymbol{\theta}), \boldsymbol{x}\right)\right\rVert{}_p, \]
How to MechML: Mechanistic penalty
Use a mechanistic penalty in the ML loss function. The model takes the following form:
\[ Loss(\boldsymbol{\omega}) \equiv \left\lVert \boldsymbol{y}_{obs} - f\left(\boldsymbol{\omega}, \boldsymbol{x}\right)\right\rVert{}_p + \lambda \left\lVert f\left(\boldsymbol{\omega}, \boldsymbol{x}\right) - \mathcal{M}(\theta)\right\rVert{}_p, \]
where \(\lambda\) is a hyperparameter that controls the weight of the mechanistic penalty.
Calibration
Susceptible-Infected Recovered model:
- Agent Based Model with
epiworldR. - Fully connected graph (everyone is connected to everyone).
- We ran 20,000 simulations
- Parameters of interest:
- Initial state
- Contact rate
- Probability of transmission
- Probability of recovery
We trained a convolutional neural network to build our calibrator:
\[ \mathcal{C}: \text{Epicurves} \to \boldsymbol{\theta} \]
Where \(\boldsymbol{\theta} = \left[\text{Init. state}\text{Contact Rate}, \text{P(transmit)}, \text{P(recover)}\right]\)
Calibration : CNN Architecture
Calibration : Model fit
Prediction: Model of gene functions (0/1)
Example model by speaker using mechanistic machine learning to improve accuracy in prediction of gene functions.
- Used a mechanistic model of evolution of gene functions to predict presence/absence of them.
- ML model: logistic regression featuring a large number of parameters.
- MechML: use the mechanistic predictions as features in the logit.
Machine Learning Sandwich for ABMs (repeat)