Make Bricks with a Little Straw: Large-Scale Spatio-Temporal Graph Learning with Restricted GPU-Memory Capacity

1 Introduction

Traffic forecasting plays a vital role in transportation planning, management, and decision-making processes [Miao et al., 2022a; Liu et al., 2024; Liu et al., 2021a; Zhang et al., 2022; Wang et al., 2023d]. Accurate predictions enable efficient resource allocation, optimal route planning, and effective traffic management strategies [Wang et al., 2020; Wang et al., 2020]. In recent years, the advancement of data collection and processing capabilities has resulted in a substantial amount of high-quality traffic data. This has sparked significant interest in the research community to develop state-of-the-art spatio-temporal learning models that can effectively extract spatio-temporal features. Currently, the prevailing approach in this field is based on spatio-temporal graph convolutional networks (STGNN). These models typically consist of spatial and temporal components, which are used to model spatial and temporal correlations separately. For example, DCRNN [Li et al., 2017] combines diffusion graph convolutional networks to represent spatial dependencies between nodes and utilizes recurrent neural networks (RNN) to capture temporal dynamics. The generated node representations are then fed into a predictor, often implemented as fully connected layers or a specialized decoder, to forecast future traffic flows.

∗Yang Wang and Pengkun Wang are corresponding authors.

The accuracy of traffic prediction has shown improvement thanks to the utilization of complex model structures and large parameter volumes. However, it is important to note that these models have primarily been developed and evaluated on small-scale datasets that possess limited road network sizes. For example, the widely used PeMSX datasets (denoted by X as 03, 04, 07, or 08) have a maximum of 883 nodes. In reality, real-world traffic networks can be much larger, often consisting of thousands of nodes. Consequently, the scalability of these models in effectively handling such large networks remains a significant concern [Liu et al., 2023b].

Existing spatio-temporal graph learning models typically take a complete spatio-temporal graph across multiple time segments as input. These models rely on graph convolution operators in GCN to perform spatio-temporal learning. However, training these models presents a significant challenge due to the dependency of loss items on a large number of neighboring nodes [Chiang et al., 2019; Hamilton et al., 2017]. Unlike other neural networks that can decompose training losses into individual items for each sample, STGNNs require storing the embeddings of all nodes within the computed graph from each STGNN layer during the backpropagation process. This storage requirement applies to various architectures, including MLP-based models [Shao et al., 2022a; Zhang et al., 2023] and Transformerbased models [Xu et al., 2020]. Consequently, when dealing with large-scale graphs, these models demand significant memory allocation that increases with the complexity of model parameters, posing a notable memory burden. While some research efforts aim to improve the training efficiency of STGNNs for large-scale graphs [Wang et al., 2023c; Wang et al., 2023e], we argue that the memory constraint is a decisive factor that directly affects the availability of these advanced models. The memory limitation not only hampers researchers but also poses challenges for Chinese researchers who may face restrictions imposed by manufacturers, limiting their access to high-end GPU devices with substantial memory capacities.

In this paper, we take a step to address the challenge of large-scale spatio-temporal graph learning with limited GPU memory. An instinct is that we can divide a large spatiotemporal graph into multiple small-scale subgraphs, conquer each subgraph individually, and then merge them incrementally. In this way, GPUs only need to store the representation of the subgraph for saving memory. In response, multiple subgraphs are used to train the model sequentially in a task streaming manner. This involves a continual learning paradigm, enabling graph learning from a static large graph to streaming subgraphs. However, the special challenge posed by continual learning is known as the tradeoff between the plasticity of new knowledge and the stability of old knowledge. In such a setting, achieving the balance between shared information among subgraphs and maintaining personalized information for each subgraph is an open question.

