A Bayesian approach to quantifying uncertainties and improving generalizability in traffic prediction models

Feed-forward neural network architectures are not explicitly designed to handle sequential data. A class of DL approaches, recurrent neural network (RNN), uses a feedback mechanism where the output from a previous time step is fed as an input to the current time step such that information from the past can propagate into future states. This feedback mechanism preserves the temporal correlation and makes it suitable to capture the temporal evolution of traffic parameters. However, RNNs are incapable of handling the long-term dependencies in temporal data due to the vanishing gradient problem (Hochreiter 1998). Long short-term memory (LSTM) (Hochreiter and Schmidhuber 1997b), a type of RNN, consists of memory cells in its hidden layers and several gating mechanisms, which control information flow within a cell state (or, memory) to selectively preserve long-term information.

The objective is to update the cell, $C_{t}$, over time using the input $x_{t}$ and the previous time step’s hidden state, $h_{t-1}$. This process involves several key operations. First, a forget gate, $f_{t}$, selectively filters information from the past. Then, an input gate, $i_{t}$, regulates the amount of information from the candidate memory cell, $\tilde{C}_{t}$, that should be incorporated into the current cell state, $C_{t}$. Finally, an output gate, $o_{t}$, governs the update of the hidden state, $h_{t}$. See Figure 1. The computations are represented as follows:

The outputs from the forget gate, $f_{t}$, input gate, $i_{t}$, and output gate, $o_{t}$ are computed as shown below:

Here, $\sigma$ and $\tanh$ represent non-linear activation functions, while $W_{f}$, $W_{i}$, $W_{o}$, and $W_{c}$ denote weight matrices corresponding to the forget gate, input gate, output gate, and candidate memory cell, respectively. Similarly, $b_{f}$, $b_{i}$, $b_{o}$, and $b_{c}$ represent the corresponding bias vectors.

Bayesian DL captures epistemic or aleatoric uncertainty by fitting probability distributions over the model parameters or model outputs, respectively. Epistemic uncertainty is modeled by placing a prior distribution over a model’s weights, and then calculating the posterior distribution of the weights given the training data, i.e. $P(w|\mathcal{D})$. During inference, the predictive probability of $\hat{y}$ for a test data $\hat{x}$ is computed by $P(\hat{y}|\hat{x})=\mathbb{E}_{P(w|\mathcal{D})}P(\hat{y}|\hat{x},w)$. However, solving this is intractable. Popular methods use variational inference (Blei, Kucukelbir, and McAuliffe 2017) to obtain the posterior distribution, however these models are computationally expensive. One such method is ‘Backpropagation by Bayes’ (Blundell et al. 2015) which approximates the posterior distribution $P(w|\mathcal{D})$ with a simple distribution $q(w|\theta)$ parameterized by $\theta$ by minimizing their Kullback- Leibler (KL) divergence:

In this study, we use dropout – a widely used regularization technique to perform the Bayesian inference by estimating the posterior distribution of weights of a NN (Gal and Ghahramani 2016a; Kendall and Gal 2017). Dropout is effective in preventing overfitting by randomly deactivating a fraction of neurons during each training iteration. This encourages the NN to become more robust and less dependent on specific neurons, thus mitigating overfitting to the training data.

To achieve Bayesian inference, we incorporate dropout at each layer of the NN during training. Additionally, we apply dropout during inference on the test data, enabling us to sample from the approximate posterior distribution. This approach, often referred to as Monte Carlo dropout, allows us to estimate uncertainty and obtain a more comprehensive understanding of model performance. In this approach, $q(\mathbf{w}|\theta)$ is approximated as a mixture of two Gaussians with small variances and the mean of one of the Gaussians is fixed at zero. In a regression setting, the epistemic uncertainty is captured by the predictive variance of the model outputs. To the contrary, the aleatoric uncertainty is estimated by fitting a Gaussian distribution over the model output to capture the noise in data. The model is trained to predict the mean, $\hat{y}$ and variance, $\sigma^{2}$. This results in the objective function being a likelihood minimization as given below.

