SAST-GNN: A Self-Attention Based Spatio-Temporal Graph Neural Network for Traffic Prediction

Abstract

Traffic prediction, which aims at predicting future traffic conditions based on historical observations, is of considerable significance in urban management. However, such tasks are challenging due to the high complexity of traffic flow. Traditional statistical learning methods perform badly in handling such regularities. Furthermore, these methods often split spatial information and temporal information as individual features, which cause extracted features hard to be fused. In this paper, we target on efficiently capturing such highly nonlinear dependencies and fusing features from both spatial and temporal dimensions to predict future traffic flow conditions accurately. To tackle this problem, we proposed the Self-Attention based Spatio-Temporal Graph Neural Network (SAST-GNN). In SAST-GNN, we innovatively proposed to add a self-attention mechanism to more accurately extract features from the temporal dimension and the spatial dimension simultaneously. At the same time, for better fusing information from both dimensions, we improved the spatio-temporal architecture by adding channels(or heads) and residual blocks. Thus features can be fused from different aspects. Experiments on two real-world datasets illustrated that our model outperforms six baselines in traffic flow prediction tasks. Especially in short-term and mid-term prediction based on long-term dependencies, our model performs much better than other baselines.