Learning Large Graph Property Prediction via Graph Segment Training
Authors: Kaidi Cao
Phitchaya Mangpo Phothilimthana
Sami Abu-El-Haija
Dustin Zelle
Yanqi Zhou
Charith Mendis
Jure Leskovec
Bryan Perozzi
Venue: NeurIPS 2023 poster (2023)
Keywords: Graph Neural Networks
Graph Property Prediction
TL;DR: We propose Graph Segment Training (GST), which is able to train on large graphs with constant (GPU) memory footprint.

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
Kaidi Cao'; Phitchaya Mangpo Phothilimthana', Sami Abu-El-Haija', Dustin Zelle?, Yanqi Zhou', Charith Mendis Jure Leskovec', Bryan Perozzi Stannorr University, *Google, 'UIC
Abstract
Learning to predict properties of a large graph is challenging because each pre- diction requires the knowledge of an entire graph, while the amount of memory available during training is bounded. Here propose Graph Segment Training (GST), general framework that utilizes a divide-and -nd-uuqueer approach to allow learning large graph property prediction with constant memory footprint. GST first divides a large graph into segments and then backpropagates through only few segments sampled per training iteration. We refine the GST paradigm by introducing a historical embedding table to efficiently obtain embeddings for seg- ments not sampled for backpropagation. To mitigate the staleness of historical embeddings, we design two novel techniques. First, we finetune the prediction head 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 typical full graph training regime.
1 Introduction
Graph property prediction is task of predicting a certain property or characteristic of an entire graph (3) Important applications include, predicting properties of molecules [32 12 predicting
properties of programs/code [1] [2 6, 41] and predicting properties of organisms based on their
protein-protein interaction networks 27
These popular graph property prediction tasks deal with relatively small graphs, so the scalability issue arises only from 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 even billions of nodes and edges [00. Training typical Graph Neural Networks (GNNs) to classify such large graphs can be computationally infeasible, the memory needed scales 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 limited amount of memory available.
Previous efforts 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 [1] 13 38 43 21. However, there a lack of research on how train scalable models for property prediction of large graphs. Training model sampled subgraph alone is insufficient for these
'The work was partially completed during Kaidi Cao's internship at Google. The work was partially completed when Charith Mendis was visiting researcher at Google.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

(a) Full Graph Training
(b) Our Work: Graph Segment Training
g
N
N
F
Feature Extractor
Produce Embeddings
N
Graph Embedding
N
N
F'
Prediction Head
F' N Y
N
Data Flow
y
Gradient Flow
Figure I: (a) Full Graph Training: Classically, models are trained using the entire graph, meaning all nodes and edges of the graph used to compute gradients. For large graphs, this might be computationally infeasible. (b) Graph Segment Training: Our solution is to partition each large graph into smaller segments and select a random subset of segments to update the model; embeddings for the remaining segments are produced without saving their intermediate activations. The embeddings of all segments are combined generate an embedding for the original large graph, which is then used for prediction. The important benefit is that GPU memory requirement only depends on the segment size (but not the full graph size).
types of tasks, 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 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 graph property.
In this paper, we address the problem of property prediction of large graphs. We propose Graph Segment Training (GST) which is able to train large graphs with constant (GPU) memory footprint.
Our approach partitions each large graph into smaller segments with a controlled size in the pre- processing 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 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 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 historical embedding table 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 reduce this bias. Our final proposed method, called GST+EFD, both memory-efficient and fast.
We evaluate our method on the following datasets: MalNet-Tiny, MalNet-Large and TpuGraphs. typical full graph training pipeline (Full Graph Training) can only train on MalNet-Tiny. and
Source code available at htttps: htattpp:://github.com/// http:/ttthwbbccooaaaii/c///////e

---

unavoidably reaches OOM on Malget-Large and TpuGraphs dataset. On the contrary, we empirically show that the proposed GST framework successfully trains on arbitrarily large graphs using single NVIDIA-VE PUIDIA-VI00 GPU with only 16GB of memory for MalNet-L 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 simultancously being 3x faster in terms of training time.
Preliminaries
Notation. For a function f(a) R2 + RR, we use fa Rdd to denote its k-th derivative of f with respect to evaluated value ag. Let f y denote the composition of function with g.
We use O to denote entry-wise product, and let a 2 - a O a. We use + to represent aggregation. ØJs a means aggregating the set {az}ısJ. where j indexes segments, and J is the number of segments in a graph. We usually drop j subscript from 4 for brevity. a, ① b, aggregates both sets {a:} and {b,} together. + can be mean or sum operators when applying vectors. We define the input graph as = (V,E), where V {u,,,.., Um is the node set and C V
the edge set.
Let dataset D ((g(), )))) contain n examples: each label ( is associated with g())
Graph Neural Network. We consider a backbone graph neural network F that takes a graph g and generates graph embedding (1)) E RRd,, followed by a final prediction head F' that takes graph embedding h() and outputs final predictions: g(t) (F' (g()). We optimize GNN with the loss function as C((F' F) ())
Historical Embedding Table. We define an embedding table T K. + Rda, where key-space K. C Z x Z is a tuple: (graph index i S n, segment index S J). We use h T(i,j) to denote embedding for graph segment of², 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 Durain = , J a common SGD update step requires calculating the gradient:
V
Σ
C((F' o F)(C)))
(g JEB
where 9 is trainable weights in F' F, and B is sampled minibatch. Graphs can differ in size (the number of nodes [)()), with some being too large to fit into the device's 's memory. This is because the memory required store all intermediate activations for computing gradients proportional to the number of nodes and edges in the graph.
To address the above issue, propose to partition each original input graph into a collection of graph segments, i.e.,
g
(i)
for {1,2, ر(o
An example of a partition algorithm is METIS 16]. This preprocessing step will result in training { Jo g. ر Number of partitions ( can vary across graphs, but the size
set Durain
of each graph segment can be bounded by a controlled size (D(()) < mOST, v(i.j) so that batch of a fixed number of graph segments can always fit within the device's memory.
When processing graph segment (. we can obtain its segment embedding through the backbone:
h F)) 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 + h! A simple realization of this aggregation is mean pooling, Note that naïvely applying

---

the prediction head F on the aggregated graph embedding 5 = (' h would not provide any reduction in peak memory consumption, as we need to keep track of the activations of all graph segments {g] (jej to perform backpropagation.
Thus, we propose perform backpropagation only few randomly sampled graph segments SO)) S (L......J()) and generate embeddings without requiring gradients for the rest. We hereby denote h, and h2 as segment embeddings that require and do not require gradient, respectively. An
entire graph embedding is then: h(() (h (t)
(Es()) D (0) JES - (8) + h
We name this
general pipeline as GST and summarize Algorithm
Algorithm 1 General Framework of GST
Require: A preprocessed training graph dataset Duain backbone F and prediction head F".
(109)" ,3) A parameterized
1: for t= to To do
2: B SampleMiniBatch((Pme
3: for (g(+) () in B do
(g] sEst t- SampleGraphSegmens(((()
4:
5:
h + ProduceEimbedding(()" for 8()
(1)
(a) + F(g! for E S(i)
6:
7: end for
+- [8 E, L (' (3) (h,")
8: SGD on loss
9: end for
One implementation of ProduceEmbeddiing() in Algorithm iis to use the same feature encoder F forward all the segments in (g/-) Jigs0 without storing any intermediate activation (by stopping gradient).
3.2 GST with Historical Embedding Table
Calculating h by stopping gradient guarantees an upper bound on peak memory consumption. However, since we do not need gradients for segments (g Jogsus, 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 T(i,j). With embedding table T, one can implement ProduceEmbeddinng( 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 C(F'(h.) (8) Dh (1) (1) 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/ that does not require
gradient, GST needs to run an actual forward pass over ( while GST+E only needs to fetch the embedding from hash table. GST+E has a small overhead from writing the updated embedding of
(()) back into the table T, which is relatively quick and can be run in separate thread so that it does impede the main training algorithm until the next iteration.
The drawback of GST+E is that historical embeddings from may be stale;
h 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 techniques 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 Ph.in(h, y) Ph,, (1h,Ah,,))) and test distribution Pres y) = P(D h v). 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:
arg min T (h.
# arg min b,
To mitigate the distribution misalignment, we introduce the Prediction Head Finetuning technique. Concretely, at the end of training, we update each embedding h 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 multi-layer perceptron.
3.4 Stale Embedding Dropout
The finetuning technique primarily addresses the negative impact of stale embeddings 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, ie.. | h || S e, V(i,j). Given this assumption, the neural network (F' a F-). is k-Lipschitz continuous, the gradients will also be bounded and never run too far from its true estimation, i.e.,
[[VF (1) F { S 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 gets updated infrequently in the embedding table T. Suppose we iterate through every graph g( in the training set Diriin for each epoch, every time we train
graph ( R + g
only few graph G. will be updated in the table T. This implies that all the other segment embeddings of g() will be at least r--iteraiion stale, with n being the number of graphs in the training set, and the most outdated segment embedding could be approximately n.) (000/5(())eeee))) -iterattioo stale, where 5 is [S()).
segments
This staleness introduces 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 we propose the second technique, Stale Embedding Dropout (SED). Unlike a standard Dropout, which uniformly drops elements and weighs up the rest, 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 7) for each segment is defined as:
loss function estimation,
p + (1 p) gou S
for g!
for g (5) with prob. (1 - p)
0
for g( with prob. p
Please refer to the theoretical analysis in the next section. By combining the two proposed techniques, we denote our final algorithm GST+EFD, which we summarized in Algorithm 2 in Appendix B
Theoretical Analysis
We characterize the effect of stale historical embeddings by studying the difference between L(F' (h!0 ))) and ((F'(()) (2)h()) )))) Assume glo has 766 segments and we perform back-
propagation on (()) segments. We let g01) A h Oh + n,)0) be the perturbation on the graph embedding.
5

---

We apply Taylor expansion around 6()) = 0 on the final loss to analyze the effect of this perturbation.
E, ( (h
(1) - C(F
(2)
-E,C(F (8) 8( L(F" D))
1 Fn + +22 220
Σ Bgg D Dagor (C Coo
(Da goll (Co F')inn"Doy
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 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 8(0), For the first term, we have
J(M) -5(()
E
B =
C
B(h (8)
J(9
J S(()
E
SED
B
Cx
J(
B(h
(8)
h2 p
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 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 IK(h h h())) is so small that the negative effect may be neglected. In our setting, the historical embeddings can be very stale, so having a factor helps reduce bias. It worthwhile to check the limiting cases. If p = (keeping all the stale embeddings without dropping any), both E geesse [ and E E5005 [R] degrade to the result
of ET. If p- 0 (droping all the stale embeddings), then the algorithm degrades to training on only ((() segments without aggregating other segments (which we denote as GST-One, when 5 = 1). B indicates that there no stale bias in this case. However, the term E giys [R] could
become too large so that it impedes training.
Experiments
5.1 Experimental Setup
Datasets. MalNet [10] is a large-scale graph representation learning dataset, with the goal 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 types, with each graph containing at most 5,000 nodes. To evaluate our approach on the regime where the graph size is large, construct alternative split from the original MalNet dataset, which we named MallNet-Large Malget-Large also contains 5,000 graphs across balanced types. MalNet-Large 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 execution time of an XLA's HLO graph with a specific compiler configuration on a Tensor Processing Unit (TPU), XLA (28) is production backend compiler for various machine learning frameworks, including TensorFlow, PyTorch, and JAX. In this particular datasct, 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
model to select the best configuration for each graph [ TpuGraphs
Wee have made public dataset
24 that closely parallels our internal dataset.

---

Table 1: Test accuracy MalNet- Tiny and Malget-Large. We report the standard deviation five runs. GST+EFD achieves better accuracy than Full Graph Training, and GST, while being much more memory efficient and computationally faster.
Dataset
MaiNet-Tiny SAGE
MalNet-LLage SAGE
Backbone
GCN
GraphGPS
GCN
GraphGPS
Full Gruph Training
87.84:11.37
88.08 1.88
90.82-0.59
0OM
0OM
OOM
GST
88.26=0.80
88 08.422+1.03
91.03=0.81
88.35++1144
88.62++0.82
91.39::0.85
GST-One
31.62+3.88
72.64+t 4.73
77.63=3.15
60.41+6.29
57.13+7.36
66.82244.71
GST+E
86.53:111.18
86.82-10.93
89.75=0.89
48.42:16.61
43.28:17.01
62.47:3.119
GST+EF
87.67-10.78
87.83-10.81
90,52-10.11
84.83.10.96
85.26110.87
91.33-10.66
GST+ED
88. 08.18400.48
88.50=0.77
90.96=0.68
82.17444.44
71883+66.31
19,46=1.36
GST+EFD
08.78=0.455
89.24:0.53
92.46: 0.66
89.67:-0.71
89.78+0.68
92.52++0.58
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, graph together with a configuration defines one g because the configuration is featurized parts of input node features to the GNN.
Please refer to Table 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(), w
+ g but randomly select only segment G/ for each graph train every iteration. (3) GST:
following the general GST framework deseribed in Algorithm [ we replace ProduceFmedding(( by
using the same feature encoder to forward all the segments in {g) ( 1jeso) without storing any intermediate activation. We set ( = 1 in our experiments. (4) E: we introduce an embedding table h = (ii, to store the historical embedding of each graph segment, and we fetch the embedding from T if we do need 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. during training.
When these techniques are combined, we concatenate the acronyms with a "4" to GST as an abbreviation. We conduct all the experiments on MalNet with single NVIDIA-VIO0 GPU with 16GB of memory, and four NVIDIA-VION 100 GPUs (for data parallelism) with 16GB of memory for TpuGraphs. Please refer to Appendix for additional implementation details.
5.2 Empirical Results on MalNet
To demonstrate the general applicability of our proposed GST framework, we consider three back- bones, namely, GCN [9)., SAGE (T]) 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 We include MalNet-Ting 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 MalNet-Large, Full Graph Training strategy can no longer fit the large graphs on GPU, so we report OOM in the table. GST's estimation on graph segment embeddings h 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.
Naively 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.
7

---

80
100
Test
Nocuroy
80
Accurazy
50
38
30
20
100
1000
30000
50.00
7030
9000
Epooh
Max Segmeett Size
Figure 2: Accuracy curve on MalNet-Large of GST+EFD with SAGE backbone. We start Prediction Head Finetuning at epoch 600.
Figure Ablation study the keep ratio in SED. We report test accuracy GST+EFD with SAGE backbone on MalNet- Large for runs.
Figure 4: Ablation study maximum segment size. We re- port test accuracy of GST+EFD with SAGE backbone on MalNet-Large for 5 runs.
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 combating the staleness issue. The combination 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
Table 2: Training and test ordered pair accuracy Table 3: Runtime analysis (average training time per iteration ms) on MalNet-L dataset.
(OPA) on TpuGraphs Dataset.
Train OPA
Test OPA
GCN
SAGE
GraphGPS
Full Graph Training
OOM
0OM
Full Graph Training
0OM
OOM
0OM
GST
87.43
86.84
GST
720.8
706.3
1285.7
GST-One
67.19
81.03
GST-One
242.2
239.6
GST+F 74.02
82.62
GST+E 258.4
253.6
451.9
GST+EFD 71.40
89.20
GST+EFD
247.9
244.7
448.2
As mentioned Section 5.1 we care more about the ranking of configurations for each graph than
the absolute runtimes in this dataset. Thus, report the ordered pair accuracy (OPA) averaged over all computational graphs in Table OPA metric is defined as:
Σ, E 1p Σ,E,
0, 1 U/1
OPA 7
Mn
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 part F, and F' is simply summation function. Since there are no learnable weights in F', we omit Prediction Head Finetuning in this experiment, and GST+EFD Table 2 excludes the finetuning stage. We observe clear tradeoff between fitting training examples and generalization. First, GST has 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 with the embedding table yiclds lower training OPA, and consequently lower test OPA as well. SED in GST+EFD functions regularization technique. Although 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 finetuning. We visualize training/test accuracy curve of GST+EFD over time in Figure 2 The staleness introduced by historical embeddings drastically hurts generalization, shown for the first 600 epochs. We start finetuning at epoch 600, and the gap between training and test accuracy decreases by 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 and visualize the results in Figure 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+FFD 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 + 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 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 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 's and GST+EFD's runtime very close GST-One's. 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 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 certain characteristic associated with the whole graph. Standard graph neural networks produce node embeddings as outputs [99 I 34]. To create 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 [1,, PipeGCN (29 and BNS- (30 achieve distributed node-level GCN training through partitioning 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 historical embedding for each node because they consider node-level tasks. As we consider graph property prediction tasks, we record historical embedding for each graph segment rather than each node.
9

---

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.
Acknowledgements
KC and JL acknowledge the support of DARPA under Nos. HR001 12190039 (TAMI). N66001 1924033 (MCS); ARO under Nos. W911NR-16-1-0342 (MURT), W911NF-16-1-011 (DURIP); NSF under Nos. OAC-1835598 (CINES), OAC-1934578 (HDR), CCF-1918940 (Expedi- tions), NIH under No. 3U54HG010426-04S (HuBMAP). Stanford Data Science Initiative, Wu Tsai Neurosciences Institute, Amazon, Docomo, GSK, Hitachi, Intel, JPMorgan Chase, Juniper Networks, KDDI, NEC, and Toshiba. CM's contributions were partially supported by ACE, one of the seven centers JUMP 2.0, a Semiconductor Research Corporation (SRC) program sponsored by DARPA and NSF under grant CCF-2316233 The content is solely the responsibility of the authors and does not necessarily represent the official views of the funding entities.
References
[1] Miltiadis Allamanis. The adverse effects of code duplication in machine learning models of code. In Proceedings of the 2019 ACM SIGPIAN International Symposium New Ideas, New Paradigms, and Reflections on Programming and Software, pages 143-153, 2019.
[2] Aleksandar Bojchevski, Johannes Klicpera, Bryan Perozzi, Amol Kapoor, Martin Blais, Benedek Rózemberezki, Michal Lukasik, and Stephan Günnemann. Scaling graph neural networks with approximate pagerank. In Proceedings of the 26th ACM SIGKDD International Conference Knowledge Discovery & Data Mining, pages 2464-2473, 2020.
[3] Ines Chami, Sami Abu-El-Haan 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): 14,, 2022.
[4] 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. PMIR, 2018.
[5] Jie Chen, Tengfei Ma, and Cao Xiao. Fastgen: Fast learning with graph convolutional networks via importance sampling. In International Conference Learning Representations, 2018.
[6] 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, NeurlPS' 18, page 3393-3404, Red Hook, NY, USA, 2018. Curran Associates Inc.
[7] Wei-Lin Chiang, Xuanqing Liu, Si Si, Yang Li, Samy Bengio, and Cho-Jui Hsieh. Cluster-gen: 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.
[8] Vijay Prakash Dwivedi, Chaitanya K Joshi, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Benchmarking graph neural networks. arXiv preprint arXiv:2003.00982 2020.
[9] 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.
10

---

[10] Scott Freitas, Yuxiao Dong, Joshua Neil, and Duen Horng Chau. A large-scale database for graph representation learnin.. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 1), 2021.
[11] Will Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large graphs. Advances in neural information processing systems, 30, 2017.
[12) 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 18 18-11333333, 2020.
[13] Wenbing Huang, Tong Zhang, Yu Rong, and Junzhou Huang. Adaptive sampling towards fast graph representation learning. Advances in neural information processing systems, 31, 2018.
[14] 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 Leurning and Systems, 2:187-198, 2020.
[15] 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. 145/3360307
[16] George Karypis and Vipin Kumar. Metis: A software package for partitioning unstructured graphs, partitioning meshes, and computing fill-reducing orderings of sparse matrices. 1997.
[17] Samuel J. Kaufman, Phitchaya Mangpo Phothilimthana, Yanqi Zhou, Charith Mendis, Sudip Roy, Amit Sabne, and Mike Burrows. A Learndd Performance Model for Tensor Processing Units. In Proceedings of Machine Learning for Systems, 2021.
[18] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.69888 2014.
[19] Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. 2016.
[20] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 17X1v:1711.05101, 2017.
[21] Elan Sopher Markowitz, Keshav Balasubramanian, Mehrnoosh Mirtaheri, Sami Abu-El-Hai 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.
[22] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary De Vito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
[23] Trang Pham, Truyen Tran, Hoa Dam, and Svetha Venkatesh. Graph classification via deep learning with virtual nodes. arXiv preprint arXiv:/708.04337. 2017.
[24] 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.
[25] 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.
[26] 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 2203:22003.48110 2022.

---

[27] Damian Szklarezyk, Annika L Gable. David Lyon, Alexander Junge, Stefan Wyder, Jaime Huerta- Cepas. Milan Simonovic, Nadezhda T Doncheva, John Morris, Peer Bork, et al. String 111: protein-protee association networks with increased coverage, supporting functional discovery in genome-wiet wide experimental datasets. Nucleic acids research, 47(D1):D607-D611,, 2019.
[28]
TensorFlow. XLA: Optimizing Compiler for TensorFlow. htpps//www..eessrfoooooookk/
[Online; accessed 19- 9--Spteemer-20119.
[29] Cheng Wan, Youjie Li, Cameron R Wolfe, Anastasios Kyrillidis, Nam Sung Kim, and Yingyan Lin. Pipegen: Efficient full-graph training of graph convolutional networks with pipelined feature communication. In International Conference on Learning Representations, 2021.
[30] Cheng Wan, Youjie Li, Ang Li, Nam Sung Kim, and Yingyan Lin. Bns-gen: Efficient full-graph training of graph convolutional networks with partition-parallllisse and random boundary node sampling. Pruceedings of Machine Learning and Systems, 4:673-693, 2022.
[31] 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.
[32] 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, 922:513-5300, 2018.
[33] 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.
[34] Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? In International Conference Learning Representations, 2018.
[35] 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.
[36] Jiaxuan You, Zhitao Ying, and Jure Leskovec, Design space for graph neural networks, Advances in Neural Information Processing Systems, 33: 17029-1702,, 2020.
[37] 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 Machine Learning, pages 25684-25701. PMLR, 2022.
[38] 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.
[39] 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.
[40] 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.
[41] Lianmin Zheng, Ruochen Liu, Junru Shao, Tianqi Chen, Joseph E. Gonzalez, lon 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 Truck (Round 1), 2021.
[42] Marinka Zitnik, Rok Sosič, Mareus 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.
[43] Difan Zou, Ziniu Hu, Yewen Wang, Song Jiang, Yizhou Sun, and Quanquan Gu. Layer- dependent importance sampling for training deep and large graph convolutional networks. Advances in neural information processing systems, 32, 2019.
12

---

A Proofs and Derivations
Theorem A.I. Under proper condition that W h 0 and W h2 0, where W 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 5(6) A h h
( be the perturbation on the graph embedding. We use ET to
denote using the embedding table without applying SED. For GST+E, we have
g
s(0 with prob.
J(g with prob. _(0-(())
h
h (3)
The randomness above comes from the fact that each segment g. is selected for backpropagation with probability s(s)
Ju
For SED, there are two folds of randomness when training graph g();; randomly selecting S
segments train and randomly select stale embedding h to drop. Thus we can rewrite 5/:0
(1- (1-p) sty (
h
with prob. 5
with prob. (1-p)((() -s())
JG)
h 6
with prob, p( (866
J(
We apply Taylor expansion around 0/ = 0 on the final loss to analyze the effect of this perturbation.
In Section A.2 of Wei et al. B1, when W (9) = 0 and W h = 0, the perturbation of 5(9) to the loss function might not be too large, it supports the use of Taylor Expansion.
(I 1 θ D C(F
()
=C(F
(h + 1)) C(F' O (0)
eo (C F') 1 1 (i)
jo (Co (')
Note that we randomly select segments with index s during training, we can then derive an approxi- mation of the expected difference during training as
E. 1 Oh, - C(F (0)0))
(3)
1 100 Σ Bgg D (Lo F 11 +
( F')
R
We can then compare the effect of SED by substituting the two versions of 51 into Eq.
So for the first term, have
E S B = (6 () E5!(0)
(1) (1) J(H) 5(0
= (D (')
E(h,)) (10)))
B = (D (Lo F')(1 (1) 1,55(())
(0) ( 5(0
= (D (Lo (
E(h 11 x))
J(
13

---

whereas for the second term, we have
Es
R3)000(00) 5(17
E genjer [R] =(D /
(C F')
2 )
= =Dx hio ho (C [1) 1 - S())
-h
2
R = (De h (C ( (1) 1 B0 8(1
SED
= =DD R0o (C o F²) 11 (0) (J(1) ((() 10 (12 - h )2
2
(J (1 p) pJ() + ps ())
2
It is casy to check that the statement satisfies given the value calculated.
Implementation Details
B.1 Missing Algorithm
Algorithm 2 Pipeline GST+EFD
Require: A preprocessed training graph dataset Dusin ( . ") A parameterized backbone F and prediction head F*. A historical segment embedding table T.
1: for t= to To do
2: B <- SampleMiniBacch(PPie
3: for (g()) ) in B do
(g(") sESt +- SampleGrrphssennens(((()))
4:
h(00 T.LookUp for j $ S(()
5:
h! + =F for S
6:
T.InserOOUpddate((. 1)))
7:
8: end for
9: SGD on 18] 4E,C(F(( ( h Th Om, (0) - (() (SED with 17 defined in Equation
10: end for
Prediction Head Finetuning
12: T. T_InertOrUpdatee(i... (((()) for every 0."0
13: for To to T; do
for g() in SampleMiniBatchh D uain) do
(1) + T LookkUp ((t,.)) for j < j(i
15:
16: end for
17: SGD with loss + B E, C (F h())
end for
Table 4: Overview of the graph datasets used in this study.
Avg. # nodes
Min. # nodes
Max. nodes
Avg. # edges
Min. edges
Max. edges
MalNet-Tiny MalNet-Large
1,410
4,994
2,860 225,474
20,096
47,838
3,374 299
541.571
20,597 380
3,278,318
TpuGraphs
38,444
615,019
62.475
1,058,278
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 (2).
14

---

Implementation details for MalNet-Large. We consider three model variations for the MalNet- Large dataset. Please refer their hyperparameters in Table We use Adam optimizer [88] with
the base learning rate of 0.01 for GCN and SAGE. For GraphGPS, we use Adam W optimizer [20] with the cosine scheduler and the base learning rate of 0.0005. We use 1.2 regularization with weight decay of le-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
keep probability p 0.5 if not otherwise specified. We train with CrossEntropy loss.
Table 5: Detailed GNN/Graph Transformer designs used in MalNet- and MalNet-Large.
model
GCN
SAGE
GraphGPS
message passing layer type pre-process layer num.
GCNConv
SAGEConv
GatcdGCN-Porfommer
0
message passing layer num.
2
2
5
post-process layer num.
3
hidden dimension activation aggregation
300
300
64
PReLU
PReLU
ReLU
mcan
Implementation details for MalNet-Tiny, We use the same model architectures/traniing 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:
(() = ΣΣ Ti >> yi] max 1 (0: -3)))
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 MalNet- Tiny dataset in Figure We show that the convergence speed of various methods studied are quite similar. In addition, the convergence speed in terms of iterations 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, 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 partition algorithms. We incorporated ablation study focusing various graph partition algorithms. The test accuracy of GST+EFD using the SAGE backbone for 5 iterations 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 decrease in graph property prediction accuracy, leading us to further investigate Vertex-Cur partition algorithms. In theory. Vertex-Cur 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 significant impact on the final prediction accuracy.
D Broader Impact
The paper presents Graph Segment Training (GST), framework for predicting properties of large approach. This method addresses the challenge of memory
graphs using
15

---

Trur D0%
Tee: Dok (5)
Figure 5: Accuracy curve on TpuGraphs dataset.
90
80
Accuracy
50
GST+EFD Train
GST+EFD Test
40
Full Graph Training Train
Graph Training
30
100
200
300
400
500
600
700
Epoch
Figure 6: Accuracy curve on MalNet-Tiny dataset.
limitation during training and has the potential 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 more resilient infrastructure, enhancing the quality of life in many communities. As with any advancement in AI, there's risk that the benefits of this technology will unevenly distributed, potentially increasing economic disparity. Companies with access to large amounts of data and the computational resources analyze 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 several potential limitations to this approach, as highlighted below:
Segmentation Limitations: The efficacy of the Graph Segment Training (GST) approach is influ- enced by the proficiency of the graph partitioning procedure. While our empirical findings indicate that various partitioning algorithms that maintain locality tend produce satisfactory outcomes, partitions created through random edge-cut methods haven't t demonstrated the same level of success.
16

---

Table 6: Test accuracy when combined with different partition algorithms.
MalNet-Tiny
MalNet-Large
Edge-Cut
Random
85.43+00988
74.02+2.23
Edge-Cul
Louvain
88.95+10.67
89.16-00.5
Edge-Cul
METIS
89.24.10.53
89.7810.68
Vertex-Cut
Random
88.12+11.17
87.69::1.51
Vertex-Cut
DBH [33
88.79=00.74
89.28=0.93
Vertex-Cut
30
89.16+0.70
89.49-10.87
The use of a historical embedding table obtain embeddings for non-sampled segments introduces additional implementation complexity and potential for errors. If not managed properly, could lead to inefficient memory usage or even slower the training process.
Historical Embedding Table:
17


Reviews:
[{"review_id": "6na8q2WvEF", "reviewer": "Reviewer_rf6m", "summary": "This 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.", "questions": "1. How to adapt the proposed method with graph properties learning that are between segmentations?", "limitations": "The 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.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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.", "weaknesses": "1. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedback and insightful comments. We appreciate the reviewer for confirming that our paper is technically sound and easy to follow. We respectfully ask the reviewer to consider increasing the score if our clarification has addressed the concerns raised by the reviewer.\n\n**Q1. How to learn a model with link prediction between two nodes from different graph segmentations.**\n\nWe thank the reviewer for the question. We would like to clarify that we use “Graph property prediction” to denote that the entire graph should deserve one prediction (contrast with: node-level, and edge-level, respectively, where each node and edge should receive a prediction). “Graph property prediction” includes “graph classification” (both single label or multi-label setting), “regression” (e.g., in chemical molecules, a graph-level GNN could output continuous values such as “boiling temperature”), ranking, and so on. Explicitly trying to predict properties of a given link connecting two segmentations is essentially a link prediction problem, which we believe is not in the scope of this paper.\n\n**Q2. The experimental part does not compare with other large-graph learning baselines such as GraphSage.**\n\nWe thank the reviewer for the comment. Most of the prior research addressing either node-level or link-level prediction issues fails to directly translate to large-scale graph property prediction tasks. Nevertheless, we did evaluate one baseline - GST-One, in our tests. It operates by training on a single sampled segment per cycle, which, if we adjust to this context, bears notable resemblance to either GraphSAGE or Cluster-GCN. Our findings indicate that this approach resulted in subpar performance.\n"}, {"review_id": "gOnKVyAl2v", "reviewer": "Reviewer_vMFq", "summary": "This paper studies an important problem on large graphs.\nThe 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. \nExtensive experiments demonstrate the effectiveness of proposed method.", "questions": "- How efficiency does the proposed framework compare to existing methods for large graph property prediction?\n- Could the proposed method be used to other types of large graphs, such as social networks or biological networks?\n- 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.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This work addresses an important problem of property prediction on large graphs, which has applicability in many real-world settings.\n- The proposed GST framework uses a divide-and-conquer approach to enable large-scale property prediction.\n- 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.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedback and insightful comments. We appreciate the confirmation that the problem we studied is important and experimental design is well-conducted. We hope we are able to address the review’s concerns, and respectfully ask to consider increasing the score.\n\n**Q1. How efficiency does the proposed framework compare to existing methods for large graph property prediction?**\n\nWe thank the reviewer for the question. Our GST framework is tailored to train on large graphs, a task that previous methods struggled with due to memory limitations. This highlights our method's memory efficiency in contrast to earlier approaches. As depicted in Table 3, our GST framework also boasts improved runtime efficiency (3x faster). Importantly, our final algorithm, GST+EFD, introduces only a slight increase in time compared to GST-One. The latter is trained on a singular sampled segment per cycle and can be viewed as a theoretical minimum time benchmark.\n\n**Q2. Could the proposed method be used for other types of large graphs, such as social networks or biological networks?**\n\nWe appreciate the reviewer's question. Indeed, our suggested approach is not limited to any particular type of graph. Most existing social network datasets tend to emphasize node-level or link-level tasks. On the other hand, the publicly available biological networks are generally of a smaller scale. This is the reason we didn’t test on these two types of graphs.\n\n**Q3. How sensitive is the proposed framework to hyperparameters?**\n\nWe thank the reviewer for the question. We conducted an ablation study adjusting some design parameters as shown in Figures 2-4. Our findings indicate that the GST+EFD configuration is notably robust to variations in maximum segment size and partition algorithms. It's worth noting that our optimization parameters remain consistent across all methods and are adopted directly from earlier implementations for Full Graph Training.\n"}, {"review_id": "Bdcd8C9tIN", "reviewer": "Reviewer_Uc8u", "summary": "This 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.", "questions": "* 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. ", "limitations": "The authors do not discuss the limitations and potential negative societal impact of their work.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "* 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.", "weaknesses": "* I don't see any.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedback and insightful comments. We appreciate the reviewer for confirming that our paper is novel, extremely well-written and highly valuable in the graph community. \n\n**Q1. 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\nWe thank the reviewer for the question. For historical embeddings, we employ a 0 initialization. Therefore, at the onset of training, if a segment hasn't been updated previously, its representation defaults to a vector of zeroes.\n\n**Q2. 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\nWe appreciate the reviewer's feedback. Our presumption is that the introduced Staled Embedding Dropout offers supplementary regularization, fostering improved feature representation. Concurrently, Prediction Head Finetuning assists in learning alignment and thus mitigates staleness.\n"}, {"review_id": "6qxAFmktiu", "reviewer": "Reviewer_U5zV", "summary": "This 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.", "questions": "When 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)}.", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Graph 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", "weaknesses": "The 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedback and insightful comments. We appreciate the confirmation that the proposed method is important and our proposed method is successful. We hope we are able to address the review’s concerns, and respectfully ask to consider increasing the score.\n\n**Q1. 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\nWe thank the reviewer for the suggestion. We will integrate the mentioned references into the related work section and provide a discussion. It's worth mentioning that we did assess a baseline, GST-One, that trains on a singular sampled segment for each cycle and, when contextualized, shares significant similarities with Cluster-GCN [4]. Additionally, we'd like to emphasize that adapting prior methods, such as [5] which uses historical embeddings, to our scenario is impractical. The method in [5] retains a historical embedding for every node, which makes it unfeasible to store all the historical embeddings given our conditions, i.e., for the malnet-large dataset, the embedding table needs to store approximately $2 \\times 10^9$ rows. Taking into account the hidden dimension, this amounts to about 1TB of memory. The memory requirement would further surge if we raise the count of training graphs.\n\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\n**Q2. When 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.**\n\nWe thank the reviewer for the question. There's a possibility of some information loss, but our empirical studies show it doesn't greatly impact performance. Addressing the reviewer's query, we delve deeper into Vertex-Cut partition algorithms in Table 5 of Appendix. Theoretically, Vertex-Cut approaches, which distribute edges among various machines and replicate nodes as required, might experience less information loss compared to Edge-Cut techniques. Our hands-on results indicate that all partitioning algorithms which preserve local structure exhibit comparable performance. This infers that edges linking different segments don't substantially influence the ultimate prediction accuracy.\n\n**Q3. Rather 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\nWe thank the reviewer for the comment. While it is feasible to proceed in that manner, it's important to highlight that refreshing the entire segment embedding table might require more time than completing one training epoch (as the number of segments is usually an order larger than the number of segments we trained on every iteration), resulting in significant overhead. Furthermore, our results indicate that GST+EFD outperforms GST in accuracy. Hence, we believe it's improbable for this alternative to surpass our suggested algorithm in either efficiency or accuracy.\n\n**Q4. In 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\nWe thank the reviewer for the comment. Indeed, the reviewer has a valid point. We reloaded $j$ to save some space in notation. We will try the reviewer’s recommendation in the final version.\n"}]

Author rebuttal:

Meta-review: The reviewers in general give positive ratings to this paper. The author rebuttal answered many questions asked in the reviewers. The remaining weaknesses pointed by the reviewers include: (1) the lack of a more detailed discussion of the limitations and potential future directions of the proposed framework; (2) limited experimental evaluations, with missing large graph training techniques that should be included.

I am happy to recommend this paper as accept. I look forward to seeing the authors addressed the aforementioned remaining issues in the camera-ready submission.
Decision: Accept (poster)