where $D$ is the cardinality of data, $y$ and $\hat{y}$ are the true and predicted values respectively, and $\sigma^{2}$ is the prediction variance. For numerical stability, the model is trained to predict the natural logarithm of variance i.e., $s_{i}:=\log\hat{\sigma}_{i}^{2}$. Consequently, the objective function becomes,

The total predictive uncertainty for $T$ stochastic forward passes of the model can be approximated as shown in Equation 6, where at each pass, a fixed fraction of neurons in each layer are randomly dropped.

A stacked LSTM model with Bayesian inference by dropout is used for traffic flow prediction and uncertainty quantification in this study. Our model consists of three LSTM layers with 20, 20 and 10 units, followed by four dense layers with 10, 10, 6 and 2 units respectively, with LeakyReLU activation function (Maas et al. 2013) for dense layers respectively. In each layer, a fraction of hidden units are stochastically dropped during model training and inference by adding dropout. For each model, dropout rates of 2, 5 and 10 percent are used to evaluate its sensitivity on model accuracy. Specifically for the LSTM, we use recurrent dropout (Semeniuta, Severyn, and Barth 2016) which applies dropout to the cell update vector $\tilde{C}_{t}$ (see Equation 7) instead of dropping gate inputs (Gal and Ghahramani 2016b) or cell state (Moon et al. 2015), which adversely affects its long-term memory.

where $d(\boldsymbol{\cdot})$ refers to the dropout mask.

The models are trained and validated on 60% and 15% of the dataset, while the remaining 25% is reserved for evaluating model performance. During the training process, we minimize the loss function (shown in Equation 5) for 100 epochs and choose the model with the lowest validation error. We use Adadelta (Zeiler 2012) as the optimizer with the learning rate of 0.10, $\rho$ value of 0.95, and epsilon of 1e-7 to train the models. Once the models are trained, the stochastic forward pass is repeated for a finite number of runs (50 in our case) to obtain uncertainty estimates.

Normalization methods usually improve the process of NN training, particularly ensuring smoother gradients, faster training, and better generalization accuracy. The gradients of the weights in any layer are highly correlated to the outputs from the previous layer. Thus, changes in the output of one layer causes significant changes in summed inputs to the following layer. Batch normalization (Ioffe and Szegedy 2015), a technique that controls the mean and variance of inputs across mini-batches, has been shown to improve training efficiency for NNs. However, applying batch normalization to LSTM is more challenging, due the presence of a temporal dimension, where the hidden state of the network is dependent on the previous time steps. Applying batch normalization to the hidden state could result in the loss of temporal dependencies, which can negatively impact the model’s performance.

Alternatively, we investigate the use of spectral normalization as an alternative approach to enhance the generalizability of models. We also compare its performance with another normalization technique called layer normalization. A brief overview of both normalization techniques are presented below:

Transfer learning involves ’transferring’ knowledge from a source domain to perform a similar task in another domain (Pan and Yang 2009). Transfer learning strategy eliminates the requirement to evaluate a new problem from scratch, rather the network is trained using data from the source domain and suitably adapted to new settings. This significantly reduces the learning time and amount of transfer-domain data required for training. Usually the discriminating layer in the NN is retrained using a small amount of data from the transfer domain, if available. In our case, we retrain the dense layers of the NN, while preserving the LSTM layers, to adapt the model to perform traffic prediction at the transfer location.

First, we present the evolution of training and validation losses for the models during the training process corresponding to a dropout rate of 2% in Figure 4. For the regular dropout model without any normalization, we observed that both the training and validation losses initially decreased as the model was trained, indicating that the model was effectively learning from the training data and generalizing well to unseen data i.e., the validation split. However, after a certain number of epochs the validation loss starts to increase. Therefore, the model was not generalizing well to the validation data and was instead memorizing the training data - a phenomenon referred to as ‘overfitting’. This indicates that the model is becoming too complex and may be fitting the model to noise in the training data.

