Graph-based Multi-ODE Neural Networks for Spatio-Temporal Traffic Forecasting

Researchers have employed traditional statistical and machine learning methods for the task of traffic forecasting. Some prominent example models are (1) K-Nearest Neighbor (KNN) (Zhang et al., 2013) which predict traffic of a node based on its k-nearest neighbors; (2) ARIMA (Van Der Voort et al., 1996; Alghamdi et al., 2019) which integrates the autoregressive model with moving average operation; and (3) SARIMA (Williams & Hoel, 2003) which adds a specific ability to ARIMA for the recognition of seasonal patterns. Many of these machine learning models only consider the temporal dependencies and ignore the spatial information. Also, they are usually based on human-designed features and have limitations in learning informative features for the intended task. Later, deep learning methods became popular due to their ability in considering the complex and high-dimensional spatio-temporal dependencies through richer representations. Early examples of deep learning models considering the spatial and temporal relations are FC-LSTM (Shi et al., 2015), ConvLSTM (Xingjian et al., 2015), ST-ResNet (Zhang et al., 2017), ST-LSTM (Liu et al., 2016), and STCNN (He et al., 2019) which are usually based on convolutional neural networks (CNNs) and recurrent neural networks (RNNs) to account for spatial and temporal information. However, these models are developed for the grid-based data and disregard the graph structure of traffic road networks. Due to this limitation, researchers moved into the application of graph-based deep learning models like graph neural networks.

Graphs provide vital knowledge about the spatial, temporal, and spatio-temporal relationships that can potentially improve performance of the final model. Recently, researchers employed the graph convolution network (GCN) (Kipf & Welling, 2017) to model the spatio-temporal interactions. DCRNN (Veeriah et al., 2015) is an early example of it that utilizes diffusion GCN with bi-directional random walks to model the spatial correlations as well as a GRU (Gated Recurrent Unit) based network to model the temporal correlations. GRU is an RNN-based method which is not effective and efficient in modeling long-range temporal dependencies. To address this limitation, works such as STGCN (Yu et al., 2018) and GraphWaveNet (Wu et al., 2021) utilize convolutional operations for both spatial and temporal domains. After the rise of attention-based models in deep learning, researchers realized that these two models still have some limitations in learning spatial and temporal correlations due to the limited capability of convolutional operations in capturing high-dimensional correlations. Therefore, two attention layers are later employed in ASTGCN (Guo et al., 2019) to capture the spatial and temporal correlations. However, these models are limited in capturing local relations and may lose local information due to the sensitivity of representations to the dilation operation. STSGCN (Song et al., 2020) used localized spatio-temporal subgraphs to enhance prior models in terms of capturing the local correlations. However, since the model formulation does not incorporate global information, this model still had limitations when it came to long-term forecasting and dealing with data that included missing entries or noise. In addition to the spatial graphs from predefined road networks, STFGNN (Li & Zhu, 2021) later introduced the use of Dynamic Time Warping (DTW) for data-driven spatial networks which helped the model to learn representations from different data-driven and domain-driven views. STFGNN also utilized a new fusion module to capture spatial and temporal graphs in parallel. However, due to the over-smoothing problem of deep GNNs (which happens when a GNN model learns similar node representations for all nodes in case of adopting deeper architecture), current GNN-based models incur some difficulties in learning rich spatio-temporal representations with deep layers.

