Title: Learning Large Graph Property Prediction via Graph Segment Training
Venue: NeurIPS 2023
Accepted: yes
Keywords: Graph Neural Networks

Graph Property Prediction
Abstract: Learning to predict properties of large graphs is challenging because each prediction requires the knowledge of an entire graph, while the amount of memory available during training is bounded. Here we propose Graph Segment Training (GST), a general framework that utilizes a divide-and-conquer approach to allow learning large graph property prediction with a constant memory footprint. GST first divides a large graph into segments and then backpropagates through only a few segments sampled per training iteration. We refine the GST paradigm by introducing a historical embedding table to efficiently obtain embeddings for segments not sampled for backpropagation. To mitigate the staleness of historical embeddings, we design two novel techniques. First, we finetune the prediction head to fix the input distribution shift. Second, we introduce Stale Embedding Dropout to drop some stale embeddings during training to reduce bias. We evaluate our complete method GST-EFD (with all the techniques together) on two large graph property prediction benchmarks: MalNet and TpuGraphs. Our experiments show that GST-EFD is both memory-efficient and fast, while offering a slight boost on test accuracy over a typical full graph training regime.

Large Graph Property Prediction via Graph Segment Training

Abstract

Learning to predict properties of a large graph is challenging because each prediction requires the knowledge of an entire graph, while the amount of memory available during training is bounded. Here we propose Graph Segment Training (GST), a general framework that utilizes a divide-and-conquer approach to allow learning large graph property prediction with a constant memory footprint. GST first divides a large graph into segments and then backpropagates through only a few segments sampled per training iteration. We refine the GST paradigm by introducing a historical embedding table to efficiently obtain embeddings for segments not sampled for backpropagation. To mitigate the staleness of historical embeddings, we design two novel techniques. First, we finetune the prediction head to fix the input distribution shift. Second, we introduce Stale Embedding Dropout to drop some stale embeddings during training to reduce bias. We evaluate our complete method GST+EFD (with all the techniques together) on two large graph property prediction benchmarks: MalNet and TpuGraphs. Our experiments show that GST+EFD is both memory-efficient and fast, while offering a slight boost on test accuracy over a typical full graph training regime.

1 Introduction

Graph property prediction is a task of predicting a certain property or a characteristic of an entire graph [3]. Important applications include, predicting properties of molecules [32, 12], predicting properties of programs/code [1, 12, 6, 41] and predicting properties of organisms based on their protein-protein interaction networks [27, 42].

These popular graph property prediction tasks deal with relatively small graphs, so the scalability issue arises only from a large number of (small) graphs.

However, graph property prediction tasks also face another scalability challenge, which arises due to the large size of each individual graph, as some graphs can have millions or even billions of nodes and edges [10]. Training typical Graph Neural Networks (GNNs) to classify such large graphs can be computationally infeasible, as the memory needed scales at least linearly with the size of the graph [40]. This presents a challenge as even most powerful GPUs, which are optimized for handling large amounts of data, only have a limited amount of memory available.

Previous efforts to improve scalability of GNNs have mostly focused on developing methods for node-level and link-level prediction tasks, which can be performed using sampled subgraphs [11, 5, 13, 38, 43, 2, 21]. However, there is a lack of research on how to train scalable models for property prediction of large graphs . Training a model on a sampled subgraph alone is insufficient for these

(a) Full Graph Training

(b) Our Work: Graph Segment Training

𝑭′

𝑭

𝑭′

Gradient Flow

𝓖

𝒚 $

𝒉

𝒚 $

𝑭

types of tasks, as the subgraph sampled may not contain all the necessary information to make accurate predictions for the entire graph. For example, if the task is to predict the diameter of the graph, it is unlikely that a fixed-size subgraph would contain sufficient features for the GNN to make an accurate and reliable prediction. Thus, it is essential to aggregate information from the entire graph to predict a graph property.

In this paper, we address the problem of property prediction of large graphs. We propose Graph Segment Training (GST) 3 , which is able to train on large graphs with constant (GPU) memory footprint.

Our approach partitions each large graph into smaller segments with a controlled size in the preprocessing phase. During the training process, a random subset of segments is selected to update the model at each step, rather than using the entire graph. This way, we need to maintain intermediate activations for only a few segments for backpropagation; embeddings for the remaining segments are created without saving their intermediate activations. The embeddings of all segments are then combined to generate an embedding for the original large graph, which is used for prediction. Therefore, each large graph has an upper bound on memory consumption during training regardless of its original size. This allows us to train the model on large graphs without running into an out-of-memory (OOM) issue, even with limited computational resources.

To accelerate the training process further, we introduce a historical embedding table to efficiently produce embeddings for graph segments that do not require gradients, as historical embeddings eliminate additional computation on such segments. However, the historical embeddings induce staleness issues during training, so we design two techniques to mitigate such issues in practice. First, we characterize the input distribution mismatch issue between training and test stages of the prediction head, and propose to finetune only the prediction head at the end of training to close the input distribution gap. Second, we identify bias in the loss function due to stale historical embeddings, and introduce Stale Embedding Dropout to drop some stale embeddings during training to reduce this bias. Our final proposed method, called GST+EFD, is both memory-efficient and fast.

We evaluate our method on the following datasets: MalNet-Tiny, MalNet-Large and TpuGraphs. A typical full graph training pipeline (Full Graph Training) can only train on MalNet-Tiny, and

unavoidably reaches OOM on MalNet-Large and TpuGraphs dataset. On the contrary, we empirically show that the proposed GST framework successfully trains on arbitrarily large graphs using a single NVIDIA-V100 GPU with only 16GB of memory for MalNet-Large and four GPUs for TpuGraphs dataset, while maintaining comparable performance with Full Graph Training. We finally demonstrate that our complete method GST+EFD slightly outperforms GST by another 1-2% in terms of final evaluation metric, and simultaneously being 3x faster in terms of training time.

2 Preliminaries

Notation. For a function f ( a ) : R d 0 → R d 1 , we use D k a f [ a 0 ] ∈ R d 0 × d 1 to denote its k -th derivative of f with respect to a evaluated at value a 0 . Let f ◦ g to denote the composition of function f with g . We use ⊙ to denote entry-wise product, and let a ⊙ 2 ∆ = a ⊙ a . We use ⊕ to represent aggregation. ⊕ j ≤ J a j means aggregating the set { a j } j ≤ J , where j indexes segments, and J is the number of segments in a graph. We usually drop j subscript from ⊕ for brevity. a i ⊕ b j aggregates both sets { a i } and { b j } together. ⊕ can be mean or sum operators when applying to vectors. We define the input graph as G = {V , E} , where V = { v 1 , ..., v m } is the node set and E ⊆ V × V is the edge set.

