Adaptive County Level COVID-19 Forecast Models: Analysis and Improvement

In this section, we adapt the TDEFSI model proposed in Wang et al. (2020) to the COVID-19 forecasting problem. The the output of each step in the TDEFSI model has a dimension for each county, allowing the network to learn spatial relationships, without requiring knowledge of specific relationships beforehand. Let $\mathbf{y}=(y_{1},y_{2},\ldots,y_{T},\ldots)$ denote the sequence of the natural logs of daily nationwide incidence, where $y_{t}\in\mathbb{R}$. Let $\mathbf{y}^{C}=(y_{1}^{C},y_{2}^{C},\ldots,y_{T}^{C},\ldots)$ denote the sequence of daily incidence for a particular county $C\in\mathcal{D}$, where $\lvert\mathcal{D}\lvert=K$ is the number of counties. Let $\mathbf{y}^{\prime}_{t}=(y^{C}_{t}\forall C\in\mathcal{D})$. The objective is defined as predicting both nationwide and county-level incidence at day $t$, where $t=T+1$, denoted as $\mathbf{z}_{t}=(y_{t},\mathbf{y}^{\prime}_{t})$. The TDEFSI loss function is given as:

where $\phi(\widehat{\mathbf{z}}_{t})$ and $\delta(\widehat{\mathbf{z}}_{t})$ are activity regularizers added to the outputs for spatial and non-negative consistency respectively, and $\lambda$ and $\mu$ are penalty parameters. Since RNNs are not affected by the temporal scale of its input sequence, (1) does not need to be adjusted to account for the fact that COVID-19 data is updated daily as opposed to the weekly ILI data. Wang et al. (2020) only considered ILI data for individual states and their counties, whereas for COVID-19 we consider all counties in the nation. Min-max normalization is used on the county-level data so that $\mathbf{y}^{\prime}\in[0,1]$. Whereas Wang et al. (2020) use the sum of $\mathbf{y}_{t}^{\prime}$ for $y_{t}$, here $y_{t}$ is the natural log of the sum of $\mathbf{y}_{t}^{\prime}$. This is done to normalize $y_{t}$ relative to $\mathbf{y}_{t}^{\prime}$ while preserving the dependency of $y_{t}$ on $\mathbf{y}_{t}^{\prime}$. Without this normalization, $y_{t}$ would be significantly larger than any $y_{t}^{C}\in\mathbf{y}_{t}^{\prime}$ for a large $K$, and would dominate the gradients used to optimize (1). This modification to $y_{t}$ requires that the regularization term $\phi(\widehat{\mathbf{z}}_{t})$ be modified to be

In Section 4.1, we use simulated data from the county level adapted SEIR model Section 3.2 to supplement available real data and analyze the adapted TDEFSI model. We provide more details on the architecture of the model in Supplementary Material.

The model is trained with features from the 1/22/20 to 5/2/20 period, targets from 1/23/20 to 5/16/20 with validation features from 5/3/20 to 5/4/20 and targets from 5/3/20 to 5/17/20. The trained model is tested by making a single 14 day forecast over 5/18/20 to 5/31/20 for all counties. The county level temporal patterns are supplemented with standard normalized county level $latitude$, $longitude$, $population$, and $population\ density$ Killeen et al. (2020), as well as $\log(population)$ and $\log(population\ density)$ prior to the application of the county distributed dense layer. We represent $n_{TF}$ features in the LSTM backbone and use 3 layers of $n_{D}$ units in the county distributed branch taking $\sigma$ to be the ReLU activation. Both training and inference use batch size of 1, with batches taken in sequential order, and previous encoding cell state shared with the next batch. Each batch uses the vector of the previous end of day’s recorded confirmed cases to forecast a target matrix of size $(n_{C},n_{F})$. Dropout is applied to the targets at a rate of $0.25$ during training to mitigate overfitting.