To address challenges, we propose a novel framework for large spatio-temporal graph learning, namely LarSTL. Specifically, we first employ the METIS algorithm as a solution to partition a large spatio-temporal graph into multiple subgraphs. Subsequently, we adopt a task streaming approach, treating each subgraph as an independent task, to sequentially train the spatio-temporal graph convolution model. To mitigate the issue of catastrophic forgetting during sequential training, we introduce a continuous learning strategy based on experience replay. This strategy allows us to consolidate the knowledge gained from previous subgraphs by saving a few nodes. Additionally, we utilize feature adaptors to extract personalized features for each subgraph. Through extensive experiments conducted on various large-scale traffic network datasets containing thousands of nodes, we demonstrate that our proposed method achieves competitive predictive performance while maintaining high training efficiency. Our contributions are as follows:

• We take a step to explore large-scale spatio-temporal graph representation learning and propose a divideand-conquer strategy based on the continuous learning paradigm, namely LarSTL.
• LarSTL first divides a large spatio-temporal graph into small-scale subgraphs and trains a model sequentially in the task streaming manner. In order to effectively consolidate the previously learned knowledge of subgraphs and their individual features, we propose an experimentreplay strategy and feature adaptors from the data and

parametric perspectives.

• We evaluate the proposed LarSTL on three large-scale traffic datasets with thousands of nodes, and the experimental results show that our model can achieve competitive performance and high efficiency.

2 RelatedWork

Traffic forecasting. Recently, the accuracy of spatialtemporal graph prediction has significantly improved with the emergence of deep learning [Zhou et al., 2023; Wang et al., 2023b; Xia et al., 2023; Huang et al., 2024b; Zhou et al., 2020b]. Spatio-Temporal Graph Neural Networks(STGNNs) are the most representative approaches to capture spatiotemporal dependencies in this field [Miao et al., 2022b; Miao et al., 2024]. For example, D2STGNN [Shao et al., 2022b] employs diffusion graph convolutional networks with RNNs to capture temporal patterns. STAEformer [Liu et al., 2023a] proposes Spatio-Temporal Adaptive Embedding transformer. ST-GAT [Zhang et al., 2019] employs graph attention networks with LSTM for efficient capture of time dependencies to improve the prediction performance of the Transformer. SSTBAN [Guo et al., 2023] follows a multitask framework by incorporating a self-supervised learner to produce robust latent representations for historical traffic data. Some works move to solving the problem of low training efficiency, for example, STHMLP [Qin et al., 2023] and ST-MLP [Wang et al., 2023f] are committed to introducing an all-MLP architecture. PFNet [Wang et al., 2023c] captures temporal correlations using deep multi-view sequence encoders and spatial correlations using graph embedding technologies. CANet [Wang et al., 2023e] effectively captures spatial dependencies by mining the implicit spatial centers of variables. Overall, when dealing with large road networks, these models are difficult to scale to large network structures due to GPU memory requirements.

Large-scale graph learning. Large-scale graph learning has garnered significant interest in various domains [Zhou et al., 2020a; Wu et al., 2020], such as recommended tasks and graph classification tasks. For example, GraphSage [Hamilton et al., 2017] employs a sampling strategy where a tree is sampled with each node as the root, progressively expanding the neighborhood by K steps with a fixed sample size. This method computes the hidden representation of the root node by hierarchically aggregating hidden representations of nodes from the bottom to the top of the sampled tree. PinSage [Ying et al., 2018] proposes short random walks to generate subgraphs for convolution operation to generate graph representation for web-scale recommendation systems. Additionally, Cluster-GCN [Chiang et al., 2019] leverages a graph pooling algorithm to sample a few nodes and subsequently performs graph convolutions on this node. SSE [Dai et al., 2018] updates node hidden states recurrently in a stochastic and asynchronous fashion. It alternatively samples a batch of nodes for state update and a batch of nodes for gradient computation. However, these models specifically designed for recommendation and node classification tasks cannot model spatio-temporal dependencies and cannot be directly applied to spatio-temporal graph learning.

3 ProblemDefinitions