Let dataset D = { ( G ( i ) , y ( i ) ) } i ≤ n contain n examples: each label y ( i ) is associated with G ( i ) .

Graph Neural Network. We consider a backbone graph neural network F that takes a graph G ( i ) and generates graph embedding h ( i ) ∈ R d h , followed by a final prediction head F ′ that takes graph embedding h ( i ) and outputs final predictions: ̂ y ( i ) = ( F ′ ◦ F )( G ( i ) ) . We optimize GNN with the loss function as L (( F ′ ◦ F )( G ( i ) ) , y ( i ) ) .

Historical Embedding Table. We define an embedding table T : K → R d h , where key-space K ⊂ Z × Z is a tuple: (graph index i ≤ n , segment index j ≤ J ). We use ˜ h ( i ) j = T ( i, j ) to denote embedding for graph segment G ( i ) j : an embedding not up to date with the current backbone F .

3 Our Method: GST+EFD

3.1 Graph Segment Training (GST)

Given a training graph dataset D train = { ( G ( i ) , y ( i ) ) } n i =1 , a common SGD update step requires calculating the gradient:

where θ is trainable weights in F ′ ◦ F , and B is a sampled minibatch. Graphs can differ in size (the number of nodes |V ( i ) | ), with some being too large to fit into the device's memory. This is because the memory required to store all intermediate activations for computing gradients is proportional to the number of nodes and edges in the graph.

To address the above issue, we propose to partition each original input graph into a collection of graph segments, i.e. ,

An example of a partition algorithm is METIS [16]. This preprocessing step will result in a training set D train = { ( ⊕ j ≤ J ( i ) G ( i ) j , y ( i ) ) } n i =1 . Number of partitions J ( i ) can vary across graphs, but the size of each graph segment can be bounded by a controlled size ( |V ( i ) j | < m GST , ∀ ( i, j ) ) so that a batch of a fixed number of graph segments can always fit within the device's memory.

When processing graph segment G ( i ) j , we can obtain its segment embedding through the backbone: h ( i ) j = F ( G ( i ) j ) . The prediction head F ′ requires information from the whole graph to make the prediction, thus we propose to aggregate all segment embeddings to recover the full graph embedding: h ( i ) = ⊕ h ( i ) j . A simple realization of this aggregation is mean pooling. Note that naïvely applying

the prediction head F ′ on the aggregated graph embedding -̂ y ( i ) = F ′ ( ⊕ h ( i ) j ) -would not provide any reduction in peak memory consumption, as we need to keep track of the activations of all graph segments {G ( i ) j } j ∈ J ( i ) to perform backpropagation.

Thus, we propose to perform backpropagation on only a few randomly sampled graph segments S ( i ) ⊆ { 1 , . . . , J ( i ) } and generate embeddings without requiring gradients for the rest. We hereby denote h s and ¯ h j as segment embeddings that require and do not require gradient, respectively. An entire graph embedding is then: h ( i ) ≈ { h ( i ) s } s ∈S ( i ) ⊕ { ¯ h ( i ) j } j / ∈S ( i ) ≜ h ( i ) s ⊕ ¯ h ( i ) j . We name this general pipeline as GST and summarize it in Algorithm 1.

Algorithm 1 General Framework of GST

Require: A preprocessed training graph dataset D train = { ( ⊕ G ( i ) j , y ( i ) ) } n i =1 . A parameterized backbone F and a prediction head F ′ .

1: for t = 1 to T 0 do

2: B ← SampleMiniBatch ( D train )

3: for ( G ( i ) , ̂ y ( i ) ) in B do

4: {G ( i ) s } s ∈S ( i ) ← SampleGraphSegments ( G ( i ) )

5: ¯ h ( i ) j ← ProduceEmbedding ( G ( i ) j ) for j / ∈ S ( i )

6: h ( i ) s ← F ( G ( i ) s ) for s ∈ S ( i )

7: end for

8: SGD on loss ← 1 |B| ∑ i L ( F ′ ( h ( i ) s ⊕ ¯ h ( i ) j ) , ̂ y ( i ) )

9: end for

One implementation of ProduceEmbedding ( · ) in Algorithm 1 is to use the same feature encoder F to forward all the segments in {G ( i ) j } j / ∈S ( i ) without storing any intermediate activation (by stopping gradient).

3.2 GST with Historical Embedding Table

Calculating ¯ h j by stopping gradient guarantees an upper bound on peak memory consumption. However, since we do not need gradients for segments {G ( i ) j } j / ∈S ( i ) , computing forward pass on these segments can be avoided to make training faster. To achieve this, we use historical embeddings acquired in previous training iterations ˜ h ( i ) j = T ( i, j ) . With an embedding table T , one can implement ProduceEmbedding ( · ) by fetching the corresponding embedding from the table without any computation. We update the embedding table after conducting the forward pass on a graph segment. We optimize the following loss L ( F ′ ( h ( i ) s ⊕ ˜ h ( i ) j ) , y ( i ) ) during training. We name the embedding version of our algorithm as GST+E.

Note that GST+E has runtime advantages over GST. For each segment G ( i ) j that does not require gradient, GST needs to run an actual forward pass over G ( i ) j , while GST+E only needs to fetch the embedding from a hash table. GST+E has a small overhead from writing the updated embedding of G ( i ) s back into the table T , which is relatively quick and can be run in a separate thread so that it does not impede the main training algorithm until the next iteration.

The drawback of GST+E is that historical embeddings from T may be stale;

˜ h ( i ) j can be the result of an out-dated feature extractor F . This type of staleness can hurt the training in various ways. Below, we provide two techniques to mitigate the staleness issue.

3.3 Prediction Head Finetuning

Let's compare the input-output distribution of the prediction head F ′ during the training and inference stage. We have training distribution P train ( h , y ) = P ( h s ⊕ ˜ h j , y ) and test distribution P test ( h , y ) = P ( ⊕ h j , y ) . Regardless of the innate distribution shift between the training and test stage of the dataset, we note that stale historical embeddings can further widen the gap between the training and

test distributions. In this case, the minimizer of the expected training loss does not minimize the expected test loss:

̸

To mitigate the distribution misalignment, we introduce the Prediction Head Finetuning technique. Concretely, at the end of training, we update each embedding h ( i ) j in the embedding table T by forwarding each graph segment in the training set with the most current feature encoder F . We then finetune only the prediction head F ′ with all the input embeddings up-to-date. We use GST+EF to denote GST+E refined with the Prediction Head Finetuning technique.

