Abstract
Accurate long-term traffic prediction is vital for intelligent transportation systems but remains challenging due to the complex entanglement of global trends and local fluctuations. We present DST-Mamba (Decomposed Spatio-Temporal Mamba), a novel framework that disentangles traffic signals into trend and seasonal components. By employing a bi-directional Mamba architecture with linear complexity, DST-Mamba effectively captures spatial dependencies across road networks while overcoming the scalability issues of traditional Transformers. Evaluations on five real-world benchmarks demonstrate that DST-Mamba achieves state-of-the-art performance in long-term forecasting.
Introduction
Predicting urban traffic flow is a cornerstone of smart city infrastructure, influencing congestion management, emergency response, and autonomous route planning. While Transformer-based models have advanced the field through global dependency modeling, they face two critical limitations in traffic scenarios. First, traffic data is physically constrained by road networks, leading to a complex entanglement of long-term global trends and high-frequency "seasonal" fluctuations. Second, the quadratic complexity ($O(L^2)$) of self-attention mechanisms makes Transformers computationally expensive for long-term forecasting on large-scale networks.
Methodology
DST-Mamba addresses these challenges through a "divide and conquer" strategy, decomposing the input signal into distinct components for specialized processing.
Temporal Decomposition and Trend Modeling
The model first splits the input signal into a Trend component (smooth, global patterns) and a Seasonal component (high-frequency fluctuations) using moving averages. The Trend module utilizes a multi-scale linear prediction strategy. It employs "top-down mixing" to down-sample the trend into multiple scales, capturing macroscopic patterns while effectively filtering out high-frequency noise.
Spatio-Temporal Mamba for Seasonal Dynamics
The core innovation lies in the Seasonal Component module, which shifts the modeling perspective from time to space:
- Spatial Tokenization: Instead of temporal steps, road network nodes (sensors) are treated as tokens. Graph aggregation, incorporating the road network's adjacency matrix, embeds spatial relationships directly into these tokens.
- Bi-directional Mamba: Unlike uni-directional sequence models, DST-Mamba uses a bi-directional approach—processing node sequences in both forward and backward orders. This is crucial for spatial networks where influence flows in multiple directions, ensuring a comprehensive global context.
- Linear Complexity: By leveraging the Selective State Space Model (Mamba), the architecture achieves near-linear complexity ($O(N)$) relative to the number of nodes, significantly improving scalability compared to Transformers.
Experimental Evaluation
DST-Mamba was evaluated on five widely-used traffic benchmarks: Traffic (862 nodes) and four PEMS datasets (PEMS03, 04, 07, 08) from California highway data. The model consistently outperformed eight major baselines, including state-of-the-art Transformers like iTransformer and PatchTST.
Key Results and Insights
Experimental results highlight several critical advantages of the DST-Mamba framework:
- Decomposition Efficacy: Ablation studies confirm that modeling the entangled raw signal yields inferior results. Specialized architectures for trend (Linear) and seasonal (Mamba) components are essential for accuracy.
- Spatial Bi-directionality: The bi-directional Mamba block is vital; uni-directional variants lose significant semantic information regarding spatial correlations.
- Predictive Robustness: Visualization shows the model accurately tracks both smooth daily cycles and sudden traffic spikes.
On the PEMS08 dataset, DST-Mamba achieved the lowest Mean Squared Error (MSE) and Mean Absolute Error (MAE) across all horizons (12 to 96 steps). Its linear complexity allows it to maintain a favorable computational profile even as network size increases.
Conclusion and Future Work
DST-Mamba provides a robust and efficient solution for long-term traffic forecasting by combining signal decomposition with the efficiency of state-space models. While highly effective, the model's reliance on graph quality (adjacency matrices) suggests that future work could explore more dynamic or learned graph structures. Additionally, further refining the integration of node-specific attributes into the Mamba blocks could enhance the representation of heterogeneous urban environments.