Algorithm and System Co-design for Efficient Subgraph-based Graph Representation Learning

Now, we formally formulate the GRL and SGRL problems.

Next, we define SGRL where for a particular query $Q$, the predictions are made based on the local subgraph around $Q$.

We list a few important examples of SGRL tasks. Link prediction seeks to estimate the likelihood of a link between two endpoints in a given graph. Additionally, it can be generalized to predict the type of links, such as relation prediction for heterogeneous graphs. In this case, the set $Q$ corresponds to a pair of nodes. The network scientific community has identified the importance of leveraging the local induced subgraphs for link prediction (Liben-Nowell and Kleinberg, 2007). For example, the number of common friends (shown as neighbors in a social network) implies how likely two individuals may become friends in the future. Another generalized form of link prediction is higher-order pattern prediction, where the set $Q$ consists of three or more nodes. The goal is to predict whether the set of nodes in $Q$ will foster a covered edge (termed hyperedge).

Graph neural networks (GNNs). Canonical GNNs associate each node $v$ with a vector representation $\boldsymbol{\mathrm{h}}$, which is learned and updated by aggregating messages from $v$’s neighbors, as

Here, UPDATE is implemented by neural networks while AGGREGATE is a pooling operation invariant to the order of the neighbors. By unfolding the neighborhood around each node, the computation graph to get each node representation forms a tree structure. According to Def. 2.4, canonical GNNs seem also able to perform SGRL by encoding the local subtree rooted at each node into a node representation (See Fig. 1a). Nevertheless, by this way, each node representation only separately reflects the subgraph around each node but cannot jointly represent the subgraph around multiple nodes, which yields the problem in Fig. 1. However, the SGRL framework considered in this work is able to learn the representation of the joint subgraph around a queried node set.

Without exception, previous works focus on improving the scalability of canonical GNNs and their system support, but some of their techniques inspire the design of SUREL.

To overcome the memory bottleneck of GPU when processing large-scaled graphs, sub-sampling the graph structure is a widely adopted strategy. GraphSAGE (Hamilton et al., 2017) and VR-GCN (Chen et al., 2018b) use uniform sampling schema and variance reduction technique respectively to restrict the size of node neighbors; PIN-SAGE (Ying et al., 2018) exploits Personalized PageRank (PPR) scores to sample neighbors. FastGCN (Chen et al., 2018a) and ASGCN (Huang et al., 2018) perform independent layer-wise node sampling to allow neighborhood sharing. Cluster-GCN (Chiang et al., 2019) and GraphSAINT (Zeng et al., 2020) study subgraph-based mini-batching approaches to reduce the size of training graphs. Note that the subgraphs in our setting are substantially different from theirs, since our subgraphs work as features for queries while their subgraphs are a compensatory choice to achieve better scalability.

Many works better the system support for GNNs. DGL (Wang et al., 2019) and PyG (Fey and Lenssen, 2019) are designed for scalable single-machine GNN training. Marius (Mohoney et al., 2021) is proposed to efficiently learn large-scale graph embeddings on a single machine. There are several distributed systems dedicated to GNNs: AliGraph (Yang, 2019) addresses the storage issue of applying GNNs on massive industrial graphs; AGL (Zhang et al., 2020) employs a subgraph-based system for GRL; ROC (Jia et al., 2020) builds a multi-GPU framework for deeper and larger GNN models; Dorylus (Thorpe et al., 2021) designs a CPU-based distributed system for GNN training. $\text{G}^{3}$ (Liu et al., 2020a) speedups GNN training via supporting parallel graph-structured operations. Zhou et al. (2021) uses feature dimension pruning to accelerate large-scale GNN inference. However, all these systems only support canonical GNNs so they all suffer from the intrinsic modeling limitations of GNNs.

Existing SGRL frameworks that extract a subgraph per query do not support efficient training and inference. $m$-hop subgraph extraction faces the size “explosion” issue as many nodes have significantly large degrees in real-world networks. Moreover, subgraphs of different sizes cause workload fluctuation, hindering load balancing and memory management.