The overhead from the finetuning is minimal, as we need to update the embedding table T only once. The rest of the finetuning stage does not involve the notoriously slow graph convolution because the prediction head F ′ is simply a multi-layer perceptron.

3.4 Stale Embedding Dropout

The finetuning technique primarily addresses the negative impact of stale embeddings on prediction head F ′ . However, the staleness also impacts the backbone F . Prior works studying the effects of stale historical embeddings commonly assume that historical embeddings do not become too stale, i.e. , ∥ ˜ h ( i ) j -¯ h ( i ) j ∥ ≤ ϵ, ∀ ( i, j ) . Given this assumption, if the neural network ( F ′ ◦ F )( · ) is k -Lipschitz continuous, the gradients will also be bounded and never run too far from its true estimation, i.e. , ∥∇ F ′ ( ˜ h ( i ) j ) -∇ F ′ ( ¯ h ( i ) j ) ∥ ≤ k ′ · ϵ . Thus, the network can often converge to the similar local minima even when using historical embeddings.

The above assumption does not hold under our GST+E framework. The rationale is that ˜ h ( i ) j gets updated infrequently in the embedding table T . Suppose we iterate through every graph G ( i ) in the training set D train for each epoch, every time we train on a graph G ( i ) ≈ ⊕ G ( i ) j , only a few graph segments G ( i ) s will be updated in the table T . This implies that all the other segment embeddings of G ( i ) will be at least n -iteration stale, with n being the number of graphs in the training set, and the most outdated segment embedding could be approximately nJ ( i ) /S ( i ) -iteration stale, where S ( i ) is |S ( i ) | .

This staleness introduces an additional source of bias and variance to the stochastic optimization; the loss function calculated with historical embeddings is no longer an unbiased estimation of its true value. To mitigate the negative impact of historical embeddings on loss function estimation, we propose the second technique, Stale Embedding Dropout (SED). Unlike a standard Dropout, which uniformly drops elements and weighs up the rest, we propose to drop only stale segment embeddings and weigh up only segment embeddings that are up-to-date. Concretely, assume with the keep probability p , the weight η for each segment is defined as:

Please refer to the theoretical analysis in the next section. By combining the two proposed techniques, we denote our final algorithm as GST+EFD, which we summarized in Algorithm 2 in Appendix B.

4 Theoretical Analysis

We characterize the effect of stale historical embeddings by studying the difference between L ( F ′ ( h ( i ) s ⊕ ˜ h ( i ) j )) and L ( F ′ ( ⊕ h ( i ) j )) . Assume G ( i ) has J ( i ) segments and we perform backpropagation on S ( i ) segments. We let δ ( i ) ≜ h ( i ) s ⊕ ˜ h ( i ) j -⊕ h ( i ) j be the perturbation on the graph embedding.

We apply Taylor expansion around δ ( i ) = 0 on the final loss to analyze the effect of this perturbation.

In the equation above, the first-order term acts as a bias term introduced by the stale historical embedding, and the second-order term acts as an additional regularizer.

Let ET denote using the embedding table without applying SED. We analyze the effect of ET and SED under the Taylor Expansion in Eq. 2 by substituting different δ ( i ) . For the first term, we have

where C is a constant matrix.

We extend the above analysis to the following theorem. Please find the complete proof in Appendix A.

Theorem 4.1. Under proper conditions, SED with a keep ratio p ensures to reduce bias term introduced by historical embeddings by a factor of p , while introducing another regularization term.

Theorem 4.1 indicates that SED can reduce the bias in the loss function introduced by the stale historical embeddings, at the cost of another regularization term, which might potentially increase total regularization. Prior works commonly make a hidden assumption that E ( ˜ h ( i ) j -h ( i ) j ) is so small that the negative effect may be neglected. In our setting, the historical embeddings can be very stale, so having a p factor helps reduce bias. It is worthwhile to check the limiting cases. If p = 1 (keeping all the stale embeddings without dropping any), both E ( i ) SED [ B ] and E ( i ) SED [ R ] degrade to the result

δ j δ j of ET. If p = 0 (droping all the stale embeddings), then the algorithm degrades to training on only S ( i ) segments without aggregating other segments (which we denote as GST-One, when S ( i ) = 1 ). E δ ( i ) SED j [ B ] = 0 indicates that there is no stale bias in this case. However, the term E δ ( i ) SED j [ R ] could become too large so that it impedes training.

5 Experiments

5.1 Experimental Setup

Datasets. MalNet [10] is a large-scale graph representation learning dataset, with the goal to predict the category of a function call graph. MalNet is the largest public graph database constructed to date in terms of average graph size. Its widely-used split is called MalNet-Tiny , containing 5,000 graphs across balanced 5 types, with each graph containing at most 5,000 nodes. To evaluate our approach on the regime where the graph size is large, we construct an alternative split from the original MalNet dataset, which we named MalNet-Large . MalNet-Large also contains 5,000 graphs across balanced 5 types. MalNet-Large 's average graph size reaches 47k with the largest graph containing 541k nodes. We will release our experimental split for MalNet-Large to promote future research.

TpuGraphs is an internal large scale graph regression dataset, whose goal is to predict an execution time of an XLA's HLO graph with a specific compiler configuration on a Tensor Processing Unit (TPU). XLA [28] is a production backend compiler for various machine learning frameworks, including TensorFlow, PyTorch, and JAX. In this particular dataset, the compiler configuration controls physical layouts of tensors in the graph, and the runtime is measured on TPU v3 [15]. This dataset cares more about the ranking of the configurations for each graph than the absolute runtimes, since the ultimate goal is to use a model to select the best configuration for each graph. 4 TpuGraphs

is similar to the dataset used in [17], but the runtime prediction is at the entire graph level rather than at the kernel (subgraph) level. TpuGraphs contains 5,153 HLO graphs and a total of 757,375 unique pairs of graphs and configurations. The average graph size is 38k, and the maximum is 615k. From the perspective of GST, a graph together with a configuration defines one G ( i ) because the configuration is featurized as parts of input node features to the GNN.

Please refer to Table 4 in Appendix for detailed statistics.