We use a graph G = (V, E, A) to represent a large-scale traffic network, where V denotes the set of nodes, and E represents the weight between the nodes. A ∈ RN×N represents the adjacency matrix of G, where N means the number of nodes and is very large for a large-scale graph. Additionally, let Xt = {x1, x2, ..., xN} represent the historical graph data at t-th time step, and xi ∈ Rd represents the state of i-th node, where d means the number of features. Given G and observations of past T time steps X ∈ RT×N×d = {X1, ·, XT }, The spatio-temporal graph prediction task aims to learn a spatio-temporal prediction model with parameter ω, which can effectively predict graph signals next Tp time steps Y ∈ RTP×N×d = {XT+1, ·,XT+TP }. The parameters ω can be decoupled into two parts: spatio-temporal feature extractor ϕ and a predictor φ which can be optimized in an end-to-end manner:

where Y and Y represent the prediction and ground-truth values, respectively.

4 Method

In this section, for clarity, we first introduce a simple spatio-temporal graph model as the backbone, namely BasSTG. Subsequently, we describe the proposed divisionand-conquest strategy, which is based on the continuous learning paradigm to facilitate large-scale spatio-temporal graph learning with BacSTG under limited GPU memory. The details and pseudocode of our proposed method LarSTL are shown in Figure 1 and Algorithm 1.

4.1 Spatio-Temporal GraphModel

GCNs have demonstrated their effectiveness for graph structures across various tasks [Huang et al., 2024a; Wang et al., 2024; Du et al., 2023]. We introduce a simple spatio-temporal graph convolution network as a backbone, BasSTG, which includes two spatio-temporal (ST) layers. Each ST layer includes a diffusion graph convolution module and a TCN module. Specifically, given an adjacency matrix Aτ ∈ RNτ×Nτ and the representation of the l-th layer Hl ∈ RNτ×Tl×Cl as input, where Nτ means the number of nodes in the input graph, Tl means the length of time steps, and Cl means the number of channels, a diffusion graph convolution layer is as follows:

where W1,W2 ∈ RCl×Cl+1 are the learnable parameters. Pkf and Pkb means forward and backward transfer matrices with k-power series, respectively, and they can computed by normalizing Aτ . The output Hl+1 is fed into a TCN layer to capture the dynamics, which is composed of causal convolution networks. After two ST layers, the output representation is fed into a forecaster to generate future traffic flow.

4.2 Continuous Learning for Large Graph

The restricted GPU memory capacity imposes limitations on the size of the input graph. To address this challenge, we propose a divide-and-conquer strategy. Initially, we divide a large spatio-temporal graph into small-scale subgraphs using a graph partitioning strategy. Subsequently, we train the model sequentially by treating these subgraphs as separate task streams. To enable this sequential learning process, we introduce a continuous learning strategy that facilitates the model's ability to acquire new knowledge from the current subgraph while consolidating the knowledge learned from the previous subgraph.

Graph Partition

Several methods for graph clustering and community detection have been developed, such as spectral clustering [Von Luxburg, 2007] and METIS [Karypis and Kumar, 1998]. We choose METIS due to its excellent balance between accuracy and speed. It partitions a graph into a predefined number of clusters such that the number of withincluster links is much higher than between-cluster links to better capture good community structure [He et al., 2023]. However, METIS is limited to finding non-overlapping clusters, which could imply losing important edge information. Especially for edge nodes at the boundary of any partition, the neighboring correlated node will be lost, and this issue will affect the prediction accuracy of the models.

To tackle this challenge, we use an overlapping subgraph partition approach. In this approach, for each partition, we identify and include spatially correlated nodes from other partitions. As a result, nodes located near the boundary of a partition will be present in multiple partitions. This strategy allows for the incorporation of nodes that exhibit spatial correlation across partitions, enabling a more comprehensive representation of the graph's connectivity and spatial relationships. Expanding all neighbors of the edge nodes into the partition is effective, but it also increases the number of nodes within this partition, leading to a high memory burden. Therefore, we downsample the spatially correlated nodes based on the shortest driving distance, which can be computed by Open Source Routing Machine (OSRM) [Luxen and Vetter, 2011].