Neural Differential Equations (NDEs) (Kidger, 2022) provide a new perspective of optimizing neural networks in a continuous manner using differential equations (DEs). In other words, DEs help to create a continuous depth generation mechanism for neural networks, provide a high-capacity function approximation, and offer strong priors on model space. There are a few types of NDEs: (1) neural ordinary differential equation (NODE) such as ODE-based RNN models (Habiba & Pearlmutter, 2020; Lechner & Hasani, 2022) (e.g., ODE-RNN, ODE-GRU, ODE-LSTM) and latent ODE models (Rubanova et al., 2019); (2) neural controlled differential equation (NCDE) (Choi et al., 2022; Kidger et al., 2020; Morrill et al., 2021) which is usually used to learn functions for the analysis of irregular time-series data; and (3) neural stochastic differential equation (NSDE) (Kidger et al., 2021b; a) which is usually employed for generative models that can represent complex stochastic dynamics. More recently, NODEs have been employed for the traffic forecasting task (Fang et al., 2021; Su et al., 2022; Pu et al., 2022). For example, STGODE (Fang et al., 2021) is proposed to address the aforementioned over-smoothing problem of GNN-based models. STGODE provides a new perspective of continuous depth generation by employing Neural Graph ODEs to capture the spatio-temporal dependencies. However, complex spatio-temporal correlations such as the lack of local temporal dependencies and dynamic node-edge communications are not captured by this model. In this study, we employ the idea of Neural Graph ODE for multiple coupled ODE-GNN modules which take into account the local temporal patterns and dynamic spatio-temporal dependencies, and as a result, can provide better function approximation for the forecasting in the presence of complex spatio-temporal correlations.

Definition 1: Traffic Graphs. We consider the traffic network as a graph $\mathcal{G}=(V,E,A)$, where $V$ is the set of nodes, $E$ is the set of edges and $A$ is the adjacency matrix such that $A\in\mathbb{R}^{N\times N}$ if $|V|=N$. In this paper, we use two types of graphs, the connection map graph and the DTW (dynamic time warping) graph. In the connection map graph, the adjacency matrix, denoted by $A^{C}$, represents the roads and actual connectivity among traffic nodes with binary value.

where $v_{i}$ and $v_{j}$ refer to traffic nodes $i$ and $j$ in the graph. In the DTW graph, the adjacency matrix, indicated by $A^{SE}$, is generated from the Dynamic Time Warping (DTW) algorithm (Berndt & Clifford, 1994) which calculates the distance between two time-series corresponding to a pair of nodes.

where $X^{i}$ and $X^{j}$ refers to the time-series data of nodes $i$ and $j$, respectively; $\epsilon$ also identifies the sparsity ratio of adjacency matrix. Notably, DTW is more effective than other point-wise similarity metrics (such as Euclidean distance) due to its sensitivity to shape and pattern similarities.

Definition 2: Graph Normalization. Adjacency matrix $A\in\{A^{C},A^{SE}\}\in\mathbb{R}^{N\times N}$ is normalized with $D^{-\frac{1}{2}}AD^{-\frac{1}{2}}$ where $D$ is the degree matrix of $A$. As shown in Eq. (3), the self-loop identity matrix $I$ is incorporated in normalization to avoid the negative eigenvalues.

where $\hat{A}$ is the normalized adjacency matrix; and $\alpha\in(0,1)$ identifies the eigenvalue range to be in $[0,\alpha]$.

Definition 3: Graph-based Neural ODE. The standard formulation of GNN-based continuous-time ODE function is defined in Eq. (4). It takes the initial value $\mathcal{H}(0)$, temporal integrals from $0$ to given time $t$, traffic graph $\mathcal{G}$, and the Neural ODE’s network parameter $\theta$.

where $f$ is the process to generate semantic message of Neural ODE parameterized by $\theta$ to model the hidden dynamics of graph $\mathcal{G}$. Since the structure of our proposed method consists of multiple ODE-GNN blocks, we have different types of Eq. (4) graph neural ODEs which are explained with more details in Section 4.

Definition 4: Tensor $n$-mode multiplication. We use subscript to identify the tensor-matrix multiplication on the specific dimension as shown below.

where $\mathcal{B}\in\mathbb{R}^{N_{1}\times N_{2}\times N_{3}}$, $\mathcal{C}\in\mathbb{R}^{N_{2}\times N_{2}^{\prime}}$, so, $\mathcal{B}\times_{2}\mathcal{C}\in\mathbb{R}^{N_{1}\times N_{2}^{\prime}\times N_{3}}$. The $n$-mode tensor-matrix multiplication is denoted as $\times_{n}$ with the $n$-th subscript.