Methods. We test combinations of the following proposed techniques and some baselines. (1) Full Graph Training: we train on all graphs in their original scale without applying any partitioning beforehand. (2) GST-One: we partition the original graph into a collection of graph segments G ( i ) ≈ ⊕ G ( i ) j , but we randomly select only one segment G ( i ) j for each graph to train every iteration. (3) GST: following the general GST framework described in Algorithm 1, we replace ProduceEmedding ( · ) by using the same feature encoder F to forward all the segments in {G ( i ) j } j / ∈S ( ⟩ ) without storing any intermediate activation. We set S ( i ) = 1 in our experiments. (4) E: we introduce an embedding table ˜ h ( i ) j = T ( i, j ) to store the historical embedding of each graph segment, and we fetch the embedding from T if we do not need to calculate gradient for the corresponding segment. (5) F: in addition to introducing the embedding table T , we finetune the prediction head F ′ with all up-to-date segment embeddings at the end of training. (6) D: we apply SED defined in Eq. 1 during training.

When these techniques are combined, we concatenate the acronyms with a '+' to GST as an abbreviation. We conduct all the experiments on MalNet with a single NVIDIA-V100 GPU with 16GB of memory, and four NVIDIA-V100 GPUs (for data parallelism) with 16GB of memory for TpuGraphs. Please refer to Appendix B for additional implementation details.

5.2 Empirical Results on MalNet

To demonstrate the general applicability of our proposed GST framework, we consider three backbones, namely, GCN [19], SAGE [11], and GraphGPS [25]. GCN and SAGE are two popular GNN architectures. GraphGPS is a Graph Transformer that recently achieves state-of-the-art performance on many graph-level tasks, but is well-known for its issue on scalability. We report the top-1 test accuracy of various methods on MalNet-Tiny and MalNet-Large in Table 1. We include MalNet-Tiny in this study because its graphs are relatively small so that it is still possible to run Full Graph Training.

Notably, we observe that GST slightly outperforms Full Graph Training in terms of test accuracy on MalNet-Tiny. GST has exactly the same number of weight parameters with Full Graph Training. This implies that GST potentially has a better hierarchical graph pooling mechanism that leads to better generalization. As we step from MalNet-Tiny to MalNet-Large, Full Graph Training strategy can no longer fit the large graphs on a GPU, so we report OOM in the table. GST's estimation on graph segment embeddings ¯ h ( i ) j that do not require gradients is accurate, and thus does not suffer from staleness issues. Therefore, we use GST as an estimation for the performance of Full Graph Training on MalNet-Large.

Naïvely training on only one graph segment (GST-One) yields inferior performance than Full Graph Training and GST, showing that it is essential to aggregate embeddings from all graph segments.

0

200

400

600

Epoch

20

30

40

50

60

70

80

90

Train

Test

0.0

0.25

0.5

0.75

1.0

p

60

65

70

75

80

85

90

95

100

100

1000

3000

5000

7000

9000

50

60

70

80

90

100

Solely introducing the historical embedding table (GST+E) significantly deteriorates the optimization due to the staleness issue. Each of the proposed techniques (Prediction Head Finetuning and SED) individually is beneficial in combating the staleness issue. The combination of our two techniques (GST+EFD) achieves the best performance, slightly outperforming GST by another 1-2% in terms of final evaluation metric.

5.3 Empirical results on TpuGraphs dataset

As mentioned in Section 5.1, we care more about the ranking of configurations for each graph than the absolute runtimes in this dataset. Thus, we report the ordered pair accuracy (OPA) averaged over all computational graphs in Table 2. OPA metric is defined as:

With some compiler domain knowledge, we found it better to predict each graph segment's runtime first and then aggregate them using sum pooling. This means the prediction head is part of F , and F ′ is simply a summation function. Since there are no learnable weights in F ′ , we omit Prediction Head Finetuning in this experiment, and GST+EFD in Table 2 excludes the finetuning stage. We observe a clear tradeoff between fitting training examples and generalization. First, GST has a much higher training OPA than the other methods, indicating accurate estimation on segments that do not require gradient is essential for training OPA. Training with one segment only or with the embedding table yields lower training OPA, and consequently lower test OPA as well. SED in GST+EFD functions as a regularization technique. Although it slightly lowers the training OPA compared to GST+E, it achieves better test OPA, even than GST, due to bias mitigation.

5.4 Ablation Studies

Effect of finetuning. We visualize the training/test accuracy curve of GST+EFD over time in Figure 2. The staleness introduced by historical embeddings drastically hurts generalization, as shown for the first 600 epochs. We start finetuning at epoch 600, and the gap between training and test accuracy decreases by a large margin instantly.

Ablation study on segment dropout ratio. To analyze the effect of the keep ratio p in SED, we vary its value from 0 to 1 and visualize the results in Figure 3. When p = 1 , GST+EFD degrades back to using the historical embedding table without SED, as the performance decreases due to staleness. When p = 0 , GST+EFD becomes GST-One, where we drop all the stale historical embeddings. This extreme case introduces too heavy regularization that impedes the model from fitting the training data, leading to a decrease in test performance ultimately. We found that p = 0 . 5 achieves a satisfactory tradeoff between fitting the training data and adding a proper amount of regularization.

Ablation study on segment size. We also alter the maximum segment size and visualize the results in Figure 4. A smaller maximum segment size will result in much more number of segments. Interestingly, we found that the proposed GST+EFD is very robust to the choice of the maximum segment size, as long as the segment size is reasonally large.

Ablation study on partition algorithms. Please refer to Appendix C.

5.5 Runtime Analysis

Next, we empirically compare runtime of different variants under the proposed GST framework. We summarize an average time for one forward-backward pass during training on MalNet-Large dataset in Table 3. Since GST runs inference for the graph segments that do not require gradients, the runtime of GST is significantly higher than others'. We also found that GST+E's and GST+EFD's runtime are very close to GST-One's; this means the overhead of fetching embeddings from the embedding table T is minimal. Moreover, GST+EFD's runtime is slightly lower than GST+E's because in the implementation, we can skip the fetching process if an embedding is set to be dropped. This result demonstrates that our proposed GST+EFD not only is efficient in terms of memory usage but also reduces training time significantly.

6 Related Works

Graph property prediction. In the context of graph property prediction, a model must predict a certain characteristic associated with the whole graph. Standard graph neural networks produce node embeddings as outputs [19, 11, 34]. To create a graph embedding (a vector representing the entire graph) out of the node embeddings, pooling methods are usually applied at the end. Common approaches to this problem include simply summing up or averaging all the node embeddings [8], or introducing a 'virtual node' connected to all the other nodes [23]. Fully-connected Graph Transformer was recently proposed with an outstanding success on existing graph property prediction benchmarks [35, 25]. However, the fully-connected attention matrix limits the applicability of Graph Transformer to only small graphs with limited number of nodes [26].