Continuous Training Strategy

After dividing graph G into multiple overlapping subgraphs {g1, · · · , gK}, where K means the number of subgraphs. This series of subgraphs is treated as a task streaming to train the model sequentially. The model can continuously absorb new knowledge from the task streaming to adapt to different subgraphs. After such training, the model is expected to extract shared features among subgraphs, while maintaining the personalized features of each subgraph, so that it can effectively make predictions for the entire large graph.

However, when we fine-tune the model with subgraph Gτ , the model may forget previously learned knowledge from the previous subgraph Gτ-1, i.e., catastrophic forgetting. This notorious phenomenon causes the decline of the model's prediction performance for previous tasks. To mitigate catastrophic forgetting, continual graph learning has been actively investigated in various fields [Wang et al., 2023a; Liu et al., 2021b]. In this paper, we propose two strategies based on experience replay and parameter adaptor to consolidate old knowledge from data and parameter perspectives.

Experience-replay strategy. In the realm of continuous learning, the effectiveness of replaying data samples from previous tasks in consolidating old knowledge has been demonstrated [Wang et al., 2023a]. Taking inspiration from this approach, we adopt a similar strategy in our framework. After each training iteration, we select a subset of nodes from the subgraph and add them to the node container B. In the subsequent training process, the nodes in the container B, along with their N-hop neighbors, are utilized to fine-tune the saved model alongside nodes from a new subgraph.

In order to improve the consolidation effect, these replay nodes should representatively cover what the model has learned in the previous task. Naturally, the model is compatible with these nodes. Based on the above analysis, during τ -1-th learning step, we sort the model's prediction performance for all nodes only within Gτ-1 in the validation dataset. Then we select top Kr% nodes with small prediction errors as the replay nodes. However, some nodes with large prediction errors need to be re-learned, which can improve the prediction accuracy. Therefore, Kc% nodes with large errors are also selected. After training, these nodes are added into B. When learning subgraph Gτ , these nodes are trained together with the new subgraph. In fact, these cross-subgraph nodes facilitate information interaction between different subgraphs.

Feature adaptor. Striking a balance between extracting new knowledge and consolidating previous knowledge poses a challenge in continuous learning. To tackle this issue, we propose the development of individualized feature adaptors for each subgraph. These adaptors are designed to extract personalized features from each subgraph, empowering the model to effectively capture and comprehend the specific spatio-temporal information within each subgraph.

Specifically, we construct an individual adaptor, θi, for i-th subgraph, which is a learnable parameter matrix. The shapes of these matrices are consistent with the ones of each parameter matrix in the spatio-temporal feature extractor ϕ. Then this adaptor θi is used to scale the parameters ϕ by dot product: θi ⊗ ϕ. In essence, the adaptor can be viewed as a learnable attention weight that characterizes the importance of spatio-temporal features for this subgraph, i.e. personalized spatio-temporal features. And adaptors are trained end-

Algorithm 1 Large-scale spatio-temporal graph learning

Input: A large-scale spatio-temporal graph G, observation historical graph data X.

Parameter: spatio-temporal feature extractor ϕ, feature adaptor θ, Predictor φ.

Output: A spatio-temporal prediction function with optimal parameters ω (ϕ, θ, φ).

• 1: Using METIS algorithm to divide G into k overlapping subgraphs {g1, ..., gk}.
• 2: while {g1, ..., gk} do
• 3: if τ == 1 then
• 4: Random initialization parameters ω1 (ϕ, θ1, φ1).
• 5: Initial replay node container B.
• 6: Train the model with the data of the subgraph g1.
• 7: else
• 8: Random initialization parameters θτ and φτ .
• 9: Load the saved parameters ϕ in ωτ-1.
• 10: Load replay-nodes in container B.
• 11: Use replay-nodes with their N-hop neighbors and nodes in gτ to train the model ωτ (ϕ, θτ , φτ ).
• 12: end if
• 13: Add replay nodes in gτ into the container B.
• 14: end while
• #########Make Prediction##########
• 15: for {g1, ..., gk} do
• 16: Make predictions for the subgraph gτ with ωτ (ϕ, θτ , φτ ).
• 17: end for

