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

1 Introduction

In recent years, advances in data collection and processing capabilities have generated large volumes of high-quality spatio-temporal data (e.g., weather conditions, traffic status) [Wang et al., 2024b; Zeng et al., 2025; Shen et al., 2025; He et al., 2025], which have been widely applied in related research fields. Due to its spatio-temporal correlations can be effectively captured by Spatio-Temporal Graph Neural Networks (STGNNs) [Wu et al., 2019; Jiang et al., 2023b; Deng et al., 2024], these learning algorithms have emerged as the dominant approach for spatio-temporal data modeling tasks.

∗ corresponding Author.

Existing STGNNs typically take spatio-temporal sequences collected by nodes (e.g., traffic segments, weather stations) as input and model them as spatio-temporal graphs to extract features (e.g., through graph convolution or attention mechanisms) [Shao et al., 2022b; Jiang et al., 2023a; Guo et al., 2023]. Early STGNNs utilized geographic information to construct graph structure [Li et al., 2018], but these methods neglected the hidden spatio-temporal relationships in the data, leading to suboptimal performance. As a result, recent studies generally focus on capturing relationships through dynamic construction of spatio-temporal graphs, followed by feature extraction from the data [Weng et al., 2023; Li et al., 2023; Dong et al., 2024]. However, these approaches have high computational complexity, resulting in substantial GPU memory overhead for STGNNs [Liu et al., 2023]. In addition, as the number of nodes increases, the memory overhead of STGNNs grows significantly, which limits the realworld applications [Wang et al., 2024a; Liu et al., 2024].

In this paper, we aim to address the challenge of spatiotemporal graph modeling under limited GPU memory constraints. One intuitive solution is to partition the spatiotemporal graph into smaller subgraphs, extract features from these subgraphs, and then aggregate node features across the subgraphs to learn the features of the entire spatio-temporal graph. This approach effectively reduces the originally high computational complexity to a manageable level, thereby saving GPU memory. However, the partitioning of the spatiotemporal graph presents a challenge. Existing partitioning methods typically rely on static partitioning, without considering the dynamic nature of the spatio-temporal graph. For example, PatchSTG [Fang et al., 2024] uses the geographical coordinates of nodes to partition subgraphs, while FCGCN [Xia et al., 2022] employs the Louvain algorithm to partition subgraphs based on the topological relationships among nodes. However, all of these methods fail to capture complex spatio-temporal dependencies.

For the above challenges, we propose a plug-and-play SubGraph Learning (SGL) method. Specifically, we design a SubGraph Partition Module (SGPM), which can be directly integrated into existing spatio-temporal graph models to partition the graph and significantly reduce the memory overhead required by the model. SGPM employs a set of learnable memory vectors as anchors and selects the most similar group of nodes in the spatiot-emporal graph as subgraphs by calculating the similarity between node features and memory vector features. This approach partitions the spatio-temporal graph into multiple smaller subgraphs and uses these subgraphs to extract node features, thus drastically reducing computational complexity and significantly lowering the GPU memory overhead.

At the same time, the way processing each divided subgraph separately solves the memory overhead and complexity problem. However, it brings a potential risk of feature redundancy. Therefore, we subsequently design a SubGraph Feature Aggregation Module (SGFAM) to eliminate it. Specifically, SGAM firstly extracts features of the node from different subgraphs it has appeared, and then averages these features. According to these operations, the node feature is more balanced and completeness. Overall, our contributions are as follows:

• We propose a plug-and-play SubGraph Learning (SGL) method to reduce the memory overhead required by existing models. SGL uses a SubGraph Partition Module (SGPM), which partitions the spatio-temporal graph into multiple subgraphs for feature extraction, effectively reducing the model's computational complexity and significantly lowering the computational cost.
• We introduce a SubGraph Feature Aggregation Module (SGFAM). This module efficiently integrates the feature of nodes across various subgraphs, avoiding redundancy in node feature and improving the completeness of spatio-temporal graph feature representation.
• We evaluate the effectiveness of SGL on traffic datasets of different scales. Experimental results show that our proposed method can effectively reduce the computational cost while maintaining model performance, with a maximum average GPU memory overhead reduction of 56.4%.

2 RelatedWork

