SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Large Graph Property Prediction via Graph Segment Training
Kaidi Cao $^{1}$ ,* Phitchaya Mangpo Phothilimthana $^{2}$ , Sami Abu-El-Haija $^{2}$ , Dustin Zelle $^{2}$ , Yanqi Zhou $^{2}$ , Charith Mendis $^{3\dagger}$ , Jure Leskovec $^{1}$ , Bryan Perozzi $^{2}$
<sup>1</sup>Stanford University, <sup>2</sup>Google, <sup>3</sup>UIUC
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.
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
*The work was partially completed during Kaidi Cao's internship at Google.
†The work was partially completed when Charith Mendis was a visiting researcher at Google.
Figure 1: (a) Full Graph Training: Classically, models are trained using the entire graph, meaning all nodes and edges of the graph are 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 to 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, 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
<sup>3</sup>Source code available at https://github.com/kaidic/GST.
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.
Notation. For a function  $f(\pmb{a}) : \mathbb{R}^{d_0} \to \mathbb{R}^{d_1}$ , we use  $D_{\pmb{a}}^k f[\pmb{a}_0] \in \mathbb{R}^{d_0 \times d_1}$  to denote its  $k$ -th derivative of  $f$  with respect to  $\pmb{a}$  evaluated at value  $\pmb{a}_0$ . Let  $f \circ g$  to denote the composition of function  $f$  with  $g$ . We use  $\odot$  to denote entry-wise product, and let  $\pmb{a} \odot 2 \triangleq \pmb{a} \odot \pmb{a}$ . We use  $\bigoplus$  to represent aggregation.  $\bigoplus_{j \leq J} \pmb{a}_j$  means aggregating the set  $\{\pmb{a}_j\}_{j \leq J}$ , where  $j$  indexes segments, and  $J$  is the number of segments in a graph. We usually drop  $j$  subscript from  $\bigoplus$  for brevity.  $\pmb{a}_i \oplus \pmb{b}_j$  aggregates both sets  $\{\pmb{a}_i\}$  and  $\{\pmb{b}_j\}$  together.  $\bigoplus$  can be mean or sum operators when applying to vectors. We define the input graph as  $\mathcal{G} = \{\mathcal{V}, \mathcal{E}\}$ , where  $\mathcal{V} = \{v_1, \dots, v_m\}$  is the node set and  $\mathcal{E} \subseteq \mathcal{V} \times \mathcal{V}$  is the edge set.
Let dataset  $\mathcal{D} = \{(\mathcal{G}^{(i)},y^{(i)})\}_{i\leq n}$  contain  $n$  examples: each label  $y^{(i)}$  is associated with  $\mathcal{G}^{(i)}$ .
Graph Neural Network. We consider a backbone graph neural network  $F$  that takes a graph  $\mathcal{G}^{(i)}$  and generates graph embedding  $h^{(i)}\in \mathbb{R}^{d_h}$ , followed by a final prediction head  $F^{\prime}$  that takes graph embedding  $h^{(i)}$  and outputs final predictions:  $\widehat{y}^{(i)} = (F^{\prime}\circ F)(\mathcal{G}^{(i)})$ . We optimize GNN with the loss function as  $\mathcal{L}((F'\circ F)(\mathcal{G}^{(i)}),y^{(i)})$ .
Historical Embedding Table. We define an embedding table  $\mathcal{T}:\mathcal{K}\to \mathbb{R}^{d_h}$ , where key-space  $\mathcal{K}\subset \mathbb{Z}\times \mathbb{Z}$  is a tuple: (graph index  $i\leq n$ , segment index  $j\leq J$ ). We use  $\tilde{\pmb{h}}_j^{(i)} = \mathcal{T}(i,j)$  to denote embedding for graph segment  $\mathcal{G}_j^{(i)}$ : 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  $\mathcal{D}_{\mathrm{train}} = \{(\mathcal{G}^{(i)},y^{(i)})\}_{i = 1}^{n}$ , a common SGD update step requires calculating the gradient:
$$
\nabla_ {\theta} \sum_ {(\mathcal {G} ^ {(i)}, y ^ {(i)}) \in \mathcal {B}} \mathcal {L} ((F ^ {\prime} \circ F) (\mathcal {G} ^ {(i)}), y ^ {(i)})
$$
where  $\theta$  is trainable weights in  $F' \circ F$ , and  $\mathcal{B}$  is a sampled minibatch. Graphs can differ in size (the number of nodes  $|\mathcal{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.,
$$
\mathcal {G} ^ {(i)} \approx \bigoplus \mathcal {G} _ {j} ^ {(i)} \quad \text {f o r} j \in \{1, 2, \dots , J ^ {(i)} \}
$$
An example of a partition algorithm is METIS [16]. This preprocessing step will result in a training set  $\mathcal{D}_{\mathrm{train}} = \{(\bigoplus_{j\leq J^{(i)}}\mathcal{G}_j^{(i)},y^{(i)})\}_{i = 1}^n$ . Number of partitions  $J^{(i)}$  can vary across graphs, but the size of each graph segment can be bounded by a controlled size  $(|\mathcal{V}_j^{(i)}| < m_{\mathrm{GST}},\forall (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  $\mathcal{G}_j^{(i)}$ , we can obtain its segment embedding through the backbone:  $h_j^{(i)} = F(\mathcal{G}_j^{(i)})$ . 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)} = \bigoplus h_{j}^{(i)}$ . A simple realization of this aggregation is mean pooling. Note that naively applying
the prediction head  $F'$  on the aggregated graph embedding  $-\widehat{y}^{(i)} = F'(\bigoplus h_j^{(i)})$  — would not provide any reduction in peak memory consumption, as we need to keep track of the activations of all graph segments  $\{\mathcal{G}_j^{(i)}\}_{j \in J^{(i)}}$  to perform backpropagation.
Thus, we propose to perform backpropagation on only a few randomly sampled graph segments  $S^{(i)} \subseteq \{1, \ldots, J^{(i)}\}$  and generate embeddings without requiring gradients for the rest. We hereby denote  $h_s$  and  $\bar{h}_j$  as segment embeddings that require and do not require gradient, respectively. An entire graph embedding is then:  $h^{(i)} \approx \{h_s^{(i)}\}_{s \in S^{(i)}} \oplus \{\bar{h}_j^{(i)}\}_{j \notin S^{(i)}} \triangleq h_s^{(i)} \oplus \bar{h}_j^{(i)}$ . 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  $\mathcal{D}_{\mathrm{train}} = \{(\bigoplus \mathcal{G}_j^{(i)},y^{(i)})\}_{i = 1}^n$  . A parameterized backbone  $F$  and a prediction head  $F^{\prime}$    
1: for  $t = 1$  to  $T_{0}$  do   
2:  $\mathcal{B}\gets$  SampleMiniBatch  $(\mathcal{D}_{\mathrm{train}})$    
3: for  $(\mathcal{G}^{(i)},\widehat{y}^{(i)})$  in  $\mathcal{B}$  do   
4:  $\{\mathcal{G}_s^{(i)}\}_{s\in S^{(i)}}\gets$  SampleGraphSegments  $(\mathcal{G}^{(i)})$    
5:  $\bar{\pmb{h}}_j^{(i)}\gets$  ProduceEmbedding  $(\mathcal{G}_j^{(i)})$  for  $j\notin S^{(i)}$    
6:  $\pmb {h}_s^{(i)}\gets F(\mathcal{G}_s^{(i)})$  for  $s\in S^{(i)}$    
7: end for   
8: SGD on loss  $\leftarrow \frac{1}{|\mathcal{B}|}\sum_i\mathcal{L}\left(F'(h_s^{(i)}\oplus \bar{h}_j^{(i)}),\widehat{y}^{(i)}\right)$    
9: end for
One implementation of ProduceEmbedding(\cdot) in Algorithm 1 is to use the same feature encoder  $F$  to forward all the segments in  $\{\mathcal{G}_j^{(i)}\}_{j\notin S^{(i)}}$  without storing any intermediate activation (by stopping gradient).
3.2 GST with Historical Embedding Table
Calculating  $\tilde{h}_j$  by stopping gradient guarantees an upper bound on peak memory consumption. However, since we do not need gradients for segments  $\{\mathcal{G}_j^{(i)}\}_{j\notin 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  $\tilde{h}_j^{(i)} = \mathcal{T}(i,j)$ . With an embedding table  $\mathcal{T}$ , one can implement ProduceEmbedding(\cdot) 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  $\mathcal{L}(F'(h_s^{(i)}\oplus \tilde{h}_j^{(i)}),y^{(i)})$  during training. We name the embedding version of our algorithm as GST+E.
Note that  $\mathrm{GST + E}$  has runtime advantages over GST. For each segment  $\mathcal{G}_j^{(i)}$  that does not require gradient, GST needs to run an actual forward pass over  $\mathcal{G}_j^{(i)}$ , while  $\mathrm{GST + E}$  only needs to fetch the embedding from a hash table.  $\mathrm{GST + E}$  has a small overhead from writing the updated embedding of  $\mathcal{G}_s^{(i)}$  back into the table  $\mathcal{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  $\mathcal{T}$  may be stale;
$\tilde{h}_j^{(i)}$  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  $\mathcal{P}_{\mathrm{train}}(\pmb{h}, y) = \mathcal{P}(\pmb{h}_s \oplus \tilde{\pmb{h}}_j, y)$  and test distribution  $\mathcal{P}_{\mathrm{test}}(\pmb{h}, y) = \mathcal{P}(\bigoplus \pmb{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:
$$
\underset {\theta} {\arg \min} \mathbb {E} _ {(\pmb {h}, y) \sim \mathcal {P} (\pmb {h} _ {s} \oplus \tilde {\pmb {h}} _ {j}, y)} \mathcal {L} (F ^ {\prime} (\pmb {h}), y) \neq \underset {\theta} {\arg \min} \mathbb {E} _ {(\pmb {h}, y) \sim \mathcal {P} (\oplus \pmb {h} _ {j}, y)} \mathcal {L} (F ^ {\prime} (\pmb {h}), y)
$$
To mitigate the distribution misalignment, we introduce the Prediction Head Finetuning technique. Concretely, at the end of training, we update each embedding  $\pmb{h}_j^{(i)}$  in the embedding table  $\mathcal{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  $\mathcal{T}$  only once. The rest of the finetuning stage does not involve the notoriously slow graph convolution because the prediction head  $F^{\prime}$  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.,  $\| \tilde{\pmb{h}}_j^{(i)} - \bar{\pmb{h}}_j^{(i)} \| \leq \epsilon, \forall (i,j)$ . Given this assumption, if the neural network  $(F' \circ F)(\cdot)$  is  $k$ -Lipschitz continuous, the gradients will also be bounded and never run too far from its true estimation, i.e.,  $\| \nabla F'(\tilde{\pmb{h}}_j^{(i)}) - \nabla F'(\bar{\pmb{h}}_j^{(i)}) \| \leq k' \cdot \epsilon$ . 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  $\tilde{h}_j^{(i)}$  gets updated infrequently in the embedding table  $\mathcal{T}$ . Suppose we iterate through every graph  $\mathcal{G}^{(i)}$  in the training set  $\mathcal{D}_{\mathrm{train}}$  for each epoch, every time we train on a graph  $\mathcal{G}^{(i)} \approx \bigoplus \mathcal{G}_j^{(i)}$ , only a few graph segments  $\mathcal{G}_s^{(i)}$  will be updated in the table  $\mathcal{T}$ . This implies that all the other segment embeddings of  $\mathcal{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  $\eta$  for each segment is defined as:
$$
\eta^ {(i)} = \left\{ \begin{array}{l l} p + (1 - p) \frac {J ^ {(i)}}{S ^ {(i)}} & \text {f o r} \mathcal {G} _ {s} ^ {(i)} \\ 0 & \text {f o r} \mathcal {G} _ {j} ^ {(i)}, \text {w i t h} \text {p r o b .} (1 - p) \\ 1 & \text {f o r} \mathcal {G} _ {j} ^ {(i)}, \text {w i t h} \text {p r o b .} p \end{array} \right. \tag {1}
$$
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.
We characterize the effect of stale historical embeddings by studying the difference between  $\mathcal{L}(F^{\prime}(\pmb{h}_s^{(i)}\oplus \tilde{\pmb{h}}_j^{(i)}))$  and  $\mathcal{L}(F^{\prime}(\bigoplus \pmb{h}_{j}^{(i)}))$ . Assume  $\mathcal{G}^{(i)}$  has  $J^{(i)}$  segments and we perform backpropagation on  $S^{(i)}$  segments. We let  $\delta^{(i)}\triangleq \pmb{h}_s^{(i)}\oplus \tilde{\pmb{h}}_j^{(i)} - \bigoplus \pmb{h}_j^{(i)}$  be the perturbation on the graph embedding.
We apply Taylor expansion around  $\delta^{(i)} = 0$  on the final loss to analyze the effect of this perturbation.
$$
\begin{array}{l} \mathbb {E} _ {s} \mathcal {L} \left(F ^ {\prime} \left(\boldsymbol {h} _ {s} ^ {(i)} \bigoplus \tilde {\boldsymbol {h}} _ {j} ^ {(i)}\right)\right) - \mathcal {L} \left(F ^ {\prime} \left(\bigoplus \boldsymbol {h} _ {j} ^ {(i)}\right)\right) \tag {2} \\ = \mathbb {E} _ {s} \mathcal {L} \left(F ^ {\prime} \left(\bigoplus \boldsymbol {h} _ {j} ^ {(i)} + \delta^ {(i)}\right)\right) - \mathcal {L} \left(F ^ {\prime} \left(\bigoplus \boldsymbol {h} _ {j} ^ {(i)}\right)\right) \\ \approx \sum_ {j} \mathbb {E} _ {\delta_ {j} ^ {(i)}} \underbrace {D _ {\boldsymbol {h} _ {j} ^ {(i)}} (\mathcal {L} \circ F ^ {\prime}) [ \boldsymbol {h} _ {j} ^ {(i)} ] \delta_ {j} ^ {(i)}} _ {B} + \underbrace {\frac {1}{2} \delta_ {j} ^ {(i) \top} (D _ {\boldsymbol {h} _ {j} ^ {(i)}} ^ {2} (\mathcal {L} \circ F ^ {\prime}) [ \boldsymbol {h} _ {j} ^ {(i)} ]) \delta_ {j} ^ {(i)}} _ {R} \\ \end{array}
$$
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  $\delta^{(i)}$ . For the first term, we have
$$
\mathbb {E} _ {\delta_ {j} ^ {(i)} ^ {\mathrm {E T}}} [ B ] = C \times \frac {J ^ {(i)} - S ^ {(i)}}{J ^ {(i)}} \mathbb {E} \left(\tilde {\boldsymbol {h}} _ {j} ^ {(i)} - \boldsymbol {h} _ {j} ^ {(i)}\right)
$$
$$
\mathbb {E} _ {\delta_ {j} ^ {(i) \mathrm {S E D}}} [ B ] = C \times \frac {J ^ {(i)} - S ^ {(i)}}{J ^ {(i)}} \mathbb {E} \left(\tilde {\boldsymbol {h}} _ {j} ^ {(i)} - \boldsymbol {h} _ {j} ^ {(i)}\right) 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,  $\text{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  $\mathbb{E}(\tilde{\pmb{h}}_j^{(i)} - \pmb{h}_j^{(i)})$  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  $\mathbb{E}_{\delta_j^{(i)\mathrm{SED}}} [B]$  and  $\mathbb{E}_{\delta_j^{(i)\mathrm{SED}}} [R]$  degrade to the result of ET. If  $p = 0$  (dropping 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$ ).  $\mathbb{E}_{\delta_j^{(i)\mathrm{SED}}} [B] = 0$  indicates that there is no stale bias in this case. However, the term  $\mathbb{E}_{\delta_j^{(i)\mathrm{SED}}} [R]$  could become too large so that it impedes training.
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.