Efficient graph neural network sub-graph sampling method based on reinforcement learning
By adaptively adjusting the sampling probability of nodes or edges through the reinforcement learning-based GraphSARM strategy and dynamically generating subgraphs, the problems of high memory consumption, low precision and slow speed in graph neural network training are solved, achieving a significant improvement in training speed and accuracy.
Patent Information
- Application Number
- CN202510641682.2
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2025-05-19
- Publication Date
- 2025-09-05
AI Technical Summary
Existing graph neural network training has problems such as high memory requirements caused by large-scale graph data, low training accuracy caused by static sampling schemes, and slow training speed caused by data transmission, resulting in low training efficiency and low accuracy.
The reinforcement learning-based GraphSARM strategy is adopted. The sampling probability of nodes or edges is adaptively adjusted through the GraphSARM-E and GraphSARM-N methods to dynamically generate subgraphs. Efficient subgraph sampling is achieved by combining the empirical degree probability and the predicted probability of the sampling agent model.
In single-machine single-card and parallel training scenarios, the training speed and accuracy of graph neural networks have been significantly improved, the convergence trend has become more stable, and an accuracy improvement of 0.3%-5% and a speed increase of 1.19x-16.14x have been achieved.
Smart Images

Figure CN120597929A_ABST
Abstract
Description
Technical Field
[0001] The present invention relates to artificial intelligence technology, and in particular to an efficient graph neural network subgraph sampling method based on reinforcement learning. Background Art
[0002] Graphs are a typical non-Euclidean data structure with excellent relational expression capabilities. Graph Neural Networks (GNNs) are designed specifically for graph data, achieving feature aggregation through message passing and are widely used in tasks such as node classification and link prediction. Although GNNs perform well in graph data processing, their training still faces three key technical bottlenecks: (1) Large-scale graph data leads to high memory requirements. (2) Static sampling schemes lead to low training accuracy. (3) Large amounts of data transmission lead to slow training speed. These three key technical bottlenecks result in low efficiency and insufficient training accuracy in large-scale graph neural network training, which are urgent issues that need to be addressed in current research on graph neural networks.
[0003] In GNN training based on subgraph sampling, computational complexity can be effectively reduced by sampling one or more subgraphs for each mini-batch. This can be achieved in two main ways. The first approach uses a graph partitioning algorithm to directly divide the graph into several independent subgraphs. During training, one or more subgraphs are randomly selected for processing, significantly reducing memory usage and computational overhead. The second approach uses a specific strategy to sample sets of nodes or edges and then induce subgraphs based on these sets. These two approaches each have their own emphasis: the former focuses on structured graph partitioning, while the latter emphasizes dynamic subgraph generation. However, both approaches use static sampling strategies to generate subgraphs.
[0004] File 1 "Cluster-GCN: An EfficientAlgorithm for TrainingDeep and Large GraphConvolutional Networks, Wei-LinChiang, Xuanqing Liu, Si Si, Yang Li, Samy Bengio, Cho-Jui Hsieh, KDD'19, August 4–8, 2019, Anchorage, AK, USA"; File 2 "GraphSAINT: GRAPHSAMPLINGBASED INDUCTIVE LEARNING METHOD, HanqingZeng, HongkuanZhou, AjiteshSrivastava, Rajgopal Kannan, Viktor Prasanna, arXiv:1907.04931v4[cs.LG]16Feb2020" and document 3 "Feature-Oriented Sampling for Fast and Scalable GNNTraining, XinZhang, Yanyan Shen, Lei Chen, 2022IEEE International Conference onDataMining(ICDM)|978-1-6654-5099-7 / 22 / $31.00 2022IEEE|DOI:10.1109 / ICDM54844.2022.00083” are three typical subgraph sampling algorithms.
[0005] File 1 uses the Metis algorithm to partition a graph into several subgraphs and randomly selects one subgraph for training in each training batch, representing the first seed graph generation method. Subgraph generation is based on graph partitioning, and the sampled subgraphs are independent and non-overlapping.
[0006] Files 2 and 3, on the other hand, sample nodes and edges using a specific probability distribution and generate subgraphs for training based on the sampling results, demonstrating the second subgraph generation method. File 2 obtains a set of subgraph nodes based on the sampling probability, then uses these subgraph nodes to determine the topology of the subgraph nodes from the full graph, ultimately generating the sampled subgraph. File 3 randomly selects a continuous segment of nodes and constructs a subgraph from these nodes.
[0007] Files 1, 2, and 3 all generate subgraphs based on static sampling strategies, and the quality of the generated subgraphs is directly determined by the sampling strategy. This static sampling strategy makes it difficult to effectively and adaptively adjust the sampling probability during the sampling process, resulting in poor training accuracy. Furthermore, the sampling steps in Files 1, 2, and 3 are all performed on the CPU, which requires significant data transfer time during training. This makes these methods very slow and inefficient for large-scale graph neural network training. Summary of the Invention
[0008] The purpose of the present invention is to provide an efficient graph neural network subgraph sampling method based on reinforcement learning to solve the problems existing in the above-mentioned prior art.
[0009] The present invention describes an efficient graph neural network subgraph sampling method based on reinforcement learning, which is applied in the training process of graph neural networks and performs GNN training based on the GraphSARM strategy.
[0010] It is characterized by:
[0011] The GraphSARM strategy includes two strategies: GraphSARM-E and GraphSARM-N;
[0012] The GraphSARM-E converts the empirical degree probability P degree and the probability P predicted by the sampling surrogate model agent Combined with the weighting factor α weight Adjust the ratio between the two;
[0013] The GraphSARM-N only considers the probability P predicted by the sampled surrogate model agent .
[0014] The advantage of the efficient graph neural network subgraph sampling method based on reinforcement learning described in the present invention is that the graph neural network training is based on the GraphSARM strategy, and its speed and accuracy are greatly improved in the single-machine single-card and single-machine multi-card parallel training scenarios, and the convergence trend is more stable. BRIEF DESCRIPTION OF THE DRAWINGS
[0015] Figure 1 It is a schematic diagram of the structure of the efficient graph neural network described in the present invention.
[0016] Figure 2 It is a schematic diagram of the training process of the efficient graph neural network described in the present invention.
[0017] Figure 3 This is a comparison chart of the convergence curves of the efficient graph neural network described in the present invention and other algorithms on the Flickr platform.
[0018] Figure 4 This is a comparison chart of the convergence curves of the efficient graph neural network described in this invention and other algorithms on the Reddit platform.
[0019] Figure 5 This is a comparison chart of the convergence curves of the efficient graph neural network described in this invention and other algorithms on the Yelp platform.
[0020] Figure 6 This is a comparison chart of the convergence curves of the efficient graph neural network described in this invention and other algorithms on the Amazon platform. DETAILED DESCRIPTION
[0021] like Figure 1 、 Figure 2 As shown in the figure, the present invention describes an efficient reinforcement learning-based graph neural network subgraph sampling method applied to the training process of a graph neural network. GNN training is performed based on the GraphSARM strategy. The input of the process is a graph G, a training node train_nid, and a sampling agent model Agent, and the output is a trained model. The main difference between GraphSARM-E and GraphSARM-N lies in the way the sampling probability is generated.
[0022] GraphSARM-E converts the empirical degree probability P degree and the probability P predicted by the sampling surrogate model agent Combined with the weighting factor α weight Adjust the ratio of the two. GraphSARM-N only considers the probability P predicted by the sampling agent model agent .
[0023] The following describes in detail the main processes of training using the GraphSARM-N and GraphSARM-E sampling strategies respectively.
[0024] (1) For GraphSARM-N, the sampling probability of each node is learned through the reinforcement learning network, and then the sampling probability of each node is adaptively adjusted. The implementation process is as follows:
[0025] S11. Pretreatment;
[0026] In the preprocessing stage, we first extract the subgraph G induced by the training node train_nid from the full graph G train , so that subsequent sampling only targets topological structures related to training. Next, the sampling agent model Agent is initialized, which is a multi-layer neural network used to predict the sampling probability of the node. In addition, the exponential moving average EMA (ExponentialMovingAverage) reward r is also initialized emaand a smoothing factor α, which is used to stabilize the reward signal during training.
[0027] S12. Subgraph sampling using the GraphSARM-N strategy;
[0028] Feature splicing: G train The node features of the graph are concatenated with the node degree to form the input of the sampling proxy model. As a lightweight topological feature, the node degree helps the model capture the structural pattern of the graph.
[0029] The sampling agent model Agent is a multi-layer neural network that processes input and outputs the original sampling probability P of each node agent .
[0030] Noise injection: To enhance exploration and prevent premature convergence to suboptimal solutions, this strategy injects noise into P agent Add random noise Noise that obeys normal distribution, and control the noise amplitude through the hyperparameter β to obtain the final sampling probability P sampling =P agent +β·Noise.
[0031] Subgraph sampling: Based on P sampling , using a sampler from G train Select a subset of nodes nodes and generate a subgraph Gsubg through these node subsets.
[0032] S13. Train on the sampled subgraph Gsubg, calculate the training loss and reward, and feed the reward back to the sampled proxy model;
[0033] GNN training and reward calculation: Train the graph neural network on Gsubg and calculate the reward r based on the training loss. Typically, the reward is inversely proportional to the loss, that is, the lower the loss, the higher the reward. This paper uses the negative loss value as the reward r.
[0034] EMA reward update: Update the EMA reward using the smoothing factor α: r ema =α·r ema +(1-α)·r to smooth out short-term fluctuations and provide a stable learning signal.
[0035] Policy gradient loss: Calculate the loss of the sampled proxy model L = -mean(P·r ema ), where P is the logarithmic probability of the sampling node, P = log(P sampling [nodes]), r ema is the smoothed reward. This loss guides the sampling agent model to favor high-reward sampling decisions and dynamically adjusts the sampling probability of the node.
[0036] S14. Sampling agent model update;
[0037] The parameters of the sampling agent model are optimized through back propagation and gradient descent to learn the sampling probability of each node and achieve adaptive adjustment of the sampling strategy.
[0038] S15. Iterate steps S12-S14 until the training converges and the model training is completed.
[0039] (2) For GraphSARM-E, the sampling probability of each edge is learned through the reinforcement learning network, and then the sampling probability of each edge is adaptively adjusted. The implementation process is as follows:
[0040] S21. Pretreatment;
[0041] First, extract the subgraph G induced by the training node train_nid from the full graph G train Next, initialize the sampling agent model Agent to predict the sampling probability of the edge. In addition, initialize the exponential moving average reward r ema and a smoothing factor α, which is used to stabilize the reward signal during training.
[0042] S22. Use GraphSARM-E strategy to perform subgraph sampling;
[0043] Feature splicing: G train The source node features and source node degrees are concatenated, the target node features and target node degrees are concatenated, and these two concatenated features are used as the input of the sampling proxy model.
[0044] The sampling agent model Agent is a multi-layer neural network that processes input and outputs the original sampling probability P of each node agent .
[0045] Calculate the empirical degree probability: define the empirical degree probability as the sum of the inverse of the source node degree and the inverse of the target node degree, that is, the empirical degree probability P degree =src_ degree s-1+dst_ degree s-1. Using parameter α weight Adjust its size.
[0046] Noise injection: agent Add random noise that follows a normal distribution to the ensemble and control the noise amplitude through the hyperparameter β. The final fusion sampling probability is:
[0047] P sampling =(1-α weight )·P degree +α weight ·P agent +β·Noise
[0048] Subgraph sampling: Based on P sampling , using a sampler from G train Select a subset of nodes nodes and generate a subgraph Gsubg through these node subsets.
[0049] S23. Train on the sampled subgraph Gsubg, calculate the training loss and reward, and feed the reward back to the sampled proxy model.
[0050] GNN training and reward calculation: Train the graph neural network on Gsubg and calculate the reward r based on the training loss. This paper uses the negative loss value as the reward r.
[0051] EMA reward update: Update the EMA reward using the smoothing factor α: r ema =α·r ema +(1-α)·r to smooth out short-term fluctuations and provide a stable learning signal.
[0052] Policy gradient loss: Calculate the loss of the sampled proxy model L = -mean(P·r ema ), where P is the logarithmic probability of the sampling node, P = log(P sampling [nodes]), r ema is the smoothed reward. This loss guides the sampling agent model to favor high-reward sampling decisions.
[0053] S24. Sampling agent model update;
[0054] The parameters of the sampling agent model are optimized through back propagation and gradient descent to learn the sampling probability of each node and achieve adaptive adjustment of the sampling strategy.
[0055] S25. Iterate steps S22-S24 until the training converges and the model training is completed.
[0056] The two independent methods proposed in this paper, GraphSARM-N and GraphSARM-E, are both dynamic subgraph sampling algorithms that can perform subgraph sampling well in GNN training. They gradually adjust the sampling proxy model through each iteration, complete the automatic adjustment of the sampling strategy, and assign appropriate sampling probabilities to nodes and edges, thereby improving the accuracy of the GNN model.
[0057] like Figures 3 to 6 As shown, the efficient graph neural network subgraph sampling method based on reinforcement learning described in the present invention has better convergence effect than other algorithms in the prior art.
[0058] Accuracy performance in a single-machine, single-GPU scenario: Utilizing the GraphSARM strategy for subgraph sampling, the accuracy of graph neural network training has been significantly improved. Experiments on Flickr, Reddit, Yelp, and Amazon datasets have demonstrated an accuracy improvement of 0.3%-5%.
[0059] Speed performance in a single-machine, single-GPU scenario: Utilizing the GraphSARM strategy for subgraph sampling, the accuracy of graph neural network training has been significantly improved. Experiments on Flickr, Reddit, Yelp, and Amazon datasets have demonstrated speed improvements of 1.19x-5.77x.
[0060] Accuracy performance in parallel training scenarios: Using the GraphSARM strategy for subgraph sampling, the accuracy of graph neural network training has been significantly improved. Experiments on Flickr, Reddit, Yelp, and Amazon datasets have achieved an accuracy improvement of 0.3%-3.5%.
[0061] Speed performance in parallel training scenarios: Using the GraphSARM strategy for subgraph sampling, the accuracy of graph neural network training has been significantly improved. Experiments on Flickr, Reddit, Yelp, and Amazon datasets have achieved speed improvements of 1.85x-16.14x.
[0062] In addition, when training graph neural networks based on the GraphSARM strategy, its convergence speed is faster and the convergence stability is better.
[0063] Those skilled in the art can make various other corresponding changes and deformations based on the technical solutions and concepts described above, and all of these changes and deformations should fall within the scope of protection of the claims of the present invention.
Claims
1. An efficient graph neural network subgraph sampling method based on reinforcement learning, applied in the training process of graph neural networks, and GNN training based on the GraphSARM strategy; It is characterized by: The GraphSARM strategy includes two strategies: GraphSARM-E and GraphSARM-N; The GraphSARM-E converts the empirical degree probability P degree and the probability P predicted by the sampling surrogate model agent Combined with the weighting factor α weight Adjust the ratio between the two; The GraphSARM-N only considers the probability P predicted by the sampled surrogate model agent .
2. The efficient graph neural network subgraph sampling method based on reinforcement learning according to claim 1 is characterized in that: The GraphSARM-N learns the sampling probability of each node through a reinforcement learning network, and then adaptively adjusts the sampling probability of each node, including the following steps: S11. Pretreatment; S12. Subgraph sampling using the GraphSARM-N strategy; S13. Train on the sampled subgraph Gsubg, calculate the training loss and reward, and feed the reward back to the sampled proxy model; S14. Update of sampling agent model; S15. Iterate steps S12-S14 until the training converges and the model training is completed.
3. The efficient graph neural network subgraph sampling method based on reinforcement learning according to claim 2 is characterized in that: In the step S11, in the preprocessing stage, first extract the subgraph Gtrain induced by the training node train_nid from the full graph G; then initialize the sampling agent model Agent to predict the sampling probability of the node; initialize the exponential moving average EMA, and reward r ema and a smoothing factor α, which is used to stabilize the reward signal during training.
4. The efficient graph neural network subgraph sampling method based on reinforcement learning according to claim 3 is characterized in that: The step S12 includes the following sub-steps: G train The node features and node degrees are spliced together to form the input of the sampling agent model; the node degree is a lightweight topological feature; the sampling agent model Agent is a multi-layer neural network that processes the input and outputs the original sampling probability P of each node agent ; To enhance exploration and prevent premature convergence to suboptimal solutions, agent Add random noise Noise that obeys normal distribution, and control the noise amplitude through the hyperparameter β to obtain the final sampling probability P sampling =P agent +β·Noise; based on P sampling , using a sampler from G train Select a node subset nodes and generate a subgraph Gsubg through the node subset.
5. The efficient graph neural network subgraph sampling method based on reinforcement learning according to claim 4 is characterized in that: The step S13 includes the following sub-steps: Train the graph neural network on the subgraph Gsubg and calculate the reward r based on the training loss; update the EMA reward using the smoothing factor α: r ema =α·r ema +(1-α)·r, to smooth short-term fluctuations and provide a stable learning signal; Calculate the loss of the sampling proxy model L = -mean(P·r ema ), where P is the logarithmic probability of the sampling node, P = log(P sampling [nodes]), r ema is the smoothed reward.
6. The efficient graph neural network subgraph sampling method based on reinforcement learning according to claim 1 is characterized in that: The GraphSARM-E learns the sampling probability of each edge through a reinforcement learning network, and then adaptively adjusts the sampling probability of each edge, including the following steps: S21. Pretreatment; S22. Use GraphSARM-E strategy to perform subgraph sampling; S23. Train on the sampled subgraph Gsubg, calculate the training loss and reward, and feed the reward back to the sampled proxy model; S24. Sampling agent model update; S25. Iterate steps S22-S24 until the training converges and the model training is completed.
7. The efficient graph neural network subgraph sampling method based on reinforcement learning according to claim 6 is characterized in that: In step S21, first extract the subgraph G induced by the training node train_nid from the full graph G train ; Initialize the sampling agent model Agent to predict the sampling probability of the edge; initialize the exponential moving average reward r ema and a smoothing factor α, which is used to stabilize the reward signal during training.
8. The efficient graph neural network subgraph sampling method based on reinforcement learning according to claim 7 is characterized in that: The step S22 includes the following sub-steps: G train The source node features and source node degree are spliced together, and the target node features and target node degree are spliced together, and these two spliced features are used as the input of the sampling agent model; the sampling agent model Agent is a multi-layer neural network that processes the input and outputs the original sampling probability P of each node agent ; Define the empirical degree probability as the sum of the inverse of the source node degree and the inverse of the target node degree, that is, the empirical degree probability P degree =src_ degree s-1+dst_ degree s-1; using parameter α weight Adjust its proportion; agent Add random noise that obeys the normal distribution, and control the noise amplitude through the hyperparameter β; the final fusion sampling probability is: P sampling =(1-α weight )·P degree +α weight ·P agent +β·Noise; based on P sampling , using a sampler from G train Select a node subset nodes and generate a subgraph Gsubg through the node subset.
9. The efficient graph neural network subgraph sampling method based on reinforcement learning according to claim 8 is characterized in that: The step S23 includes the following sub-steps: Train the graph neural network on Gsubg and calculate the reward r based on the training loss; update the EMA reward r using the smoothing factor α ema =α·r ema +(1-α)·r to smooth short-term fluctuations and provide a stable learning signal to calculate the loss of the sampling agent model L = -mean(P·r ema ), where P is the logarithmic probability of the sampling node, P = log(P sampling [nodes]), r ema is the smoothed reward.