Asymmetric Learning for Graph Neural Network based Link Prediction

TL;DR

This paper introduces AML, a novel asymmetric learning approach for graph neural network-based link prediction, significantly improving training speed on large-scale graphs without sacrificing accuracy.

Contribution

It proposes the first asymmetric learning strategy for GNN-LP, combining GNN and MLP models for node representations, and introduces a row-wise sampling method for efficient training.

Findings

AML is 1.7X to 7.3X faster than symmetric GNN-LP methods.

AML maintains comparable accuracy to existing methods.

The approach effectively scales GNN-based link prediction to large graphs.

Abstract

Link prediction is a fundamental problem in many graph based applications, such as protein-protein interaction prediction. Graph neural network (GNN) has recently been widely used for link prediction. However, existing GNN based link prediction (GNN-LP) methods suffer from scalability problem during training for large-scale graphs, which has received little attention by researchers. In this paper, we first give computation complexity analysis of existing GNN-LP methods, which reveals that the scalability problem stems from their symmetric learning strategy adopting the same class of GNN models to learn representation for both head and tail nodes. Then we propose a novel method, called asymmetric learning (AML), for GNN-LP. The main idea of AML is to adopt a GNN model for learning head node representation while using a multi-layer perceptron (MLP) model for learning tail node representation. Furthermore, AML proposes a row-wise sampling strategy to generate mini-batch for training, which is a necessary component to make the asymmetric learning strategy work for training speedup. To the best of our knowledge, AML is the first GNN-LP method adopting an asymmetric learning strategy for node representation learning. Experiments on three real large-scale datasets show that AML is 1.7X~7.3X faster in training than baselines with a symmetric learning strategy, while having almost no accuracy loss.

Tables

Table 1. Table 1 : Complexity analysis. L 𝐿 L is the number of model layers. s = ∥ 𝐀 ∥ 0 / N 𝑠 subscript delimited-∥∥ 𝐀 0 𝑁 s{=}\lVert{\bf A}\rVert_{0}/N is the average number of neighbors for a node in 𝒢 𝒢 \mathcal{G} . r 𝑟 r is the dimension of node representation. | ℰ | ℰ |\mathcal{E}| is the number of links in the training set ℰ ℰ \mathcal{E} . N 𝑁 N is the number of nodes in graph 𝒢 𝒢 \mathcal{G} . k 𝑘 k in local GNN-LP is the number of hops for the enclosed subgraphs. “RWS” denotes row-wise sampling. “w/ RWS” in local GNN-LP and global GNN-LP means generating mini-batch with RWS. “w/o RWS” in AML means generating mini-batch without RWS but with an edge-wise sampling strategy.

Method	Computation complexity
Local GNN-LP	$𝒪\left(2L{s}^k{r}^2\cdot \left|ℰ\right|\right)$
Local GNN-LP (w/ RWS)	$𝒪\left(2L{s}^k{r}^2\cdot \left|ℰ\right|\right)$
Global GNN-LP	$𝒪\left(2s^L{r}^2\cdot \left|ℰ\right|\right)$
Global GNN-LP (w/ RWS)	$𝒪\left(s^L{r}^2\cdot \left(\left|ℰ\right|+N\right)\right)$
AML (w/o RWS)	$𝒪\left(\left(s^L{r}^2+L{r}^2\right)\cdot \left|ℰ\right|\right)$
AML	$𝒪\left(s^L{r}^2\cdot N+L{r}^2\cdot \left|ℰ\right|\right)$

Table 2. Table 2: Statistics of datasets.

Datasets	ogbl-collab	ogbl-ppa	ogbl-citation2
#Nodes	235,868	576,289	2,927,963
#Edges	1,285,465	30,326,273	30,561,187
Features/Node	128	128	128
#Training links	1,179,052	21,231,931	30,387,995
#Validation links	160,084	9,062,562	86,682,596
#Test links	146,329	6,031,780	86,682,596
Metric	Hits@50	Hits@100	MRR

Table 3. Table 3: Comparison with global GNN-LP baselines. “Time” denotes the whole time to complete training. “Gap” denotes the accuracy of AML minus that of baselines.

Methods	ogbl-collab			ogbl-ppa
	Hits@50 (%) $\uparrow$	Gap	Time (s) $\downarrow$	Hits@100 (%) $\uparrow$	Gap	Time (s) $\downarrow$
SAGE	54.57 $\pm$ 0.82	-	$7.9\times {10}^3$	50.13 $\pm$ 0.55	-	$1.0\times {10}^5$
AML	57.26 $\pm$ 1.25	$+2.69$	$4.4\times {\mathrm{𝟏𝟎}}^{\mathrm{𝟑}}$	49.73 $\pm$ 0.89	$-0.40$	$3.0\times {\mathrm{𝟏𝟎}}^{\mathrm{𝟒}}$

Table 4. (a) SAGE as the GNN model.

Methods	ogbl-collab			ogbl-ppa
	Hits@50 (%) $\uparrow$	Gap	Time (s) $\downarrow$	Hits@100 (%) $\uparrow$	Gap	Time (s) $\downarrow$
SAGE	54.57 $\pm$ 0.82	-	$7.9\times {10}^3$	50.13 $\pm$ 0.55	-	$1.0\times {10}^5$
AML	57.26 $\pm$ 1.25	$+2.69$	$4.4\times {\mathrm{𝟏𝟎}}^{\mathrm{𝟑}}$	49.73 $\pm$ 0.89	$-0.40$	$3.0\times {\mathrm{𝟏𝟎}}^{\mathrm{𝟒}}$

