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.

CN120597929APending Publication Date: 2025-09-05SOUTH CHINA UNIV OF TECH
View PDF 0 Cites 0 Cited by

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

Technical Problem

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.

Method used

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.

Benefits of technology

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.

✦ Generated by Eureka AI based on patent content.

Smart Images

  • Figure CN120597929A_ABST
    Figure CN120597929A_ABST
Patent Text Reader

Abstract

The invention discloses an efficient graph neural network sub-graph sampling method based on reinforcement learning, relates to an artificial intelligence technology, and provides a scheme for solving the problems of low efficiency and the like in the prior art. The method is applied to the training process of a graph neural network, and GNN training is carried out based on a GraphSARM strategy. The GraphSARM strategy comprises a GraphSARM-E strategy and a GraphSARM-N strategy, and the GraphSARM strategy comprises a GraphSARM-E strategy and a GraphSARM-N strategy; the GraphSARM-E combines an empirical degree probability P < degree > and a probability Pagent predicted by a sampling agent model, and adjusts the proportion of the Pagree and the Pagent through a weighting factor alpha < weight >; the GraphSARM-N only considers the probability Pagent predicted by the sampling agent model. The method has the advantages that graph neural network training is carried out based on the GraphSARM strategy, the speed and precision are greatly improved in single-machine single-card and single-machine multi-card parallel training scenes, and the convergence trend is more stable.
Need to check novelty before this filing date? Find Prior Art

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.