The spatio-temporal forecasting problem is described as learning a mapping function $\mathcal{F}$ that transforms the current historical spatio-temporal data $\mathcal{X}=(X_{t-L+1},X_{t-L+2},...,X_{t})$ into the future spatio-temporal data $\mathcal{Y}=(X_{t+1},X_{t+2},...,X_{t+L^{\prime}})$, where $L$ and $L^{\prime}$ denote the length of historical and target time-series to be predicted, respectively. In the traffic forecasting problem, we have a historical tensor $\mathcal{X}\in\mathbb{R}^{B\times N\times L\times C}$ and traffic graph $\mathcal{G}$, where $B$ is the batch size; $N$ is the number of nodes; $L$ is the input temporal length; and $C$ is the number of input features (e.g., traffic speed, flow, density). The goal is to find $\hat{\mathcal{Y}}=\mathcal{F}(\mathcal{X},f,\mathcal{G})$ in which $\mathcal{F}$ is the overall forecasting network and $f$ corresponds to different graph-based Neural ODE processes.

The overview of our proposed GRAM-ODE is shown in Fig. 2, which is composed of two streams of operations: ($i$) DTW-based graph (top row) which is constructed based on semantical similarities, and ($ii$) connection map graph (bottom row) which is constructed based on geographical spatial connectivities. These two types of adjacency matrices are fed into the model separately to capture spatial correlations from both data-driven and geographical domain-driven views. They are later integrated with the multi-head attention mechanism in the final layer. As shown in Fig. 2, we have three parallel channels for each graph, each of which has two consecutive GRAM-ODE layers with a multi ODE-GNN block sandwiched between two blocks of temporal dilated convolution (TCN). The final features of each graph from different channels are then concatenated and imported into the attention module ($AM$) which is designed to effectively fuse these two separate sets of operations by taking into account all the high-dimensional relations towards the intended forecasting task. We provide further details about each of these components in the subsections below.

Each GRAM-ODE layer consists of a multi ODE-GNN block placed between two TCNs. Fig. 3 illustrates details of our proposed Multi ODE-GNN block in which we combine multiple ODE modules to take into account all the essential temporal and spatial patterns from different viewpoints and extract informative features from their fusion. In the multi ODE-GNN block, we design three types of message passing operations, a message filtering constraint as well as a new aggregation module to consider high-dimensional correlations.

We use the attention mechanism in the last layer of GRAM-ODE to effectively aggregate the final learned embeddings of two traffic graphs (i.e., DTW-based graph and the connection map graph) in a way that is better aligned with the forecasting objective. The attention module (AM) is designed to replace the previous fully connected layers while capturing the correlations of high-dimensional information. In this module, we first concatenate the embeddings of two graphs and then compute the attention scores between them. Eqs. (25) and (26) mathematically describe this attention operation.

where $W_{q}(b_{q}),W_{k}(b_{k})$, and $W_{v}(b_{v})$ represent the attention query, key, and value weight matrix (bias vector), respectively; $X$ is the input of attention module; $h$ is the number of heads; and $\sqrt{\frac{h}{C^{\prime}}}$ is the normalization factor with $C^{\prime}$ representing the embedding dimension. Therefore, $Q$:$\{Q_{1},Q_{2},...,Q_{h}\}$, $K$:$\{K_{1},K_{2},...,K_{h}\}$, $V$:$\{V_{1},V_{2},...,V_{h}\}$ shows the query, key, and value set for multiple attention heads. Finally, output of the attention module $X_{i}^{\prime}$ can be mapped to the feature space of original data for prediction with a linear dense layer.

In regression problem settings (such as the traffic flow forecasting), Huber loss function between the real value and predicted value is widely used in the literature. It combines the merits of $L1$ and $L2$ loss functions. Huber loss is a piece-wise function which consists of two parts: (1) a squared term with small and smooth gradient when the difference between the true and predicted values is small (i.e., less than a $\delta$ threshold value) and (2) a restricted gradient term when the true and predicted values are far from each other. Eq. (27) represents the standard form of Huber loss function.

