Automated Dilated Spatio-Temporal Synchronous Graph Modeling for Traffic Prediction

In recent years, spatio-temporal graph modeling has become a mainstream method for traffic prediction. Most of these works combine spatial graph convolution networks (GCNs) with some temporal encoder to capture the complex spatio-temporal dependencies. STGCN [2] first integrate TCNs and GCNs for spatio-temporal modeling. Based on this, ASTGCN [8] adopts attention mechanism to enhance the representation learning capability of GCNs and TCNs, Graph WaveNet [6] involves the adaptive graph to incorporate with dilated temporal convolution networks. RNN is a widely used model for sequence leaning, whose variant gated recurrent unit (GRU) is utilized in DCRNN [5] and T-GCN [33] to capture temporal correlations of hidden representations from GCNs. In addition, self-attention mechanism is powerful tool not only for spatial dynamic learning, but also for capturing temporal dependencies [24]. Both STGNN [3] and GMAN [34] adopt self-attention mechanism in spatial GCNs and temporal dependencies learning. In some most recent works, some novel frameworks are introduced in this field. STSGCN [25] first proposes the framework of spatio-temporal synchronous modeling. Based on this, STFGNN [26] presents an informative fusion graph and parallel TCNs for further improving spatio-temporal dependencies learning. AutoSTG [35] first combines neural architecture search approach with spatio-temporal graph framework to improve the adaptability for different data. However, these existing works are not only difficult to flexibly capture the long-short term complex spatio-temporal dependencies, but also hard to characterize the spatio-temporal relations in different periods. Different from them, our model can balance short-term and long-term spatio-temporal correlations by dilation mechanism, and the graph search operation in our model can construct diverse spatio-temporal relations automatically for different datasets.

Automated Machine Learning (AutoML) aims to obtain appropriate features or models for various downstream tasks [36]. This field can be roughly divided into two main categories: automated feature engineering and automated model design, while automated feature engineering aims to synthesize or select informative features for model training [37, 38, 39]. Automated model design aims to select appropriate models or construct reasonable model architecture. Neural architecture search (NAS) is one of the most important direction in automated model design, which is widely used in deep learning. There are three mainstream types of methods in NAS, reinforcement learning-based [40, 41, 42], evolutionary learning-based [43, 44, 45] and gradient-based [46, 47, 48] respectively. Among these three types, gradient-based methods are relatively more efficient. Since efficiency is important in traffic prediction, we adopt gradient-based framework DARTS [46] in this paper. Different from these previous works, we employ DARTS to search graph structure rather than neural architecture in this paper.

Given the traffic sensor set with $N$ nodes $V(|V|=N)$, the sensor network can be defined as a graph $\mathcal{G}=(V,E,A)$. $E$ denotes the set of edges, whose relations between different nodes is characterize by the adjacency matrix $A$. The graph signal at time step $t$ contains $d$-dimensional original traffic features (e.g., the speed, volume), which is defined as ${X}_{\mathcal{G}}^{(t)}\in\mathbb{R}^{N\times d}$. The aim of traffic prediction task in sensor network is to learn a non-linear function $f(\cdot)$ from historical $T$-step graph signals for forecasting next $T^{{}^{\prime}}$-step graph signals. The mathematical form is defined as follows:

In this paper, the graph that characterizes the spatio-temporal relations in one time step is called as the meta graph, which is the basic unit of STSGs. To take both geographical proximity and pattern similarity into account, we introduce two types of meta graph: spatial graph $A_{SG}\in\mathbb{R}^{N\times N}$ and temporal graph $A_{TG}\in\mathbb{R}^{N\times N}$. The adjacency matrix of spatial graph can be formulated as:

Dynamic Time Warping (DTW) algorithm is adopted to calculate the similarity of two time series [49], which can characterize the pattern similarity between different nodes. For example, given two time series $X=(x_{1},x_{2},\cdots,x_{m})$ and $Y=(y_{1},y_{2},\cdots,y_{n})$, DTW is a dynamic programming algorithm defined as:

where $D(i,j)$ denotes the shortest distance between sub-sequence $X_{s}=(x_{1},x_{2},\cdots,x_{i})$ and $Y_{s}=(y_{1},y_{2},\cdots,y_{j})$. As a result, we can obtain $DTW(X,Y)=D(m,n)$ as the final distance between $X$ and $Y$. This method not only does not need to limit the length of the input sequences, but also better reveals the similarity of two time series compared with the Euclidean distance. Thus, we can define the adjacency matrix of temporal graph through the DTW distance as following:

where $X^{i}$ and $X^{j}$ are speed data series attached to node $i$ and node $j$ respectively. $\epsilon$ is a threshold to control the sparsity of $A_{TG}$, the setting of which is the same as that in [26].

Spatio-temporal synchronous graph (STSG) is a special architecture to establish the unified spatio-temporal correlations by graph structures, which is first proposed in [25]. Given a spatial graph with $N$ nodes, the adjacency matrix of STSG in [25] is designed as a $3N\times 3N$ expanded matrix, as shown in Fig. 2(a). In [26], the STSG is extended to $4N\times 4N$ and improved by involving more informative graphs such as temporal graph, which is shown in Fig. 2(b). Although this modeling approach can characterize more complex spatio-temporal relations, the larger expanded STSG could bring the problem of high computational overhead.

Recall that a normal spatial GCN layer [50] has the form:

where $\mathbf{X}\in R^{N\times D}$ and $\mathbf{Z}\in R^{N\times D^{\prime}}$ respectively denote the input and output node embedding of the GCN layer. $\mathbf{\Theta}\in R^{D\times D^{\prime}}$ denotes the shared weight for nodes’ feature mapping and $\sigma(\cdot)$ denote the activation function. $\hat{\mathbf{A}}\in R^{N\times N}$ denotes the normalized adjacency matrix which represents the message passing between one-hop neighbors.

To capture the complex spatio-temporal correlations, the normal spatial GCN can be extended to the spatio-temporal scale. A special GCN-based framework called spatio-temporal synchronous graph convolution network is proposed [25], but there are two main limitations of this modeling approach. The first one is that the fixed spatio-temporal receptive field limits the learning capability for long-term dependencies. To capture the longer-term spatio-temporal dependencies by this framework, the receptive field should become larger. Whenever the receptive field is expanded by one unit, the calculation amount of matrix multiplication will be expanded exponentially. The second limitation is that the adjacency matrix of STSG is shared for each spatio-temporal synchronous graph convolution module, which can not reflect the diverse spatio-temporal correlations in different periods. In addition, the deterministic adjacency matrix designed manually is hard to characterize the complex relations between different spatio-temporal nodes for different datasets.

To overcome these problems, we design the Auto-DSTSG module with multiple parallel blocks. The form of the Auto-DSTSG block is defined as:

where $\mathbf{X}(t,k,d)\in R^{kN\times D}$ and $\mathbf{Z}\in R^{N\times D^{\prime}}$ respectively denote the input and output representation. $k$ denotes the kernel size and $d$ denotes the dilation factor to control the skipping distance in the sequence input. $[\cdot]$ denotes the concentration operation. Compared with the normal GCN, the receptive field has been expended to $d\times k$ times when the range of spatio-temporal synchronous graph covers $k$ time steps. Here $\mathcal{F}_{G}(\cdot)$ represents the graph structure search operation, whose output is a mixed adjacency matrix $\mathcal{M}_{A}\in R^{kN\times kN}$ during searching phase. This operation can adjust the structure of STSGs according to different data, which can characterize complex spatio-temporal relations in different periods and even different datasets. $\mathcal{C}(\cdot)$ denotes cropping operation to convert the range of data from $kN$ to $N$. The common method is to select the graph data of the middle or last time step as the final output, which is described in [3]. When $\mathcal{F}_{G}(\cdot)$ is a fixed adjacency matrix with $k=1,d=1$ and $\mathcal{C}(\cdot)$ is an identity mapping, eq (6) is equivalent to eq (5).

The weight sharing mechanism in temporal convolution is conducive to learning global temporal dependencies for individual nodes [26, 52]. Although Auto-DSTSG module can flexibly capture the complex spatio-temporal correlations in a unified framework, its weight non-sharing mechanism for different blocks is more conducive to capturing diverse local spatio-temporal dependencies for different time periods rather than global dependencies for each node.

The dilated temporal convolution operation is defined as:

where $\mathbf{x}\in\mathbf{R}^{T}$ denotes the given 1D sequence input, $\mathbf{f}\in\mathbf{R}^{K}$ denotes a temporal convolutional filter at step $t$, $\Theta$ denotes the learnable weights of the filter and $d$ denotes the dilation factor. Similar to [51], we obtain the larger receptive field by expanding the dilation factor in temporal convolution when the layer goes deeper. To keep the consistency of the receptive field in each layer, the kernel size $k$ of temporal convolution is fixed as 2 with a sequence of dilation factors $[1,2,4,4]$.