Table 5. (b) GAT as the GNN model.

Methods	ogbl-collab			ogbl-ppa
	Hits@50 (%) $\uparrow$	Gap	Time (s) $\downarrow$	Hits@100 (%) $\uparrow$	Gap	Time (s) $\downarrow$
GAT	56.43 $\pm$ 0.86	-	$7.8\times {10}^3$	49.71 $\pm$ 0.48	-	$1.0\times {10}^5$
AML	57.60 $\pm$ 0.71	$+1.17$	$4.5\times {\mathrm{𝟏𝟎}}^{\mathrm{𝟑}}$	50.23 $\pm$ 0.78	$+0.52$	$3.2\times {\mathrm{𝟏𝟎}}^{\mathrm{𝟒}}$

Table 6. Table 4 : Comparison with non-GNN and local GNN-LP baselines. “(S)” in AML means using SAGE as the GNN model, and “(G)” in AML means using GAT as the GNN model. Accuracy of non-GNN and local GNN-LP baselines on ogbl-ppa and ogbl-citation2 is from  [ ZLX + 21 ] . “*” denotes the results achieved by rerunning the authors’ code on the clean ogbl-collab.

Category	Methods	ogbl-collab		ogbl-ppa		ogbl-citation2
		Hits@50(%)$\uparrow$	Time(s)$\downarrow$	Hits@100(%)$\uparrow$	Time(s)$\downarrow$	MRR(%)$\uparrow$	Time(s)$\downarrow$
Non-GNN	CN	49.96$\pm$0.00*	-	27.60$\pm$0.00	-	51.47$\pm$0.00	-
	AA	56.49$\pm$0.00*	-	32.45$\pm$0.00	-	51.89$\pm$0.00	-
	Node2vec	49.29$\pm$0.64*	-	22.26$\pm$0.88	-	61.41$\pm$0.11	-
	MF	37.93$\pm$0.76*	-	32.29$\pm$0.94	-	51.86$\pm$4.43	-
Local GNN-LP	DE-GNN	57.87$\pm$0.79*	$6.1\times {10}^4$	45.70$\pm$3.46	$2.0\times {10}^6$	78.85$\pm$0.17	$1.1\times {10}^6$
	SEAL	57.55$\pm$0.72*	$6.5\times {10}^4$	48.80$\pm$3.16	$2.0\times {10}^6$	87.67$\pm$0.32	$1.1\times {10}^6$
Global GNN-LP	AML (S)	57.26$\pm$1.25	$4.4\times {\mathrm{𝟏𝟎}}^{\mathrm{𝟑}}$	49.73$\pm$0.89	$3.0\times {\mathrm{𝟏𝟎}}^{\mathrm{𝟒}}$	86.55$\pm$0.06	$1.0\times {\mathrm{𝟏𝟎}}^{\mathrm{𝟒}}$
	AML (G)	57.60$\pm$0.71	$4.5\times {10}^3$	50.23$\pm$0.78	$3.2\times {10}^4$	86.70$\pm$0.05	$1.0\times {\mathrm{𝟏𝟎}}^{\mathrm{𝟒}}$

Table 7. Table 5 : Experiment to verify the necessity of GNN in AML. SMLP applies MLP with pre-encoding to learn representation for both head nodes and tail nodes. “Gap” denotes the accuracy of AML minus that of SMLP.

Methods	ogbl-collab		ogbl-ppa		ogbl-citation2
	Hits@50 (%) $\uparrow$	Gap	Hits@100 (%) $\uparrow$	Gap	MRR (%) $\uparrow$	Gap
SMLP	47.25 $\pm$ 0.89	-	47.42 $\pm$ 1.37	-	69.82 $\pm$ 0.05	-
AML (S)	57.26 $\pm$ 1.25	$+10.01$	49.73 $\pm$ 0.89	$+2.31$	86.55 $\pm$ 0.06	$+16.73$
AML (G)	57.60 $\pm$ 0.71	$+10.35$	50.23 $\pm$ 0.78	$+2.81$	86.70 $\pm$ 0.05	$+16.88$

Table 8. Table 6 : Reversed asymmetric learning. AML-R denotes the reversed case of AML, which learns representation for head nodes with an MLP model while learning representation for tail nodes with a GNN model. “(S)” means using SAGE as the GNN model, and “(G)” means using GAT as the GNN model.

Methods	ogbl-collab	ogbl-ppa	ogbl-citation2
	Hits@50 (%) $\uparrow$	Hits@100 (%) $\uparrow$	MRR (%) $\uparrow$
AML-R (S)	57.15 $\pm$ 0.32	49.73 $\pm$ 0.45	85.70 $\pm$ 0.10
AML (S)	57.26 $\pm$ 1.25	49.73 $\pm$ 0.89	86.55 $\pm$ 0.06
AML-R (G)	57.08 $\pm$ 1.19	50.30 $\pm$ 0.61	85.91 $\pm$ 0.04
AML (G)	57.60 $\pm$ 0.71	50.23 $\pm$ 0.78	86.70 $\pm$ 0.05

Table 9. Table 7 : Ablation study. “KT” represents knowledge transfer. “ Δ ( L ) superscript Δ 𝐿 \Delta^{(L)} ” indicates the residual variable defined in (6). “PE” means pre-encoded information of graph structure. “HO” represents homophily. “Gap” denotes the accuracy of variants of AML minus that of AML.

