A Memory-Efficient SubGraph Learning Method for Urban Spatio-Temporal Forecasting

Abstract

Spatio-temporal graph learning has been widely applied to urban spatio-temporal forecasting tasks, analyzing the relationships between data to achieve accurate predictions. However, in recent years, model structures have become increasingly complex in pursuit of higher performance, resulting in greater computational complexity and GPU memory consumption during execution. In addition, urban expansion requires models to cover larger areas, further increasing memory demands. These factors collectively limit the applicability of existing models in real-world scenarios. To address these challenges, we propose a plug-and-play SubGraph Learning (SGL) method to reduce the memory overhead without compromising performance. Specifically, we introduce a SubGraph Partition Module (SGPM), which leverages a set of learnable memory vectors to select node groups with similar features from the graph, effectively partitioning the graph into smaller subgraphs. Noting that partitioning the graph may lead to feature redundancy, as overlapping information across subgraphs can occur. To overcome this, we design a SubGraph Feature Aggregation Module (SGFAM), which mitigates redundancy by averaging node features from different subgraphs. In addition, we propose a SubGraph Distributed Computation (SGDC) method that divides subgraphs into multiple groups and distributes them across different GPUs for feature extraction, thereby reducing the load on any single GPU during the subgraph processing phase. Experiments on four traffic network datasets of various scales demonstrate that SGL can significantly reduce memory overhead, achieving up to a 56.4% reduction in average GPU memory overhead, while maintaining robust prediction performance.