STCGAT: A Spatio-temporal Causal Graph Attention Network for traffic flow prediction in Intelligent Transportation Systems

Graph Convolutional Networks (GCN)[22] has received much interest in recent years. They are now frequently utilized for various tasks involving graph-structured data, including graph classification, node classification, and link prediction. The clustering method divides GCN into two basic categories: Spectral-based and Spatial-based[8]. Spectral-based techniques implement filters to define the graph convolution from the perspective of graph signal processing and extend the graph convolution to the spectral domain by locating the corresponding Fourier bases. The primary approaches are Chebyshev Spectral CNN (ChebNet)[23] and Adaptive Graph Convolution Network (AGCN)[24]. The Spatial-based approach represents graph convolution as aggregating feature information from the neighborhood and defining the graph convolution by the spatial relationship of the nodes. The most prominent methods are GAT and Gated Attention Network (GAAN)[25].

Earlier traffic forecasting projects were often multivariate time series analysis problems utilizing time series modeling techniques such as History Average (HA) Model[26], Vector Autoregression (VAR)[27] and Support Vector Regressor (SVR)[28]. However, these algorithms rely too heavily on the assumption of ideal smoothness, which is incompatible with the nonlinear correlation of traffic data. With its well-established data modeling and autonomous learning capabilities, deep learning has steadily taken over time series forecasting tasks in recent years. Most studies have used RNN-based models to capture temporal dependencies, but the complex long-loop structure is time-consuming. It may also be accompanied by gradient disappearance and explosion phenomena as the long-loop links grow. To address this problem, some work uses Temporal Convolutional Network (TCN)[29][30] to process time series in parallel, learning more helpful information through a larger field of perception. Recently, several researchers have concentrated on the robust modeling capabilities of Transform[31] for time series data and have presented many variant models[32][33][34] based on Transform with excellent performance on long time series prediction tasks.

To model traffic data’s spatial and temporal dependencies, some studies[16][17][18] modeled traffic networks as regular two-dimensional grids, then process the two-dimensional data using CNN to capture the spatial dependence, and finally further model the temporal dependence of traffic time series using RNN or CNN. However, CNN’s are not always applicable in traffic road networks of non-Euclidean nature. To solve the problem, more and more researchers are turning to GCN-based models for Spatio-temporal data prediction. For example, DCRNN[9] models spatial dependencies by wandering in both directions on the traffic road topology graph and then uses GRU to capture temporal dependencies. ASTGCN[10] models the spatial-temporal dependence of traffic data with spatial attention and temporal attention, respectively. STSGCN[35] captures the local Spatio-temporal correlation of traffic data by combining Spatio-temporal blocks through a local Spatio-temporal synchronized graph convolution module. STFGNN[19] learns both local and global Spatio-temporal dependencies by processing data-driven spatial and temporal graphs at different moments in parallel. Following STFGNN, STGODE[36] draws on the Dynamic Time Warping (DTW)[37] used by STFGNN to generate semantic adjacency matrices for traffic road topology maps to capture deeper Spatio-temporal dependencies. In recent works, Z-GCNETs[38] introduce the concept of zigzag persistence in the traffic network diagram structure and integrate it into GCNs to enhance the stability of the model. STG-NCDE[39] designs two neural control differential equations dealing with temporal and spatial dependencies, respectively, and integrates both to capture Spatio-temporal dependencies simultaneously.

Traffic flow forecasting uses historical traffic information on the road to predict traffic condition information in the future period. In our model, traffic speed is chosen as the traffic condition information of the road.

Definition 1: We use the undirected graph $G=(V,E)$ to represent the topological information of the traffic roads. Where $V=\{v_{1},v_{2},\cdots,v_{N}\}$ is the set of all nodes of graph $G$, $N$ denotes the number of all road nodes, and $E$ denotes the set of connected edges of all nodes of graph $G$.