Subgraph extraction can be replaced with efficient walk-based sampling, which sidesteps all above issues via regulating the number and the length of sampled walks. The number and the length of these walks are small constants, so the space and time complexity here is only linear w.r.t the number of nodes. Specifically, during preprocessing, SUREL reduces the subgraph around each node in a given graph to a set of random walks originated from it. To compensate for the loss of structural information after breaking subgraphs into walks, an intra-node distance feature termed relative positional encoding (RPE) is proposed, which enables locating each node in the sampled subgraph. The collected set of walks paired with its RPEs is hosted in the walk-based subgraph storage, with a dedicated data structure designed to support rapid and intensive access. The preprocessing flow is presented in the upper part of Fig. 2.

For training and testing, given a query (set of nodes), SUREL employs subgraph joining to implicitly construct a subgraph around the entire query in full parallel. First, all the walks originated from the queried node set are grouped. Then, the precomputed node-level RPEs are joined into query-level RPEs. SUREL further adopts neural networks to encode the grouped walks paired with query-level RPEs, and makes final predictions based on the obtained subgraph representation. A mini-batching strategy is designed to maximize data reuse during training by exploiting the query overlaps.

The bottleneck of current SGRL frameworks is how to cheaply acquire the $m$-hop neighbors for each queried set of nodes. SUREL proposes to decompose the $m$-hop subgraph into a set of $m$-length walks that start from the queried set of nodes. As the walks are regular, their storage and access are extremely efficient. This also resolves the computational problem caused by the long-tailed distribution of node degrees. More importantly, the collected walks grouped by their starting nodes can be shared and reused among different queries. Our design decouples SGRL from redundant subgraph extraction and enables the reusability of preprocessed data. We summarize the preprocessing routines with the support of hash-indexed storage in Algorithm 1 and introduce the specifics next.

Walk Sampling. During preprocessing, SUREL samples $M$-many $m$-step walks for every node in a given graph. As Fig. 3 (upper left) shows, the sampled walks are grouped in a set $\mathcal{W}_{u}$, where $u$ denotes the starting node of these walks. Walk sampling can be easily divided into parallelizable pieces. The parallelization is implemented based on NumPy and OpenMP framework in C. Moreover, to further accelerate walk sampling, we use compressed sparse row (CSR) to represent the graph. The CSR format consists of two arrays, idxptr of length $|\mathcal{V}|+1$ used to record the degrees of nodes, and indices of size $|\mathcal{E}|$, each row of which corresponds to the neighbor list per node. CSR allows intensive fast access to the neighbors of a node while keeping the memory cost low, which is vital for walk sampling in large-scale graphs.

Relative Positional Encoding (RPE). Structural information gets lost after breaking subgraphs into walks. SUREL compensates such loss via RPE to locate the relative position of a node in each sampled subgraph, which characterizes the structural contribution of the node to its corresponding subgraph.

For each set of walks $\mathcal{W}_{u}$, we first establish a set $\mathcal{V}_{u}$ that contains distinct nodes appearing in $\mathcal{W}_{u}$. Define node-level RPE $\mathcal{X}_{u}:\mathcal{V}_{u}\rightarrow\mathbb{R}^{m+1}$ as follows: for each node $x\in\mathcal{V}_{u}$, a vector $\mathcal{X}_{u,x}\in\mathbb{R}^{m+1}$ is assigned, where $\mathcal{X}_{u,x}[i]$ is the landing counts of node $x$ at position $i$ in all walks of $\mathcal{W}_{u}$. In SUREL, RPE can be computed on the fly as walks get sampled, thus resulting in nearly zero extra computational cost. The set of walks $\mathcal{W}_{u}$ paired with the RPE $\mathcal{X}_{u}$ essentially characterize a sub-sampled subgraph around the node $u$. Next, we present a dedicated data structure to host $\mathcal{W}_{u}$ and $\mathcal{X}_{u}$ altogether.