GNN with graph partitioning. ClusterGCN [7] is designed for node-level tasks by training on graph segments. It partitions a graph into graph segments and randomly selects graph segments to form a minibatch during training. ROC [14], PipeGCN [29] and BNS-GCN [30] achieve distributed node-level GCN training through partitioning a graph into small segments such that each could be fitted into a single GPU memory, and then training multiple segments in parallel. All the above graph partitioning techniques for GNNs rely on the fact that an ego-subgraph (a subgraph centered around a node) contains sufficient information to make a prediction for the centered node. This is not true for graph property prediction tasks where we need to aggregate information from the whole graph to make an accurate prediction.

GNN with historical embeddings. The idea of historical embeddings was first introduced in VR-GCN [4], which uses historical embeddings to control the variance of neighborhood sampling. GNNAutoScale [9] incorporates historical embeddings to recover a more accurate neighborhood estimation in a scalable fashion. Developed upon GNNAutoScale, Yu et al. [37] uses a momentum step to incorporate historical embeddings when updating feature representations to further alleviate the staleness issue. These prior works maintain a historical embedding for each node because they consider node-level tasks. As we consider graph property prediction tasks, we record a historical embedding for each graph segment rather than each node.

7 Conclusion

We study how to train a GNN model for large graph property prediction tasks. We propose Graph Segment Training (GST), a general framework for learning large graph property prediction tasks with a constant memory footprint. We further introduce a historical embedding table to efficiently produce embeddings for graph segments that do not require gradients, and design two novel techniques Prediction Head Finetuning and Stale Embedding Dropout - to mitigate the staleness issue. In conclusion, we believe that our proposed method is a step toward making graph property prediction learning more practical and scalable.

References

Miltiadis Allamanis. The adverse effects of code duplication in machine learning models of code. In Proceedings of the 2019 ACM SIGPLAN International Symposium on New Ideas, New Paradigms, and Reflections on Programming and Software , pages 143-153, 2019.

Aleksandar Bojchevski, Johannes Klicpera, Bryan Perozzi, Amol Kapoor, Martin Blais, Benedek Rózemberczki, Michal Lukasik, and Stephan Günnemann. Scaling graph neural networks with approximate pagerank. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining , pages 2464-2473, 2020.

Ines Chami, Sami Abu-El-Haija, Bryan Perozzi, Christopher Ré, and Kevin Murphy. Machine learning on graphs: A model and comprehensive taxonomy. Journal of Machine Learning Research , 23(89):1-64, 2022.

Jianfei Chen, Jun Zhu, and Le Song. Stochastic training of graph convolutional networks with variance reduction. In International Conference on Machine Learning , pages 942-950. PMLR, 2018.

Jie Chen, Tengfei Ma, and Cao Xiao. Fastgcn: Fast learning with graph convolutional networks via importance sampling. In International Conference on Learning Representations , 2018.

Tianqi Chen, Lianmin Zheng, Eddie Yan, Ziheng Jiang, Thierry Moreau, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. Learning to Optimize Tensor Programs. In Proceedings of the 32nd International Conference on Neural Information Processing Systems , NeurIPS'18, page 3393-3404, Red Hook, NY, USA, 2018. Curran Associates Inc.

Wei-Lin Chiang, Xuanqing Liu, Si Si, Yang Li, Samy Bengio, and Cho-Jui Hsieh. Cluster-gcn: An efficient algorithm for training deep and large graph convolutional networks. In Proceedings of the 25th ACM SIGKDD international conference on knowledge discovery & data mining , pages 257-266, 2019.

Vijay Prakash Dwivedi, Chaitanya K Joshi, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Benchmarking graph neural networks. arXiv preprint arXiv:2003.00982 , 2020.

Matthias Fey, Jan E Lenssen, Frank Weichert, and Jure Leskovec. Gnnautoscale: Scalable and expressive graph neural networks via historical embeddings. In International Conference on Machine Learning , pages 3294-3304. PMLR, 2021.

Scott Freitas, Yuxiao Dong, Joshua Neil, and Duen Horng Chau. A large-scale database for graph representation learning. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 1) , 2021.

Will Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large graphs. Advances in neural information processing systems , 30, 2017.

Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. Advances in neural information processing systems , 33:22118-22133, 2020.

Wenbing Huang, Tong Zhang, Yu Rong, and Junzhou Huang. Adaptive sampling towards fast graph representation learning. Advances in neural information processing systems , 31, 2018.

Zhihao Jia, Sina Lin, Mingyu Gao, Matei Zaharia, and Alex Aiken. Improving the accuracy, scalability, and performance of graph neural networks with roc. Proceedings of Machine Learning and Systems , 2:187-198, 2020.

Norman P. Jouppi, Doe Hyun Yoon, George Kurian, Sheng Li, Nishant Patil, James Laudon, Cliff Young, and David Patterson. A Domain-Specific Supercomputer for Training Deep Neural Networks. Commun. ACM , 63(7):67-78, June 2020. ISSN 0001-0782. doi: 10.1145/3360307.

George Karypis and Vipin Kumar. Metis: A software package for partitioning unstructured graphs, partitioning meshes, and computing fill-reducing orderings of sparse matrices. 1997.

Samuel J. Kaufman, Phitchaya Mangpo Phothilimthana, Yanqi Zhou, Charith Mendis, Sudip Roy, Amit Sabne, and Mike Burrows. A Learned Performance Model for Tensor Processing Units. In Proceedings of Machine Learning for Systems , 2021.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. 2016.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101 , 2017.

Elan Sopher Markowitz, Keshav Balasubramanian, Mehrnoosh Mirtaheri, Sami Abu-El-Haija, Bryan Perozzi, Greg Ver Steeg, and Aram Galstyan. Graph traversal with tensor functionals: A meta-algorithm for scalable learning. In International Conference on Learning Representations , 2021.

Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.

Trang Pham, Truyen Tran, Hoa Dam, and Svetha Venkatesh. Graph classification via deep learning with virtual nodes. arXiv preprint arXiv:1708.04357 , 2017.

Phitchaya Mangpo Phothilimthana, Sami Abu-El-Haija, Kaidi Cao, Bahare Fatemi, Charith Mendis, and Bryan Perozzi. Tpugraphs: A performance prediction dataset on large tensor computational graphs. arXiv preprint arXiv:2308.13490 , 2023.

Ladislav Rampášek, Mikhail Galkin, Vijay Prakash Dwivedi, Anh Tuan Luu, Guy Wolf, and Dominique Beaini. Recipe for a general, powerful, scalable graph transformer. arXiv preprint arXiv:2205.12454 , 2022.

