A Lightweight Spatio-Temporal Cosine Transformer (STCformer) for Non-Stationary Traffic Flow Prediction

Traffic flow prediction faces challenges such as prediction difficulties caused by data non-stationarity and the high computational complexity of existing attention mechanisms. To overcome these challenges, this paper proposes a traffic flow prediction model named STCformer based on non-stationary normalization and cosine attention. The model first eliminates sequence distribution drift through a non-stationary normalization module; second, it employs a GRU residual learning framework to efficiently capture dynamic residual information and introduces cosine similarity to replace the Softmax-based attention mechanism, reducing the computational complexity to linear. Experimental results on four real-world datasets show that STCformer outperforms the compared baseline models in prediction performance. In particular, in the 12-step prediction task on the PEMS08 dataset, compared with the second-best baseline model, MAE, RMSE, and MAPE are reduced by approximately 11.62%, 5.77%, and 8.33%, respectively, demonstrating excellent long-term prediction capability. Computational cost experiments show that on the PeMS07 dataset, compared with traditional attention variants, STCformer reduces GPU memory usage by 66.4% and training and inference time by 36.6% and 37.8%, respectively, meeting the real-time prediction requirements of large-scale road networks.