Unified Spatio-Temporal Tokens are Bases for Generalizable Traffic Forecasting

Abstract

Traffic forecasting plays a crucial role in real-world applications such as traffic management and urban planning. Recent studies have mainly focused on spatio-temporal graph neural networks (STGNNs) and attention-based methods, which have shown promising results. Nevertheless, both approaches model spatial information implicitly, which limits their ability to generalize across different traffic networks. In this paper, we propose Spatio-Temporal Unified Network (STUNet), a framework to explicitly encode spatial features into unified representations and integrate them with temporal information effectively. To obtain spatial representations explicitly, we design a spatial tokenizer that segments the adjacency matrix of the relation graph into patches to serve as spatial tokens. Furthermore, to effectively integrate spatial and temporal representations, we introduce query-aggregate attention, which simulates the process of tracing upstream and downstream nodes and aggregating their information, thereby capturing complex spatio-temporal dependencies. Extensive experiments on traffic benchmarks demonstrate that STUNet achieves generalization across different traffic networks with competitive performance. Code is available at https://github.com/JimmyChen6/STUNet.