To better control the information flow and reserve the useful information, we adopt the gating mechanism in this case. Similar to [53], a simple gated temporal convolution network only contains an output gate, which is expressed as:

where $X\in R^{N\times T\times D}$ is the given input, $X_{\mathbf{f}}\in R^{N\times(T-k-d+2)\times D^{\prime}}$ is the output, $\mathbf{\Theta_{1}}$, $\mathbf{\Theta_{2}}$, $\mathbf{b_{1}}$ and $\mathbf{b_{2}}$ are model parameters, $\odot$ is the element-wise product, $tanh(\cdot)$ is the tanh activation function of the outputs, and $\sigma(\cdot)$ is the sigmoid function which controls the ratio of information flow put forward to the next layer.

To enhance the global temporal correlations for individual nodes, we aggregate the output of the dilated temporal convolution module and the output of Auto-DSTSG module in each layer. The output of each layer in Auto-DSTSGN can be defined as:

where $X_{\mathbf{out}}$ denotes the output from each layer, $X_{\mathbf{g}}$ and $X_{\mathbf{f}}$ are respectively the output of Auto-DSTSG module and dilated temporal convolution module. $\mathbf{Agg}(\cdot)$ denotes the aggregation function, which is set as sum function in our model.

In graph structure search operation, all computations are differentiable. Similar to DARTS framework, a bi-level gradient-based optimization algorithm can be employed to update the weight parameters $\theta$ of the network (including the parameters in TCNs, GLUs, FCs and MLP) and the architecture parameters $\omega$ (including the scores of candidate matrices) alternately. As shown in Algorithm 1, the weight parameters (Line 4-5) and architecture parameters (Line 6-7) are alternately updated based on the training and validation sets, until the stopping criteria is met. Then, the structures of STSGs can be obtained by selecting the candidate operations with the highest operation scores.

We evaluate our model on PEMS03, PEMS04, PEMS07 and PEMS08 which are collected from Caltrans Performance Measurement System (PeMS). The time granularity of all datasets is set to 5 minutes. The spatial graph for each dataset are constructed based on road network topology. Z-score normalization is applied to the traffic flow data. Detailed statistics of datasets are shown in Table II.

Each dataset is chronologically split with 60% for training, 20% for validation and 20% for testing. We use the historical traffic flow in the last one hour to forecast the future traffic flow in the next one hour. The spatial graph and temporal graph are designed inspired by [26]. Our model is implemented by Pytorch 1.5 with NVIDIA TESLA V100 GPU. The max hop of graph convolution in each Auto-DSTSG block is set as 2 by default. The dimension of hidden representations is set as 40. The optimizer of our model is set as Adam. The batch size is 64 and the learning rate is 0.001. Our model is evaluated five times on each dataset. During search process, we utilize the early stopping strategy for graph structure search with tolerance 15 for 60 epochs. During training process, we reinitialize the optimizer and employ early stopping with tolerance 30 for 200 epochs.

We compare our model with the following nine state-of-art baselines in recent years:

• •

  FC-LSTM: This is a variant of Long Short-Term Memory Network, which adopts fully connected hidden units [22]. We set the number of hidden layer as 1 and the hidden units as 64.

• •

  DCRNN: Diffusion Convolution Recurrent Neural Network, which integrates GCNs into encoder-decoder RNNs [5]. The hop of diffusion graph convolution is set as 2 and the hidden dimension is set as 64.

• •

  STGCN: Spatio-Temporal Graph Convolution Network, which employs GCNs and TCNs for spatio-temporal learning [2]. Each spatio-temporal cell in this model contains two TCNs and one GCN. The number of spatio-temporal cell is set as 2 and the hidden dimension is set as 64.

• •

  ASTGCN: Attention based Spatial Temporal Graph Convolution Network, which utilizes spatial and temporal attention mechanisms [8]. Similar to STGCN, there are two spatio-temporal cells in this model and the hidden dimension is set as 64.

• •

  Graph WaveNet (GWN): Graph WaveNet adopts the GCNs with adaptive adjacency matrix and 1D dilated TCNs for spatio-temporal modeling [6]. Each layer in this model contains a gated TCN and a spatial GCN. The number of stacked layers in this model is set as 8 with the dilation rate [1, 2, 1, 2, 1, 2, 1, 2, 1, 2] and the hidden dimension is set as 64.