Spatio-temporal graph construction. As the key component of spatio-temporal graph models, spatio-temporal graph construction has attracted significant attention from researchers. DCRNN [Li et al., 2018] uses predefined matrix based on road segment distances as spatio-temporal graphs. AGCRN [Bai et al., 2020] generates adaptive graphs using adaptive embeddings to learn the correlations between nodes. STFGNN [Li and Zhu, 2021] constructs a Spatial-Temporal Fusion Graph using the FastDTW algorithm to extract local correlations. Attention mechanisms also play a crucial role in spatio-temporal graph learning. For example, PDformer [Jiang et al., 2023a] uses DTW and predefined matrix as mask matrix to calculate the weights of the graph structure. STWave [Fang et al., 2023] constructs a local graph attention network using predefined matrix to efficiently model spatial correlations. LarSTL [Wang et al., 2024a] uses the METIS

algorithm [Karypis and Kumar, 1998] and geographic coordinates to partition the traffic road network into multiple subgraphs for continuous learning. However, these methods still model spatio-temporal graphs with static relationships, without accounting for the dynamic nature of spatio-temporal graphs.

Spatio-temporal graph neural networks. In recent years, with the rise of deep learning, Spatio-Temporal Graph Neural Networks (STGNNs) have emerged as one of the most representative methods for capturing spatio-temporal dependencies , becoming a rising star in the field of spatio-temporal graph modeling. For example, GMAN [Zheng et al., 2020] uses a transformed attention mechanism to reduce biases between the past and the future. STAEformer [Liu et al., 2023] introduces spatio-temporal adaptive embedding to enhance the performance of Transformers. DyHSL [Zhao et al., 2023] designs a dynamic hypergraph learning method to capture the dynamic and complex relationships in traffic networks. PDG2Seq [Fan et al., 2024] leverages periodic information to generate dynamic graphs and utilizes an RNN-based iterative prediction method to capture future trend changes. Additionally, some studies have approached model performance enhancement from other perspectives. For instance, STEP [Shao et al., 2022a] and STD-MAE [Gao et al., 2024] enhance performance by designing pre-training models to generate rich contextual representations that can seamlessly integrate with any architecture. EXPERT [Lee and Ko, 2024] proposes a novel Mixture of Experts (MoE) framework, which generates new graphs based on evolving environmental conditions to address the issue of spatial distribution drift during testing. However, these methods come with substantial computational overhead despite enhancing accuracy.

3 Method

In this section, we first provide the mathematical definition of the problem studied in this paper. Next, we introduce the feature extraction process of spatial feature extraction modules commonly used in current spatio-temporal graph models. Finally, we explain in detail the working principles of SGPM and SGFAM and how they integrate into existing models.

3.1 Spatio-Temporal GraphModeling

A spatio-temporal graph is represented by G = (V, E), where V is the set of nodes and E is the set of edges representing the adjacency relationships between road segments. At time point t, the historical feature data of graph G can be represented as Xt = {x1t , x2t , ..., xNt } ∈ RN×d , where xit represents the traffic conditions of node i at time t, d represents the number of features, N is the number of nodes.

We aim to learn a spatio-temporal prediction model g() that can effectively predict the feature data for the next Q time steps using the historical feature data from the past P time steps, X = {X1, ..., XP } ∈ RP×N×d . The goal is to predict the future data Y = {XP+1, ..., XP+Q} ∈ RQ×N×d :

3.2 Spatio-Temporal Graph Feature Extraction

In existing STGNNs, spatio-temporal graph features are typically extracted using graph convolution or attention mechanisms [Guo et al., 2021]. As shown in Figure 1, the spatiotemporal graph feature extraction process can be divided into three steps.

First, feature extraction is applied to each node, mapping its original features into the latent space and enhancing the node representations. It can be represented as:

where Xin ∈ RN×C represents the input spatio-temporal data, H ∈ RN×D represents the extracted hidden features, and Fin is the feature extraction function used to extract the hidden features from the spatio-temporal data. This function can be implemented using methods like MLP, among others.

According to spatio-temporal characteristics of each node, the correlation matrix construction module then represents the relationships among nodes in the graph. This can be expressed as:

where A ∈ RN×N represents the correlation matrix between nodes, and G() is the method for constructing the correlation matrix. This can be achieved using techniques such as dynamic graph construction or attention score matrix construction.

Finally, the feature weighting module enhances the representations of each node through the spatiotemporal correlation matrix extracted by the previous step. This process can be expressed as:

where Hout ∈ RN×C represents the extracted spatiotemporal graph features, Fout() transforms the extracted features into a format suitable for input to the next layer.