It is also important to analyze the rate at which the validation loss increases because a faster rate of increase indicates that the model is overfitting more quickly, which means that the model’s generalizability to unseen data may be severely limited. Bayesian LSTM with layer-normalization still shows some signs of overfitting, however, this happens at a much lower rate compared to the case when no normalization were used. By controlling the mean and standard deviations of inputs to different layers of NN, the layer-normalization allowed the model to focus more on important features and patterns in the data, rather than fitting noise in the training data. Hence, the model’s performance on unseen data can be improved with layer-normalization by reducing overfitting, while still maintaining a high level of performance on the training data.

In contrast, when applying the spectral normalization method, both the training and validation losses exhibited decreasing trends. This indicates that the model was not overfitting and suggests that the normalization technique effectively reduced the complexity of the model By controlling the spectral norms of the weight matrices in the NN, this method encouraged the model to learn simpler and more generalized representations of the data. This, in turn, can lead to improved generalization performance on new data. However, due to over-regularization, both the training and validation losses for the spectral normalized model were higher compared to the regular model and the layer-normalized model. Specifically, the validation losses for the regular, layer-normalized, and spectral normalized models were -1.3382, -1.3990, and 1.3122 for 2% dropout. Additionally, the experiments were performed for dropout rates of 5% and 10%, however the results are only presented corresponding to 2% dropout for brevity. The results and conclusions look similar for the other dropout rates, however, the prediction accuracy decreases and overfitting increases for larger dropout rates.

In the context of evaluating the prediction performance of each model on the test data, the mean prediction was calculated. The results are presented in Table 1, which includes the RMSE, $R^{2}$, and MAPE metrics. Specifically, the regular model exhibited the lowest RMSE, however the layer- and spectral normalization did not increase the RMSE much. The general predictions are fairly accurate – the mean prediction (red curve) accurately follows the real traffic flows (black curve), see Figure 5. Further, the RMSE and $R^{2}$ values follow similar trends as the validation loss obtained for each model. This behavior can be attributed to the fact that the validation loss, as defined in Equation 5, also takes into account the squared deviations between the predicted and true values.

Next, the uncertainty estimates are compared. Figure 5 shows the 95th percentile confidence intervals of the estimates of the aleatoric and epistemic uncertainties using grey and red bands in each subplot. As depicted, there are significant differences in the uncertainty estimates among the three models. To better understand how the overall aleatoric and epistemic uncertainty of each model compares, Figure 6 plots the distribution of the magnitude of the uncertainties estimated over the test data. In these box-plots, the colored boxes represent the boundaries for the 25th and 75th percentile of the standard deviations, with the notch inside the box representing the median. The whiskers extend from the box by 1.5 times the inter-quartile range and flier points are those past the end of the whiskers. The figures illustrate that, although the aleatoric uncertainty estimates in layer normalized model exhibit a high degree of dispersion, the median standard deviation is lower than the regular model. Similarly, the layer normalized model outperforms the regular model in terms of epistemic uncertainty. Nevertheless, in both cases, the spectral normalization approach effectively controls the standard deviation.

The variation in the magnitude of the estimated aleatoric and epistemic uncertainties over time are shown in Figure 7. The aleatoric uncertainty estimates follow the general trend of the data for all models, with lower uncertainty observed during periods characterized by lower flow rates and and greater uncertainty observed during periods of higher flows, such as the peak congestion. This observation can be attributed to the chaotic and unpredictable behavior of the system during congested states, which are commonly associated with higher flows.

However, minor differences are observed for regular and layer normalized models, where they show a slight increase in aleatoric uncertainty during low flow regions (0 to 6 hr and 24 to 30 hr) and morning peak congestion period (  8 hr), compared to other flow regimes. This may be attributed to the limited knowledge of the models about the low flow and peak congestion regimes, as the training data only consisted of approximately 20% and 6% of data belonging to these regimes.