We perform a capacity study 2 by randomly sampling 20 configurations from $\{n_{D},n_{TF}\in\mathbb{Z}|\quad 10\leq n_{D}\geq 50,\ 1\leq n_{TF}\geq 5\}$ with equal probability. While there is no clear trend in forecast mean squared error; we observe that a network with $n_{TF}=5,\ n_{D}=20$ is reproducible and lightweight. Anecdotally, training proceeds best on the edge of instability, encouraging annealing of the network towards more stable configurations. Different size networks learn in different regimes, and capacity must be tuned in conjunction with dropout rate. We note that sharing cell state between sequential batches is essential for network convergence to an effective forecasting configuration. Using both the time distributed and county distributed dense layers is substantially more effective than using only a time distributed, or a time distributed, county distributed layer. Additionally, predictions from all 20 model configurations are ensembled by averaging and scored over the forecast horizon. The ensembled predictions perform better than any single model’s predictions, even in the presence of poor scoring component models. Due to the lightweight nature of of the architecture, ensembled predictions are easily obtained, mitigating the instability inherent in dropout and common in RNN architectures.

The best and worst performing model forecasts are collected in Figure 4. Figure 6 shows the second through fourth worst and the forth through second best forecasts made using CLEIR-Net model. Additional forecasts are provided in the Supplementary Material. Best predicted by CLEIR-Net are large, connected, counties. In particular, the top five best scoring state forecasts were the New York and four adjacent states: Pennsylvania, Massachusetts, New Jersey, and Connecticut. Worst predicted are small or isolated counties, such as Hawaii, and Montana. Also, poorly predicted are states with peculiar dynamics relative to the rest of the nation. For example, the five best predicted states exhibit the early stages of curve flattening, while Arkansas, and Wisconsin, the fourth and fifth worst predictions have an accelerating number of cases. Figure 5 illustrates this via a map showing the MSE for the contiguous states.

A summary of the forecast performance and size of all the models explored in this paper is presented in Table 2. All methods were used to forecast for 5/18/20 to 5/31/20. We observe that CLEIR-Net requires significantly low numbers of parameter to train versus the adapted TDFESI model.

The SEIR model assigns one of four status to a proportion of the total population in a space with constant population. It then builds a set of ordinary differential equations (ODEs) which describe rate of change of each status population. We use a version of the SEIR model that neglects birth and death effects and assumes a lack of any vaccination though it assumes recovered patients gain permanent resistance. The model is identical to one found in Y. Fang [2020] which is recreated below.

We denote $p_{i}$ to be the total population of county $i$, and $p_{k,i}$ to be the population with status $k$ in county $i$, where $k\in\{S,\ E,\ I,\ R\}$. $x_{k,i}:=\displaystyle p_{k,i}\big{/}p_{i}$ is the proportion of the population with status $k$. The parameters, $\beta$, $\sigma$, and $\gamma$, along with initial conditions govern the dynamics of the epidemic. This model does not account for any interaction between the given county’s population and all other U.S. counties.

A conservation law for the population of the county requires the rate of population accumulation to equal the rate of population inflow minus the rate of population outflow plus the rate of population generation. The population generation term is given by our SEIR equations. We assume no change in the county population requiring the population flow and generation terms to add to zero. Additionally, intuitively we know the SEIR equations which govern the transition of status should not induce population accumulation or depreciation. This requires our population flow terms to balance.

We decided on population mixing terms that would account for the change in a status’ proportion of the population without changing the total county population. To do this, the inflow term for given $k$ population in county $i$ is a foreign county’s population flow into county $i$ multiplied by the foreign county’s proportion of population with status $k$ summed over all foreign counties. Symmetrically, the outflow term is the same population flow variable, but this time multiplied by proportion of population with status $k$ in county $i$ summed over each foreign county. The multiplication of the proportions, which must sum to one, and the flow terms, which are equivalent across ODEs for a county pair, ensure the county population does not change.

Let $f_{i,j}$ be the total population flow of people from county $i$ to $j$ such that $f_{i,j}=f_{j,i}$ and $f_{i,i}=0$. These conditions guarantee each county has no net change in population. The derivatives of $p_{k,i}$, $k\in\{S,E,I,R\}$ can be written as