It is easy to manage the collected set of walks due to its regularity. An $m*M$-sized chunk is allocated to each set of walks, which assists to speed up data fetching. How to organize node-level RPE presents a real challenge because the cardinality of the set $|\mathcal{V}_{u}|$ varies from node to node. One naïve way to avoid such irregularity is to directly scatter these RPEs back to nodes in previously collected walks. But, this gives an $m*M*(m+1)$ tensor, resulting in an unrealizable memory need. Moreover, it loses track of node IDs in walks that are needed for joining subgraphs later.

We use an associative array $\mathcal{A}$ to organize all walk-based subgraphs as shown in the upper part of Fig. 3. For each node $u\in\mathcal{V}$, its corresponding entry in $\mathcal{A}$ is a node-level subgraph formed as a tuple $(\mathcal{W}_{u},\mathcal{H}_{u})$. Here, $\mathcal{W}_{u}$ is a set of walks starting from $u$, and $\mathcal{H}_{u}$ is a dictionary that maps the unique node set $\mathcal{V}_{u}$ of $\mathcal{W}_{u}$ to its corresponding node-level RPE $\mathcal{X}_{u}$. The use of dictionary resolves irregularities in $\mathcal{V}_{u}$ mentioned above, while maintaining the connection between node IDs and their RPEs. In addition, array $\mathcal{T}$ is introduced to store RPE values centrally, rather than scattered across dictionaries. As Fig. 3 (upper right) shows, the value of $\mathcal{H}_{u}(x)$ is now replaced with the index of the RPE value $\mathcal{X}_{u,x}$ stored in $\mathcal{T}$ accordingly, noted as $\text{RPE-ID}_{u,x}$. This design overall guarantees the access of RPE in $O(1)$ time.

The above $\mathcal{A}$ and $\mathcal{H}_{u}$ are built on top of uthash’s macros ¹¹1https://troydhanson.github.io/uthash/, with extended support for arbitrary insertions and deletions of key–value pairs. It offers data access and search in $O(1)$ time on average, which is about as good as the direct address table but greatly reduces the space wastage. In particular, it has no dependency or need for communication between multiple hash queries, thus can be pleasingly executed in parallel. Both $\mathcal{A}$ and $\mathcal{H}_{u}$ are stored in RAM on the CPU side. As we observed in Fig. 3, there are many repeated RPE values. Once all nodes are sampled, the array $\mathcal{T}$ can be pruned to remove duplicates. RPE-IDs will be updated synchronously when $\mathcal{T}$ is reindexed. For example, both node $a$ and $v$ have the RPE value of $[0,0,1]$, whose index in $\mathcal{T}$ is (1) after pruning. Thus, both $\mathcal{H}_{u}(a)$ and $\mathcal{H}_{u}(v)$ are assigned to the new RPE-ID as (1). The shape of $\mathcal{T}$ is regular and its size is usually small after pruning, which can be fully loaded in GPU. In practice, we found that pining RPEs in GPU memory is critical, as it can significantly reduce the communication cost of moving data back and forth between RAM and SDRAM.

The storage designed above records the downsampled subgraph around each node. As SGRL is mostly useful for making predictions over a set of nodes $Q$, here we further illustrate how to get the joined subgraph around all the nodes $u\in Q$.

The idea is to concatenate all set of walks $[...,\mathcal{W}_{u},...]$ for $u\in Q$, since each set of walks $\mathcal{W}_{u}$ can be viewed as a subgraph around $u$. Besides, each node $x$ in the walks will be paired with a query-level RPE $\mathcal{X}_{Q,x}$ that characterizes the relative position of node $x$ in the joint subgraph around the queried set $Q$. Specifically, $\mathcal{X}_{Q,x}$ is defined by joining all RPEs $\mathcal{X}_{u,x}$ for $u\in Q$, i.e., $\mathcal{X}_{Q,x}=\uplus_{u\in Q}\mathcal{X}_{u,x}(\triangleq[...,\mathcal{X}_{u,x},...])\in\mathbb{R}^{(m+1)\times|Q|}$. There will be some $u\in Q$ such that $x\notin\mathcal{V}_{u}$, for which $\mathcal{X}_{u,x}$ is set to all zeros. Through this procedure, the joined subgraph with query-level RPEs is sent to GPU for representation learning and then model inference.