Methods	ogbl-collab		ogbl-ppa		ogbl-citation2
	Hits@50 (%) $\uparrow$	Gap	Hits@100 (%) $\uparrow$	Gap	MRR (%) $\uparrow$	Gap
AML (w/o KT)	54.50$\pm$0.90	$-2.76$	49.16$\pm$1.07	$-0.57$	86.24$\pm$0.05	$-0.31$
AML (w/o ${\Delta }^{\left(L\right)}$)	54.61$\pm$0.62	$-2.65$	49.00$\pm$0.26	$-0.73$	86.26$\pm$0.02	$-0.29$
AML (w/o HO)	55.70$\pm$1.75	$-1.56$	47.74$\pm$0.66	$-1.99$	85.75$\pm$0.10	$-0.80$
AML (w/o PE)	43.90$\pm$0.78	$-13.36$	1.97$\pm$0.27	$-47.76$	60.82$\pm$0.14	$-25.73$
AML	57.26$\pm$1.25	-	49.73$\pm$0.89	-	86.55$\pm$0.06	-

Table 10. (a) SAGE as the base GNN model.

Methods	ogbl-collab		ogbl-ppa		ogbl-citation2
	Hits@50 (%) $\uparrow$	Gap	Hits@100 (%) $\uparrow$	Gap	MRR (%) $\uparrow$	Gap
AML (w/o KT)	54.50$\pm$0.90	$-2.76$	49.16$\pm$1.07	$-0.57$	86.24$\pm$0.05	$-0.31$
AML (w/o ${\Delta }^{\left(L\right)}$)	54.61$\pm$0.62	$-2.65$	49.00$\pm$0.26	$-0.73$	86.26$\pm$0.02	$-0.29$
AML (w/o HO)	55.70$\pm$1.75	$-1.56$	47.74$\pm$0.66	$-1.99$	85.75$\pm$0.10	$-0.80$
AML (w/o PE)	43.90$\pm$0.78	$-13.36$	1.97$\pm$0.27	$-47.76$	60.82$\pm$0.14	$-25.73$
AML	57.26$\pm$1.25	-	49.73$\pm$0.89	-	86.55$\pm$0.06	-

Table 11. (b) GAT as the base GNN model.

Methods	ogbl-collab		ogbl-ppa		ogbl-citation2
	Hits@50 (%) $\uparrow$	Gap	Hits@100 (%) $\uparrow$	Gap	MRR (%) $\uparrow$	Gap
AML (w/o KT)	54.60$\pm$0.60	$-3.00$	47.38$\pm$0.28	$-2.85$	86.28$\pm$0.05	$-0.42$
AML (w/o ${\Delta }^{\left(L\right)}$)	54.56$\pm$1.43	$-3.04$	47.04$\pm$1.03	$-3.19$	86.30$\pm$0.01	$-0.40$
AML (w/o HO)	55.73$\pm$0.34	$-1.87$	47.86$\pm$1.28	$-2.37$	85.83$\pm$0.04	$-0.87$
AML (w/o PE)	43.95$\pm$1.01	$-13.65$	2.95$\pm$0.35	$-47.28$	60.89$\pm$0.22	$-25.81$
AML	57.60$\pm$0.71	-	50.23$\pm$0.78	-	86.70$\pm$0.05	-

Equations

$${\mathbf{U}}_{i\ast }^{(\ell )}=f\left({\hat{\mathbf{A}}}_{i\ast }{\mathbf{U}}^{(\ell -1)}{\mathbf{W}}_1^{(\ell )}+{\mathbf{U}}_{i\ast }^{(\ell -1)}{\mathbf{W}}_2^{(\ell )}\right),$$

$${\tilde{\mathbf{V}}}^{(0)}={\hat{\mathbf{A}}}^L\mathbf{X}.$$

$${\tilde{\mathbf{V}}}^{(\ell )}=f\left({\tilde{\mathbf{V}}}^{(\ell -1)}{\mathbf{W}}_1^{(\ell )}+{\tilde{\mathbf{V}}}^{(\ell -1)}{\mathbf{W}}_2^{(\ell )}\right),$$

$${\Delta }^{(0)}=\mathbf{X}-{\hat{\mathbf{A}}}^L\mathbf{X},$$

$${\Delta }^{(\ell )}=f\left({\Delta }^{(\ell -1)}{\mathbf{W}}^{(\ell )}\right),$$

$${\mathbf{V}}^{(L)}={\tilde{\mathbf{V}}}^{(L)}+{\Delta }^{(L)},$$

$$\mathbf{U}={\mathbf{U}}^{(L)}+{\mathbf{V}}^{(L)},\mathbf{V}={\mathbf{V}}^{(L)},$$

$$S_{ij}=f_{\mathbf{\Theta }}\left({\mathbf{U}}_{i\ast }\odot {\mathbf{V}}_{j\ast }\right),$$

$$\underset{\mathcal{W}}{\min }\frac{1}{|\mathcal{E}|}\sum _{(i,j)\in \mathcal{E}}{f}_{loss}(S_{ij},Y_{ij})+\frac{\lambda }{2}\sum _{\mathbf{W}\in \mathcal{W}}\Vert \mathbf{W}{\Vert }_F^2,$$

$$\underset{\mathcal{W}}{\min }\frac{1}{|{\mathcal{E}}_{\Omega }|}\sum _{(i,j)\in {\mathcal{E}}_{\Omega }}{f}_{loss}(S_{ij},Y_{ij})+\frac{\lambda }{2}\sum _{\mathbf{W}\in \mathcal{W}}\Vert \mathbf{W}{\Vert }_F^2.$$

