Spatio-Temporal Attention-based Graph Convolution Networks for Traffic Prediction

Yiling Chen, Linjiang Zheng, Weining Liu

Abstract: Accurate traffic prediction is critical to the effectiveness of intelligent transportation systems. However, traffic data are highly nonlinear with complicated dynamic spatio-temporal correlations, accurate traffic forecasting, particularly long-term forecasting, remains a difficulty. Existing models generally learn a fixed graph to capture spatial correlations, which makes them difficult to effectively capture changing spatial dynamics over time, resulting in poor prediction outcomes. To tackle these challenges, we propose a new spatio-temporal graph convolution network model, named Spatio-Temporal Attention-based Graph Convolution Network (STAGCN), which jointly fuses dynamic evolving spatial correlations and long-term temporal correlations to improve the performance. First, we design a spatio-temporal multi-head self-attention module to capture both spatial heterogeneity and temporal correlations. Second, we propose an adaptive evolving graph convolution module that can learn a new graph at each time step and make STAGCN evolve dynamically. Meanwhile, self-attention is used to construct a dynamic adjacency matrix to further capture spatial correlations. In addition, we optimize the original Transformer by employing relation-aware attention mechanism to make it better suitable for time series prediction, thereby improving the long-term prediction performance of the model. Extensive experiments are conducted on two real-world datasets, demonstrating that our proposed model achieves state-of-the-art performance and consistently outperforms other baselines.