The data structure described in Sec. 3.3 enables a highly parallel implementation of subgraph joining along with optimized memory management. On the CPU side, $\mathcal{X}_{Q,x}$ is not directly used to assemble walks. Instead, we use a query-level RPE-ID that joins node-level RPE indices in $\mathcal{T}$, i.e. use $\text{RPE-ID}_{Q,x}=[...,\text{RPE-ID}_{u,x},...]\in\mathbb{R}^{|Q|}$ for $u\in Q$, which reduces the memory cost from $(m+1)*|Q|$ to $|Q|$. For instance, in Fig. 3 (bottom right), $\mathcal{X}_{Q,u}=([2,0,0],[0,0,1])$ can be substituted by $\text{RPE-ID}_{Q,u}=(3,1)$, as their RPE values locate at the entry (3) and (1) of $\mathcal{T}$. As follows, SUREL pre-allocates an array with the fixed-size $[m*M*|Q|,|Q|]$, where $m*M*|Q|$ is the size of walks around $Q$. Then, SUREL fills the index array with $\text{RPE-ID}_{u,x}$ by multithreads. Note that $\text{RPE-ID}_{u,x}$ can be rapidly retrieved via the dictionary operation $\mathcal{H}_{u}(x)$. Lastly, assembling RPE values to walks is performed on GPU via the indexing operation $\mathcal{X}_{u,x}=\mathcal{T}(\text{RPE-ID}_{u,x})$, where $\mathcal{T}$ is pinned in GPU memory earlier. SUREL incorporates a Python/C hybrid API for subgraph joining, building on top of NumPy, PyTorch, OpenMP and uthash.

Some remarks can be made here. First, the above algorithm contains some redundancy to compute the query-level RPE-ID for the nodes that appear multiple times in the walks. In practice, we find that about half of the nodes appear only once, thus doubling the computation time at most. To avoid such redundancy, one can first compute the set union $\mathcal{V}_{Q}=\cup_{u\in Q}\mathcal{V}_{u}$, and then compute the query-level RPE-ID by traversing all nodes in $\mathcal{V}_{Q}$. However, parallel set union is difficult to implement efficiently. When multithreading is enabled, we observe a significant increase in the efficiency of SUREL, as opposed to the union operation. Also, by dynamically adjusting the number of threads, the workload between CPU and GPU can be well balanced. Second, we have empirically found that using RPE-ID instead of RPE to assemble walks provides an observable performance boost (speed up by $2\times$ or more), otherwise data communication between CPU and GPU would the main bottleneck.

After subgraph joining for each query, the obtained subgraph is represented by a concatenated set of walks on which nodes are paired with query-level RPEs (See Fig. 3). Next, we introduce neural networks to encode these walks into a subgraph representation $h_{Q}$.

Due to its regularity, any sequential models, e.g., MLP, CNN, RNN, and transformers can be adopted for sampled walks. We test RNN and MLP for neural encoding, both of which achieve similar results. Next, we take the RNN as an example. We encode each walk $W=(w_{0},w_{1},...,w_{m})\in\mathcal{W}$ as $\text{enc}(W)=\text{RNN}(\{f\left(\mathcal{X}_{Q,w_{i}}\right)\}_{i=0,1,\dots,m})$, where $w_{i}$’s denote the node at step $i$ in one sampled walk. Here, $f$ is to encode the query-level RPE. Node or edge attributes for each step $w_{k}\in W$ can be supported by attaching those attributes after its RPE. To obtain the final subgraph representation of $Q$, we aggregate the encoding of all the associated walks through a mean pooling, i.e., $h_{Q}=\text{mean}(\{\text{enc}(W)|W\;\text{starts from some $u\in Q$}\}).$ In the end, a two-layer classifier is used to make prediction by taking $h_{Q}$ as input. In our experiments, all the tasks can be formulated as binary classification, and thus we adopt Binary Cross Entropy as the loss function.