Definition 2: We represent the historical traffic speed information of time length $T$ by a feature matrix $X=\{X_{1},X_{2},\cdots,X_{N}\}\in R^{N\times T\times F}$, where $F$ denotes the feature dimension, $X_{t}=\{\vec{x_{t:1}},\vec{x_{t:2}},\cdots,\vec{x_{t:N}}\}\in R^{N\times F}$ denotes the set of traffic information of all road nodes at any $t$ time, and $\vec{x_{t:i}}\in\mathbb{R}^{F}$ then denotes all the feature vectors of node $v_{i}$.

Traffic prediction has a substantial spatial and temporal correlation, considering the spatial correlation and temporal dependence of individual road nodes in the road network. Therefore, we aim to use the road network topology $G$ and the traffic information feature matrix $X$ to predict the traffic speed information $Y^{\prime}=[X_{t+1}^{\prime},X_{t+2}^{\prime},\cdots,X_{t+T}^{\prime}]$ in the next $T$ moments by learning a function $f(\cdot)$.

The traffic condition data of distinct road nodes are strongly and dynamically connected in the spatial dimension. However, traditional graph neural networks construct predefined adjacency matrices to perform graph convolution operations based on information such as connectivity or distance between graph nodes. Although the predefined adjacency matrix intuitively represents the position relationship between nodes, it cannot reflect the dynamic spatial dependence between road nodes at different moments dynamically. Therefore, as shown in (b) in Fig. 2, we use a node-adaptive learning mechanism to learn the dynamic correlation information among road nodes at different moments in a fine-grained manner to generate the traffic subgraph $G_{ad}$ for the corresponding moment. As shown in Eq. 2, this mechanism generates the adjacency matrix $\tilde{A^{t}}\in R^{N\times N}$ corresponding to that moment at any moment $t$.

Where $E_{At}\in R^{N\times d}$ is the embedding dictionary encoding each traffic road node, $d$ is the embedding dimension, and $E_{At}^{T}$ is the transpose matrix of $E_{At}$. $softmax$ is the normalization function, and $ReLU$ is the nonlinear activation function.

Meanwhile, to capture spatial dependencies among nodes in the spatial dimension adaptively and dynamically. We propose a NAL-GAT to extract the spatial features of traffic roads by integrating the node-adaptive learning mechanism with GAT. Specifically, at any moment $t$, GAT computes the attention coefficients of its neighboring nodes vertex by vertex for the node correlation information generated by the node adaptive learning mechanism. Finally, it aggregates the spatial dependencies among the nodes on the graph at moment $t$. As shown in Eq. 3, the computation process of attention coefficient $e_{ij}^{t}$ between node $v_{i}$ and its arbitrary neighbor node $v_{j}$ is demonstrated.

Where $a$ is the computational function of the attention mechanism, and $W\in R^{F\times F^{\prime}}$ is the graph’s weight matrix of all nodes. The attention coefficient $\alpha_{ij}^{t}$ of the graph’s attention layer is then generated by normalizing the attention coefficients of node $v_{i}$’s neighbors.

Where $\tilde{A}_{i}^{t}$ denotes all the neighbor nodes of $v_{i}$.

In addition, we note that all nodes in the GAT share the same parameter space $W\in R^{F\times F^{\prime}}$, which may lead to too large graph $W$ with more nodes making the model difficult to optimize. To solve this problem, we construct a shared weight pool $W_{p}\in R^{d\times F\times F^{\prime}}$, which can get the weight matrix $W^{\prime}=E_{At}\cdot W_{p}\in R^{F\times F^{\prime}}$ of each node according to the node’s embedding dictionary $E_{At}$. Then the complete attention coefficient calculation process is shown in Eq. 5.

Where $\vec{a}\in\mathbb{R}^{2F^{\prime}}$ is the weight matrix, $\parallel$ denotes the connection operation, and $LeakReLu$ is the nonlinear activation function.

To capture deeper feature information, this paper further uses the multi-head attention mechanism to model spatial dependence. As shown in Eq. 6, $Q$ sets of mutually independent attention mechanisms are invoked.