to-end along with the model. For smooth learning, we freeze the model parameters trained on the first subgraph and only update feature adapters in the subsequent training process.

Furthermore, different subgraphs can exhibit distinct traffic patterns, for example, the traffic trends in shopping malls and residential areas can vary significantly. Recognizing this, we propose the development of dedicated predictors tailored for specific subgraphs. These specialized predictors are designed to effectively capture the unique features and dynamics of each traffic subgraph.

5 Experiment

In this chapter, we evaluate the effectiveness of LarSTL on a large-scale network dataset with thousands of nodes and answer the following underlying concerns: Q1. What is the prediction performance and efficiency of the model? Please refer to Section 5.2. Q2. Is every component of the method valid? Please refer to Section 5.3. Q3. How sensitive is the model to the number of subgraphs and the graph partition algorithm? Please refer to Section 5.4.

5.1 Experiment Setting

Data. We evaluate the validity of the model on three large traffic datasets, GLA, GBA, and CA, proposed by from [Liu et al., 2023b]. These three emerging datasets, with thousands of nodes, are far larger than those previously used in the traffic forecasting community [Liu et al., 2023b]. All datasets are collected from the California Department of Transportation (CalTrans) Performance Measurement System (PeMS) and aggregate the traffic readings from 5-minute intervals into 15-minute windows. The details of traffic networks in each dataset are shown in Table.1.

Restricted GPU device. All models are trained on a V100 GPU with 16G memory. We give priority to running a model with the optimal hyperparameters from their official codes. And we specify a maximum batch size of 64. If a model is unable to run with this setting, we would reduce the hyperparameter and the batch size of the models until it fully occupies the memory on this V100 GPU. Eventually, if the model still doesn't work, we consider this model to be out of memory.

Experiment setting. We chronologically split the data into train, validation, and test sets, with a ratio of 6:2:2 for all datasets. AdamW [Loshchilov and Hutter, 2017] is used as the optimizer with an initial learning rate of 0.001. we use an early stopping strategy to accelerate the training process. We divide the entire graph into three subgraphs. To evaluate the performance of our framework, we adopt Mean Absolute Errors (MAE), Root Mean Squared Errors (RMSE), and Mean Absolute Percentage Errors (MAPE) as the metrics for effectiveness. The replay node numbers Kr and Kc are set to 1%. We report the overall efficiency by considering the total training time and evaluation time of all models, along with the average training time per epoch.

Baseline. We compare some advanced spatio-temporal prediction models as follows:

• LSTM is an RNN variant to model long-term temporal dependencies.
• TCN is a lightweight temporal modeling framework using causal convolution.
• DCRNN [Li et al., 2017] introduces graph diffusion graph convolutional networks and LSTM to capture long-term dependencies.
• Z-GCNETs [Chen et al., 2021] propose to enhance DL architectures with the most salient time-conditioned topological information of the data.
• AGCRN [Bai et al., 2020] is also a RNN-based sptiotemporal graph learning model.
• ASTGCN [Guo et al., 2019] focuses on adaptive modeling of spatial dependencies using attention mechanism.
• HGCN [Guo et al., 2021] designs hierarchical GCN to capture both the spatial relationships of microscopic nodes and relationships between macroscopic regions.

• DGCRN [Li et al., 2023] leverages and extracts dynamic characteristics from node attributes, while the parameters of dynamic filters are generated at each time step.
• STGODE [Fang et al., 2021] uses neural ordinary differential equations to effectively model the continuous changes of traffic signals.
• DSTAGNN [Lan et al., 2022] uses a dynamic spatialtemporal aware graph to replace the predefined static graphs used in traditional graph convolution.
• D2STGNN [Shao et al., 2022b] proposes a unique estimation gate with residual decomposition mechanism to decouple traffic signals into diffusion and inherent traffic information.