SUREL organically incorporates the storage designed in Sec. 3.3 and the subgraph-joining operation described in Sec. 3.4 to achieve efficient training and model serving.

Subgraph queries $Q$’s are sets of nodes, which often come from a common ambient on a large graph. There might be many overlaps between different queries and their $m$-hop induced subgraphs. If the queried subgraphs are known in prior, we may put these queries with high node overlap into the same batch to improve data reuse. Here, queries of each given task are assumed to have the same size, e.g. $|Q|=2$ for link prediction. In practice, test queries are usually given online while the training ones can be prepared in advance. Hence, we propose to accelerate the training pipeline by mini-batching the overlapping queries. Practitioners can choose the appropriate pipeline according to the specific situation. Algorithm 2 summarizes the overall training procedure of SUREL.

Mini-batching for Training. We first randomly sample a seed-set of nodes $\bar{\mathcal{V}}$ from the union of queried node sets $\cup Q$. Then, we run breadth-first search (BFS) to expand the seed-set $\bar{\mathcal{V}}$. Neighbor fetching of the BFS here is based on the grouped queries instead of the original graph: a neighbor of node $u$ is defined as the node that shares at least one query with it. During BFS, the reached queries will be added to a set $\mathcal{Q}$. The expansion stops once the size of either the seed-set $\bar{\mathcal{V}}$ or the mini-batch $\mathcal{Q}$ reaches some pre-defined limits. Since the data structure for each query in SUREL after subgraph joining is regular, it is easy to decide the size limits of seed-set and mini-batch based on resource availability (i.e. GPU memory). In practice, this BFS procedure improves reusability of data within each mini-batch, and may significantly decrease the communication cost between CPU and GPU. If the training set only contains positive queries (often in link/motif prediction tasks), we design an efficient sampling strategy for negative queries by the same principle that randomly pairs them within the same batch.

We conduct extensive experiments to evaluate the proposed framework with three kinds of graphs (homogeneous, heterogeneous, and higher-order homogeneous) on three corresponding types of tasks, namely, link prediction, relation prediction and higher-order pattern prediction. Homogeneous graphs are the graphs without node/link types. Heterogeneous graphs include node/link types. Higher-order graphs contain higher-order links that may connect more than 2 nodes. The dataset statistics are summarized in Table 1, most of which are larger than the datasets used in (Zhang et al., 2021b; Zhou et al., 2021), not to mention that our node-set prediction task is much more complex than the node classification task considered in the previous works.

Open Graph Benchmark (OGB). We use three link prediction and one relation prediction datasets (Hu et al., 2020): ppa - a protein interaction network, collab - a collaboration network, and citation2 - a citation network; and one heterogeneous network ogb-mag, which contains four types of nodes (paper, author, institution and field) and their relations extracted from MAG (Wang et al., 2020).

Higher-order Graph Dataset. DBLP-coauthor is a temporal higher-order network that records co-authorship of papers as timestamped higher-order links. tags-math contains sets of tags that are applied to questions on the website math.stackexchange.com as higher-order links. For the two higher-order graphs, SUREL and all the baselines will treat them as standard graphs by projecting higher-order links into cliques. However, the training and test queries are generated based on higher-order links detailed next.

Settings. For Link Prediction, we follow the data split as OGB requires to isolating the validation and test links (queries) from the graphs. For Relation Prediction, the relations of paper-author (P-A) and paper-citation (P-P) are selected. The dataset is split based on timestamps. 0.5% of existing edges of each target relation type are selected from ogb-mag. For each paper, two authors/citations are picked from its P-A/P-P relations respectively, one for validation and the other for testing. The remaining links are used for training. For Higher-order Pattern Prediction, we focus on predicting whether two nodes will be connected to a third node concurrently via a higher-order link in the future. Specifically, positive queries are node triplets, where two nodes are linked before the timestamp $t$ and the third node establishes connection to the pair via a higher-order link after $t$. The split ratio of positive node triplets is 60/20/20 for training/validation/testing. For Relation Prediction and Higher-order Pattern Prediction, each positive query is paired with 1000 randomly sampled negative queries (except tags-math uses 100) in testing. For fair comparison, all baselines are tested with the same set of negative queries sampled individually for each dataset. All experiments are run 10 times independently, and we report the mean performance and standard deviations.

