Attention Mechanism Based Spatial-Temporal Graph Convolution Network for Traffic Prediction

Attention Mechanism Based Spatial-Temporal Graph Convolution Network for Traffic Prediction

Considering the complexity of traffic systems and the challenges brought by various factors in traffic prediction, we propose a spatial-temporal graph convolutional neural network based on attention mechanism (AMSTGCN) to adapt to these dynamic changes and improve prediction accuracy. The model combines the spatial feature extraction capability of graph attention network (GAT) and the dynamic correlation learning capability of attention mechanism. By introducing the attention mechanism, the network can adaptively focus on the dependencies between different time steps and different nodes, effectively mining the dynamic spatial-temporal relationships in the traffic data. Specifically, we adopt an improved version of graph attention network (GAT_v2) in the spatial dimension, which allows the model to capture more complex dynamic spatial correlations. Furthermore, in the temporal dimension, we combine gated recurrent unit (GRU) structure with an attention mechanism to enhance the model’s ability to process sequential data and predict traffic flow changes over prolonged periods. To validate the effectiveness of the proposed method, extensive experiments were conducted on public traffic datasets, where AMSTGCN was compared with five different benchmark models. Experimental results demonstrate that AMSTGCN exhibits superior performance on both short-term and long-term prediction tasks and outperforms other models on multiple evaluation metrics, validating its potential and practical value in the field of traffic prediction.