Where $\alpha_{ik}^{t,q}$ is the weight coefficient computed by the attention mechanism of the qth group at time $t$, $E_{At}^{q}\cdot W_{p}^{q}$ is the weight matrix of the corresponding group, and $\vec{x_{i}}^{\prime}\in\mathbb{R}^{QF^{\prime}}$ is the new feature representation obtained by passing node $v_{i}$ through the attention layer of the multi-headed graph.

When using multi-headed attention, the network’s last layer should not output too many features, so we use a separate self-attentive mechanism to limit the node’s output feature length. Specifically, we use a new weight pool $W_{p}^{\prime}\in R^{d\times QF^{\prime}\times F^{\prime\prime}}$ to map the output dimension of nodes from $\mathbb{R}^{QF^{\prime}}$ to $\mathbb{R}^{F^{\prime\prime}}$ to obtain the final output result $\vec{x}_{t:i}^{\prime\prime}\in\mathbb{R}^{F^{\prime\prime}}$.

When all nodes in the graph have completed the above graph attention layer operation, we can obtain the output features $X_{t}^{\prime\prime}=\{\vec{x}_{t:1}^{\prime\prime},\vec{x}_{t:2}^{\prime\prime},\cdots,\vec{x}_{t:N}^{\prime\prime}\}\in R^{N\times F^{\prime\prime}}$. For the convenience of presentation, we express this process in Eq. 8.

where $\sigma(\cdot)$ is the computation function for the graph attention layer.

There is a correlation between traffic conditions in the time dimension at various times. As shown in (C) in Fig. 2, we replace the gating unit of GRU with NAL-GAT to further capture the Spatio-temporal dependence of the traffic time series data. Specifically, the spatially dependent time series data $X_{t}^{\prime\prime}$ at any moment $t$ is used as the input data of the GRU.

Where $\overrightarrow{h}_{t-1}$ is the output at the previous moment, $[\cdot]$ denotes the concat operation in the feature dimension, $\widetilde{h_{t}}$ is the candidate hidden layer state, $\odot$ denotes the multiplication by elements, and $\overrightarrow{h_{t}}\in R^{N\times F^{\prime\prime}}$ is the output at the current moment.

It is important to note that as the input time length rises, so does the network depth of the model. However, the deep network may lead to issues such as gradient disappearance and overfitting in the model. Therefore, we use the residual module to connect the network’s layers to improve the model’s capacity for long-term capture.

Where $\omega_{1}$ and $\omega_{2}$ are both one-dimensional convolution kernels, $\varepsilon$ is the nonlinear activation function, $\otimes$ denotes the convolution operation, and $\overrightarrow{h}_{t}^{\prime}\in R^{N\times F^{\prime\prime}}$ is the output of residual concatenation. Until the completion of the above operations at the Tth time step, we can obtain the sequence data containing the Spatio-temporal dependence $\overrightarrow{H}^{\prime}=(\overrightarrow{h}_{t-T}^{\prime},\overrightarrow{h}_{t-(T-1)}^{\prime},\cdots,\overrightarrow{h}_{t}^{\prime})$.

In addition, traffic data are not always sequentially correlated, and there are complex causal correlations between traffic events. Therefore, we use bidirectional GRU to capture the local causal, temporal relationships. The reverse operation is similar to the above operation, and the output results are finally stitched to obtain the output $H\in R^{N\times T\times 2F^{\prime\prime}}$.

From the above procedure, it is evident that GRU is processed by progressively unfolding along the timeline, which causes the output at the present time to depend on the state at the previous time and so lacks the capacity to capture long-term temporal dependence. We deploy a parallel temporal convolutional network along the time axis to improve the performance of extracting long-term temporal dependencies. The output of Equation 11’s data deformation $H^{\prime}\in R^{N\times(T*2F^{\prime\prime}})$ is utilized as the TCN’s input data. In the time series convolution process, the time series data $H_{i:}^{\prime}\in\mathbb{R}^{(T*F^{\prime\prime})}$ of any node $v_{i}$ and a filter $f:\{0,\cdots,l-1\}\ \ \overrightarrow{}\ \ \mathbb{R}$ are first extended for the elements $s$.

