The ongoing coronavirus disease 2019 (COVID-19) global pandemic poses immense threats not only to public health, but also to the stability of healthcare infrastructures and economies around the world. In an attempt to slow down the outbreak, many countries have imposed unprecedented lockdown and social distancing measures that were eventually proven to be effective in downscaling the volume of COVID-19 fatalities [1, 2] -however, the instigation of such measures gave rise to various "What if?" questions that have become central to public discourse, e.g., what the number of deaths would have been had the government acted earlier? Would lifting the lockdown cause a second wave of infections? Data-driven models that predict the effects of (lockdown) policies on the trajectory of COVID-19 fatalities are crucial for answering these questions; they are also necessary for informing governments and decision-makers on what policy directions to follow at different stages of the pandemic [3, 4] .
In response to calls for data-driven policy-informing models, the academic community has produced various models for forecasting COVID-19 fatalities [5, 6, 7 ] -these models have been used by public health organizations, such as the WHO and the CDC [22] , to advise local authorities on what policy directions to follow. (For instance, the IHME model developed by the University of Washington [5] has been repeatedly cited by the White House in press conferences on COVID-19. 2 ) Despite a plethora of epidemiological models for COVID-19 spread, models that can predict the effect of "counterfactual" lockdown policies on COVID-19 fatalities -which is crucial for conducting scenario analyses and policy planning -are still lacking. Moreover, despite the prospects of machine learning (ML) as a key tool for developing such models [8] , existing models have been based primarily on epidemiological approaches, with ML techniques being used merely for parameter optimization [9, 10] .
In this paper, we develop one of the first ML models for learning lockdown policy effects on COVID-19 fatalities in a global context -i.e., we treat each country hit by the pandemic as a distinct data point, and exploit the variations in policies followed by different countries to learn country-specific policy effects. We characterize each country with a feature vector that comprises economic, social, demographic, environmental and public health indicators; "counterfactual" fatality curves under a hypothetical policy for a given country are thus inferred from "factual" fatality curves of countries "similar" in features wherein this policy was actually implemented. (For example, our model would use data from Sweden to predict what the number of fatalities in Norway would have been under a "herd immunity" policy [11] .) We envision that our model would be used by policy-makers to conduct scenario analyses by assessing the volume of COVID-19 fatalities under different possible policiesthis is especially timely as governments seek policies for gradual lifting of lockdown measures [12] . How does our model work? How is it different? A high-level pictorial illustration of our model is provided in Fig. 1 . We follow a Bayesian approach to jointly model COVID-19 fatalities across many countries through the following hierarchical, two-layer Gaussian process (GP) prior:
Upper-layer GP: f U = R 0 (Country features, Policy indicators), Lower-layer GP: f L = COVID-19 fatalities over time at a given R 0 .
The lower-layer GP f L (.) models the COVID-19 fatality curve over time within each country -it uses a compartmental SEIR (Susceptible, Exposed, Infectious, Recovered) model [13] as its prior mean function, parameterized with the reproduction number R 0 which characterizes the rate by which the pandemic spreads, i.e., how many people (on average) does each new patient infect [14] . The upperlayer GP f U (.) models the R 0 parameter as a function of country features and policy indicators, allowing the model to share data and parameters across different countries that experimented with different policies (e.g., different lockdown timing). Our model captures uncertainty in both the inferred parameters and the predicted COVID-19 fatalities through the posterior variance across the two layers.
Because of the rarity of pandemics in our modern condition, little related work has been done within the machine learning community to address this problem. In what follows, we provide a brief overview of previous works -an elaborate discussion of the related literature is provided in Appendix A. Previous works prior to the current pandemic have been primarily focused on learning contagion (diffusion) processes on networks, e.g., [15, 16] ; unfortunately, these models do not apply to the pandemic as information on social network structures underlying disease spread is hard to obtain. In response to the COVID-19 pandemic, two strands of research work have emerged: (a) methods for devising optimal control (lockdown) policies to contain disease spread [17, 18, 19] , and (b) models for forecasting the disease spread and its expected fatalities [5, 6, 7, 9, 10, 20, 21] . Our paper belongs to category (b) in that the developed model is trained to forecast the future number of fatalities. But unlike existing models in category (b), our model predicts fatalities under arbitrary lockdown policy choices, hence it can be used for research in category (a) to derive optimal policies.
The most prominent models in category (b) -developed by various academic institutions -have been recognized by the CDC and used to issue national forecasts of COVID-19 fatalities within the United States [22] . Most of these models (e.g., the "UCLA" model [9] , the "MIT" model [21] and the IHME model [5] ) are SEIR models fit for individual countries, with ideas from "machine learning" being used only to optimize the basic SEIR parameters via gradient descent. Our model differs from these models in that it: (1) jointly models fatalities across all of the 170 countries affected by the pandemic, (2) incorporates individual country features to learn how disease dynamics and policy effects vary across countries, (3) enables evaluating future projections under alternative policy scenarios, and (4) combines both mechanistic SEIR models and data-driven machine learning models. Hierarchical modeling has been previously used in the "Imperial" model developed in [20] (which was fit across 11 European countries). This model can be viewed as a special case of ours as it assumes policy effects to be fixed across all countries in its upper layer with no machine learning components to model heterogeneity, and its lower layer uses a serial interval distribution to predict short-term deaths only. A detailed table of comparison between all models is provided in Appendix A.
Setup and background. Let Y i (t) ∈ N ∪ {0} be the number of reported COVID-19 deaths in a given geographical area i on the t th day since the beginning of the outbreak. Throughout this paper, we assume that a geographical area corresponds to a country, and consider a set of N countries hit by the pandemic. 3 Each country i is characterized by a feature vector X i ∈ R d comprising its economic, social, demographic, environmental and public health indicators. Because the number of confirmed daily COVID-19 cases depends greatly on the testing capabilities and strategies within each country [23] , we use the reported daily deaths as a more concrete indicator of disease spread.
We model the COVID-19 lockdown policy of each country i over days t to t through the sequence
where p i (t) ∈ {0, 1} K is a K-dimensional policy indicator variable, reflecting whether country i applies each of K different COVID-19 containment and social distancing measures (e.g., school closure, sheltering in place, travel bans, etc) on day t. To build our model, we used data for N = 170 countries, each with d = 35 features and K = 9 policy indicators (See Section 4 and Appendix C for details).
Our key objective is to address the following question: "How would different lockdown policies affect future COVID-19 fatalities?" To this end, we build a data-driven model to learn lockdown policy effects on COVID-19 deaths. We envision this model to be trained and applied in a global context, where data capturing variations in lockdown policies applied in different countries enables us to learn counterfactual policy effects within each country. Predictions made by our model can be used to conduct scenario analyses that inform government policy.
To build our model, we consider a data set D N,t for N countries covering a period of t days, i.e.,
For each country i, we forecast the trajectory of the number of COVID-19 deaths within a future time horizon of T days under a given lockdown policy, i.e.,
In addition to point predictions, we also estimate a sequence of uncertainty intervals C i [t : t + T ] that cover the true trajectory of future COVID-19 fatalities, i.e., Y i [t : t + T ], with high probability.
In this Section, we develop a Bayesian model for the effects of lockdown policies on future COVID-19 fatalities. We provide our model's specification in Section 3.1, and then develop the corresponding model training and parameter learning algorithms in Section 3.2.
Two-layer GP prior. We assume a (hierarchical) two-layer Gaussian process (GP) [24, 25] prior on the COVID-19 fatality curves {Y i (t)} i for the N countries under consideration as follows:
is a transformation function (described later this Section), α and θ i are the hyper-parameter sets for the upper-and lower-layer (associated with country i) GPs, respectively. The model in (4) assumes that the fatality curve Y i (t) for each country i is a noisy draw from a GP prior with a mean function D θi and kernel function K θi . The prior on lower-layer parameters θ i is specified via the upper-layer GP -the function f U maps country i's features X i and its adopted policy p i to a parameter θ i . The mean function m α for the upper-layer GP is assumed to be a constant, and a Matérn kernel is selected for both layers. Because the upper-layer GP shares its parameters α across all countries, countries with "similar" features and policies will share similar parameters, and hence similar fatality profiles. The graphical model for (4) is provided in Fig. 2 (a) . Compartmental priors. With little data available at the early stages of the pandemic, our prior knowledge on the expected COVID-19 fatality curves are limited to mathematical models for the spread of infectious diseases. Thus, we model the prior mean function D θi (t) in (4) trough a compartmental SEIR (Susceptible, Exposed, Infectious, Recovered) model [26] , with parameters
where β i is the contact rate (average number of contacts per person per unit time), σ is the virus's incubation rate, γ i is the patient recovery rate, and µ i is the mortality rate. With the exception of σ, all of the other SEIR parameters are country-specific since they depend on population demographics (e.g., contact rates depend on social mobility and mortality rates depend on age distribution [27] ). The parameter set θ i for the lower-layer GP comprises the SEIR parameters, in addition to the length-scale L of the kernel function in (4), i.e., the lower-layer parameter set is θ i = ( L , β i (t), σ, γ i , µ i ).
In an SEIR model, individuals within the population of country i are assigned to 5 compartments: Susceptible S i (t), Exposed E i (t), Infectious I i (t), Recovered R i (t) or Deceased D i (t) -each individual progresses through these compartments as dictated by the following differential equations [28, 29] :
with dR i /dt = γ i I i − µ i R i , and dD i /dt = µ i I i , where n i is the population size of country i. In (6), we assume that each country's contact rate β i (t) is time-dependent in order to account for the impact of (lockdown) policy changes over time. The death compartment D i obtained by solving (6) is set as the prior mean function D θi (t) for the lower-layer GP in (4) . Our choice of the SEIR model as the prior fatality curve is motivated by the dynamics of the SARS-CoV-2 virus that causes COVID-19; since the SEIR model captures incubation rates [29] , it more accurately represents SARS-CoV-2 dynamics -compared to other SIR variants -which are known to exhibit significant incubation periods [30] . Note that while we select an SEIR model as our prior, other compartmental models can be used as well. A block diagram describing the compartmental prior mean function is given in Fig. 2 (b) .
The Bayesian nature of our model enables combining the rigorous mechanistic foundation of compartmental models with the data-driven (nonparameteric) nature of GPs. That is, at the early stages of the pandemic, the early fatality forecasts would be dominated by the prior mean function D θi (t) derived from the SEIR prior -as more data on COVID-19 fatalities are collected over time, the GP posterior will refine the SEIR forecasts based on observed patterns in the data. Because of its hybrid nature, we call our model a compartmental Gaussian process (CGP).
Modeling policy effects. Governments decide on (lockdown) policies over time in order to control the basic reproduction number (R 0 ) of COVID-19, i.e., the rate by which the pandemic spreads [14] . The reproduction number within country i as represented by the SEIR prior in (6) is given by [28] :
where R 0,i (t) is the reproduction number in country i at time t. An R 0,i (t) of a value less than 1 means that country i's fatality curve Y i (t) is "flattening" -the policy p i (t) aims at driving R 0,i (t) below 1 by imposing social distancing measures that would minimize the contact rate β i (t) in (7).
Because different countries with comparable features apply different policies, the upper-layer GP function f U (X, p) in (4) will learn counterfactual fatality curves for each country under alternative policies that have been tried in other countries (See Fig 2(c) ). For instance, we can learn the fatality curves for Scandinavian countries (such as Norway and Denmark) under a hypothetically less restrictive lockdown policy using data from Sweden, which adopted less stringent policy measures [11] .
The lower-layer parameters are linked to the upper layer as follows: f U is specified as a multi-output GP; each output corresponding to a distinct parameter in θ i = ( L , β i (t), σ, γ i , µ i ). The transformation function v(.) in (4) is an identity map for all outputs corresponding to parameters ( L , σ, γ i , µ i ), except for the contact rate parameter β i (t) where it performs the following mapping:
whereβ is a reference value for the contact rate obtained using early data from Wuhan province [31] .
Scenario analysis via posterior inference. What is the expected number of COVID-19 fatalities in country i given a future lockdown policy scenario P i [t : t + T ]? To answer this question, we compute the posterior distribution of Y i [t : t + T ] given the data D N,t and the policy scenario P i [t : t + T ], i.e.,
Upper layer , where conditioning on D N,t in the left hand side was omitted for brevity. Here, the lower-and upperlayer posteriors are computed analytically using the closed-form expression for GP posteriors [24] , and the integral is evaluated via a Monte Carlo approximation. Point estimates and uncertainty intervals are obtained by evaluating the mean and variance of the resulting distribution.
Accurate posterior inferences of Y i [t : t+T ] requires training the CGP model by optimizing the parameter α of the upper-layer GP using the observed data D N,t by maximizing the model's log-likelihood:
with α * = arg max α L(D N,t | α). Because the integral in (9) is intractable, we resort to a variational inference approach for optimizing the model's likelihood [32, 33, 34, 35] . That is, to train our model, we minimize the Evidence lower bound (ELBO) on (9) given by:
where Q(.) is the variational distribution with parameters φ, and conditioning on X i and P i [1 : t] is suppressed for notational brevity. We choose a normal distribution for Q(.) -this renders analytic evaluation of the ELBO objective and its gradients possible. We use stochastic gradient descent via ADAM algorithm to optimize the ELBO objective [36] , and update the lower-layer parameters in each gradient iteration by solving the SEIR differential equations in (6) use Euler's method [37] . A pseudo-code for our training algorithm is provided in Appendix B. 
In this Section, we use our CGP model to forecast COVID-19 fatalities in various countries around the world, taking into account the lockdown policies within these countries. We compare the projections of our model with other models listed by the Center for Disease Control (CDC) (Section 4.2), and show how our model can be used to analyze counterfactual policy scenarios (Section 4.3).
Data description. We collated the data set D N,
for N = 170 countries affected by the COVID-19 pandemic using three data sources: (1) published World Bank reports 4 were used to extract a set of d = 35 features X i for each country, (2) the COVID-19 CSSE data repository at Johns Hopkins University [38] was used to extract each country's fatality time-series Y i , and (3) the Oxford COVID-19 Government Response Tracker (OxCGRT) -curated by the Blavatnik School of Government at Oxford University [39] -was used to extract K = 9 policy indicators P i for each country over time. Our data set covered the period spanning from January 22, 2020 to May 8, 2020.
Each country's features included a comprehensive set of economic, demographic, environmental, social and health indicators (e.g., population density, prevalence of obesity, air transport frequency, median age, prevalence of hand-washing facilities, health expenditure, etc). Policy indicators included: information on school and workplace closure, public events' cancellation, travel restrictions, public transport closure, etc. A complete list of variables included in our model is provided in Appendix C.
Implementation. We implemented our CGP model using Pyro [40] , a auniversal probabilistic programming language in Python supported by a PyTorch backend. The variational inference algorithm in Section 3.2 was implemented using ADAM with 1000 iterations and a learning rate of 0.01. Further details on hyper-parameter tuning is provided in Appendix C. Projections from all baselines involved in our comparisons were obtained from the official CDC website [22] .
Baselines. We considered the most prominent baseline models listed by the CDC [22]: the "UCLA" model [9] , the "MIT-DELPHI" model [21] , the Los Alamos National Laboratory "LANL" model [ the "Imperial" model [20] , the IHME model [5] , the "YYG" model [10] which was found to be the best performing CDC-listed model in recent weeks [42] , in addition to the CDC-ensemble, which issues a national-level forecast for the United States by combining the predictions of 16 individual models [22] . The UCLA, MIT-DELPHI, LANL, YYG and IHME models are all variants of the SEIR model -the IHME has been the most influential of these models, having been often cited during White House press briefings on COVID-19 modeling efforts [44] . In addition to these models, we consider a vanilla SEIR model and a standard Gompertz curve fitting approach as baseline models for fatality curves [41] .
Forecasting Accuracy for the United States. We start off by comparing the accuracy of fatality projections by our model with the baselines for the United States. In Table 1 , we evaluate the accuracy of the 7-day and 14-day projections issued at three stages of the pandemic: before the peak of infections (March 28), in the midst of the peak (April 11) [38] , and in the "plateauing" stage (April 25). Accuracy was evaluated by computing the difference between true and predicted cumulative deaths throughout the forecasting horizon, i.e., T k=1 (Yi(t + k) − Yi(t + k)) for T = 7 and T = 14. For each forecast, only data preceding the forecast date was used for training our model. The results are summarized in Table 1 , with the best performing model in each forecast highlighted via an underlined bold font.
As we can see in Table 1 , our model consistently outperforms the CDC-ensemble and the prominent IHME models at all forecasts; most notably the CGP predictive accuracy is an order of magnitude better in the current plateau stage of the pandemic. Our model also performs consistently well across all forecasts -its performance is comparable to the best model for each forecast. Note that, as mentioned earlier in Section 1, the Imperial model provides a competitive performance on short-term predictions because it is tailored to inferences based on short-term transmission dynamics, but it falls short when it comes to long-term forecasts, which can be crucial for anticipating the timing of the peak.
The benefits of joint (global) modeling across multiple countries are demonstrated through an ablated version of our model that uses United States (US) data only for training. The global version of our model consistently achieves a significant performance gains compared to the version that uses US data only; similar patterns were found in other countries (See Appendix C). This shows the value of using machine learning to capture country-specific disease spread parameters based on country features.
Accuracy of global projections. In Table 2 , we compare the forecasting accuracy of our model with the CDC-listed baselines that offer projections for countries other than the US (IHME, Imperial and YYG). We evaluated the performance of all baselines in 10 countries that were significantly affected by the pandemic in Europe, Asia, Africa, and the Americas. (Note that ours is the only model that covers all countries and spans all continents.) Since these countries were at different stages of the pandemic at any given time, we evaluated the 7-day and 14-day forecasts on April 25, 2020, when all countries have had a significant number of infections. As we can see in Table 2 , our model outperforms the baselines in almost all countries on both forecasting horizons. Central to the accuracy of our model is its hierarchical GP structure which shares data across countries based on their " feature similarity", enabling accurate predictions even when little data is available for individual countries.
Our model does not only learn the fatality curves within each country, but also the country-specific effects of lockdown. Since the same policy would naturally yield different effects in different countries, this is a major source of performance gain for our model compared to others (e.g., Imperial model) which assume fixed policy effects. To demonstrate the heterogeneity of policy effects, Fig. 3 (a) shows the inferred R 0 in the United Kingdom (UK) against the stringency of the policy applied at different points of time. We quantify the policy stringency through the stringency index defined in [39] , which collapses all policy indicators into a single number between 0 and 100 -a higher index corresponds to a stricter policy. The regression slope in Fig. 3 (a) reflects the policy effects, i.e., how much a unit increase in policy stringency reduces R 0 . As we can see, the lockdown effects learned by our model differs among countries based on their features; here we compare policy effects for the UK and Italy as an exemplar. We specify the country features most relevant to policy effects in Appendix C.
As we have seen in Section 4.2, our model performs favorably compared to existing models in terms of predicting COVID-19 fatalities, but how can it be used to inform decision-making? Here, we use our model to analyze the lockdown policy of the UK in the early stages of the pandemic, which sparked controversies within the British society [45, 46, 47], and provide projections for future COVID-19 fatalities under the gradual re-opening plan announced by the UK government [48].
In Fig. 3 (b) we display the predictions and counterfactual inferences of our model from the policymaker's perspective in the period spanning from May 8 until the end of May, when a lockdown lifting policy was being planned for. At this point of time, the policy-maker can be presented with counterfactual fatality curves of what would have happened before May 8 had lockdown been implemented 1 week earlier (blue curve) or 1 week later (red curve). Our model predicts that 13,827 lives would have been saved with an earlier lockdown, and 22,405 more deaths would have occured under a later one. Looking into the future, our model predicts that under the current UK government re-opening plan (black curve), daily deaths would stabilize around 200. Maintaining the current lockdown (blue curve) would lead daily deaths to fall under 100 in Augusts, which would save 6,215 more lives compared to the current re-opening plan. On the other hand, a complete re-opening in June (red curve) would result in a second peak in August-September although there is large uncertainty about the volume of the second peak. Similar analyses for other countries are provided in Appendix C.
We envision our model to be used by governments to make measured decisions that might impact the lives of millions of people all over the world. We hope that our model exemplifies the importance of machine learning-based decision-making for public health in the post-coronavirus world.