We can write this system of equations more concisely if we vectorize each equation. This can be done by creating a flow matrix, $\mathbf{F}$. The elements are made up of the aforementioned $f_{i,j}$ terms. Again, using the constraints $f_{i,j}=f_{j,i}$ and $f_{i,i}=0$ we create a symmetric flow matrix with a diagonal of zeros. Here, the $i$th column represents population outflows from county $i$ and the $i$th row represents population inflows to county $i$. Lastly, we need our matrix to satisfy our constant population constraint. In other words, we want the diagonals of the matrix to be the negative sum of the other elements in their row. This can be done by subtracting the diagonal of one transpose $\mathbf{F}=\left\{f_{ij}\right\}_{ij}$, written as $\mathbf{F}^{bal}=\mathbf{F}-diag(\mathbf{1}^{T}\mathbf{F})$. This matrix exhibits the property $\sum\mathbf{F}^{bal}\times\mathbf{1}=0$, guaranteeing the conservation laws are satisfied.

The network is trained over a sequence of inputs where the target is to forecast the number of recorded infections for each county for each of the next $n_{F}$ days in the future, given the days elapsed since the first national recorded infection, $t$, the vector of the current day’s county level recorded infections, $I(t)$, and, crucially, the time invariant county level features, latitude, longitude, population, population density, log of population and population density. A batch size of one is used where each sample consists of a single given day of the pandemic and the corresponding targets are the days in it’s forecast horizon. That is, a sequence of overlapping samples, taken in sequential order. Additionally, each batch’s encoder and remember cells use the cell and hidden states from the previous batch’s encoder and remember cells as their initial states and provide their output states to the next sample’s corresponding cells. Sharing states between batches allows the LSTM cell to effectively utilize its memory property over the entire sequence but limit gradient exposure to short sequences. An initial state of zero is used for the first sample in the sequence.

After training, learned weights are used to forecast future infection trajectories. Since the model learns to use states from past batches, effective inference requires the model be initialized, post learning and prior to forecasting, by making sequential predictions over the entire training data.

The network’s basic definitions, inputs, targets, and outputs for a single batch are summarized here.

Definitions

• •

  $n_{F}:=n_{Forecast}$ Number of days in the forecast horizon

• •

  $n_{TF}:=n_{Time\ Features}$ Number of time features to model in the LSTM backbone

• •

  $n_{C}:=n_{Counties}$ Number of counties in the prediction space

• •

  $n_{X}:=n_{CountyLevelFeatures}$ Number of additional county level features used by the county distributed dense layer

• •

  $n_{D}:=n_{Dense\ Units}$ Number of units in the county distributed dense layer

Inputs:

• •

  $t$: Time elapsed in days since first nationally recorded infection

• •

  $I(t)$: Current day’s county level infections

• •

  $X$: Time invariant county level features

• •

  $C^{(E)}_{t-1}$: State of the previous batch’s encoder cell

• •

  $h^{(E)}_{t-1}$: State of the current batch’s encoder cell

• •

  $C^{(R)}_{t-1}$: State of the previous batch’s remember cell

• •

  $h^{(R)}_{t-1}$: State of the current batch’s remember cell

Targets:

The goal of the network is to forecast the end of day recorded infections for each US county for each day in the forecast horizon, therefore our targets are the following.

• •

  $I(t+i)\quad\forall\ i\in 1...n_{F}$: Actual end of day recorded county level infections for the $i^{th}$ day in the forecast horizon.

Outputs: The corresponding model predicted infections are then denoted by the following.

• •

  $\hat{I}(t+i)\quad\forall\ i\in 1...n_{F}$: Predicted end of day county level infections for the $i^{th}$ day in the forecast horizon.

Additionally, each batch share’s the hidden and cell states of it’s encoder and remember LSTMs with the following batch.

• •

  $C^{(E)}_{t}$: Cell state of the current batch’s encoder cell

• •

  $h^{(E)}_{t-1}$: Hidden state of the current batch’s encoder cell

• •

  $C^{(R)}_{t}$: Cell state of the current batch’s remember cell

• •

  $h^{(R)}_{t-1}$: Hidden state of the current batch’s remember cell

The LSTM backbone of the model consists of an encoder cell and remember cell which translate the national total recorded infections and days elapsed into a low dimensional time varying pattern, and a repeatable forecast cell which propagates the time pattern into the future.

Encoder Cell:

• •

  Purpose: The encoder cell, given the previous batch’s encoder state, transforms the total national recorded infections, $\sum I(t)$, and days elapsed since the first nationally recorded infection into a low dimensional time varying pattern.

• •

  Inputs: $t$, $\sum I(t)$, $C^{(E)}_{t-1}$, $h^{(E)}_{t-1}$

