M2STFormer: A Mixture-of-Experts Multi-level Spatio-Temporal Transformer for Traffic Flow Prediction

Traffic flow prediction is a fundamental task in intelligent transportation systems. A key challenge lies in effectively capturing the complex and dynamic spatio-temporal dependencies in traffic data. In recent years, Graph Neural Networks (GNN) and Transformer-based models have shown great promise in addressing this challenge. However, the inherent spatio-temporal heterogeneity of traffic flow data remains a significant obstacle. Existing models often struggle to model the rapidly changing spatio-temporal dependencies among road segments and differentiate the importance of extracted features at multiple spatial and temporal scales. To this end, we propose M2STFormer, a Mixture-of-Experts Multi-Level Spatio-Temporal Transformer (M2STFormer) for traffic flow prediction. Our model incorporates a Mixture-of-Experts (MoE) module to learn unique patterns of each spatio-temporal location, and a hierarchical selection mechanism to adaptively fuse multi-level spatio-temporal features. We validate the effectiveness of M2STFormer through extensive experiments on two public datasets and two newly collected real-world datasets. The results show that our method achieves state-of-the-art performance across multiple datasets compared to single-stage baselines. Furthermore, while greatly improving computational efficiency, our approach matches or surpasses the performance of representative two-stage pretraining methods on several key metrics.