Differentiable agent-based epidemiological modeling for end-to-end learning
Abstract
Mechanistic simulators are an indispensable tool for epidemiology to explore the behavior of complex, dynamic infections under varying conditions and navigate uncertain environments.
ODE-based models are the dominant paradigm that enable fast simulations and are tractable to gradient-based optimization, but make simplifying assumptions about population homogeneity.
Agent-based models (ABMs) are an increasingly popular alternative paradigm that can represent the heterogeneity of contact interactions with granular detail and agency of individual behavior.
However, conventional ABM frameworks are not differentiable and present challenges in scalability; due to which it is non-trivial to connect them to auxiliary data sources easily.
In this paper we introduce G RADABM which is a new scalable, fast and differentiable design for ABMs.
G RAD ABM runs simulations in few seconds on commodity hardware and enables fast forward and differentiable inverse simulations.
This makes it amenable to be merged with deep neural networks and seamlessly integrate heterogeneous data sources to help with calibration, forecasting and policy evaluation.
We demonstrate the efficacy of G RAD ABM via extensive experiments with real COVID-19 and influenza datasets.
We are optimistic this work will bring ABM and AI communities closer together.
1. Introduction
Mechanistic simulation models explicitly model mechanisms and causal connections to explore empirically observed variables.
They are an indispensable tool for epidemi- ology to explore the behavior of complex, dynamic infections and guide forecasting and decision making.
Here, compartmental ODE-based models are the dominant paradigm which use a system of differential equations to describe disease transmission by modeling the flow rate of individuals from one disease stage (compartment) to another.
Classical examples include the ‘so-called’ S-I-R (susceptible-infectedrecovered) model with 3 compartments [30, 23]; with some works which further stratify the population via separate compartmental models for each sub-population (eg: AgeSEIR [38], Geo-SEIR [9]).
These SIR-like models are simple to design, very fast to execute and amenable to gradientbased optimization of simulation parameters [41, 8].
Hence they have also been used in hybrid models (with deep neural networks) for forecasting influenza [24], Ebola [29] and also COVID pandemic [27].
This tractability is achieved through simplifying assumptions about homogeneous transmission and perfect mixing (each agent interacts with every other agent) in the population, which may often not be desirable.
Specifically, these SIR-like models do not represent the underlying contact graphs which prevents them from recreating heterogeneity of the interacting populations and modeling the adaptability of individual behavior, in the real world [10, 38].
Agent-based models (ABMs) have emerged as an alternate paradigm which can alleviate these concerns.
ABMs are discrete simulators which comprise a collection of agents which can act and interact within a computational world [32, 15, 39, 59, 48].
They can explicitly represent heterogeneity via the underlying contact networks and also model the adaptability of individual behavior to enable more realistic simulations.
However, they are conventionally slow, difficult to scale to large population sizes [40] and tough to calibrate with real-world data [26].
This is a challenge since simulation results (emergent behavior) can be highly sensitive to the scale of the input population and calibration of the input parameters.
In addition, incorporating novel sources of data that could inform calibration and other downstream tasks (e.g., forecasting) is often laborious and adds overhead complexity to the ABM (e.g., incorporating digital exposure data to ABMs [31]).
In this paper, we introduce G RADABM to alleviate these concerns and realize the potential of ABMs for practical decision making in epidemiology.
The key insight of G RAD ABM is to enable fast forward and differentiable inverse simulations.
We realise this through a differentiable sparse tensor-calculus based implementation.
G RAD ABM can simulate millions-size networked populations in few seconds on commodity hardware, run on both CPUs and GPUs and is end-to-end differentiable.
Conventional ABM frameworks follow an object-oriented (agent-centered) design where the agents are modeled as objects.
While conceptually appealing and extensible, this is often inefficient to represent inter-agent infection transmission (over contact networks) and within-agent disease progression as the size of agent populations and interactions scale.
Our framework G RAD ABM follows a networkcentric design [20] where agents are represented as tensors, their interaction networks as (sparse) adjacency matrices and a continuous relaxation of the stochastic disease transmission model is used to produce gradient estimates with automatic differentiation.
We use this scalable, fast and differentiable design of G RAD ABM to merge with deep neural networks (DNNs) and seamlessly integrate heterogeneous data sources to help with calibration, forecasting and policy evaluation.
Particularly, we use COVID-19 and Influenza modeling as examples to show the benefit of G RAD ABM for (time-space) personalized and robust forecasting; and evaluating efficacy of pharmaceutical interventions for decision making.
We believe that differentiable agent-based epidemiology can open new venues for epidemic understanding and aid practical decision making by learning from millions of agents in few seconds.
Our Contributions : (i)
We present G RAD ABM, a differentiable ABM that can simulate million-scale populations in a few seconds on commodity hardware and be merged with
DNNs for end-to-end learning.
The key idea of G RAD ABM is a general-purpose differentiable sparse tensor-calculus based implementation which we validate for epidemiology here.
(ii)
We demonstrate the utility of G RAD ABM for robust forecasting and analysis of COVID-19 and Influenza.
(iii)
We show the use of G RAD ABM in evaluating pharmaceutical interventions for policy decision making.
2. Background and Related Work
First, we briefly mention work in Agent-based Models (ABMs) and with relevance to epidemiology.
Then, we summarize relevant work on learning in ABMs by integrating with deep neural networks, methods to calibrate ABM parameters and scalability of simulation parameters.
We conclude a brief mention of a useful insight we leverage for our solution (permutation invariance).
We refer the reader to [45, 23, 30, 37] for a more extensive discussion.
ABMs define systems [32, 15] as collections of agents with can act and interact with each other within a computational world.
They have been used to model cells in a tumor micro-environment to diagnose cancers [39], humans in a physical environment to study economic policies [59], infectious diseases [48, 5] and avatars in a digital world to study misinformation [16].
ABMs in Epidemiology (EpiABM) are used to understand how disease spreads and evaluate efficacy of health interventions for mitigating them.
EpiABMs build simulations where agents host and transmit virus within their multiple contact networks (there is a within-host disease model e.g. SIR); and [40, 25] discuss methods and assumptions for designing such simulators.
Briefly, population contact networks are assumed to resemble a random graph with a specific degree distribution; and an epidemiological model [31] is used to describe betweenhost transmission and within-host progression of infection.
Recently, such EpiABMs are used to evaluate the benefit of delaying 2nd dose of the vaccine [48], deploying mobile apps for digital contact tracing [5], prioritizing test speed over accuracy [36] to contain the spread of COVID19.
The utility of such simulators for practical decision making depends upon their ability to recreate the population with great detail, integrate with real-world data streams and analyze the sensitivity of results.
Learning from data in
ABMs by integrating with DNNs is being explored across multiple fields.
The dominant method is to treat it as a multi-agent reinforcement learning task where the ABM defines the environment and a DNN is used to parameter- ize agent action policies.
Similar ideas has been used by [59] to learn equitable economic policies, [51] to analyze societal segregation dynamics, [42] to learn oil and gas macro strategies.
An alternate strategy has focused on using
DNNs to emulate agent-based models [11, 7, 43] by using the ABM to generate training dataset for DNNs.
Recently, [11] used an epidemiological ABM to generate synthetic datasets and trained DNN models to predict infectiousness for proactive contact tracing.
[43] used a similar strategy (of ABM as surrogate) to identify what simulation parameters to calibrate.
In contrast, we propose to make the ABM differentiable instead so that it can be trained end-to-end with the DNN modules using gradient-based optimization.
This makes it possible to integrate heterogeneous data to infer latent micro variables and improve forecasting power of ABMs through hybrid DNN-ABM pipelines.
This also has interesting implications for calibrating ABMs, as we discuss next.
Calibration of ABMs involves identifying appropriate values for micro variables and is essential to ensure reliability of results.
Generally, a hybrid strategy is used where some parameters are sourced from control trials conducted offline by clinical experts and then others are calibrated in-silico by achieving a goodness-of-fit between emergent ABM output and real-world macro data (e.g. number of deaths).
The conventional technique is to use grid search [18, 48] or beam search [55] by running several forward simulations and finding the best-fit.
This is often slow and tough to scale to many parameters.
Recent works is using optimization methods for faster and better calibration.
[43] uses an ML model to identify what parameters to calibrate but values are identified with random search.
[5] used compartmental (SEIR-like) models where parameters are learned by gradient-based optimization in the
SEIR-like model and are then used in the ABM.
However,
SEIR-like models have few tunable parameters which limits the extent and correctness of calibration with real-world data for the ABM.
In contrast, the differentiable design of
G RAD ABM allows calibrating parameters with gradientbased optimization inside the ABM itself.
This helps learn personalized parameters resulting in better (and robust) forecasting and policy decision making.
Such gradient-based calibration is iterative and hence requires fast forward simulations, as is reviewed next.
Scalability of ABMs is a key consideration since modeling granular population details is computationally expensive.
Conventional frameworks like
Mesa follow an object-oriented design and while ease to use (Python-API) are prohibitively slow to scale (require few hours for a single iteration).
There have been attempts to reduce this burden and simulate realistic scale with distributed
HPC systems and GPU-optimized implementations.
EpiFast [14] demonstrated simulating large contact networks in a few minutes on distributed systems.
However, the required compute for such executions is expensive and nontrivial to access.
OpenABM [31] presents an optimized
C++/Cuda API to build an epidemiological ABM for fast execution on commodity hardware but is tough to extend and is not amenable to gradient-based learning.
Motivated by recent work in molecular dynamics [50], G RAD ABM is implemented using highly optimized sparse tensor APIs of auto-diff packages (e.g.,: pytorch, JAX) to accelerate and scale simulations while preserving ease of use.
Invariances, Computation and ABMs Efficacy of DNN architectures results from overcoming the curse of dimensionality by leveraging the pre-defined invariances arising from the underlying low-dimensionality and structure of the physical world [17].
This is key to efficient computation on grids with CNNs (translation invariance) and graphs with
GNNs via neural message passing (permutation invariance).
We posit that while useful for learning DNNs, utilizing these invariances can make computation tractable in physical systems, where they exist naturally.
Specifically, we observe that epidemiological models [20, 31, 5, 48] also adhere to permutation invariance wherein the order of infectious interactions (pair-wise message passing) within a step does not matter when estimating the probability of infection from those interactions.
We use this to make G RAD ABM differentiable where the transmission model can be implemented as a variant of message passing parameterized with physical equations instead of neural networks.
More details in appendix.
3. Differentiable Agent-based Modeling for Epidemiology
The pipeline is summarized in Fig.
2, where the inner loop is comprised of the epidemiological model G RAD ABM and the outer loop of the calibration model C ALIB NN.
Training is executed over 4 stages: 1) generating epidemiological input parameters using a parameterized neural network (C ALIB NN), 2) epidemiological forward simulation with the generated parameters (inner loop), 3) computing a loss/error between output of the epidemiological model and ground truth data, and 4) compute the gradient of this loss using back-propagation through all the previous 3 stages and then update parameters of C ALIB NN using gradient descent.
In the rest of this section, we first describe the epidemiological model, followed by the calibration and training pipeline.
3.1. Epidemiological Model
Here, we describe the micro loop shown in Fig.
We consider a K step discrete-event simulation with a networked population of n stateful agents, where A t is the underlying contact network at given step t .
We build upon a standard epidemiological model used in several clinical papers [20, 31, 48], which is comprised of a Transmission
Model and a Progression Model.
These two components recursively update the state of all agents, which is described next.
Agent State is denoted by a 3-D tensor X tj for an agent j ∈ { 1 , . . . , n } at time step t ∈ { 0 , . .
. , K } .
X tj = { a j , d tj , e tj } where a j ∈ { 0 − 10 , 11 − 20 , 21 − 30 , . . .
, 71 − 80 , 80+ } is the age of the agent, which can take one of nine discrete values; d tj is the current disease stage of agent j and for example in COVID-19, it can take values in ∈ { S, E, I, R, M } where S, E, I, R, M denote susceptible, exposed, infected, recovered and dead (mortality) stages respectively, e tj ∈ {− 1 , . .
, t − 1 } denotes the time step of last exposure.
In some simulations, we augment the state using two additional controlled dynamic variables, which are: v tj ∈ { 0 , 1 , 2 } that denotes the current vaccination status and q tj ∈ { Y, N } that denotes the current quarantine status.
Transmission Model is a parameterized model that computes the probability of infection transmission as a result of an interaction between susceptible and infected (or exposed) agents.
Interaction Networks are the sites for contact between agents that transmit infection.
Following [48, 5], we separately model contact graphs within a county and represent agent interactions in multiple scenarios: household, random and work related interactions.
The graph is generated using demographic-stratified mobility information obtained from [49], evolves at every time step (different at each step).
Consider an interaction at time t , between susceptible and infected agents with states X ti =
[ a i , d ti , e ti ] X tj = [ a j , d tj , e tj ] , where d ti = E , e ti = − 1 , d tj =
I and d tj ∈ { 0 , . . .
t − 1 } respectively.
The rate of transmission at any step t primarily depends upon: (i) infectiousness of the pathogen at time t
( R t ), (ii) susceptibility of infectee i to transmission ( S i ), and (iii) transmissibility of the infector j
( T j ) and (iv) time since exposure for infector ( ∆ E tj =
t − e tj ).
The infectiousness varies over time, starting at zero when the agent is infected, peaks at an intermediate time and eventually tends to zero and is modeled using a gamma distribution [5].
The susceptibility of the infectee is age-stratified as described in [31].
The probability of transmission ( q ( . ; . ) )
from the interaction is represented as: q ( d ti , d tj ) =
1 − e − λ
( R,S i ,T j , ∆ E tj ) , where,
[Equation]
models the rate of transmission in the interaction,
ˆ I i is the expected number of interactions for agent i , q ( d ti , d tj ) is then aggregated over all interactions for any agent i and then the aggregated value is used to parameterize a Bernoulli distribution and sample a discrete transmission event.
A successful exposure then causes an update in X ti ( d ti and e tj ).
Progression Model Once successfully exposed (as explained in the Transmission Model) or when in an infected state, an agent j enters a hierarchy of disease stage progression which triggers subsequent changes in the state of the agent.
We follow a slightly modified SEIRM [57] progression model.
Susceptible ( S ) agents may become exposed ( E ) as explained in the Transmission model above, and exposed agents progress to infection ( I ) and eventually end in a stage of recovery or mortality ( R, M ).
The Progres- sion Model is parameterized by stage transition times and mortality rate.
The key difference between our model and the SEIRM-like models is that we track states of individual agents rather than just the aggregate numbers of agents in various disease stages.
Furthermore, we model the S → E exposure step, using interaction networks which take into account local agent interactions, which is different from the standard SEIRM-like models.
This model can be parameterized to study multiple infectious diseases, and through experiments we validate generality for COVID-19 and influenza, with possible changes in the type and number of disease stages.
Forward Simulation with G RAD ABM
We combine these components of the epidemiological model to simulate the disease dynamics over a horizon of K steps.
The input to the simulator are the time-dependent parameters which govern the transmission and progression models, denoted as θ t T =
[ R t , S, T ] , θ t P =
[ m, τ EI , τ IR , τ IM ] respectively, where t ∈ { 0 , . . .
, K } .
Here R t is the time-dependent disease reproduction number (measure of infectiousness),
S is the age-stratified susceptibility, T is the age-stratified infection transmissibility, m is the mortality rate (fraction of expired agents amongst all agents in the { R, M } disease stage),
τ EI , τ IR , and τ IM are the stage transition times.
At each step t , every agent i interacts with a set of neighbors { j : j ∈ N i } and may accumulate or transmit infection.
This disease stage evolution, at any step t , is given by:
[Equation]
where Update ( X ti , N i , ( X tj ) j ∈N i , θ tT , θ tP ) is = (cid:40)
Transmit ( X ti , N i , ( X tj )
j ∈N i , θ tT ) , if d ti = S,
Progress ( X ti , θ tP
[Equation]
The transmission model is
Transmit ( X ti , N i , ( X tj ) j ∈N i , θ T ) where
R t , S i , T j are obtained from θ T and e tj from X jt is written as:
[Equation]
where, λ denotes the rate of transmission from a single interaction, (cid:83) is an aggregation function which accumulates transmission over multiple interactions and ˆ Q =
Bernoulli(1 − e q ( t ) ) maps this transmission to a probability of infection.
Following the invariance of infection transmission noted in section 2 (more in appendix), (cid:83) is the summation ( (cid:80) ) function in our implementation.
The progression model Progress ( X ti , θ P ) updates the disease stage from E → I or from I → { R, M } at times determined by the stage transition time parameters.
Finally, the aggregate cumulative deaths is determined as:
[Equation]
where m ∈ θ
[Equation]
3.2. Calibration and Forecasting
Training In order to use G RAD ABM for real-world applications, we need to be able to calibrate its parameters ( θ tT , θ tP ) in such a way that the aggregate predicted quantities from the model such as cumulative deaths match the observed values from real-world data.
This is done using the calibration and training setup (the outer loop in Fig. 2).
This is a W step loop, with each step w comprising the following four stages as shown in the figure.
Stage 1 : Generating epidemiological parameters ( θ tT , θ tP ) wt ∈{ 0 ,...,K } using C ALIB NN ( f ) with neural network parameters ϕ w at step w and input as the auxiliary data ( D ) from sources as described in section 4.
For ease of exposition, we drop the t ∈ { 0 , . . .
, K } subscript and refer to this input as ( θ tT , θ tP ) w .
Then, we have ( θ t T , θ tP ) w = f
( D ; ϕ w ) Stage 2 Performing a K -step simulation using G RAD ABM (the inner loop in Fig. 2) with ( θ tT , θ tP ) w from C ALIB NN in the previous stage.
Stage 3 :
The output of G RAD ABM is aggregated to compute cumulative deaths.
This macro value is compared with corresponding ground truth values and the following loss is computed:
L (ˆ y w , y ; ( θ tT , θ tP ) w ) =
MSE (ˆ y w , y ) , where y is the ground truth data representing cumulative deaths, ˆ y w is the cumulative deaths predicted by the epidemiological model at the the w th training step with input parameters ( θ tT , θ tP ) w ) and MSE denotes the mean-squared error function.
Stage 4 : Once this loss is computed, we use automatic differentiation to estimate gradient for minimizing the loss.
The gradient is computed over the K -steps of the epidemiological model
G RAD ABM (which recursively calls Update ( · ; · ) ) and calibration model C ALIB NN which is parameterized with a neural network.
In practice, auto-diff packages such as Pytorch implement back-propagation through such a ‘for loop’ by unrolling this loop in the internal computational graph that the library tracks to enable end-to-end differentiation.
The only requirement is that all steps within the loop be differentiable, which we establish by showing that the components of our epidemiological model - Transmission model and Progression model are differentiable.
The Transmission Model is differentiable as the functions comprising it, which are: λ is a smooth function, (cid:83) is a permutation invariant function ( (cid:80) ) and is linear.
In order to make the sampling from ˆ Q differentiable, we utilize a continuous relaxation with Gumbel-Softmax reparametrization [33].
Next, the Progression Model is differentiable as it is effectively a linear deterministic model, with the parameters being the transition times.
Thus, the entire K step epidemiological model is differentiable.
The final step of this gradient computation involves backpropagation through the parameters of the C ALIB NN neural network ( ϕ ) as well.
Once the gradient is thus determined, the parameters of C ALIB NN are then updated as given below:
[Equation]
where α is the learning rate, ( θ tT , θ tP ) w = f
( D ; ϕ w ) and ˆ y w is computed by calling Update K times with same parameters ( θ tT , θ tP ) w .
This training is continued till an acceptable level of aggregate prediction error is achieved.
This process of learning via gradient descent requires multiple runs of the K -step epidemiological model.
Hence, for this training to be feasible, our epidemiological model should be fast and also scalable to enable real-world relevant simulations.
Thus, we note that the two key properties of
G RAD ABM which make such iterative gradient-based optimization tractable are: i) fast simulations at realistic scalewith sparse tensor calculus; ii) differentiable simulations for gradient-based optimization.
Inference Once trained, the hybrid pipeline with C ALIB NN and G RAD ABM can be used to forecast in unseen scenarios as well as aid in policy decision making.
These unseen scenarios may include predicting future the evolution of multiple infections for a county/region, or it may involve making predictions for a new county/region for which historical data for calibration is not available, or it may involve making predictions even when the data (inputs to C ALIB NN) for a county/region is noisy.
In each of these scenarios, the data available for the county/region is fed to C ALIB NN, which then outputs time-space personalized parameters to be used in the G RAD ABM.
G RAD ABM is then run for K steps using these parameters and the simulation results are used to make the desired predictions.
4. Experimental Setup
We conduct experiments for COVID-19 and influenza on multiple counties of the state of Massachusetts, USA, and we learn personalized parameters for each county.
To predict county-specific parameters, we design C ALIB NN (Sec. 3.2) which ingests diverse data sources and is optimized using gradients from the prediction error of G RAD ABM.
Next, we describe the data sources used to build the simulator and training pipeline.
• Heterogeneous Data Sources G RADABM allow bringing together data across multiple granularities (individual, census-block, county, city, country) and modalities (mobility, policy, clinical, digital surveys) due to flexibility of agent-based simulators.
We collected important publicly available 1 datasets from a variety of trusted sources.
Here we introduce our data sources in the order of our training pipeline in Fig.
i)
We use county-level time- varying features that are ingested by C ALIB NN.
Specifically, for COVID-
we use 5 signals including insurance claims data, online symptoms surveys from Facebook, and line-list data.
For flu, we use the 14 signals from the Google symptom dataset.
Flu data is reported at state-level, so we utilize the same data for all counties.
ii) G RAD ABM leverages individual-level data to set agents’ states created based on demographic information from the US Census [2].
iii)
Our interaction networks are created using census block level mobility data from Safegraph.
Specifically, we use MCMC sampling to generate synthetic contact graphs inside each census-block using demographic information (e.g., age).
iv)
To set some of the parameters of our transmission and progression models, we utilize clinical data from reliable sources such as CDC [3] and clinical papers [31, 28].
The target variable for COVID-19 is COVID-associated mortality, while in flu we use influenza-like-illness (ILI) counts, which is collected by the CDC.
The ground truth data for the target variables is obtained from JHU and CDC.
Additional details are in the appendix.
• Baselines We compare
G RAD ABM against multiple popular approaches for calibrating an ABM - Vanilla-ABM: Following [31, 48], it uses parameters set with expert and/or authoritative information with randomized search.
For COVID-19, we get Rand case-fatality rate from [13, 4] and R 0 for flu from [21].
PC-ABM:
Following [5], it uses R 0 and initial infections learned using a compartmental model (SEIRM and SIRS for
COVID-19 and flu, respectively).
We also perform studies to compare different configurations of G RAD ABM.
G RADABM (w/o TL, w/o C ALIB NN): Removes C ALIB NN and its data inputs, and instead uses scalar learnable parameters.
G RAD ABM (w/o TL): trains with only data from one county and one ABM; instead we use transfer learning in
G RAD ABM to learn a single C ALIB NN that can predict for all counties.
• Metrics For forecasting error metrics, we adopt standard metrics used by previous work for evaluating epidemic predictions [54, 6].
Specifically, we use normal deviation (ND), root mean squared error (RMSE) and mean absolute error (MAE) which are detailed in our appendix.
Following CDC forecasting guidelines [12, 22], we make weekly predictions for 1 to 4 weeks ahead in the future.
Our evaluation for both diseases is of at leastmonths in 10 counties.
Specific counties and evaluation periods are in the appendix.
• Implementation Details Data preparation: We process the time series data into padded sequences which are normalized per region.
Architecture details : C ALIB NN is an encoder-decoder neural architecture based on GRUs [19] and self-attention [56].
ABM parameters ( θ P , θ T ) :
Our time-varying ABM parameters have the flexibility to be changed on each time step of the simulator (daily), but we found this to be challenging for learning.
Thus, we opted to use the same set of ABM parameters for every 7 days (i.e., we learn a new set of parameters for every week).
Details can be found in the appendix.
Our code and data are publicly available 2 .
5. Results
First , we show that G RAD ABM integrates with CalibNN to improve ABM forecasting performance on multiple infectious diseases.
Next , we show that G RAD ABM is robust to noisy data, can seamlessly scale to simulate large populations and be used to evaluate policy interventions for decision making.
5.1. Forecasting multiple infectious diseases
Results presented in Table 1 shows that G RAD ABM consistently outperforms typical ABM calibration methods when forecasting infection spread of COVID-19 and Influenza, across all metrics (rows 1-3 in Table 1) 3 .
For instance,
G RAD ABM improves ND to 0.97 from 8.75 in VanillaABM and 2.21 in PC-ABM on COVID-19; and to 0.41 from 0.57 in Vanilla-ABM and 0.59 in PC-ABM on Influenza.
Specifically, we observe that G RAD ABM achieves a consistent gain of 8x to 12x over Vanilla-ABM on all metrics in
COVID-19.
This performance benefit is attributed to three key features of G RAD ABM which allow us to learn personalized time-varying parameters: gradient-based calibration, integration of heterogeneous data sources via C ALIB NN, and transfer learning from joint county training.
We also conduct ablation studies to analyse these features (rows 3in Table 1) and note some observations.
First , using C ALIB NN is significant to performance gain as it allows: i) integration of heterogeneous data sources for learning parameters and ii) rich and flexible optimization by coupling the ABM with a neural network.
Second , training counties independently result in marginal fall in performance (e.g.:
ND goes from 0.97 to 1.26 in COVID) when compared to joint training which utilizes transfer learning.
This indicates that G RAD ABM can automatically take advantage of what is being learned in one county can be useful for another county, but, in contrast to other methods that directly copy the learnt values, our method personalizes appropriately.
Finally, we note that transfer learning does not help in Influenza like in the COVID-19 case.
As noted in Sec. 4, this results from shared state-level influenza data for all counties and hence C ALIB NN cannot personalize parameters across counties.
5.2. Robustness, Scalability and Decision Making
Here, we study the effectiveness of deploying G RAD ABM to real-world scenarios by focusing on three key consid-
We performed the unpaired t- test ( overruns
G RAD gains w.r calibration methods are statistically significant erations: i) scalability to a large real-world population, ii) robustness to noisy data, iii) utility for practical decision making by evaluating policy interventions.
G RAD ABM runs fast forward simulations on large populations Figure 3 shows that the run-time for G RAD ABM scales linearly with number of interactions (and agents) in the population and executes very quickly (even as the adjacency matrix scales quadratically).
For instance, G RADABM executes a simulation with 800,000 agents (5 million interactions) over 133 steps in 4 seconds on a GPU (andseconds on CPU).
This is roughly 300x faster than the equivalent Mesa implementation.
This performance efficiency can be attributed to the design of a sparse tensor-calculus based implementation of G RAD ABM.
We note that the fast execution of forward simulations provides the following benefits for practical use: i) makes it possible to conduct sensitivity analyses of results through multiple runs with different hyperparameters, ii) essential for integrating with
DNN pipelines for parameter calibration with requires a simulation runs in each optimization step.
G RAD ABM is robust to noisy data Figure 4 shows that
G RAD ABM outperforms baselines even when trained with noise data.
Data quality issues are very common in public health data [35, 47] and we evaluate performance of
G RAD ABM in such scenarios.
To study this, we introduce
Gaussian noise to the ground truth target with µ noise = 0 and several scales of standard deviation s noise .
For each county, we first obtain the standard deviation of the target and multiply it by the factor λ to get the s noise .
We utilize 4 different values of λ .
Specifically, we observe that even with large amount of noise ( λ = 4 ) our error on all metrics is lower than the other baselines that were trained without noise.
This is attributed to our differentiable design which allows integrating with C ALIB NN such that: i) heterogeneous data sources are used to learn simulation parameters, ii) simulation parameters are not represented as scalar variables but the outputs of a neural network.
GradABM helps in evaluating policy interventions:
What if we delay the second dose of COVID-19 vaccine?
Here, we investigate the utility of GradABM for policy de-
[Table]
G RAD ABM achieves lower forecasting error than it λ the
This is achieved due to design which allows encoding simulation parameters cision making and focus on the question of delaying the second dose.
We follow the clinical setup as in [48] and study sensitivity of the decision to protection of first dose ( 50% , 60% , 70% , 80% ) and calibration mechanism (VanillaABM vs G RAD ABM).
We compare cumulative deaths for
Franklin County in MA under the two prioritization schedules — S1: Prioritized second dose, S2: Delayed second dose.
The sensitivity of this prioritization decision is shown in Fig.
5, where the Y-axis shows the ratio of cumulative deaths under strategy S1 to S2.
Whenever this ratio is below 1 (referenced by the green line in the figure), S1 is recommended over S2 and vice-versa.
It is evident from our experiment that at vaccine first dose protection of 60%,
G RAD ABM and Vanilla-ABM provide different recommendations, where the main difference is the calibration approach.
Since, G RAD ABM learns personalized calibration parameters which provide better forecasts, it also helps in better policy decision making.
More details in the appendix.
6. Discussion
We introduce G RAD ABM, a differentiable ABM that can simulate million-scale populations in a few seconds on commodity hardware and be merged with DNNs for end-to-end learning.
The key idea of G RAD ABM is a general-purpose differentiable sparse tensor-calculus based implementation which we validate for epidemiology here.
We demonstrate the utility of G RAD ABM to learn personalized time-varying parameters using heterogeneous data sources to better inform forecasting and policy decision making.
Future work could explore other benefits like the incorporation of multiple hierarchies of data (also known as macro-micro predictions [58]).
In our experiments, we found that even with sparse optimization, G RAD ABM utilizes a high amount of
CPU/GPU memory (especially in the backward computation) which could be an issue for simulations of 6+ months.
Addressing this limitation could be an interesting future direction.
Our ABM used a linear deterministic model for the disease progression.
Future work could explore how to incorporate more complex and stochastic disease progression models (e.g., [48]).
Acknowledgements
This work was supported in part by the NSF (Expeditions CCF-1918770, CAREER IIS-2028586, RAPID IIS2027862, Medium IIS-1955883, Medium IIS-2106961,
CCF-2115126), CDC MInD program, ORNL, faculty research award from Facebook and funds/computing resources from Georgia Tech and MIT.
AC was supported by Adobe
Data Science Research Award.
A. Extended related work Permutation invariance: key insights and assumptions
We begin by highlighting some key assumptions that are followed by epidemiologists and then leverage them to motivate the specific implementation of G RAD ABM.
The key insight of G RAD ABM is to utilize differentiable sparse-tensor computation in interaction networks and symmetries of the transmission model for fast and endto-end differentiable simulations.
Furthermore, we note that: i) clinical data shows that the natural infection takes multiple days to incubate, ii) CDC policy decisions are guided by 4-week ahead forecasting and analysis.
These considerations motivate the following standard assumptions in epidemiological simulation work: Assumptions : Several works in epidemiological ABMs [20, 31, 5, 48] ratify the following assumptions:
A1 : the granularity of each step of the simulator is one day, A2 : infectious interactions ”do not collide” i.e. on any given day, an agent can not also infect other agents the day it is exposed, A3 : an agent cannot be reinfected while already infected, A4 : an agent cannot be infected when it is in the RM , i.e. recovered or mortality stage, i.e., there are no reinfections (this assumption can be relaxed based on the nature of the disease being modelled)
A5 : cumulative transmission over multiple interactions
can cause infection, A6 : infection does not accumulate over days.
It is important to note that real-world interaction networks routinely have low degree and high centrality and are modeled as small-world graphs as in [5, 48].
This allows use of sparse tensors to represent and process interaction networks; and scale to realistic populations (larger n ) which improves performance when outputs are scale dependent.
Also, Assumptions A2-A5 together imply that the order of interactions within a step does not matter.
This enables us to invoke permutation invariance in interactions and makes the transmission model differentiable.
B. C ALIB NN details
For C ALIB NN, we encode the feature time series until t N by passing it through a Gated Recurrent Unit (GRU) [19] to obtain a condensed representation for each time step:
[Equation]
where h t is the hidden state of the GRU for time step t .
To capture long-term relations and prevent over-emphasis on last terms of sequence we use self-attention layer:
[Equation]
where Self-Atten [56] involves passing the embeddings into linear layers to extract meaningful similarities before normalizing the similarities using Softmax.
Then, we use the attention weights to combine the latent representations and obtain a single embedding representing the time series of features from t 0 to t N : u
[Equation]
[Equation]
therefore, our vector u t 0 : t N summarizes the input sequence and represents the context to be given to the decoder.
Next, we use another GRU as a decoder, which takes the context u t 0 : t N and a positional encoding τ k that informs the model of how many days in the future we want to predict.
We simply use τ k to be a float between 0 and 1.
[Equation]
We want to unroll our decoder to predict not only the values of the simulator, but also the future ones.
For this, we pass the decoder output through a feedforward network FFN ( o t ) .
We found the optimization can be challenging if we directly use the output of the neural network as parameters of our simulator.
Therefore, we bound the output in a similar manner as [8]: θ t = θ
L
+ ( θ U − θ L ) · σ ( FFN ( o t )) , where θ L and θ U are the lower and upper bounds of θ t for all t , and σ is the Sigmoid function.
C. More details on experimental setup
Code and data are publicly available 4 .
Please refer to the
README in code repository for more details.
Computational setup.
All experiments were conducted using a 4 Xeon E7-4850 CPU with 512GB of 1066 Mhz main memory and 4 GPUs Tesla V100 DGXS 32GB.
Our method implemented in PyTorch trains on GPU in 15 mins.
Inference is takes only a few seconds.
Real-time forecasting.
We follow the literature on evaluating epidemic forecasting methodologies [52, 34, 6] and use the real-time forecasting setup.
We simulate real-time forecasting by making models train only using data available until each of the prediction weeks and make predictions for 1 to 4 weeks ahead in the future.
Data revisions in public health data are large and may affect evaluation and conclusions [35, 22], therefore, we utilize fully revised data following previous papers on methodological advances [6, 46].
Evaluation.
As evaluation is on weekly predictions, we aggregate daily predictions to obtain weekly predictions (sum for COVID-19 and average for flu).
As stated in the main paper, we opted to use the same set of ABM parameters for every 7 days (i.e., we learn a new set of parameters for every week).
This is a reasonable assumption as it has been often found that the parameters of mechanistic models do not change much from one day to the other [44].
In our evaluation, we work with the following counties in
Massachusetts: 25001, 25003, 25005, 25009, 25011, 25013, 25015, 25021, 25023, 25027.
The specific evaluation period is determined with epidemic weeks 5 which is the standard in CDC’s epidemic prediction initiatives 6 .
For COVID19 these are 202014, 202016, 202018, 202020, 202022, 202024, 202026, 202028, 202030.
For flu, we evaluate in epidemic weeks 201746, 201748, 201750, 201752, 201802, 201804, 201806, 201808, 201810.
As noted before, this means that for each epidemic week, we make 4 predictions in the future.
Metrics.
Let ˆ y w,τ be the prediction week w for τ -weeks ahead in the future, y w,τ be the corresponding ground truth value, and e w,τ
= ˆ y w,τ − y w,k .
In this paper τ takes fromto T = 4 .
Then, we define our metrics as follows:
ND =
(cid:80) | e w,τ | / (cid:80) | y w,τ | , RMSE = (cid:80) e 2 w,τ /WT and MAE =
(cid:80) | e w,τ | /WT .
Implementation details.
• Data preparation : All our time series data that enter to C ALIB NN are padded with a minimum sequence length of 20 for COVID and 5 for influenza.
We also normalize each features with mean 0 and variance 1.
To inform C ALIB NN of the county to predict, we use a one hot encoding for counties.
• Architecture details :
In C ALIB NN, the encoder is a 2-layer bidirectional GRU and decoder is a 1-layer bidirectional GRU, both with hidden size of 32.
Output layer has two linear layers of size 32x16x D with ReLU activation function and D is ABM parameter dimensions dimensions:
D = 3 for COVIDand D = 2 for flu.
• Hyperparameters : We found a learning rate of 10 − 3 , Xavier initialization, and the Adam optimization algorithm work best.
• ABM parameters ( θ P , θ T ) :
For COVID-19, we have three parameters: R 0 , mortality rate, and initial infections percentage.
These are bounded with θ L = [1 . 0 , 0 . 001 , 0 . 01] and θ U =
For flu, we have two parameters: R 0 and initial infections percentage.
These are bounded with θ L = [1 . 05 , 0 . 1] and θ U =
The initial infections percentage is the percentage of the population that is infected at time step t = 0 of the simulation.
• ABM clinical parameters : To set some of the parameters of our transmission and progression models, we utilize clinical data from reliable sources.
Specifically, for COVID-19 we use age-stratified susceptibility and parameters of a scaled gamma distribution to represent infectiousness as a function of time as per [31, 28].
For influenza, we set those parameters based on CDC flu facts [3].
Target variables.
The target variable for COVID-19 is
COVID-associated mortality, while in flu we use influenzalike-illness (ILI) counts, which is collected by the CDC.
ILI measures the percentage of healthcare seekers who exhibit influenza-like-illness symptoms, defined as ”fever (temperature of 100°F/37.8°C or greater) and a cough and/or a sore throat without a known cause other than influenza” [1].
The ground truth data for the target variables is obtained from JHU CSSE COVID-19 data repository and ILI from CDC influenza dashboard.
Details on baseline implementation.
Vanilla-ABM:
As noted in the main paper, R 0 and case-fatality rate are obtained from authoritative sources.
To set the initial infections percentage for this baseline, we set it to the mean value of the search range.
PC-ABM: We present details on each of the ODE models we used for this baseline.
(COVID-19) SEIRM [57]: The SEIRM model consists of five compartments: Susceptible ( S ), Exposed ( E ), Infected ( I ), Recovered ( R ), and Mortality ( M ).
It is parameterized by four variables Ω = { β, α, γ, µ } , where β is the infectivity rate, 1 /α is the mean latent period for the disease, 1 /γ is the mean infectious period, and µ is the mortality rate.
The basic reproductive number R
= β/
( γ + µ ) .
[Equation]
(Flu) SIRS [53]: This model consists of three compartments:
Susceptible ( S t ), Infected ( I t ), and Recovered ( R t ).
It is parameterized by three variables Ω = { β, D, L } , where β is the infectivity rate, D is the mean duration of immunity, and L is the mean duration of the immunity period.
The basic reproductive number R 0 = βD .
[Equation]
D. More details on evaluating policy interventions
We reproduce the experimental setup in [48].
For interventions, we simulate standard COVID-19 vaccination versus delayed second dose vaccination prioritizing the first dose.
Sensitivity analyses included first dose vaccine efficacy of 50%, 60%, 70%, 80%, and 90% after daypost-vaccination with a vaccination rate of 0.3% population per day; assuming the vaccine prevents only symptoms but not asymptomatic spread (that is, non-sterilizing vaccine).
We measure cumulative COVID-19 mortality, cumulative SARS-CoV-2 infections, and cumulative hospital admissions due to COVID-19 over 74 days.
We explicitly modeled the confirmation of infections with polymerase chain reaction testing and quarantining of known infected agents
[Table]
with imperfect compliance over time.
To simulate a natural pattern of infection at the point vaccinations begin, we started our simulation with 10 agents infected and ran the simulation for 20 days before starting vaccinations, which corresponds to a cumulative infection rate of 1%, similar to the one in the US, UK, and most of Europe when vaccinations were started and corresponding to the time horizon used in our analysis.
In both our vaccination strategies, we started administering vaccines on the basis of age, starting with people over 75, then those over 65, and so on.
We compare different vaccine regimens and prioritization schedules.
Consider the following six hypothetical individuals with their age and dose eligibility on any given step of the simulations: i) Adam - first dose eligible 78 yr old; ii) Betty - second dose eligible 78 yr old; iii)
Charlie - first dose eligible 68 year old; iv) David - second dose eligible 68 year old; v) Eleanor - first dose eligible 40 year old; vi)
Frank - second dose eligible 40 year old.
Table 2 shows the order in which these individuals are prioritized for vaccine administration under the strategies we compare.