Baselines. We consider three classes of baselines. Graph Embedding methods for transductive learning: Node2vec (Grover and Leskovec, 2016) and DeepWalk (Perozzi et al., 2014), which learns a single embedding for each node and may suffer from the information over-squashing issue; Canonical GNNs: GCN (Kipf and Welling, 2017), GraphSAGE (Hamilton et al., 2017), GraphSAINT (Zeng et al., 2020), Cluster-GCN (Chiang et al., 2019), Relational GCN (R-GCN) (Schlichtkrull et al., 2018), Relation-aware Heterogeneous Graph Neural Network (R-HGNN) (Yu et al., 2022); SGRL models: SEAL(Zhang and Chen, 2018; Zhang et al., 2021a), DE-GNN(Li et al., 2020). SEAL supports both offline and online subgraph extraction per query. However, it takes SEAL 2+ hours and 102GB RAM to offline extract 2% training subgraphs on citation2. Thus, we only keep the online setting for SEAL. DE-GNN only supports offline subgraph extraction. Table 2 compares subgraph sampling for different SGRL methods. We adopt official implementations of above baselines with tuned parameters that match reported results. SUREL uses an 2-layer MLP for embeddings of RPEs and an 2-layer RNN to encode query-level joined walks. The obtained subgraph embeddings are fed into an MLP classifier for final prediction. Default training parameters are: learning rate lr=1e-3 with early stopping of 5-epoch patience, dropout p=0.1, Adam (Kingma and Ba, 2015) as the optimizer, batch capacity $B_{1}=1500$, and batch size $B_{2}=32$. Hidden dimension $d$ and walk parameters $M,m$ are investigated in Sec. 4.4.

Metric. The evaluation metrics include Hits@K and Mean Reciprocal Rank (MRR). Hit@K counts the percentage of positive samples ranked at the top-K place against all the negative ones. MRR firstly calculates the inverse of the ranking of the first correct prediction against the given number of paired negative samples, and then an average is taken over the total queries.

Environment. We use a server with four Intel Xeon Gold 6248R CPUs, 1TB DRAM, and eight NVIDIA RTX 6000 (24GB) GPUs.

Table 3 shows results of three prediction tasks. Apparently, for these three link prediction benchmarks, the performance of SGRL models is significantly better than transductive graph embedding models and canonical GNNs, particularly for the challenging tasks over ppa and collab. Within SGRL models, SUREL sets two SOTA results on ppa and citation2, and gets comparable performance on collab against SEAL, which validates the modeling effectiveness of our proposed walk-based framework. For relation prediction and higher-order pattern prediction, we observe a large gap (up to 60%) between canonical GNNs and SUREL-based models, especially in higher-order cases. This again demonstrates the inherent modeling limitation of canonical GNNs to predict over a set of nodes. DE-GNN suffers from serious scalability issues when employing subgraph extraction for higher-order pattern prediction. Our best attempt is to deploy DE-GNN on tags-math by using 10% training samples, while the other three graphs failed. DE-GNN spends more than 300 hours preprocessing just $5\%$ training queries of DBLP-coauthor.

Table 4 reports the runtime, memory consumption comparison on a single machine (using one GPU) between canonical GNNs and SGRL models. SUREL offers a reasonable total runtime on these benchmarks compared with canonical GNNs. Meanwhile, its preprocessing overhead is negligible as showed in Table 4 under the term ‘Prep.’, and the higher-order case can be efficiently handled as well. SEAL adopts online extraction, and thus the cost is not counted in preprocessing, while its training suffers from the computation bottleneck. DE-GNN uses offline extraction, and it takes 15+ hours and 98GB RAM to process training queries in tags-math, which is obviously incapable of scaling to DBLP-coauthor (so not present in Table 4). Overall, SUREL substantially accelerates the subgraph extraction and makes it feasible for SGRL on large-scale graphs.