This effect is particularly evident in the epistemic uncertainty, where we observe a more pronounced higher standard deviation specifically for low flow regimes in these two models. In contrast, the spectral normalized model was able to capture trends despite the limited representation of low flow data in the training dataset. This suggests that the spectral normalization approach may be more effective at capturing and generalizing to unseen datasets, compared to the other models. However, for regimes with higher data representation, the layer normalized model performs better than spectral norm model - therefore provides lower epistemic uncertainty estimates.

Here, we aim to investigate the impact of normalization on the generalizability of the models by evaluating their performance, both in terms of accuracy and uncertainty, on out-of-distribution datasets, namely VDS 1114211 and VDS 1118170. Initially, we assess the performance of the trained models on these datasets without any additional training, providing insights into how well the models can generalize to data that differs from the training distribution. Subsequently, we proceed to retrain the models using partial data from these out-of-distribution datasets. Finally, we re-evaluate the models’ performance to determine the impact of retraining on their ability to handle these datasets.

The prediction accuracy of the models in transfer learning without and with retraining on out-of-distribution datasets is shown in Table 2 corresponding to dropout rate of 2%. We observe that the performance of models at a new station is not the same as that at the training station. This is expected since the models are being tested on samples that were not encountered during training. We observe the regular model to exhibit the highest $R^{2}$ and least RMSE value compared to both layer and spectral-norm models. Similar trends hold true for other dropout values, and hence omitted for brevity. However, it was observed that as the dropout rate increases, models exhibits higher regularization, resulting in a decrease in prediction accuracy.

Surprisingly, even when faced with significantly different temporal patterns in these out-of-distribution datasets, we observe acceptable levels of prediction performance – RMSE values are not observed to increase. Interestingly, the model’s performance is contingent upon the underlying data distribution, rather than the temporal pattern itself. For instance, encountering previously unseen values can lead to errors or uncertainties in the model’s predictions. To further comprehend this phenomenon, we can refer to Figure 8, which illustrates a comparison between the distribution of the normalized flow values from out-of-distribution datasets and the training data.

A sample of model predictions over a 48 hour period on one of the out-of-distribution datasets, namely VDS 1118170 without retraining is shown in Figure 9. It can be seen that the models’ predictions in low flow regimes tend to be characterized by errors and uncertainty. This is likely due to the data distribution of VDS 1118170 exhibiting a high skewness with an abundance of low flow values that were not encountered during the model’s training phase. Interestingly, the spectral normalized model provides a constant prediction for the low flow regimes, unlike the regular or layer normalized models, see Figure 9. The spectral normalized model treats extreme low flows as perturbations to the low flow values observed during training and therefore restricts its predictions to a constant value. However, it is worth noting that all models performs remarkably well in other flow regimes, resulting in low RMSE values within the overlapping regions of the distributions. Note that from a traffic prediction perspective, errors in the low flow regime are less critical since congestion management strategies are targeted for high flow regimes.

In contrast, the distribution of VDS 1114211 closely resembles the training distribution, suggesting that the model has encountered similar flow conditions during training, resulting in accurate predictions and low RMSE values.

When the dense layers (and layer normalization parameters for the layer normalized model) of the models are retrained using 20% of the data from each of these locations, the RMSE values tend to decrease in majority of cases. However, it is important to note that the MAPE often increases in few cases. The reason behind this behavior lies in the objective of transfer learning, which aims to minimize the loss function described in Equation 5, primarily based on square error loss, and the sensitivity of MAPE to errors in low flow regions, which might be affected during the retraining process.

Normalization is also observed to have a significant impact on uncertainty estimation. The uncertainty plots presented in Figures 10 and 11 demonstrate the aleatoric and epistemic uncertainty at the two test locations, both before and after retraining. Consistently with the trends observed in the training data, we find that normalization reduces both aleatoric and epistemic uncertainty in these plots. Notably, spectral normalization exhibits a greater reduction in uncertainty compared to layer normalization. This advantage becomes more pronounced as the dissimilarity between the out-of-distribution datasets increases. Despite the marginally higher RMSE obtained with spectral normalization, one might still prefer to utilize it due to its ability to provide more confident predictions. The reduced uncertainty associated with spectral normalization enhances the model’s reliability and provides a higher level of confidence in its predictions, which can be crucial in certain applications.