Where $d$ is the dilation factor, $l$ is the filter size, and $s-d\cdot i$ denotes the direction of the timeline past. When $d=0$, the dilation convolution becomes a regular convolution. In addition, as shown in (d) in Fig. 2, the TCN needs to perform a series of transformations such as Weight Norm and Dropout and use the residual join to obtain the output $o\in R^{N\times(T*2F^{\prime\prime})}$ as in Eq. 13.

Finally, we perform a dimension-specific linear transformation of the output sequence by a two-layer fully connected neural network.

Where $W_{1}\in R^{F^{\prime\prime\prime}\times(T*2F^{\prime\prime})}$ and $W_{2}\in R^{(T\times F)\times F^{\prime\prime\prime}}$ are the weight matrices, $b_{1}$ and $b_{2}$ are the corresponding bias terms, and $Y^{\prime}\in R^{N\times T\times F}$ is the final prediction result.

During model training, we optimize the model using the $L1$ loss function and the Adam optimizer to make the error between the predicted $Y^{\prime}=[X_{t+1}^{\prime},X_{t+2}^{\prime},\cdots,X_{t+T}^{\prime}]$ and labeled values $Y=[X_{t+1},X_{t+2},\cdots,X_{t+T}]$ as small as possible.

We have done extensive experiments on the following four real public transportation datasets, PeMS03, PeMS04, PeMS07, and PeMS08[40], to illustrate the effectiveness of the proposed Spatio-temporal prediction framework. These datasets were obtained on Caltrans’ Performance Measurement System (PeMS). The system has over 39,000 traffic detectors deployed on California freeways to collect real-time traffic flow data and geographic information on the roadways. The collected data are aggregated every five minutes. Table I summarizes the critical statistics for the four datasets.

STCGAT was compared to some of the most advanced baseline models. The following is a summary of these baselines.

• •

  HA[26]: The model uses the average of historical traffic data as the forecast value.

• •

  VAR[27]: The model captures the dependencies between temporal data.

• •

  FC-LSTM[14]: The model is based on the traditional RNN model to model the time dependence of historical traffic data.

• •

  TCN[29]: The model allows for processing long-time series data at a fraction of the cost.

• •

  DCRNN[9]: The model is based on GCN to capture spatial dependencies and uses an encoder-decoder architecture to capture temporal correlations.

• •

  ASTGCN[10]: The model uses spatial and temporal attention mechanisms to capture underlying Spatio-temporal patterns.

• •

  STSGCN[35]: The model captures local Spatio-temporal dependencies using spatial graph convolution and one-dimensional temporal convolution, respectively.

• •

  AGCRN[11]: The model proposes an adaptive graph convolution module that can autonomously learn the adjacency matrix of a traffic road network.

• •

  STFGNN[19]: The model effectively fuses multiple spatial and temporal graphs to learn the Spatio-temporal relationships hidden by traffic data.

• •

  STGODE[36]: The model captures deeper Spatio-temporal dependencies by expanding the perceptual field of the GCN.

• •

  Z-GCNETs[38]: The model is designed with a time-aware GCN to capture the complex Spatio-temporal dependencies in traffic data.

• •

  STG-NCDE[39]: The model is designed with two independent neural control differential equations for modeling spatial and temporal dependence.

We divide the used dataset into 60% training set, 20% validation set, and the remaining 20% test set with Z-score standardization. The partitioned dataset is then processed through a sliding window of length $2T$, where the first $T$ time lengths of serial data are used as historical data, and the last $T$ time lengths of data are used as labeled values. Here, we set the size of $T$ to $12$. That is, one hour’s historical traffic data is used to predict traffic data for the next hour.

