Let’s Group: A Plug-and-Play SubGraph Learning Method for Memory-Efficient Spatio-Temporal Graph Modeling

Let’s Group: A Plug-and-Play SubGraph Learning Method for Memory-Efficient Spatio-Temporal Graph Modeling

Wenchao Weng, Hanyu Jiang, Mei Wu, Xiao Han, Haidong Gao, Guojiang Shen, Xiangjie Kong

Proceedings of the Thirty-Fourth International Joint Conference on Artificial Intelligence
Main Track. Pages 3471-3479. https://doi.org/10.24963/ijcai.2025/386

Spatio-temporal graph modeling is widely applied to spatio-temporal data, analyzing the relationships between data to achieve accurate predictions. However, despite the excellent predictive performance of increasingly complex models, their intricate architectures result in significant memory overhead and computational complexity when handling spatio-temporal data, which limits their practical applications. 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. 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. The source code is available at https://github.com/wengwenchao123/SubGraph-Learning.
Keywords:
Data Mining: DM: Mining spatial and/or temporal data
Machine Learning: ML: Applications
Machine Learning: ML: Time series and data streams
Multidisciplinary Topics and Applications: MTA: Sensor networks and smart cities