Trans-Mamba: The Cross-network of Transformer and Mamba for Traffic Flow Prediction

Traffic flow prediction is a crucial aspect of intelligent transportation systems. It holds significant importance in optimizing traffic management and enhancing road safety. However, the task of traffic flow forecasting remains a significant challenge due to its inherent complexity, dynamic nature, and the presence of long-range dependencies. Existing graph neural network methods exhibit two primary limitations. Firstly, when static graphs are employed to capture spatial dependencies, they fail to account for the dynamics in these dependencies. Secondly, these methods struggle to effectively capture long-range spatial dependencies. Furthermore, attention-based methods, while potentially addressing some of these issues, require substantial computational and memory resources, leading to inefficiencies. We design a cross-network model that combines Transformer and Mamba architectures for traffic flow prediction, which efficiently captures dynamic long-range spatio-temporal dependencies. The model architecture is composed of three main components: a learnable spatio-temporal embedding layer that adaptively captures the unique features of different nodes and time; stacked spatio-temporal blocks, which include two temporal blocks and one spatial block arranged between them to effectively capture spatio-temporal dependencies; and a fully connected output layer that generates the final predictions. Experiments on three public real-world datasets demonstrate the effectiveness of TransMamba. Specifically, the model improves the mean absolute error (MAE) and root mean square error (RMSE) by 2.59% and 4.14%, respectively, on the PeMS08 dataset. Additionally, ablation studies confirm the effectiveness and efficiency of each component of the model.