• •

  STSGCN: Spatial-Temporal Synchronous Graph Convolution Network, which utilizes multiple STSG modules for localized spatio-temporal joint dependencies modeling [25]. The size of spatial-temporal synchronous graph is set as $3N\times 3N$, the number of STSG layers is set as 3, and the hidden dimension is set as 64 in this model.

• •

  STFGNN: Spatial-Temporal Fusion Graph Convolution Network, which utilizes spatio-temporal fusion graph convolution and parallel TCNs for learning localized and global spatio-temporal dependencies respecively [26]. The size of spatial-temporal fusion graph is set as $4N\times 4N$, the number of layers is set as 3, and the hidden dimension is set as 64 in this model.

• •

  STGODE: Spatial-Temporal Graph Ordinary Differential Equation Network, which integrates the tensor-based ordinary differential equation into GCN modules [54]. The number of graph ODE layers is set as 6 and the hidden dimension is set as 64 in this model.

• •

  AutoSTG: Automated Spatio-Temporal Graph Network, which integrates NAS with spatio-temporal graph learning modules [35]. The number of spatio-temporal graph learning cells is set as 5, the number of mixed operations is set as 6 in each cells, and the the hidden dimension is set as 64 in this model.

The evaluation metrics are mean absolute errors (MAE), root mean squared errors (RMSE) and mean absolute percentage errors (MAPE) averaged over five times for one hour ahead prediction. For these three metrics, smaller values means better performance and the formula of these three metrics are defined as follows:

where $\widehat{y}_{i}$ denotes the prediction results and ${y}_{i}$ denotes the ground-truths.

From the results in Table III, we can observe that our model Auto-DSTSGN consistently outperforms the sub-optimal baselines with 4.9%$\sim$10.3% improvements in terms of MAE on the four datasets, which demonstrates the superiority of our proposed method. In order to make the comparison results more intuitive visually, we also provide the box plot of Table III, as shown in Figure 5. Next, we analyze and compare the strengths of our proposed model with some well-performing baselines.

AutoSTG is the only baseline that integrates the neural architecture search to adjust its architecture corresponding to the data. From the comparison results, AutoSTG significantly outperforms most non-auto state-of-art models such as DCRNN and Graph WaveNet, since it can search the optimal neural architectures for different data scenarios. However, AutoSTG is still weaker than our model. There are two main reasons. The one is that the spatio-temporal synchronous graphs of our model can characterize more complex dynamics than the normal graphs. Another one is that AutoSTG focuses on the neural architecture search but ignores the importance of informative spatio-temporal graph construction. Both STFGNN and STSGCN adopt the spatio-temporal synchronous graphs but they can hardly capture the short-term and long-term spatio-temporal correlations flexibly. Our model presents the dilated spatio-temporal synchronous graph convolution framework to balance the short-term and long-term dependencies learning. Both STFGNN and STGODE are the baselines that adopt informative adjacency matrices for spatio-temporal dependencies learning. However, these two models design the spatio-temporal synchronous graphs or neural architectures manually and empirically, thus they are difficult to adapt to different data scenarios. In contrast, our model can construct different spatio-temporal synchronous graphs for different data scenarios based on auto machine learning, which enhances the adaptability and generalizability of our model.

In summary, the dilated spatio-temporal synchronous graph structures in our model can flexibly characterize the short and long-term spatio-temporal dependencies, the auto machine learning mechanism for graph structure search can help our model adapt to different data inputs and achieve the optimal dilated spatio-temporal synchronous graph modeling. These are why our model can outperform other baselines significantly.

We conduct ablation study on PEMS04 and PEMS08 to evaluate the effectiveness of key components in our model. As shown in Table IV, we compared Auto-DSTSGN with following variants: 1) w/o GCN, which removes all the Auto-DSTSG modules from our models 2) w/o TCN, which removes all the dilated temporal convolution modules from our model. 3) w/o Dilation, which removes the dilation mechanism in Auto-DSTSG module from our model. 4) w/o GSS, which replaces the Auto-DSTSG module with the fixed STFGNN module [26]. 5) GSS Random, which samples an complete adjacency matrix from search space in graph structure search module during the search process. 6) w/o DTW, which replaces the DTW with Pearson coefficient to construct the temporal graphs.