where $\delta$ is a hyperparameter value set for the intended threshold; $\mathcal{Y}$ is the true future spatio-temporal data; and $\hat{\mathcal{Y}}$ is the predicted future data.

Algorithm 2 provides the complete pseudocode of GRAM-ODE training. In each optimization step of this algorithm, the outputs of all GRAM-ODE layers across parallel channels are concatenated, and these concatenated outputs of different graph types are then fed into the attention module for better aggregation and prediction.

We show the performance results of our model on six widely used public benchmark traffic datasets ¹¹1Datasets are downloaded from STSGCN github repository https://github.com/Davidham3/STSGCN/ : PEMS03, PEMS04, PEMS07, and PEMS08 released by (Song et al., 2020) as well as PEMS-BAY (Li et al., 2017) and METR-LA (Jagadish et al., 2014). The first four datasets (PEMS03, PEMS04, PEMS07, PEMS08) are collected based on the California Districts they represent (District 3, District 4, District 7, and District 8, respectively). PEMS-BAY covers the San Francisco Bay Area and METR-LA focuses on the traffic data of the Los Angeles Metropolitan Area. All these datasets collect three features (flow, occupation, and speed) at each location point over a period of time (with 5-minute time intervals). The spatial connection network for each dataset is constructed using the existing road network. The details of data statistics are shown in Table. 2. To use these datasets in experiments, we pre-process the features by z-score normalization.

We use Mean Absolute Error (MAE), Mean Absolute Percentage Error (MAPE), and Root Mean Squared Error (RMSE) metrics to evaluate the spatio-temporal forecasting. These metrics are defined as follows.

We compared our proposed GRAM-ODE with the following baselines.

• •

  ARIMA (Box & Pierce, 1970): Auto-Regressive Integrated Moving Average is one of the most well-known statistical models for time-series analysis.

• •

  DCRNN (Veeriah et al., 2015): Diffusion Convolutional Recurrent Neural Network utilizes diffusion graph convolutional networks with bidirectional random walks on directed graphs, and seq2seq gated recurrent unit (GRU) to capture spatial and temporal dependencies, respectively.

• •

  STGCN (Yan et al., 2018): Spatio-Temporal Graph Convolutional Network combines graph structure convolutions with 1D temporal convolutional kernels to capture spatial dependencies and temporal correlations, respectively.

• •

  GraphWaveNet (Wu et al., 2021): GraphWaveNet integrates adaptive graph convolution with 1D dilated casual convolution to capture spatio-temporal dependencies.

• •

  STSGCN (Song et al., 2020): Spatio-Temporal Synchronous Graph Convolutional Networks decompose the problem into multiple localized spatio-temporal subgraphs, assisting the network in better capturing of spatio-temporal local correlations and consideration of various heterogeneities in spatio-temporal data.

• •

  STFGNN (Li & Zhu, 2021): Spatio-Temporal Fusion Graph Neural Networks uses Dynamic Time Warping (DTW) algorithm to gain features, and follow STSGCN (Song et al., 2020) in using sliding window to capture spatial, temporal, and spatio-temporal dependencies.

• •

  STGODE (Fang et al., 2021): Spatio-Temporal Graph ODE Networks attempt to bridge continuous differential equations to the node representations of road networks in the area of traffic forecasting.