• •

  Outputs: $C^{(E)}_{t}$, $h^{(E)}_{t}$

• •

  Parameters: $\theta_{E}:=\theta_{Encode}$, $\theta_{E}\in(\mathbb{R}^{2\times n_{TF}},\ \mathbb{R}^{n_{TF}\times n_{TF}})$

Remember Cell:

• •

  Purpose: The remember cell, given the context from it’s previous state, and the output representation from the encoder, learns to remember the current state of the LSTM backbone used by the repeated forecast cell.

• •

  Inputs: $C^{(R)}_{t-1}$, $h^{(R)}_{t-1}$, $h^{(E)}_{t}$

• •

  Outputs: $C^{(R)}_{t}$, $h^{(R)}_{t}$

• •

  Parameters: $\theta_{R}:=\theta_{Remember}$,$\theta_{R}\in(\mathbb{R}^{n_{TF}\times n_{TF}},\ \mathbb{R}^{n_{TF}\times n_{TF}})$,

Forecast Cell:

• •

  Purpose: The forecast cell is similar to the remember cell in that its function is to model the underlying low dimensional time pattern. However, while the remember cell incorporates information both from the encoder cell and its own previous state, the forecast cell uses only the LSTM backbone state. For the first day in the forecast, the cell propagates the current state from the remember cell into the future. To achieve multiple days in the forecast horizon, the forecast cell is repeated, with each consecutive cell after the first taking the state from the previous cell as input.

• •

  Inputs: $C^{(F)}_{t+i-1}$, $h^{(F)}_{t+i-1}$ when $i>1$, $C^{(R)}_{t}$, $h^{(R)}_{t}$ otherwise

• •

  Outputs: $C^{(F)}_{t+i}$, $h^{(F)}_{t+i}$

• •

  Parameters: $\theta_{F}:=\theta_{Forecast}$, $\theta_{F}\in\mathbb{R}^{n_{TF}\times n_{TF}}$,

The time and county distributed layers are used to enforce a hierarchical framework of shared factors. The time distributed layer uses the same parameters across all time steps to learn a time pattern underlying each county, given the low dimensional patterns from the backbone. Then, given the underlying time pattern for each county, the county distributed layer predicts the day to day change in infections using the same set of explanatory factors shared by each county.

Time Distributed Layer:

• •

  Purpose: If we consider the low dimensional time pattern learned by the LSTM backbone to be a mixed signal representing the underlying national dynamics, then the goal of the time distributed linear layer is to separate the national time pattern into its county level components, providing a county level time dependent condition on which the county distributed layer is applied.

• •

  Inputs: $h^{(F)}_{t+i}$

• •

  Outputs: $h^{(C)}_{t+i}$

• •

  Parameters: $\theta_{TD}:=\theta_{Time\ Distributed}$, $\theta_{TD}\in\mathbb{R}^{2\times n_{C}}$

County Distributed Layer:

• •

  Purpose: Given descriptive factors comparable across counties, latitude, longitude, population, population density, and log population and population density, and the county level time feature learned by the time distributed layer, the county distributed layer predicts the county level daily changes in infections for all counties, $\Delta I(t)$.

• •

  Inputs: $h^{(C)}_{t+i}$, $X$

• •

  Outputs: $\Delta I(t+i)$

• •

  Parameters: $\theta_{CD}:=\theta_{County\ Distributed}$, $\theta_{CD}\in(\mathbb{R}^{(n_{X}+2)\times n_{D}}$,  $\mathbb{R}^{(n_{D}+1)\times n_{D}},\ \mathbb{R}^{(n_{D}+1)\times 1}$) (For the input, hidden, and output layers of the County distributed dense layer.)

Final Prediction:

• •

  Purpose: The predicted changes in infection, $\Delta I(t+i-1)$, are added to the previous day’s predicted county level infections, $\hat{I}(t+i-1)$, to determine the current day’s predicted infections, $\hat{I}(t+i)$.

• •

  Inputs: $\Delta I(t+i-1)$, and $\hat{I}(t+i-1)$ when $i>1$, or $I(t)$ otherwise

• •

  Outputs: $\hat{I}(t+i)$

• •

  Parameters None

