Cross Space and Time: A Spatio-Temporal Unitized Model for Traffic Flow Forecasting

TL;DR

The Spatio-Temporal Unitized Model (STUM), a unified framework designed to capture both spatial and temporal dependencies while addressing spatio-temporal heterogeneity through techniques such as distribution alignment and feature fusion, is introduced.

Abstract

Predicting spatio-temporal traffic flow presents significant challenges due to complex interactions between spatial and temporal factors. Existing approaches often address these dimensions in isolation, neglecting their critical interdependencies. In this paper, we introduce the Spatio-Temporal Unitized Model (STUM), a unified framework designed to capture both spatial and temporal dependencies while addressing spatio-temporal heterogeneity through techniques such as distribution alignment and feature fusion. It also ensures both predictive accuracy and computational efficiency. Central to STUM is the Adaptive Spatio-temporal Unitized Cell (ASTUC), which utilizes low-rank matrices to seamlessly store, update, and interact with space, time, as well as their correlations. Our framework is also modular, allowing it to integrate with various spatio-temporal graph neural networks through components such as backbone models, feature extractors, residual fusion blocks, and predictive modules to collectively enhance forecasting outcomes. Experimental results across multiple real-world datasets demonstrate that STUM consistently improves prediction performance with minimal computational cost. These findings are further supported by hyperparameter optimization, pre-training analysis, and result visualization. We provide our source code for reproducibility at https://anonymous.4open.science/r/STUM-E4F0.

Figures (6)

• Figure 1: Motivation of our proposed method. (a) shows the sensor distribution of the PEMS04 dataset. (b) is a visual result of the traffic flow of a pair of residential areas over a random period. And (c) is spatio-temporal dependencies shown in traffic flow prediction tasks.
• Figure 2: The overview of our proposed method. (a) shows the architecture of the Spatio-Temporal Unitized Model (STUM), where MLP represents the model prototype and STGNN represents a way of enhancement. (b) shows the computing process of The Multi-Layer Residual Fusion (MLRF) blocks. (c) shows the construction of Adaptive Spatio-temporal Unitized Cells and how the information transmission Cross Space and Time.
• Figure 3: Comparison of traditional and our proposed Methods. The traditional method (a) separates temporal and spatial modeling, while our approach (b) integrates them using low-rank matrix factorization for better joint prediction.
• Figure 4: The Efficiency Study. The results compare the variation in training time and the reduction in MAE metrics for STUM equipped with all six baselines as backbone networks on the PEMS04 dataset.
• Figure 5: Results of the t-SNE visualization of the Spatio-temporal Unitized Model embedding on the PEMS04 dataset. The left part represents the embedding space obtained using STAEFormer as the backbone extractor, and the right side represents the embedding space enhanced by Multi-Layer Residual Fusion equipped with four Adaptive Spatio-temporal Unitized Cells.