5.2 Performance and Efficiency Analysis (Q.1)

The prediction performance of each model on the three datasets is presented in Table 2, showcasing the competitive efficiency and prediction accuracy of our model.

Comparing different methods on the GBA, we observe that temporal models, such as TCN and LSTM, perform poorly because they fail to capture spatial information efficiently. DGCRN and D2STGNN demonstrate impressive performance on the GBA datasets, highlighting the importance of considering the dynamic features of spatial topology. However, their performance can not meet expectations because the device's memory capacity limited the size of their parameters, resulting in inadequate learning of the spatio-temporal information. Meanwhile, their complex model designs limit their scalability when applied to larger datasets like GLA and CA datasets. On the other hand, RNN-based models, such as DCRNN and AGCRN, maintain good predictive performance across multiple datasets, which can be attributed to their ability to effectively capture spatio-temporal information. However, when running in the limited GPU, these models all struggle to effectively handle the large road network, resulting in a rapid deterioration of prediction performance. In terms of efficiency, TCN-based approaches such as Z-GCNETs and HGCN are generally faster than other baselines, thanks to their parallel computation of temporal convolution operations. Among the RNN-based models, DCRNN and AGCRN are slower than TCN-based methods due to the use of RNNs as decoders.

Our model achieves superior performance and efficiency, especially for large-scale datasets, thanks to the proposed divide-and-conquer strategy, which reduces the learning complexity of spatio-temporal information.

5.3 Ablation Experiment (Q.2)

In this section, we evaluate the effectiveness of each component. We make the following variants: w/o tem means that we remove the temporal module in LarST. w/o rep means that we remove the experience-replay strategy for consolidating old knowledge learned from the previous subgraphs. w/o adp means that we directly fine-tune the saved model each time instead of configuring feature adaptors for each subgraph. The prediction results of each variant are shown in Figure.2.

We can see that each variant is more or less inferior to our proposed model, proving the validity of each component of the model. The prediction performance of w/o tem is not good, this is because it does not model temporal correlations and has not learned sufficient spatiotemporal information, resulting in suboptimal prediction performance. The prediction error of w/o adp is high, which indicates that the feature adaptor can effectively help the model consolidate the old knowledge. At the same time, it shows that learning the personalized features of each subgraph improves the model's prosperity in the new subgraph.

5.4 Hyperparameter Experiment (Q.3)

We assess the impact of the number of graph partitions on clustering and its subsequent effect on model prediction efficacy. The outcomes derived from a 60-horizon analysis on the GLA dataset are presented in Table 3, revealing that optimal model performance is achieved when K is equal to 3. This because that a higher value of K results in clusters with fewer nodes, which fails to furnish adequate spatio-temporal data, leading to model underfitting. Moreover, such a setup causes the model to readily forget previously learned information when it encounters new knowledge within current subgraphs. Conversely, a small value of K implies a larger number of nodes per cluster, and the scale of model parameters is constrained by GPU memory, resulting in underfitting of the model.

6 Conclusion

We propose an efficient approach for large-scale spatiotemporal graph learning, employing a divide-and-conquer strategy. Our method partitions a large graph into smaller subgraphs, allowing individual processing and analysis. To enhance the learning process, we introduce two strategies: experience replay and feature adaptator. Experience replay strategy can consolidate knowledge from preceding subgraphs, while the feature adaptator strategy ensures scalability for incorporating new knowledge from subsequent subgraphs. We comprehensively evaluate our approach on datasets featuring large traffic networks, demonstrating its effectiveness and efficiency in large-scale spatio-temporal graph learning under limited GPUs.

Acknowledgments

This paper is supported by the National Natural Science Foundation of China (No.62072427, No.12227901), the Project of Stable Support for Youth Teamin Basic Research Field, CAS (No.YSBR-005), and Academic Leaders Cultivation Program, USTC.