Following the previous works in this domain, we perform experiments by splitting the entire dataset into 6:2:2 for train, validation, and test sets. This split follows a temporal order, using the first 60% of the time length for training, and the subsequent 20% each for validation and testing. We use the past one hour to predict the future one hour. Since the time interval of data collection is 5 minutes, $L,L^{\prime}=12$ temporal data points. Based on Table 2, we have different number of sensors among different datasets, therefore, $|V|$ is different. DTW threshold ($\epsilon$) in Eq. (2) is 0.1; number of channels ($C$) in the historical data is $3$ (i.e., flow, speed, and occupation) and in the embedding space is $64$. The shared temporal weights $W_{s1},W_{s2}\in\mathbb{R}^{12\times 12}$ are initialized randomly from normal distribution. The length of latent space for the input of local ODE block is $L^{\prime\prime}=4$, and in the final attention module, number of attention heads $h=12$. During training, we use the learning rate of $10^{-4}$, $10^{-4}$, $10^{-5}$, $10^{-5}$, $10^{-4}$, and $10^{-4}$ for PEMS03, PEMS04, PEMS07, PEMS08, PEMS-BAY, and METR-LA datasets, respectively. The optimizer is AdamW. All experiments are implemented using PyTorch (Paszke et al., 2019) and trained using Quadro RTX 8000 GPU, with 48GB of RAM.

Our proposed GRAM-ODE outperforms other baselines on all datasets in Table 3, except for the MAPE metric in PEMS07, where it is slightly greater than that of STFGNN. ARIMA performs considerably worse than other baselines, likely because it ignores the graph structure of spatio-temporal data. GraphWaveNet performs relatively poorly, possibly due to its limited capability in stacking spatio-temporal layers and expanding the receptive field learned from 1D CNN temporal kernels. DCRNN uses bi-directional random walks and a GRU to model spatial and temporal information, respectively, but its relatively low modeling efficacy and efficiency for long-range temporal information may contribute to its subpar performance. STGCN uses GCN for the spatial domain and 1D dilated convolutional operations for the temporal domain, but may lose local information due to the dilation operation, while its absence of attention-based operations and limited capability of convolutional operations in modeling high-dimensional spatial, temporal, and spatio-temporal correlations may also result in relatively poor performance. STSGCN captures local correlations with localized spatio-temporal subgraphs but may miss global information and thus perform poorly in long-range forecasting and with data containing missing entries or noise. STFGNN uses DTW for latent spatial networks and performs well, but is limited in learning comprehensive spatio-temporal correlations. STGODE uses Neural ODEs to capture spatio-temporal dynamics and achieves very good performance compared to other baselines, but still lacks the ability to capture complex spatio-temporal correlations, balance local and global patterns, and model dynamic interactions between its components.

We selected two nodes from the PEMS08 road networks to conduct our case study for the qualitative demonstration of results. As Fig.4 shows, the predicted curves by our proposed GRAM-ODE (red curve) is better (more closely) aligned with the ground truth than STGODE (grey curve). The ground truth in node 109 has more fluctuations compared to the node 17 which causes more difficulty in the forecasting task. We can also observe that our model provides faster response in case of these abrupt fluctuations. This highlights the effectiveness of local ODE-GNN block in our model architecture which helps the model to better learn the local fluctuations.

To evaluate the robustness of our model, we add noise to the historical input of training data, which can potentially mimic uncertainties and biases that can arise during the data collection process. The added noise follows zero-mean i.i.d. Gaussian distribution with fixed variance, e.g., $\mathcal{N}(0,\gamma^{2})$, where $\gamma^{2}\in\{2,4\}$. We conduct robustness analysis across different values of $n\in\{0.1,0.2,0.3,0.4\}$ representing the ratio of training data impacted by noise. In other words, $n=0$ captures the performance without any noise. Fig. 6 represents the results of robustness comparisons between GRAM-ODE and STGODE on PEMS04 dataset across all the three metrics (MAE, MAPE, and RMSE). It can be observed that GRAM-ODE performance is more robust than STGODE with different levels of added noise $n=0.1;0.2;0.3;0.4$ which is probably due to the powerful spatio-temporal feature extraction gained by multiple ODE-GNN modules. We can also notice that, when the noise levels are high (with $\gamma^{2}=4$ and $n=0.4$), GRAM-ODE can still beat many baseline models listed in Table 3, demonstrating the significant benefit of incorporating various ODE modules in our framework which can improve robustness.