Meanwhile, STCGAT is implemented using the PyTorch deep learning framework. In the hyperparameter settings of our model, the embedding feature dimension of the nodes is set to 10, the hidden layer size is 64, the number of multi-head attention mechanisms is set to 3, and the convolutional kernel size is set to 2. During the training process, the learning rate is set to 0.001, the batch size is set to 64, and the model is optimized using the Adam optimizer with a maximum number of iterations of $300$. The training environment of the model is shown in Table II.

We use the following three performance metrics to measure the model’s predictive power.

• •

  Mean Absolute Error(MAE):

• •

  Root Mean Squared Error(RMSE):

• •

  Mean Absolute Percentage Error(MAPE):

Where $L$ denotes the total number of samples. The lower the values of the three indicators above, the greater the model’s predictive accuracy. We conduct each experiment five times and then calculate the mean value as the test result.

On PeMS03, PeMS04, PeMS07, and PeMS08, our model was compared with the twelve representative baseline approaches described previously. Table III shows the prediction performance results of STCGAT and other baseline models within one hour (12 prediction steps). We can observe that 1) The GCN-based solution outperforms HA, VAR, and LSTM traditional machine learning methods, demonstrating the significance of explicitly modeling spatial correlation and the efficacy of GCN in traffic flow prediction tasks; 2) The performance metrics of our improved GAT-based method on each dataset are significantly better than other advanced baseline models, achieving significant results; 3) Our proposed method has stable short- and long-term Spatio-temporal prediction capabilities. As shown in Fig. 3, which demonstrates the comparison of the prediction performance of STCGAT with other Spatio-temporal prediction models on different Horizon, we can observe that the oscillations of the performance curve of STCGAT on each dataset are relatively small, indicating that our proposed method is insensitive to the prediction Horizon and the prediction performance is very stable; 4) As shown in Fig. 4, we use STCGAT and a representative baseline model to predict the traffic conditions at any road node for 288 consecutive time steps (24 hours). We can observe that the prediction results of STCGAT are more closely matched to the actual values, proving that STCGAT can more accurately capture the temporal and spatial correlations in the traffic flow sequence and achieve the best prediction results.

We created four STCGAT-based model variants and contrasted STCGAT with these four variables on the PeMS04 and PeMS08 datasets to understand the influence of STCGAT’s many modules. The differences between these four variants of the model are described below.

1. 1.

   w/o node embedding: The model removes the node embedding operation and uses only the traditional predefined adjacency matrix.

2. 2.

   w/o ResNet: The model removes the residual connection structure in STCGAT.

3. 3.

   w/o reverse GRU: The model one-way GRU captures the Spatio-temporal correlation of the traffic network.

4. 4.

   w/o TCN: The model removes the temporal convolutional network.

Table IV shows that we conducted experiments with these four variants and STCGAT on PeMS04 and PeMS08. Combined with the histogram of evaluation metrics for each variant model at different time steps in Fig. 5, we can obtain the following observations: 1) Each evaluation metric with TCN is at its maximum value, and the performance metrics are similar in the short-term (15 min) and long term (60 min). This illustrates the necessity of using temporal convolutional networks to extract the global temporal correlation of traffic flow; 2) The performance metrics of w/o node embedding decrease significantly when the node self-learning module is removed. This proves that learning the dynamic spatial correlation among nodes from the state information of traffic flow at different moments better expresses how the traffic flow dynamics change in reality; 3) The performance metrics values of w/o ResNet compared to STCGAT are significantly larger, i.e., removing the residual module has a more significant impact on both datasets. This indicates that the application of the residual module helps STCGAT to mitigate the problem of overfitting or gradient disappearance caused by the superposition of network layers to a certain extent; 4) The performance metrics of with or without reverse GRU also increase on both the PeMS04 and PeMS08 datasets, indicating that effectively capturing the causality of traffic flow data helps to analyze the Spatio-temporal correlation of traffic flow more comprehensively; 5) Compared with the four variants, STCGAT has the best performance. On the one hand, it shows the importance of each module in STCGAT, and on the other hand, it shows that STCGAT can extract the Spatio-temporal correlations in traffic flow series more accurately.