Yu Shi, Shuxin Zheng, Guolin Ke, Yifei Shen, Jiacheng You, Jiyan He, Shengjie Luo, Chang Liu, Di He, and Tie-Yan Liu. Benchmarking graphormer on large-scale molecular modeling datasets. arXiv preprint arXiv:2203.04810 , 2022.

Damian Szklarczyk, Annika L Gable, David Lyon, Alexander Junge, Stefan Wyder, Jaime Huerta-Cepas, Milan Simonovic, Nadezhda T Doncheva, John H Morris, Peer Bork, et al. String v11: protein-protein association networks with increased coverage, supporting functional discovery in genome-wide experimental datasets. Nucleic acids research , 47(D1):D607-D613, 2019.

TensorFlow. XLA: Optimizing Compiler for TensorFlow. https://www.tensorflow.org/xla. [Online; accessed 19-September-2019].

Cheng Wan, Youjie Li, Cameron R Wolfe, Anastasios Kyrillidis, Nam Sung Kim, and Yingyan Lin. Pipegcn: Efficient full-graph training of graph convolutional networks with pipelined feature communication. In International Conference on Learning Representations , 2021.

Cheng Wan, Youjie Li, Ang Li, Nam Sung Kim, and Yingyan Lin. Bns-gcn: Efficient full-graph training of graph convolutional networks with partition-parallelism and random boundary node sampling. Proceedings of Machine Learning and Systems , 4:673-693, 2022.

Colin Wei, Sham Kakade, and Tengyu Ma. The implicit and explicit regularization effects of dropout. In International conference on machine learning , pages 10181-10192. PMLR, 2020.

Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh S Pappu, Karl Leswing, and Vijay Pande. Moleculenet: a benchmark for molecular machine learning. Chemical science , 9(2):513-530, 2018.

Cong Xie, Ling Yan, Wu-Jun Li, and Zhihua Zhang. Distributed power-law graph computing: Theoretical and empirical analysis. Advances in neural information processing systems , 27, 2014.

Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? In International Conference on Learning Representations , 2018.

Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie-Yan Liu. Do transformers really perform badly for graph representation? Advances in Neural Information Processing Systems , 34:28877-28888, 2021.

Jiaxuan You, Zhitao Ying, and Jure Leskovec. Design space for graph neural networks. Advances in Neural Information Processing Systems , 33:17009-17021, 2020.

Haiyang Yu, Limei Wang, Bokun Wang, Meng Liu, Tianbao Yang, and Shuiwang Ji. Graphfm: Improving large-scale gnn training via feature momentum. In International Conference on Machine Learning , pages 25684-25701. PMLR, 2022.

Hanqing Zeng, Hongkuan Zhou, Ajitesh Srivastava, Rajgopal Kannan, and Viktor Prasanna. Graphsaint: Graph sampling based inductive learning method. In International Conference on Learning Representations , 2019.

Chenzi Zhang, Fan Wei, Qin Liu, Zhihao Gavin Tang, and Zhenguo Li. Graph edge partitioning via neighborhood heuristic. In Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining , pages 605-614, 2017.

Hengrui Zhang, Zhongming Yu, Guohao Dai, Guyue Huang, Yufei Ding, Yuan Xie, and Yu Wang. Understanding gnn computational graph: A coordinated computation, io, and memory perspective. Proceedings of Machine Learning and Systems , 4:467-484, 2022.

Lianmin Zheng, Ruochen Liu, Junru Shao, Tianqi Chen, Joseph E. Gonzalez, Ion Stoica, and Ameer Haj Ali. Tenset: A large-scale program performance dataset for learned tensor compilers. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 1) , 2021. URL https://openreview.net/forum?id=aIfp8kLuvc9 .

Marinka Zitnik, Rok Sosiˇ c, Marcus W Feldman, and Jure Leskovec. Evolution of resilience in protein interactomes across the tree of life. Proceedings of the National Academy of Sciences , 116(10):4426-4433, 2019.

Difan Zou, Ziniu Hu, Yewen Wang, Song Jiang, Yizhou Sun, and Quanquan Gu. Layerdependent importance sampling for training deep and large graph convolutional networks. Advances in neural information processing systems , 32, 2019.

A Proofs and Derivations

Theorem A.1. Under proper a condition that W · ˜ h ( i ) j ≈ 0 and W · h ( i ) j ≈ 0 , where W is the first linear transformation in F ′ , SED with a keep ratio p ensures to reduce bias term introduced by historical embeddings by a factor of p , while introducing another regularization term.

Proof. Let δ ( i ) ≜ h ( i ) s ⊕ ˜ h ( i ) j -⊕ h ( i ) j be the perturbation on the graph embedding. We use ET to denote using the embedding table without applying SED. For GST+E, we have

The randomness above comes from the fact that each segment G ( i ) j is selected for backpropagation with probability S ( i ) J ( i ) .

For SED, there are two folds of randomness when training on graph G ( i ) : randomly selecting S ( i ) segments to train and randomly select stale embedding ˜ h ( i ) j to drop. Thus we can rewrite δ ( i ) j as

We apply Taylor expansion around δ ( i ) j = 0 on the final loss to analyze the effect of this perturbation. In Section A.2 of Wei et al. [31], when W · ˜ h ( i ) j ≈ 0 and W · h ( i ) j ≈ 0 , the perturbation of δ ( i ) to the loss function might not be too large, so it supports the use of Taylor Expansion.

Note that we randomly select segments with index s during training, we can then derive an approximation of the expected difference during training as

We can then compare the effect of SED by substituting the two versions of δ ( i ) j into Eq. 3. So for the first term, we have

whereas for the second term, we have

It is easy to check that the statement satisfies given the value calculated.

B Implementation Details

B.1 Missing Algorithm

Algorithm 2 Pipeline of GST+EFD