Furthermore, it is important to consider that aleatoric uncertainty estimates obtained on these datasets without retraining might not accurately reflect their true representation. This is because the estimated data variance ($\sigma^{2}$) is learned solely from the training data and remains unchanged. However, after retraining, the models exhibit a more precise estimation of data uncertainty as they adapt to the characteristics of the new datasets. Notably, for the spectral normalized model, the standard deviation increased for both stations, accounting for changes in flow patterns between the training and out-of-distribution datasets. On the other hand, the other models did not demonstrate significant changes in the aleatoric uncertainty estimates. Regarding epistemic uncertainty, there were no notable changes observed across the models. This observation suggests that the regular and layer normalized models tend to be overconfident in their predictions, possibly due to their inability to generalize well to out-of-distribution datasets. In contrast, the spectral normalized model exhibits a more realistic estimation of uncertainty, indicating its ability to provide reliable and generalizable predictions.

To gain a deeper understanding of the model performances, traffic time series were labeled into four regimes: 1) low flow (0 to 6 hr), 2) increasing flow (6 to 8 hr), 3) high flow (8 - 18 hr) or congestion and 4) decreasing flow (18 to 24 hr) which are appropriately color-coded as shown below. By comparing the gradients and feature space outputs of each model across these flow regimes, we aim to discern their specific characteristics and differences.

Figure 13 shows the color-maps of gradients normalized across time for the three models corresponding to different input vectors over a period of 24 hours (i.e. 288 samples in intervals of 5 minutes). Note that the x-axis represents the feature index, which is the historical flow data from the previous time steps with the latest flow data appearing at the right-hand-side end, whereas the y-axis corresponds to different times of the day. The model gradients are compared to identify patterns that could lead to higher generalizability of the spectral normalized models.

The regular models exhibit higher gradients, indicating increased sensitivity of the model outputs to historical traffic flows, particularly during the low flow regime (0 to 6 hr) and the increasing flow regime (6 to 8 hr). The layer normalized models demonstrate slightly lower gradients compared to the regular models, while the spectral normalized models exhibit the lowest gradients. The increased sensitivity to specific data patterns in regular and layer normalized models leads to larger variations in predictions and, consequently, wider uncertainty bands. On the other hand, the presence of added regularization in spectral normalization allows the model to focus more on generic trends rather than specific data characteristics. As a result, the spectral normalized model demonstrates relatively consistent gradient characteristics throughout the day, showing less sensitivity to different inputs across traffic regimes.

Further, feature-space representations of the outputs from the penultimate layers of the models were visually compared in Figure 14 to understand their relation with the gradients explored before. Based on the model configuration, the six dimensional feature outputs were suitably reduced to two dimensions using a non-linear dimension reduction technique, $t$-Stochastic Neighbor Embedding (Van der Maaten and Hinton 2008).

We have observed that the outputs from the regular dropout model exhibit high dispersion in feature space for data within the same traffic states, indicating sensitivity to specific data and consequently, limited generalizability. Particularly, we find that the feature outputs demonstrate significant dispersion in space for low- and increasing flow states, which aligns with the higher gradients observed in these regimes. In Table 3, we compare the variances of outputs belonging to specific traffic regimes in the 6-dimensional feature space for each model. Similar trends are observed for the layer and spectral normalized models, although these model outputs show localized behavior and less sensitivity. A notable improvement in feature space localization is observed for the spectral normalized model. This enhanced localization of traffic data from different flow regimes in the feature space has important implications for the model’s learning process and contributes to the stability and improved generalizability of the spectral normalized model. This characteristic allows the model to handle data perturbations more effectively and provides increased robustness and reliability in capturing different traffic patterns.

These trends hold true across various out-of-distribution datasets, although we have omitted the detailed results for brevity. This aspect is particularly valuable in scenarios where it is not feasible to train models specifically for a given location due to limited data availability. In such cases, models trained and validated on abundant data can be effectively transferred for use in different domains, such as a station with distinct characteristics.