An efficient spatial-temporal transformer with temporal aggregation and spatial memory for traffic forecasting

Traffic forecasting technology has widespread applications in various domains, such as urban traffic planning and intelligent transportation systems. Traffic forecasting encounters challenges in effectively capturing the intricate spatial-temporal correlations in traffic data. While the latest methods have achieved satisfactory performance, they still suffer from two limitations: (i) Most methods overlook the memory of valuable traffic patterns at each traffic node, thus making it struggle to reveal dynamic spatial-temporal correlations using their inherent periodicity and trend characteristics from a broader perspective. (ii) As the research progresses, recently proposed models become increasingly complex and massive. To address these issues, we propose a Spatial-Temporal Aggregation Memory Transformer (STAMT) for traffic forecasting. Specifically, we propose a memory bank to enhance vanilla spatial attention and cache the traffic patterns of historical input. By querying these traffic patterns, which contain rich spatial-temporal semantic information, the model can optimize prediction performance by extracting trends and regularities across various periods. To reduce the computational costs, we introduce a temporal module to capture temporal correlations while reducing temporal dimension information. In addition, we leverage the random feature map and matrix multiplication associativity property, which reduce the quadratic complexity of spatial attention to linearity with regard to the number of nodes. Ultimately, our theoretical analysis concludes that a single-layer spatial attention network is sufficient to capture spatial-temporal correlations deeply without stacking. Extensive experiments on nine real-world datasets demonstrate that STAMT outperforms state-of-the-art baselines in regular, long-range, and large-scale traffic forecasting tasks while significantly reducing the computational costs. Codes are available at https://github.com/LiuAoyu1998/STAMT.