$${\mathcal{E}}_{\Omega }=\{(i,j)|i\in {\mathcal{V}}_{\Omega }\wedge (i,j)\in \mathcal{E}\}.$$

Full text

1 Introduction

Link prediction [LK07], a fundamental problem in many graph based applications, aims to predict the existence of a link that has not been observed. Link prediction problem widely exists in real applications, like drug response prediction [SCK17], protein-protein interaction prediction [QBJKS06], friendship prediction in social networks [AA03], knowledge graph completion [NMTG16, RBF*+*21, SJWT23] and product recommendation in recommender systems [KBV09]. Its increased importance in real applications also promotes a great interest in research for link prediction algorithms in the machine learning community.

Link prediction algorithms have been studied for a long time [LZ11, MBT17, SQN*+*20], while learning based algorithms are one dominant class in the past decades. The main idea of learning based algorithms is to learn a deterministic model [KBV09, ME11, ZC17] or a probabilistic model [SM07, GZFA09, GSP09] to fit the observed data. In most learning based algorithms, models learn or generate a representation for each node [ZXK*+*21], which is used to generate a score or probability of link existence. Traditional learning based algorithms typically do not adopt graph neural network (GNN) for node representation learning. Although these non-GNN based learning algorithms have achieved much progress in many applications, they are less expressive than GNN in node representation learning.

Recently, graph neural network based link prediction (GNN-LP) methods have been proposed and become one of the most popular algorithms due to their superior performance in accuracy. The key to the success of GNN-LP methods is that they learn node representation from graph structure and node features in a unified way with GNN, which is a major difference between them and traditional non-GNN based learning algorithms. Existing GNN-LP methods mainly include local methods [ZC18, LWWL20, ZLX*+*21, YSYL21] and global methods [KW16, PHL*+*18, HHN*+*19, YKL*+*21]. Local methods apply GNN to subgraphs that capture local structural information. Specifically, they first extract an enclosed $k$-hop subgraph for each link and then use various labeling tricks [ZLX*+*21] to capture the relative positions of nodes in the subgraph. After that, they learn node representation by applying a GNN model to the labeled subgraphs, and then they extract subgraph representation with a readout function [GSR*+*17] for prediction. Global methods learn node representation by directly applying a GNN model to the global graph and then make prediction based on head and tail node representation. Although there exists difference in local methods and global methods, all existing GNN-LP methods have a common characteristic that they adopt a symmetric learning strategy for node representation learning. In particular, they adopt the same class of GNN models to learn representation for both head and tail nodes. An illustration of existing representative GNN-LP methods is presented in Figure 1. Although existing GNN-LP methods have made much progress in learning expressive models, they suffer from scalability problem during training for large-scale graphs, which has attracted little attention by researchers.

In this paper, we propose a novel method, called asymmetric learning (AML), for GNN-LP. The contributions of this paper are listed as follows:

• •

We give computation complexity analysis of existing GNN-LP methods, which reveals that the scalability problem stems from their symmetric learning strategy adopting the same class of GNN models to learn representation for both head and tail nodes.

• •

AML is the first GNN-LP method adopting an asymmetric learning strategy for node representation learning.

• •

AML proposes a row-wise sampling strategy to generate mini-batch for training, which is a necessary component to make the asymmetric learning strategy work for training speedup.

• •

Experiments on three real large-scale datasets show that AML is $1.7{\times}{\sim}7.3{\times}$ faster in training than baselines with a symmetric learning strategy, while having almost no accuracy loss.

2 Preliminary

In this section, we introduce notations and some related works for link prediction.

Notations.

We use a boldface uppercase letter, such as ${\bf B}$, to denote a matrix. We use a boldface lowercase letter, such as ${\bf b}$, to denote a vector. ${\bf B}_{i*}$ and ${\bf B}_{*j}$ denote the $i$th row and the $j$th column of ${\bf B}$, respectively. ${\bf X}{\in}\mathbb{R}^{N\times u}$ denotes the node feature matrix, where $u$ is the feature dimension and $N$ is the number of nodes. ${\bf A}{\in}\{0,1\}^{N\times N}$ denotes the adjacency matrix of a graph $\mathcal{G}$. $A_{ij}{=}1$ iff there is an edge from node $i$ to node $j$, otherwise $A_{ij}{=}0$. $L$ denotes the number of layers for GNN models. $\mathcal{E}$ denotes the set of links for training. For a link $(i,j)$, we call node $i$ a head node and call node $j$ a tail node.

Graph Neural Network.

GNN [GMS05, SGT*+*09, BZSL14, KW17, HYL17, VCC*+*18, JLM*+*22] is a class of models for learning over graph data. In GNN, nodes can iteratively encode their first-order and high-order neighbor information in the graph through message passing between neighbor nodes [GSR*+*17]. Due to the iteratively dependent nature, the computation complexity for a node exponentially increases with iterations. Although some works [ZHW*+*19, ZZS*+*20, YL21, GWJ18] propose solutions for the above problem of exponential complexity, the computation complexity of GNN is still much higher than that of a multi-layer perceptron (MLP).

Graph Neural Network based Link Prediction.