Require: A preprocessed training graph dataset D train = { ( ⊕ G ( i ) j , y ( i ) ) } n i =1 . A parameterized backbone F and a prediction head F ′ . A historical segment embedding table T . 1: for t = 1 to T 0 do 2: B ← SampleMiniBatch ( D train ) 3: for ( G ( i ) , y ( i ) ) in B do 4: {G ( i ) s } s ∈S ( i ) ← SampleGraphSegments ( G ( i ) ) 5: ˜ h ( i ) j ←T . LookUp (( i, j )) for j / ∈ S ( i ) 6: h ( i ) s ← F ( G ( i ) s ) for s ∈ S ( i ) 7: T . InsertOrUpdate (( i, s ) , h ( i ) s ) 8: end for 9: SGD on 1 |B| ∑ i L ( F ′ ( η ( i ) s · h ( i ) s ⊕ ( η ( i ) j · ˜ h ( i ) j ) , y ( i ) ) {SED with η defined in Equation 1} 10: end for 11: # Prediction Head Finetuning 12: T . InsertOrUpdate (( i, j ) , F ( G ( i ) j )) for every G ( i ) j 13: for t = T 0 to T 1 do 14: for G ( i ) in SampleMiniBatch ( D train ) do 15: ¯ h ( i ) j ←T . LookUp (( i, j )) for j ≤ J ( i ) 16: end for 17: SGD with loss ← 1 |B| ∑ i L ( F ′ ( ⊕ ¯ h ( i ) j ) ) 18: end for

We follow GraphGym [36] to represent design spaces of GNN as (message passing layer type, number of pre-process layers, number of message passing layers, number of post-process layers, activation, aggregation). Our code is implemented in PyTorch [22].

Implementation details for MalNet-Large. We consider three model variations for the MalNetLarge dataset. Please refer to their hyperparameters in Table 5. We use Adam optimizer [18] with the base learning rate of 0.01 for GCN and SAGE. For GraphGPS, we use AdamW optimizer [20] with the cosine scheduler and the base learning rate of 0.0005. We use L2 regularization with a weight decay of 1e-4. We train for 600 epochs until convergence. For Prediction Head Finetuning, we finetune for another 100 epochs. We limit the maximum segment size to 5,000 nodes, and use a keep probability p = 0 . 5 if not otherwise specified. We train with CrossEntropy loss.

Implementation details for MalNet-Tiny. We use the same model architectures/training schedules as in the MalNet-Large dataset. The only difference is that as graphs in MalNet-Tiny have no more than 5000 nodes, so we limit maximum segment size to 500 here.

Implementation details for TpuGraphs. We only consider SAGE with configurations (SAGEConv, 0, 4, 3, 128, ReLU, sum) for the TpuGraphs dataset. We use Adam optimizer with the base learning rate of 0.0001. We train for 200,000 iterations until convergence. We by default limit the maximum segment size to 8,192 nodes, and use a keep probability p = 0 . 5 if not otherwise specified. Since we care more about relative ranking than the absolute runtime, we use PairwiseHinge loss within a batch during training:

C Additional Results

Convergence analysis. To study the effect on convergence for the proposed framework GST and technique SED, we visualize training/test curve per epoch on TpuGraphs dataset in Figure 5, MalNetTiny dataset in Figure 6. We show that the convergence speed of various methods studied are quite similar. In addition, the convergence speed in terms of iterations is similar between Full Graph Training and the proposed GST. Due to certain implementation overheads, e.g., on-the-fly graph segment extraction, embedding table query, we didn't observe speed up in terms of wall-clock time for our current implementation yet. Nevertheless, the implementation can be optimized further to reduce the overhead, we leave it for future work.

Ablation study on partition algorithms. We incorporated an ablation study focusing on various graph partition algorithms. The test accuracy of GST+EFD using the SAGE backbone for 5 iterations is depicted in Table 6. Our findings indicate that the Random Edge-Cut algorithm delivers subpar results due to its inability to maintain the integrity of the subgraph structure. Some attentive readers might hypothesize that ignoring edges between segments could lead to a decrease in graph property prediction accuracy, leading us to further investigate Vertex-Cut partition algorithms. In theory, Vertex-Cut techniques, which distribute edges across different machines and duplicate nodes as necessary, are likely to result in less information loss compared to Edge-Cut methods. Our empirical findings show that all partition algorithms that manage to retain the local structure have quite similar performance levels. This suggests that edges connecting different segments do not have a significant impact on the final prediction accuracy.

D Broader Impact

The paper presents Graph Segment Training (GST), a framework for predicting properties of large graphs using a divide-and-conquer approach. This method addresses the challenge of memory

25000

50000

75000

100000

125000

150000

175000

200000

Iterations

60

65

70

75

80

85

Train OPA (%)

25000

50000

75000

100000

125000

150000

175000

200000

Iterations

60

65

70

75

80

85

90

Test OPA (%)

Method

GST

GST+E

GST+EFD

GST-One

0

100

200

300

400

500

600

700

Epoch

30

40

50

60

70

80

90

GST+EFD - Train

GST+EFD - Test

limitation during training and has the potential to bring significant impact across various domains, with notable efficiency and accuracy improvements. GST can be particularly beneficial to industries that need to manage and interpret large-scale graph data. In telecommunications, it can help optimize network infrastructure, while in cybersecurity, it could improve anomaly detection in network traffic. Furthermore, enabling the capability in large network analysis can contribute to more resilient infrastructure, enhancing the quality of life in many communities. As with any advancement in AI, there's a risk that the benefits of this technology will be unevenly distributed, potentially increasing economic disparity. Companies with access to large amounts of data and the computational resources to analyze it might reap disproportionate benefits, which could further exacerbate the digital divide.

E Limitations

While Graph Segment Training (GST) represents a significant advancement in graph property prediction, there are several potential limitations to this approach, as highlighted below:

Segmentation Limitations : The efficacy of the Graph Segment Training (GST) approach is influenced by the proficiency of the graph partitioning procedure. While our empirical findings indicate that various partitioning algorithms that maintain locality tend to produce satisfactory outcomes, partitions created through random edge-cut methods haven't demonstrated the same level of success.

Historical Embedding Table : The use of a historical embedding table to obtain embeddings for non-sampled segments introduces additional implementation complexity and potential for errors. If not managed properly, it could lead to inefficient memory usage or even slower the training process.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work uses graph segment training to reduce memory requirements in order to address the scalability concerns of training with large graphs. Additionally, to reduce computation time, historic embeddings for graph segments are stored. These historic embeddings are updated once for all embeddings at the end of training and periodically dropped during training according to the Stale Embedding Dropout to reduce the bias of stale embeddings to the prediction head and loss function.\n\nSTRENGTHS:\nGraph segment training is an intuitive and seemingly successful method for achieving high graph property prediction on large graphs while reducing maximum memory requirements and total runtime.\n\nThe study of stale embeddings in this work is particularly important as embedding table lookup have been utilized in prior studies.\n\n[1] Zhang, Shichang, et al. \"Motif-driven contrastive learning of graph representations.\" arXiv preprint arXiv:2012.12533(2020).\n\n[2] Tan, Qiaoyu, Ninghao Liu, and Xia Hu. \"Deep representation learning for social network analysis.\" Frontiers in big Data 2 (2019): 2.\n\nWEAKNESSES:\nThe experiments only evaluate GST and its several variants. Some evaluation on other large graph training techniques should be included. Although some of these are used for node classification, they can easily be adapted and evaluated for graph property prediction.\n\n[3] Zou, Difan, et al. \"Layer-dependent importance sampling for training deep and large graph convolutional networks.\" Advances in neural information processing systems 32 (2019).\n[4] Wei-Lin Chiang, Xuanqing Liu, Si Si, Yang Li, Samy Bengio, and Cho-Jui Hsieh. Cluster-gcn: An efficient algorithm for training deep and large graph convolutional networks. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 257–266, 2019.\n[5] Huang, Kezhao, et al. \"ReFresh: Reducing Memory Access from Exploiting Stable Historical Embeddings for Graph Neural Network Training.\" arXiv preprint arXiv:2301.07482 (2023).\n\nIn particular, [5] also utilizes historical embeddings for training on large graphs. A comparison with this method specifically would be useful to include.\n\nQUESTIONS:\nWhen partitioning the graph, what is done with the connecting edge(s) between segments? Is this feature information lost? Would we see performance improvements if this structure information was reincorporated into the model. For example, if each graph was processed twice, using different segmentation sets each time, then the potential loss of structural information can be recovered.\n\nRather than performing stale embedding dropout, would it be reasonable to simply update the segment embedding table at fixed intervals? Or, update stale embeddings rather than dropout according to the stale embedding keep probability.\n\nIn Eq. 2, why is the effect of stale historical embeddings being characterized by the difference between the L(F’(h_s^{(i)}⊕\\hat(h)_j^{(i)})) and L(F’(⊕h_j^{(i)}))? It doesn’t seem that the sampled graph segments are being accounted for in the calculation of the second loss. To mitigate differences not caused by stale historic embeddings, would it not be better to concatenate the sampled graph segments in both loss calculations? Perhaps this is a confusion of notation, in which case, j should not be overloaded to represent either any segment in J^{(i)} or any segment in J^{(i)} excluding those in S^{(i)}.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper aims to predict properties of very large graphs, by segmenting the large graph into multiple subgraphs with the existing graph partitioning algorithm and then learning over segmented subgraphs where gradients are calculated on some of them for memory-efficient training. Also, to further efficiently train with segmented subgraphs, the authors use the embedding table that stores subgraph representations and provides embeddings for subgraphs that are not back-propagated, while calculating embeddings for remaining subgraphs. Moreover, in order to prevent the staleness issue where subgraph representations in the embedding table are outdated, the authors not only finetune only the property prediction head as the post step of training but also randomly drop some subgraph representations with regard to their stalenesses. The authors validate the proposed method, namely Graph Segment Training (GST), with its variants (GST+EFD) on multiple large-scale graph datasets, showing the efficacy of the proposed methods.\n\nSTRENGTHS:\n* The idea of training with partitioned subgraphs of the large graph for its property prediction problem is interesting, novel, and highly valuable to the graph community.\n* Each proposed ingredient (e.g., historical embedding table, prediction head fine-tuning, and stale embedding dropout), which composes the final GST+EFD architecture, has its own unique benefit in learning with segmented subgraphs of the large graph; having solid contributions.\n* The theoretical results show the benefit of stale embedding dropout, which randomly drops some subgraphs with respect to their stalenesses and results in reducing bias from stale embeddings.\n* The proposed GST and GST+EFD outperform the full graph training mechanism while being much more efficient.\n* This paper is extremely well-written. All contents are clear and easy to follow.\n\nWEAKNESSES:\n* I don't see any.\n\nQUESTIONS:\n* How to fetch the embeddings of subgraphs from the embedding table, if the embedding table does not have the representations for them. For example, at the beginning of training, the historical embedding table is empty and there may be no representations of partitioned subgraphs. \n* It is unclear why GST+EFD can outperform GST, given that GST calculates all subgraph embeddings every time, while GST+EFD sometimes uses the stable subgraph embeddings from the embedding table.\n\nLIMITATIONS:\nThe authors do not discuss the limitations and potential negative societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-03"}

