ST-Mamba: Spatial-Temporal Selective State Space Model for Traffic Flow Prediction

TL;DR

This work addresses the challenge of accurately forecasting traffic flow while balancing computational efficiency. It introduces ST-Mamba, a Spatial-Temporal Selective State Space model that forgoes graph-based representations in favor of a unified spatial-temporal processing framework via the ST-Mixer and ST-SSM block, enabling effective long-range dependency modeling with reduced complexity. The approach achieves state-of-the-art computational efficiency (61.11% faster) and competitive accuracy (0.67% improvement on targeted metrics) across multiple real-world datasets, outperforming many graph-based and transformer-based baselines. Practically, ST-Mamba supports real-time traffic management with lower resource requirements and is extensible to additional data sources and longer horizons. The paper also provides rigorous ablations and a complexity analysis, showing linear scaling with sequence length and demonstrating the potential for deployment in resource-constrained transportation systems.

Abstract

Traffic flow prediction, a critical aspect of intelligent transportation systems, has been increasingly popular in the field of artificial intelligence, driven by the availability of extensive traffic data. The current challenges of traffic flow prediction lie in integrating diverse factors while balancing the trade-off between computational complexity and the precision necessary for effective long-range and large-scale predictions. To address these challenges, we introduce a Spatial-Temporal Selective State Space (ST-Mamba) model, which is the first to leverage the power of spatial-temporal learning in traffic flow prediction without using graph modeling. The ST-Mamba model can effectively capture the long-range dependency for traffic flow data, thereby avoiding the issue of over-smoothing. The proposed ST-Mamba model incorporates an effective Spatial-Temporal Mixer (ST-Mixer) to seamlessly integrate spatial and temporal data processing into a unified framework and employs a Spatial-Temporal Selective State Space (ST-SSM) block to improve computational efficiency. The proposed ST-Mamba model, specifically designed for spatial-temporal data, simplifies processing procedure and enhances generalization capabilities, thereby significantly improving the accuracy of long-range traffic flow prediction. Compared to the previous state-of-the-art (SOTA) model, the proposed ST-Mamba model achieves a 61.11\% improvement in computational speed and increases prediction accuracy by 0.67\%. Extensive experiments with real-world traffic datasets demonstrate that the \textsf{ST-Mamba} model sets a new benchmark in traffic flow prediction, achieving SOTA performance in computational efficiency for both long- and short-range predictions and significantly improving the overall efficiency and effectiveness of traffic management.

Figures (6)

• Figure 1: An illustration of the conversion of road maps into spatial-temporal information.
• Figure 2: The framework of ST-Mamba model.
• Figure 3: The intricacies of Mamba
• Figure 4: Comparison of FLOPS vs total computational time on ST-Mamba and STAEformer in PEMS08.
• Figure 5: Comparison of RMSE, MAE, and MAPE across various layers for STAEformer and ST-Mamba.