Benefited from the powerful ability of GNN in modeling graph data, GNN-LP methods are more expressive than traditional non-GNN based learning algorithms in node representation learning. GNN-LP methods include two major classes, local methods and global methods . For local methods, different methods vary in the labeling tricks they use, which mainly include double radius node labeling (DRNL) [ZC18], distance encoding (DE) [LWWL20], partial zero-one labeling trick [YSYL21] and zero-one labeling trick [ZLX*+*21]. As shown in [ZLX*+*21], local methods with DRNL and DE perform better than other methods. However, one bottleneck for DRNL and DE is that they need to compute the shortest path distance (SPD) between target nodes and other nodes in subgraphs, and computing SPD is time-consuming during the training process [ZLX*+*21]. Although we can compute SPD in the preprocessing step, costly storage overhead for subgraphs occurs instead. For global methods [KW16, PHL*+*18, HHN*+*19, YKL*+*21], they mainly apply different GNN models on the global graphs to generate node representation. Almost all existing GNN-LP methods, including both local methods and global methods, adopt a symmetric learning strategy which utilizes the same class of GNN models to learn representation for both head and tail nodes.

Non-GNN based Methods.

Besides GNN-LP methods, WLNM [ZC17] and SUREL [YZW*+*22] also show competitive performance in accuracy. The main difference between them and GNN-LP methods is that they do not apply GNN to learn from graphs or subgraphs. For example, WLNM applies an MLP model to learn subgraph representation from the adjacency matrices of extracted subgraphs. SUREL proposes an alternative sampler for subgraph extraction and applies a sequential model, like recurrent neural networks (RNNs), to learn subgraph representation.

In general, GNN-LP methods have higher accuracy than non-GNN based methods [ZLX*+*21], but suffer from scalability problem during training for large-scale graphs. Although global GNN-LP methods are more efficient than local GNN-LP methods, they still have high computation complexity in training due to the adopted symmetric learning strategy. This motivates our work in this paper.

3 Asymmetric Learning for GNN-LP

Like most deep learning methods, GNN-LP methods are typically trained in a mini-batch manner. Suppose the number of links in the training set is $|\mathcal{E}|$. Then existing GNN-LP methods with symmetric learning need to perform $2|\mathcal{E}|$ times of GNN computation within each epoch. In particular, $|\mathcal{E}|$ times of GNN computation are for head nodes and another $|\mathcal{E}|$ times of GNN computation are for tail nodes. It is easy to verify that $|\mathcal{E}|$ times of GNN computation are inevitable for both head and tail node representation learning with a symmetric learning strategy. Since GNN is of exponential computation complexity and $|\mathcal{E}|$ is of a considerably large value, existing GNN-LP methods incur a huge computation burden for large-scale graphs.

To solve the scalability problem caused by symmetric learning, we propose AML which is illustrated in Figure 2. The main idea of AML is to adopt a GNN model for learning head node representation while using a multi-layer perceptron (MLP) model for learning tail node representation. Meanwhile, AML pre-encodes graph structure to avoid information loss for MLP. The following content of this section will present the details of AML.

3.1 Node Representation Learning with AML

We use ${\bf U}{\in}\mathbb{R}^{N\times r}$ to denote the representation of all head nodes and use ${\bf V}{\in}\mathbb{R}^{N\times r}$ to denote the representation of all tail nodes, where each row of ${\bf U}$ and ${\bf V}$ corresponds to a node. We apply a GNN model to learn representation for head nodes while using an MLP model to learn representation for tail nodes222We can also apply a GNN model to learn representation for tail nodes while using an MLP model to learn representation for head nodes. For convenience, we only present the technical details of one case. The technical details for the reversed case are the same as the presented one..

We take SAGE [HYL17], one of the most representative GNN models, as an example to describe the details. Note that other GNN models can also be used in AML. Let $\hat{{\bf A}}$ denote the normalization of ${\bf A}$, which can be row-normalization, column-normalization or symmetric normalization. Formulas for one layer in SAGE are as follows:

[TABLE]

where $\ell$ is layer number, $f(\cdot)$ is an activation function, ${\bf W}^{(\ell)}_{1}{\in}\mathbb{R}^{r\times r}$ and ${\bf W}^{(\ell)}_{2}{\in}\mathbb{R}^{r\times r}$ are learnable parameters at layer $\ell$. ${\bf U}^{(\ell)}{\in}\mathbb{R}^{N\times r}$ is the node representation at layer $\ell$ and ${\bf U}^{(0)}{=}{\bf X}$. In (1), node $i$ encodes neighbor information via $\hat{{\bf A}}_{i*}{\bf U}^{(\ell-1)}{\bf W}^{(\ell)}_{1}$ and updates its own message together with ${\bf U}^{(\ell-1)}_{i*}{\bf W}^{(\ell)}_{2}$.

Different from existing GNN-LP methods which adopt the same class of GNN models to learn representation for both head and tail nodes, AML applies an MLP model to learn representation for tail nodes. However, naively applying MLP for tail nodes will deteriorate the accuracy. Hence, as in [KWG19, WJZ*+*19, CWD*+*20], we first pre-encode graph structure information into node features in the preprocessing step. The pre-encoding step is as follows:

[TABLE]

To further improve the representation learning for tail nodes, we propose to transfer knowledge from head nodes to tail nodes by sharing parameters. The formula is as follows:

[TABLE]

where we perform knowledge transfer between head nodes and tail nodes by sharing ${\bf W}_{1}^{(\ell)}$ and ${\bf W}_{2}^{(\ell)}$. Since sharing parameters somewhat restricts the expressiveness of $\tilde{{\bf V}}^{(L)}$, we propose to apply an MLP model to learn over the residual $\Delta^{(0)}$. Adding the residual $\Delta^{(L)}$ to $\tilde{{\bf V}}^{(L)}$, we obtain the representation for tail nodes which is shown as follows:

[TABLE]

where ${\bf W}^{(\ell)}\in\mathbb{R}^{r\times r}$ is a learnable parameter at layer $\ell$.

Note that if BatchNorm [IS15] is applied to ${\bf U}^{(\ell)}$ and $\tilde{{\bf V}}^{(\ell)}$, we keep individual parameters of BatchNorm for ${\bf U}^{(\ell)}$ and $\tilde{{\bf V}}^{(\ell)}$. Since ${\bf U}^{(\ell)}$ and $\tilde{{\bf V}}^{(\ell)}$ have different scales and lie in different representation spaces, it is reasonable to keep individual parameters of BatchNorm for them. Here pre-encoding step only needs to be performed once and the resulting $\hat{{\bf A}}^{L}{\bf X}$ in (2) can be saved for the entire training process.

3.2 Learning Objective of AML

According to the definitions in (1) and (6), modeling links with ${\bf U}^{(L)}$ and ${\bf V}^{(L)}$ can capture directed relations while being a little difficult to capture undirected relations. Our solutions are twofold. Firstly, we formulate each undirected link by two directed ones with opposite directions. Secondly, motivated by the work in [LYZ11], in AML we propose to model both homophily and stochastic equivalence [Hof07]. As a result, the formulas for the prediction of a pair $(i,j)$ are as follows:

[TABLE]

where ${\bf U}$ and ${\bf V}$ are representation for head and tail nodes, respectively. ${\bf V}^{(L)}$ in ${\bf U}$ is included to model the homophily feature in graph data. $S_{ij}$ is the prediction for the pair $(i,j)$. $\odot$ denotes element-wise multiplication. $f_{\bm{\Theta}}(\cdot)$ is an MLP model with parameter $\bm{\Theta}$.

Given node representation ${\bf U}$ and ${\bf V}$, the learning objective of link prediction is as follows:

[TABLE]

where $\mathcal{E}$ denotes the training set. $\mathcal{W}$ is the set of all learnable parameters. $Y_{ij}$ is the ground-truth label for the pair $(i,j)$. $f_{loss}(\cdot,\cdot)$ is a loss function, such as cross-entropy loss. $\lambda$ is a coefficient for the regularization term of $\mathcal{W}$. $\lVert\cdot\rVert_{F}$ denotes the Frobenius norm of a matrix.

3.3 Row-wise Sampling for Generating Mini-Batch

Like most deep learning methods, existing GNN-LP methods are typically trained in a mini-batch manner. But existing GNN-LP methods adopt an edge-wise sampling strategy to generate mini-batch for training. In particular, they first randomly sample a mini-batch of edges $\mathcal{E}_{\Omega}$ from $\mathcal{E}$ at each iteration and then optimize the objective function based on $\mathcal{E}_{\Omega}$. For example, if we adopt the edge-wise sampling strategy for the objective function of AML in (9), the corresponding learning objective at each iteration will be as follows:

[TABLE]

Suppose $\mathcal{E}_{\Omega}{=}\{(i_{1},j_{1}),{\cdots},(i_{B},j_{B})\}$ with $B$ as the mini-batch size. We respectively use $C$ and $M$ to denote the computation complexity for generating a node representation by GNN and MLP. Since $\mathcal{E}_{\Omega}$ is edge-wise randomly sampled from $\mathcal{E}$ and $N$ is of a large value for large-scale graphs, $(i_{1},{\cdots},i_{B})$ will have a relative small number of repeated nodes. Then the edge-wise sampling strategy for AML has a computation complexity of $\mathcal{O}(|\mathcal{E}|{\cdot}(C{+}M))$ for each epoch, which has the same order of magnitude in computation complexity as existing GNN-LP methods and hence is undesirable.

To solve this high computation complexity problem of edge-wise sampling strategy adopted by existing GNN-LP methods, in AML we propose a row-wise sampling strategy for generating mini-batch. More specifically, we first sample a number of row indices $\mathcal{V}_{\Omega}$ (head nodes) from $\{1,2,\cdots,N\}$ for each mini-batch iteration. Then we construct the mini-batch $\mathcal{E}_{\Omega}$ as follows:

[TABLE]

By using this row-wise sampling strategy, AML has a computation complexity of $\mathcal{O}(|\mathcal{V}_{\Omega}|{\cdot}C+(|\mathcal{V}_{\Omega}|/N)|\mathcal{E}|\cdot M)$ for each mini-batch iteration. Since $\mathcal{V}_{\Omega}$ iterates over $\{1,{\cdots},N\}$ for $N/|\mathcal{V}_{\Omega}|$ times to go through the whole training set, this row-wise sampling strategy for AML has a computation complexity of $\mathcal{O}(N{\cdot}C{+}|\mathcal{E}|{\cdot}M)$ for each epoch. Therefore, the row-wise sampling strategy enables AML to decouple the factor of $|\mathcal{E}|$ from the computation complexity of GNN, leading to a complexity reduction by orders of magnitude compared to the edge-wise sampling strategy.

Algorithm 1 summarizes the whole learning algorithm for AML.

3.4 Complexity Analysis