From the above formulas, it can be seen that the complexity of spatio-temporal graph feature extraction is highly dependent on the number of nodes N.

3.3 SubGraph Partition

The limited GPU memory capacity imposes restrictions on the scale of the input graph. To address this challenge, we propose a SubGraph Partition Module (SGPM). As shown in Figure 1, SGPM is placed after the feature extraction function and partitions the spatio-temporal graph into multiple smaller subgraphs based on similarity. By extracting features from these subgraphs separately, this approach reduces the GPU memory capacity required.

As shown in Figure 2, SGPM uses a set of randomly initialized memory vectors P = [P1, P2, ..., PM] ∈ RM×D as anchors to construct subgraphs.

Specifically, we first compute the similarity between each memory vector Pi and the node features H, which can be represented as:

where wi ∈ RN×1 represents the similarity between the memory vector Pi and the node features H.

Since highly correlated nodes tend to have similar features, Pi can serve as an anchor point to select nodes with higher similarity, forming a subgraph, which can be expressed as:

where idxi ∈ RK represents the indices of the top K nodes with the highest similarity to Pi. Using idxi, we collect the corresponding nodes and their features to form a subgraph, which can be represented as:

thus, a subgraph is obtained for each memory vector, and the related subgraph indices and feature sets can be represented as: idx = {idx1, ..., idxM} ∈ RM×K , and Hsub = {Hsub 1 , ...,H sub M } ∈ RM×K×D .

3.4 Subgraph Feature Extraction

Similar to the general spatio-temporal graph feature extraction methods, each subgraph extracts the features between its nodes by constructing a correlation matrix, represented as:

where Asub i ∈ RK×K represents the correlation matrix of subgraph i, and all subgraphs share the same G(). The collection of correlation matrices for all subgraphs is denoted as Asub = {Asub 1 , ...,Asub M } ∈ RM×K×K .

Finally, the spatial features of the subgraph nodes are extracted by multiplying the node features of the subgraph with the correlation matrix:

where Hsub out ∈ RM×K×C represents the spatio-temporal graph features extracted from the subgraph.

3.5 SubGraph Feature Aggregation

As shown in Figure 1, we propose a SubGraph Feature Aggregation Module (SGFAM), placed after the feature weighting layer, to integrate node features from different subgraphs. As shown in Figure 3, we use the subgraph indices idx to collect the features of the same nodes across different subgraphs. This process can be expressed as:

where Hi represents the set of features for node i extracted from all subgraphs. To ensure a more comprehensive representation of the extracted features, we average these features to obtain the final output feature, expressed as:

where Hiout ∈ RN×D is the final hidden feature for node i, sum() represents the sum of features in Hi , and count() denotes the number of features in Hi .

Finally, the aggregated hidden features for all nodes are represented as:

where Hout ∈ RN×C represents the final output of SGFAM.

3.6 Complexity Analysis

In traditional spatio-temporal graph feature extraction, the complexity is tightly coupled with the number of nodes N due to the construction of the correlation matrix and the feature weighting process. For instance, both graph convolution and attention mechanisms have a complexity of O(N 2 ). In contrast, under the subgraph learning approach, spatiotemporal feature extraction consists mainly of three components: subgraph partitioning, correlation matrix construction, and node feature aggregation. The complexities of these components are as follows: subgraph partitioning is O(NM +MK), correlation matrix construction is O(K2 ) (since subgraphs are computed in parallel), and node feature aggregation is O(MK+N). Thus, the overall complexity of spatio-temporal feature extraction using the subgraph learning approach is O(NM + MK + K2 + N). Since M is significantly smaller than K and N, it can be considered a constant, simplifying the complexity to O(N +K2 ). When K is considerably smaller than N, the computational complexity of the subgraph learning method is significantly lower than that of the traditional spatio-temporal graph feature extraction methods, leading to a substantial reduction in memory consumption.

4 Experiment

To evaluate the performance of our proposed SGL, we selecte a series of representative spatio-temporal models that utilize dynamic graph convolution or attention mechanisms to extract spatial features. We integrate SGPM and SGFAM into their spatial feature extraction modules for subgraph partitioning and feature aggregation to assess the effectiveness of SGPM and SGFAM. In this work, we aim to address the following questions:

• Q1: How do SGL and dynamic graph convolution or attention mechanisms in different backbone networks perform in terms of spatio-temporal graph modeling tasks, efficiency, and memory overhead?
• Q2: Are SGPM and SGFAM effective?
• Q3: How sensitive is SGPM to the number of subgraphs and the number of nodes per subgraph?
• Q4: What are the characteristics of the nodes in the subgraphs divided by SGPM?

Backbones To validate the effectiveness of SGL, we select a series of models based on graph convolution and attention mechanisms to extract spatial features as backbones for the experiments. These models can be categorized into two types based on their computation methods: 1) Parallel Computation Models: GMAN [Zheng et al., 2020], STWavefull GAT [Fang et al., 2023], STWave-ESGAT [Fang et al., 2023], STAEformer [Liu et al., 2023], DGCNet(P) [Weng et al., 2024]; 2) Serial Computation Models: DGCRN [Li et al., 2023], DDGCRN [Weng et al., 2023], DGCNet(R) [Weng et al., 2024].

Dataset We select four widely-used real-world traffic datasets (PEMS03, PEMS04, PEMS07, and PEMS08) [Guo et al., 2019; Li and Zhu, 2021] for a series of experiments to investigate the performance, efficiency, and memory overhead differences of SGL across various backbone networks' spatial feature extraction modules in traffic prediction tasks. These experiments aim to evaluate SGL in scenarios with different node scales: small-scale (PEMS08), medium-scale (PEMS03, PEMS04), and large-scale (PEMS07). The detailed information for each dataset is provided in Table 1.

Experimental Setup Following the settings from previous studies [Guo et al., 2019; Li and Zhu, 2021], we divide the datasets into training, validation, and test sets in a 6:2:2 ratio based on the time order. Our experiments are conducted on a GPU server equipped with a GeForce GTX 4090 graphics card, using the PyTorch framework. All models, including their SGL variants, strictly follow the original training configurations to ensure fair comparison, including optimizers, maximum training epochs, and early stop strategies. For the PEMS07 dataset, the batch size is set to 16, while the batch size for the other three datasets is set to 64. The settings for the number of subgraphs M and the number of nodes per subgraph K for each dataset are detailed in Table 2.

To comprehensively evaluate the performance of SGL , we assess it from three aspects:

1. Performance Evaluation: We use three evaluation metrics: Mean Absolute Error (MAE), Mean Absolute Percentage Error (MAPE), and Root Mean Squared Error (RMSE).
2. Memory Overhead: We use GPU cost as a metric.
3. Running Efficiency: We evaluate using train time and inference time as metrics.

4.1 Performance and Effciency Analysis (Q1)

Table 3 shows the prediction performance, efficiency, and memory overhead of each model and its SGL variant. From the table, we observe that the models and their SGL variants achieve comparable prediction performance, but the SGL variants significantly reduce memory overhead, demonstrating the effectiveness of the SGL method.

On datasets of varying scales, as the dataset size increases, the memory savings of SGL become more pronounced. For instance, under identical conditions on the small-scale PEMS08 dataset, the memory overhead of DDGCRN-SGL is reduced by 18.2% compared to DDGCRN. On the large-scale PEMS07 dataset, DDGCRN-SGL reduces memory overhead by 60.5% compared to DDGCRN, highlighting the practicality of SGL's divide-and-conquer subgraph partitioning strategy. Although STWave-ESGAT can also reduce memory overhead while maintaining performance, it requires predefined matrix to determine the attention range, which may not be applicable in certain situations. On the other hand, STWave-SGL does not have this limitation, making it more versatile.

Regarding operational efficiency, models like GMAN, representing parallel computation frameworks, fully leverage parallel processing in their SGL variants, significantly accelerating computation speeds. In contrast, serial computation models like DGCRN, which rely on RNNs for iterative feature extraction, do not achieve parallel computation. Nevertheless, their SGL variants maintain comparable operational efficiency to their prototypes while demonstrating superior efficiency when the number of nodes is large.

4.2 Ablation Experiment (Q2)

SubGraph Partitioning Method To verify the rationality of SGPM, we design two alternative subgraph partitioning methods and one node clustering-based feature extraction method and test them on DGCNet(R) to demonstrate the effectiveness of our proposed SGPM:

1. DGCNet(R)-RD: This method randomly partitions the original traffic network into equally sized subgraphs for feature extraction.
2. DGCNet(R)-METIS: This method uses the DTW algorithm to generate a similarity matrix for the nodes and then applies the METIS algorithm to partition the subgraphs.
3. DGCNet(R)-DC: This method clusters node features into subgraph representations using a mapping matrix [Qin et al., 2023; Zhang et al., 2024] and then employs these subgraph cluster representations for spatio-temporal graph feature extraction.

Table 4 presents the performance of the three variants on PEMS08. We observe that the other two subgraph partitioning methods perform worse than SGPM. DGCNet(R)- RD generates subgraphs through random partitioning without considering node correlations, resulting in the poorest performance. DGCNet(R)-METIS constructs subgraphs where nodes exhibit high correlation by calculating similarities via dynamic programming, thus outperforming DGCNet(R)-RD. However, as the subgraphs in DGCNet(R)-METIS remain fixed and fail to account for the dynamic changes in node correlations, its performance is still inferior to DGCNet(R)- SGL. DGCNet(R)-DC utilizes a mapping matrix to transform node features into subgraph representations. However, this approach results in the loss of individual node features, which significantly hampers its performance compared to other methods.

Feature Aggregation Methods To validate the rationality of the averaging operation for feature aggregation, we design two alternative aggregation methods for comparison:

1. DGCNet(R)-sum: This method directly uses the sum of node features instead of averaging them.
2. DGCNet(R)-max: This method selects the maximum feature from the set of subgraph node features as the node's

final feature.

Figure 5 shows the performance of the three aggregation methods. It is evident that the averaging operation yields the best results. DGCNet(R)-sum suffers from feature redundancy, resulting in the poorest performance among the three variants. DGCNet(R)-max captures the most representative feature by selecting the largest feature as the node's feature but neglects feature from other subgraphs, leading to slightly worse performance than DGCNet(R)-SGL. These results demonstrate that the averaging operation avoids redundancy and proves its effectiveness in aggregating feature.

4.3 Hyperparameter Experiment (Q3)

We conduct hyperparameter experiments on the PEMS07 dataset using DDGCRN-SGL to evaluate the impact of the number of nodes per subgraph K and the number of subgraphs M on the model's predictive performance. Figure 4 illustrates the effect of K and M on the model's performance.

Figure 4: Sensitivity analysis of parameter K and M on PEMS07 dataset.

It is evident that setting K or M too low results in a sharp decline in performance. This is because an insufficient number of subgraphs or nodes per subgraph leads to inadequate extraction of spatio-temporal features, causing the model to underfit and fail to achieve the same performance as the original model. When the product of K and M approaches or exceeds N, the performance of DDGCRN-SGL matches that of DDGCRN, while requiring significantly less memory. Therefore, properly configuring K and M is crucial for balancing memory efficiency and maintaining performance.

4.4 Subgraph Visualization (Q4)

In this section, we visualize the subgraphs generated by DDGCRN-SGL during prediction on the test set to analyze the feature similarity among subgraph nodes and the distribution of overlapping nodes across subgraphs. Figure 5 presents the visualization results of subgraph features after dimensionality reduction using the T-SNE algorithm [Van der Maaten and Hinton, 2008].

The results show that the nodes in each subgraph exhibit clustering behavior, indicating a high degree of similarity among nodes selected based on their feature similarity. Additionally, the visualization reveals overlapping or closely located nodes across different subgraphs, demonstrating that some nodes are shared among multiple subgraphs. This overlap further validates the necessity of the SGFAM module to address potential redundancies in aggregated node feature.

5 Conclusion

In this paper, we propose a plug-and-play SubGraph Learning method to reduce the memory overhead of spatio-temporal graph models. Specifically, we introduce a SubGraph Partition Module that dynamically partitions the original spatiotemporal graph into smaller subgraphs for extracting spatial features, significantly reducing memory overhead. To address the issue of redundant node feature caused by nodes appearing in multiple subgraphs, we design an SubGraph Feature Aggregation Module that performs average aggregation on the multiple features obtained by nodes. We conduct a comprehensive evaluation of this method on traffic network datasets of various scales, and the results demonstrate that it can significantly reduce GPU consumption without compromising model performance, providing a practical solution for deploying spatio-temporal models in real-world scenarios.

Acknowledgments

This work was supported in part by the "Pioneer" and "Leading Goose" R&D Program of Zhejiang under Grant 2024C01214, in part by the National Natural Science Foundation of China under Grant 62072409 and 62476247, and Zhejiang Supcon Information Co., Ltd. KYY-HX-20230833.