Network parameters are optimized by minimizing the weighted mean squared error (MSE) between the forecast and target matrices with NAdam with a learning rate of 0.001. Each sample is weighted by a factor of $w_{i,j}=(\log(Population_{j}+1)\log(i+1))^{-1}$, where $i$ is the target number of days ahead in the forecast, and $j$ is the county. Training is stopped when the validation loss is unimproved for 30 epochs at which point weights from the best performing epoch are used for inference. Random dropout is applied to the targets during training at a rate of 0.25 to mitigate overfitting to any particular county’s patterns. Light L1 and L2 regularization of 0.00005 are applied to the kernel and bias weights of all dense layers and to all recurrent weights to improve conditioning of the gradient. Forecasting is performed using learned parameters by resetting cell states and making predictions with a batch size of one in sequential order over all available days in the input range; ensuring the memory property is active during inference. Predictions from the last available day are taken as the forecast.
Prior to training, 4 counties from New York are removed from the data. Bronx, Kings, Queens, and Richmond counties have no recorded cases in the JHU data and are believed to be aggregated into the New York City totals.

Table 3 compares the forecast performance of the CLEIR-Net architecture both with the county distributed layer (Variant II) and without (Variant I). In both cases the time distributed layer is used to transform the national time pattern into the county level time patterns. Variant I uses the time distributed layer to directly predict the county level daily recorded infection changes while Variant II uses the county distributed layer to make further refinements to the output of the time distributed layer to predict the the county level daily recorded infection changes. In both cases, predicted changes in infections are added to the previous days predicted infections for each day forecast into the future and the cost is computed between predicted and recorded infections over the entire forecast. We train both variants using both mean squared logarithmic error and mean squared error to forecast a 7 day horizon. Variant I slightly outperforms the naïve no-change benchmark by all measures. When trained using the mean squared error objective, Variant II significantly outperforms the mean squared error of Variant I and the naïve no-change benchmark. This yields a more meaningful predicted national change in cases but at the expense of the mean squared logarithmic error, reducing accuracy in counties with lower recorded infections.

The architecture of the model used in the TDEFSI experiments in Section 4.1 is adapted specifically from the TDEFSI-LONLY model. This consists of $k$ LSTM layers, each with a latent dimension of $H^{(i)}$, followed by a dense layer with $H$ outputs, and a final dense layer with $K+1$ outputs. The input sequence to the network is $\mathbf{y}$, and the output sequence is $\widehat{\mathbf{z}}$. All TDEFSI-LONLY experiments used the parameters $k=2$, $H^{(i)}=128$, $H=256$, $\lambda=0.01$, $\mu=0.0001$ when relevant, which Wang et al. [2020] found to be optimal for their data.

Through determining the dependency between confirmed case of counties we can predict which counties are likely to be accurately forecasted by CLEIR-Net. The model, seeking county level and national level trends, fits itself to counties which are more inter-connected with other counties since their influence on each other induces overall trends. Our approach is to determine a non-linear dependency between counties by calculating the mutual information (MI), between counties to determine the connection between counties.

To estimate the MI between two counties, $C$ and $C^{\prime}$, denoted $\rho(C,C^{\prime}):=\widehat{MI}(C,C^{\prime})$, we use the hash-based method, proposed in Noshad et al. [2018]. However, it is prohibitively expensive to compute the MI between every county pair. Furthermore, intuitively there is lower dependency between counties that are far apart. Therefore, we compute MI between a county and its neighbors and average over all neighbors. We use the US Census Bureau’s county adjacency data to determine neighbors. Our approach is described in Algorithm 1.

The dependency values and error for the 48 contiguous states are shown in Figures 7 and 8, respectively. We observe the correlation between the dependency values and error of counties after training. We show that the inverse correlation between dependency and error is significantly high.

Note that by calculating the county dependency we can get a quick estimate of which counties won’t be accurately forcasted by the CLEIR-Net model before having to train it. This information is crucial to have as soon as possible because if we want to delegate resources based on the results of the model we have to know and disclose the limitations so that they can be planned around ahead of time. There are further potential applications for the dependency values as well. Notable to this work is potentially estimating flow rates between counties in the spatial mixing extension to SEIR outlined in Section 1.

Figure 9 shows the CLEIR-Net forecasts for all states and the District of Columbia, except those already shown in Section 4.2, Figure 6.