In terms of memory management, SUREL achieves comparable RAM usage to canonical GNNs, because the number of walks $M$ and the steps $m$ are small constants in practice. The extra memory cost is linear in $|\mathcal{V}|$, so the total memory cost is still dominated by the original graph. However, SEAL induces much more RAM usage as it extracts subgraphs of long-tail sizes, and its total memory cost is often super-linear in $|\mathcal{V}|$. Both SEAL and SUREL consume much less SDRAM because they do not need GPU to load large adjacency matrices and host node representations.

We further profile the training and inference performance, and present it in Fig. 4. The upper half plots the time-to-accuracy comparison between canonical GNNs and SGRL models. Each dot indicates one training epoch for full-batch GCN, SEAL and SUREL, 10 training epochs for Cluster-GCN and GraphSAINT. As it shows, both SEAL and SUREL use 1-3 epochs to get good enough performance, and each epoch of SUREL takes around 1/10 time of SEAL on citation2. The time per epoch of full-batch GCN is comparable with SUREL, while Cluster-GCN and GraphSAINT are faster. However, these models generally take longer time to converge to even subpar performance. On ppa, the curve of SEAL is pretty oscillating, leading to longer convergence. SUREL uses large $M$ and $m$ to achieve better and more stable performance on ppa, so the training time per epoch is comparable with SEAL. The training curves of canonical GNN baselines are not plotted for ppa because of their poor performance (See Table 3).

The bottom half of Fig. 4 provides the comparison of end-to-end inference throughput between two classes of models. Canonical GNNs offer rapid inference, since they generate node representations as the intermediate computation results that are shared across all queries. But as aforementioned, sharing node representations may over-squash useful information and degenerate performance as shown in Table 3. SEAL, as SGRL, achieves good prediction performance but its inference is extremely slow, because of subgraph extraction per query. SUREL fundamentally solves this bottleneck by replacing the extraction with walk-based subgraph joining. It is $\textbf{4}-\textbf{16}\times$ faster than SEAL on inference for link prediction, and achieves even more speedup than DE-GNN in higher-order settings.

The number $M$, the step $m$ of walks and the hidden dimension $d$ effect scalability and accuracy of SUREL. To examine their impact, we evaluate SUREL on citation2, a large sparse graph, and collab, a medium dense graph, for different values of $m$, $M$, and $d$.

Prediction Performance. Fig. 5(a) and 5(d) show the prediction results. As expected, the performance consistently increases if we use a larger number of walks $M$. But for the step $m$, it is not always true that longer steps will guarantee better results, which depends on the specifics of the dataset. For instance, in network citation2, to accurately predict the link between two papers, more steps are needed as it would capture a larger group of papers which share similar semantics. While for collab, the case is different, as deeper walks would include more noise for predicting collaborations between two authors. In general, some small $m\leavevmode\nobreak\ (2\sim 5)$ and $M\leavevmode\nobreak\ (50\sim 400)$ ensure adequate performance. By adjusting $m$ and $M$, we can achieve the trade-off between accuracy and scalability, none of which is achievable through other SGRL models. Moreover, SUREL is insensitive to the hidden dimension as shown by Fig. 5(d).

Training and Inference Time Cost. As Figs. 5(b) and 5(c) demonstrated, the time of walk sampling and subgraph joining is nearly linear w.r.t. the total number of walks ($m*M$) under the same number of threads (16 by default). Here, we do not regulate $M$ based on the degree of each node in a query, which may induce certain duplication in sampled walks originated from the nodes with small degrees. Using degree-adaptive $M$ is promising to further improve the scalability of SUREL while keeping good prediction performance. We leave such investigation for future study.

To investigate the scaling performance of the parallel implementation, we examine the runtime of heavy operations in SUREL by using different numbers of threads. Fig. 6(a) shows the throughput of walk sampler and query-level RPE joining on citation2. The runtime is also compared to the estimated runtime by Amdahl’s law (Grama et al., 2003) shown in Fig. 6(b): walk sampling and RPE joining are in good agreement with the expected speedup, thus implying well parallelized implementation.