From the experimental results, we can find that Auto-DSTSGN outperforms all the ablation variants. Compared with the results of w/o Dilation, Auto-DSTSGN improves 1.7%, 4.0% in terms of MAE and MAP on PEMS04. Meanwhile, it also improves 2.4%, 7.0% in terms of MAE and MAPE on PEMS08, which illustrates the effectiveness of dilation mechanism on learning long-term spatio-temporal dependencies. Compared with the results of w/o GSS, Auto-DSTSGN improves 3.2%, 6.1% in terms of MAE and MAPE on PEMS04. In the meanwhile, it also improves 3.7%, 6.9% in terms of MAE and MAPE on PEMS08. Compared with the results of GSS Random, Auto-DSTSGN improves 4.2%, 7.7% in terms of MAE and MAPE on PEMS04. Meanwhile, it also improves 5.5%, 12.4% in terms of MAE and MAPE on PEMS08, which illustrates the effectiveness of Graph structure search module in learning diverse spatio-temporal correlations and adapting to different data. There is a significant fall of the performance without graph convolutions and temporal convolutions (w/o GCN and w/o TCN), demonstrating the effectiveness of these two parts for spatio-temporal representation learning. Moreover, when we replace DTW with Pearson coefficient to construct the temporal graphs (w/o DTW), the performance drops on both PEMS04 and PEMS08 datasets. The reason is that Pearson coefficient can only characterize the linear correlations between different time series while DTW can better characterize the non-linear time series similarity.

We select PEMS04 and PEMS08 to further investigate the relations between the attributes of meta graphs and the optimal structure of STSGs on them. For the two different meta graphs SG and TG, we choose mean degree of them as the most important attribute. The higher mean degree of a graph means stronger correlations among the nodes. Specifically, the higher mean degree of SG represents the stronger spatial correlations while the higher mean degree of TG represents the stronger temporal correlations. Thus, empirically we need more graph convolution operations on the related graphs to achieve a larger receptive field for capturing long-range spatio-temporal correlations. We count the average number of SG and TG in the learned structure of STSGs on two datasets. All the results are shown in Table V. We observe that the STSGs learned on PEMS04 contains more TGs while the STSGs learned on PEMS08 contains more SGs. This can be explained by the mean degree of TG and SG on two datasets. Since the higher mean degree characterizes stronger correlations between different nodes, our model can obtain stronger capability for message passing in spatio-temporal scale by adjusting the structure of STSGs automatically. In addition, we also visualize the learned structures of STSGs on the PEMS04 and PEMS08 in Figure 6.

Time consumption and memory occupancy on GPU are two intuitive metrics to reflect the time and space complexity of different models. They are also two important metrics that measure the efficiency and scalability of the model in industrial scenarios. We select three best baselines STFGNN, STGODE and AutoSTG to compare with our model on the two metrics. The results are shown in Figure 7. From the absolute perspective, the efficiency of our model during training phase is slightly improved compared with STFGNN and STGODE in general and the GPU occupancy of our model is significantly lower than them. We involve the dilation mechanism into the spatio-temporal synchronous graph modeling, so our model significantly reduces model complexity compared with STFGNN. STGODE employs the multiple neural ODE architectures, which greatly increase the time and space complexity of the model. From the relative perspective, the GPU occupancy of our model is almost the same during search phase and training phase, but AutoSTG significantly costs more GPU occupancy in the searching phase. This indicates that graph structure search is more lightweight and efficient than neural architecture search. This is because graph structure search only involves simple matrix calculations whereas neural architecture search involves complex computations of neural network structures with learnable parameters.

To further investigate the effectiveness of our model, we conduct parameter study on PEMS04 and PEMS08, including the dimension of hidden representations $D$ and the max hop of graph convolution $H$ in each block. The experimental results are shown in Figure 8. We can find that MAE, MAPE and RMSE on two datasets are the optimal when $D$ is equal to 48. When $D$ is too small, the learning capability of our model become worse, resulting in poor prediction performance. When $D$ is too large, the three metrics on both PEMS04 and PEMS08 become worse. This is because too large hidden dimension cause the over-fitting. For parameter $H$, we can observe that MAE, MAPE and RMSE on PEMS04 achieve the optimal results when $H$ is equal to 2. On PEMS08, MAE, MAPE and RMSE obtain the best values when $H$ is equal to 3. This implies that aggregating the neighbor information of the appropriate order in the traffic network can better learn the spatial dependencies, and an excessively large number of hops will lead to the over-smoothing phenomenon of the GCNs.