The computation complexity for different methods are summarized in Table 1. For large-scale graphs, we often have $s^{L}{<}N$ and $s^{k}{<}N$. Typically, $s^{L}r^{2}$ has the same order of magnitude as $Ls^{k}r^{2}$. According to (1), GNN has a computation complexity of $C{=}\mathcal{O}\left(s^{L}{\cdot}r^{2}\right)$ to generate a node representation while the corresponding computation complexity of MLP is $M{=}\mathcal{O}\left(L{\cdot}r^{2}\right)$ according to (3)-(6). It is easy to verify that $M{\ll}C$. Although many works [HYL17, ZHW*+*19, ZZS*+*20, YL21] have proposed solutions to reduce $C$, $C$ is still much larger than $M$. Here we suppose all methods are trained in a mini-batch manner which has been adopted by almost all deep learning models including GNN.

From Table 1, we can get the following results. Firstly, AML has a computation complexity of $\mathcal{O}(N{\cdot}C{+}|\mathcal{E}|{\cdot}M)$, which is much lower than $\mathcal{O}(2|\mathcal{E}|{\cdot}C)$ of existing GNN-LP methods. Secondly, even with our proposed row-wise sampling strategy, existing GNN-LP methods still have a computation complexity of $\mathcal{O}((|\mathcal{E}|{+}N){\cdot}C)$, without change in the order of magnitude. The reason is that they still need to perform $|\mathcal{E}|$ times of GNN computation for tail nodes within each epoch.

4 Experiments

In this section, we evaluate AML and baselines on three real datasets. All methods are implemented with Pytorch [PGM*+*19] and Pytorch-Geometric Library [FL19]. All experiments are run on an NVIDIA RTX A6000 GPU server with 48 GB of graphics memory.

4.1 Settings

Datasets

Datasets for evaluation include ogbl-collab333In ogbl-collab, there are data leakage issues in the provided graph adjacency matrix ${\bf A}$. We remove those positive links in the validation and testing set from ${\bf A}$., ogbl-ppa and ogbl-citation2444https://ogb.stanford.edu/docs/linkprop/. The first two are medium-scale datasets with hundreds of thousands of nodes. The last one is a large-scale dataset with millions of nodes. For ogbl-ppa, since the provided node features are uninformative, we apply matrix factorization [ME11] to generate new features for nodes. The first two datasets are for undirected link prediction, while the last one is for directed link prediction. The statistics of datasets are summarized in Table 2. Since most GNN-LP methods adopt the evaluation metrics provided by [HFZ*+*20], we also follow these evaluation settings.

Baselines

AML is actually a global GNN-LP method. We first compare AML with existing global GNN-LP baselines by adopting the same GNN for both AML and baselines. Since almost all existing global GNN-LP methods are developed based on the graph autoencoder framework proposed in [KW16], we mainly adopt the GNNs under the graph autoencoder framework. In particular, we adopt SAGE [HYL17] and GAT [VCC*+*18] as the GNNs for both AML and baselines, because SAGE and GAT are respectively representative non-attention based and attention based GNN models under the graph autoencoder framework.

Then we compare AML with non-GNN baselines and local GNN-LP baselines. Non-GNN baselines include common neighbors (CN) [New01], Adamic-Adar (AA) [AA03], Node2vec [GL16] and matrix factorization (MF) [ME11]. Local GNN-LP baselines inlcude DE-GNN [LWWL20] and SEAL [ZC18, ZLX*+*21]. For local GNN-LP methods, we extract enclosed subgraphs in an online way to simulate large-scale settings by following the original work [ZLX*+*21].

Hyper-parameter Settings

Hyper-parameters include $L$ (layer number), $r$ (hidden dimension), $\lambda$ (coefficient for the regularization of parameters), $T$ (maximum number of epoches), $\eta$ (learning rate) and $B$ (mini-batch size). On ogbl-collab, $L{=}3$, $r{=}256$, $\lambda{=}0$, $T{=}400$, $\eta{=}0.001$ and $B{=}65,536$. On ogbl-ppa, $L{=}3$, $r{=}256$, $\lambda{=}0$, $T{=}50$, $\eta{=}0.01$ and $B{=}65,536$. On ogbl-citation2, $L{=}3$, $r{=}256$, $\lambda{=}0$, $T{=}50$, $\eta{=}0.005$ and $B{=}65,536$. We use Adam [KB15] as the optimizer. We use GraphNorm [CLX*+*21] to accelerate the training. We adopt BNS [YL21] as the neighbor sampling strategy for large-scale training. In BNS, there are three hyper-parameters, including $\tilde{s}_{0}$, $\tilde{s}_{1}$ and $\delta$. $\tilde{s}_{0}$ denotes the number of sampled neighbors for nodes at output layer. $\tilde{s}_{1}$ denotes the number of sampled neighbors for nodes at other layers. $\delta$ denotes the ratio of blocked neighbors. On ogbl-collab, $\tilde{s}_{0}{=}7$, $\tilde{s}_{1}{=}2$, $\delta{=}1/2$. On ogbl-ppa, $\tilde{s}_{0}$ equals the number of all neighbors, $\tilde{s}_{1}{=}4$, $\delta{=}5/6$. On ogbl-citation2, $\tilde{s}_{0}{=}7$, $\tilde{s}_{1}{=}2$, $\delta{=}1/2$. We run each setting for 5 times and report the mean with standard deviation.

4.2 Comparison with Global GNN-LP Baselines