comments: SUMMARY:
This paper studies an important problem on large graphs.
The authors propose a new Graph Segment Training (GST) method for large-scale prediction of properties. The proposed method utilizes a divide-and-conquer approach to allow learning large graph property prediction with a constant memory footprint. GST divides a large graph into segments and then backpropagates through a few segments sampled per training iteration. 
Extensive experiments demonstrate the effectiveness of proposed method.

STRENGTHS:
- This work addresses an important problem of property prediction on large graphs, which has applicability in many real-world settings.
- The proposed GST framework uses a divide-and-conquer approach to enable large-scale property prediction.
- The experiments are well-designed and demonstrate the effectiveness of the proposed framework.

WEAKNESSES:
- The paper could benefit from a more detailed discussion of the limitations and potential future directions of the proposed framework.

QUESTIONS:
- How efficiency does the proposed framework compare to existing methods for large graph property prediction?
- Could the proposed method be used to other types of large graphs, such as social networks or biological networks?
- How sensitive is the proposed framework to hyperparameters?

LIMITATIONS:
The paper could benefit from a more comprehensive discussion of the limitations and broader societal impact of the proposed framework.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper deals with large graph learning tasks via graph segmentations. More specifically, in each training step, the author sample nodes from graph segmentations and only update parameters related to the selected nodes. To optimize memory consumption, the author further introduced a historical embedding table. To bridge the training and prediction gap, the author also designed a prediction head fine-tuning scheme. Besides, the author provided some theoretical analysis of the proposed method. Empirical results show the proposed method remains good memory efficiency and test accuracy.\n\nSTRENGTHS:\n1. The paper targets an important problem with high application value, which makes property predictions on Large graphs. The proposed segment training idea looks reasonable and works well practically. \n\n2. Based on the graph segmentation idea, the author has made comprehensive and detailed consideration of the memory usage, the training-testing gap, and the theoretical analysis of the historical approximation bias, which makes the proposed method more technically sound.\n\n3. The paper is well-organized and easy to follow.\n\nWEAKNESSES:\n1. The proposed method didn't discuss how to deal with the inter-segmentation training, for example, how to learn a model with link prediction between two nodes from different graph segmentations.\n\n2. About Baselines: The experimental part does not compare with other large-graph learning baselines such as GraphSage[1]. More baselines should be compared to evaluate the effectiveness of proposed GST.\n\nReference:\n[1] GraphSage: Representation Learning on Large Graphs, NIPS 2017.\n\nQUESTIONS:\n1. How to adapt the proposed method with graph properties learning that are between segmentations?\n\nLIMITATIONS:\nThe proposed method can process large-scale graphs, which might be used for user information mining on large social networks. There are potential risks to user privacy when applying such methods.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}