Comparison with global GNN-LP baselines is shown in Table 3 and Figure 3. According to the results, we can draw the following conclusions. Firstly, AML has almost no accuracy loss in all cases compared to baselines555When AML’s accuracy is within the standard deviation of baselines’ accuracy, we say that AML achieves almost no accuracy loss compared to baselines.. For example, AML’s accuracy is within the standard deviation of baselines’ accuracy on ogbl-ppa and ogbl-citation2. Instead, AML achieves an accuracy gain of $1.17\%{\sim}2.69\%$ on ogbl-collab compared to baselines. Secondly, AML is about $1.7{\times}{\sim}7.3{\times}$ faster in training than baselines when achieving almost no accuracy loss. For example, AML is $1.7{\times}$ faster than baselines on ogbl-collab, $3.1{\times}$ faster on ogbl-ppa, and $7.3{\times}$ faster on ogbl-citation2. In particular, baselines need about 5.8 days and 4.2 days to get the mean of results while AML only needs 1.7 days and 0.6 days to get the mean of results, on ogbl-ppa and ogbl-citation2 respectively. Thirdly, the speedup of AML relative to baselines increases with the size of graphs. For example, the number of nodes in ogbl-collab, ogbl-ppa and ogbl-citation2 increases in an ascending order, and the speedup of AML relative to baselines increases in a consistent order on these three graph datasets. Finally, AML has a better accuracy-time trade-off than baselines, which can be concluded from Figure 3. For example, we can find that AML is faster than baselines when achieving the same accuracy.

4.3 Comparison with non-GNN and Local GNN-LP Baselines

Comparison with non-GNN and local GNN-LP baselines is shown in Table 4. According to the results, we can draw the following conclusions. Firstly, AML is comparable with state-of-the-art local GNN-LP methods in accuracy. For example, AML has almost no accuracy loss compared to the best baseline DE-GNN on obgl-collab. On ogbl-ppa, AML achieves an accuracy gain of $3.43\%$ compared to the best baseline SEAL. On ogbl-citation2, AML gets an accuracy loss less than $1\%$ compared to the best baseline SEAL. Secondly, AML is about $13{\times}{\sim}110{\times}$ faster than local GNN-LP baselines. For example, AML is about $13{\times}$ faster than DE-GNN and SEAL on ogbl-collab, about $66{\times}$ faster on ogbl-ppa, and about $110{\times}$ faster on ogbl-citation2. In particular, DE-GNN and SEAL need 3.5 days to get the mean accuracy on ogbl-collab which is a relatively small-scale dataset. By contrast, AML only needs 1.2 hours to get the mean accuracy. Finally, GNN-LP methods can achieve better accuracy than non-GNN methods. For example, DE-GNN improves by $13\%$ in accuracy over MF on ogbl-ppa and improves by $17\%$ over Node2vec on ogbl-citation2. AML improves by $17\%$ in accuracy over MF on ogbl-ppa and improves by $16\%$ over Node2vec on ogbl-citation2.

4.4 Necessity of GNN in AML

Here we perform experiment to verify the necessity of GNN in AML. We design a method called symmetric MLP (SMLP) which applies MLP with pre-encoding to learn representation for both head nodes and tail nodes. More specifically, SMLP adopts the techniques for learning tail node representation in AML, i.e., (2) and (3), to learn representation for both head and tail nodes.

Results are shown in Table 5. We can find that SMLP is much worse than AML on all datasets. This shows that training a GNN model for node representation learning is necessary to achieve high accuracy.

4.5 Reversed Asymmetric Learning

In above experiments, AML learns representation for head nodes with a GNN model while learning representation for tail nodes with an MLP model. Here we verify whether the reversed case can also behave well. We denote the reversed case as AML-R, which learns representation for head nodes with an MLP model while learning representation for tail nodes with a GNN model. Results are shown in Table 6. Results show that AML and AML-R have similar accuracy.

4.6 Ablation Study

In this subsection, we study the effectiveness of different components in AML, including knowledge transfer, residual term $\Delta^{(L)}$, pre-encoding graph structure and modeling the homophily. Results are shown in Table 7. We can find the following phenomenons. Firstly, knowledge transfer between head and tail nodes effectively improves the accuracy of AML. For example, knowedge transfer can improve the accuracy of AML by $3.00\%$ on ogbl-collab, by $2.85\%$ on ogbl-ppa and by $0.42\%$ on ogbl-citation2. Secondly, $\Delta^{(L)}$ is beneficial for AML. For example, including $\Delta^{(L)}$ in AML can improve accuracy by $3.04\%$ on ogbl-collab, by $3.19\%$ on ogbl-ppa and by $0.40\%$ on ogbl-citation2. Thirdly, modeling homophily is helpful for AML. For example, modeling homophily in AML can improve accuracy by $1.87\%$ on ogbl-collab, by $2.37\%$ on ogbl-ppa and by $0.87\%$ on ogbl-citation2. Finally, pre-encoding graph structure plays a crucial role in AML. For example, AML has an accuracy loss of about $13\%$ on ogbl-collab, $47\%$ on ogbl-ppa and $25\%$ on ogbl-citation2 without pre-encoding graph structure.

5 Conclusions

Graph neural network based link prediction (GNN-LP) methods have achieved better accuracy than non-GNN based link prediction methods, but suffer from scalability problem for large-scale graphs. Our computation complexity analysis reveals that the scalability problem of existing GNN-LP methods stems from their symmetric learning strategy for node representation learning. Motivated by this finding, we propose a novel method called AML for GNN-LP. To the best of our knowledge, AML is the first GNN-LP method adopting an asymmetric learning strategy for node representation learning. Extensive experiments show that AML is significantly faster than baselines with a symmetric learning strategy while having almost no accuracy loss.