Link Prediction with Contextualized Self-Supervision

Traditional link prediction methods typically infer link existence by calculating some heuristic metrics on graphs, such as Common Neighbors [5], Jaccard Index [16], Preferential Attachment Index [17], Adamic-Adar Index [6], Katz Index [7], Rooted PageRank [18], and SimRank [19]. However, heuristic methods simply define rigid and one-sided structural measures to infer the link existence between node pairs, failing to capture the richer structural context and the essential information conveyed by node attributes. Recently, learning based methods have been proposed to advance link prediction under two categories.

Self-supervised learning (SSL) is a promising learning paradigm for deep neural networks (DNNs) in the domains of computer vision [35, 36] and natural language processing [37, 38]. The central theme of SSL is focused on defining self-supervised pretext tasks to assist DNNs in learning more generalized and robust data representations from the unlabeled data. The training strategies fall into two categories. 1) Pretraining: A DNN model is first pretrained with self-supervised pretext tasks and then finetuned with the target supervision task. 2) Joint training: The self-supervised tasks and the target supervision task are jointly trained with a multi-task learning objective.

Recently, SSL has been explored to boost the performance of machine learning tasks on graph-structured data through defining various pretext tasks [39, 40]. Inspired by image inpainting in computer vision [41], node attribute reconstruction has been utilized to enhance the effectiveness of node embeddings learned by GCNs [12, 40, 42, 43]. On the other hand, node structural properties have also been leveraged to design pretext tasks, such as node degree [40], structural distance between nodes [40], graph partition affiliations of nodes [44, 12], node subgraph motifs [45], and node centrality [46]. The above-mentioned self-supervised pretext tasks are designed with the pipeline of predicting attribute/structure properties from the learned node embeddings, which mainly favor the node classification task. Pioneered by DIG [47], another line of graph self-supervised learning focuses on maximizing the mutual information between node embeddings and the embeddings of their high-level graph summaries. DIG [47] and MVGRL [48] achieve this objective by contrasting node embeddings and the global graph representation. Extensions are also explored to contrast node embeddings and the embeddings of their surrounding subgraphs [49, 42, 50, 51, 52]. Moreover, the idea has been generalized to improve graph representation learning by contrasting local subgraph embeddings and the global graph representation [53, 54, 55].

To date, the current graph self-supervised learning algorithms are primarily designed to favor node-level and graph-level classification tasks. However, how to leverage self-supervision to facilitate link prediction is still under-explored. To fill this research gap, our work thoroughly analyses the roles of structural context in link-forming mechanisms on graphs and explicitly exploits structural context prediction as self-supervision to boost the link prediction performance.

Following traditional learning based algorithms [30, 32], we cast the link prediction problem as a binary classification task on edge embeddings constructed from the attributes of paired nodes. With the carefully designed architecture, the constructed edge embeddings well capture node attribute difference/similarity, exerting the power of the homophily property [15] towards high-quality link prediction. The overall pipeline from node attributes to the link existence likelihood prediction is end-to-end, which can be taken as a generalization of the existing learning based attribute-driven link prediction methods [30, 32].

For nodes $v_{i}$ and $v_{j}$, their node embeddings are respectively constructed from their attribute vectors $\bm{x}_{i},\bm{x}_{j}\in\mathbb{R}^{m}$:

where $\bm{\Phi}(v_{i})$ and $\bm{\Phi}(v_{j})\in\mathbb{R}^{d}$ are the embeddings of node $v_{i}$ and node $v_{j}$ with $d$ dimensions, $\bm{W}_{emb}\in\mathbb{R}^{d\times m}$ is the weight matrix for constructing node embeddings from node attributes, and $\sigma(\cdot)$ is the sigmoid function. From node embeddings $\bm{\Phi}(v_{i})$ and $\bm{\Phi}(v_{j})$, we construct edge embedding $\bm{\Phi}(\langle v_{i},v_{j}\rangle)\in\mathbb{R}^{d}$ for edge $\langle v_{i},v_{j}\rangle$ with the following aggregation operators [21]:

where $\bm{\Phi}(\langle v_{i},v_{j}\rangle)_{r}$, $\bm{\Phi}(v_{i})_{r}$ and $\bm{\Phi}(v_{j})_{r}$ indicate the $r$-th dimension of the embedding vectors $\bm{\Phi}(\langle v_{i},v_{j}\rangle)$, $\bm{\Phi}(v_{i})$ and $\bm{\Phi}(v_{j})$, respectively.

On the constructed edge embedding $\bm{\Phi}(\langle v_{i},v_{j}\rangle)$, the link existence probability $\hat{y}_{v_{i},v_{j}}$ between nodes $v_{i}$ and $v_{j}$ is predicted as

where $\bm{W}_{link}\in\mathbb{R}^{1\times d}$ is the weight matrix for predicting the link existence probability. The link prediction task aims to minimize the following loss:

where $\ell(\cdot,\cdot)$ is the binary cross entropy loss and $y_{v_{i},v_{j}}\in\{0,1\}$ is the ground-truth edge existence label, where 1 indicates nodes $v_{i}$ and $v_{j}$ are connected and 0 indicates otherwise. To enable the model to learn non-linkage patterns, we augment the training set with a set of negative edges $\mathcal{E}_{neg}$ with size $|\mathcal{E}_{tr}|$, where each element is the sampled node pair with no observed links between them in $\mathcal{G}_{tr}$, and is labeled with 0.

To enable the learned node and edge embeddings to better capture specific contextual information beneficial to link prediction, we incorporate the contextualized self-supervised learning task to predict structural context. Specifically, we use node embeddings $\bm{\Phi}(v_{i})$ and $\bm{\Phi}(v_{j})$ to respectively predict their structural contexts. For nodes $v_{i}$ and $v_{j}$, we respectively sample their structural contexts as $c(v_{i})$ and $c(v_{j})$. The existence probability of $c(v_{i})$ as node $v_{i}$’s structural context, $\hat{z}_{v_{i},c(v_{i})}\in[0,1]$, is modeled as

where $\bm{\Psi}(c(v_{i}))$ is the embedding of structural context $c(v_{i})$. Similarly, the existence probability of structural context $c(v_{j})$ as node $v_{j}$’s structural context, $\hat{z}_{v_{j},c(v_{j})}\in[0,1]$, is modeled as

where $\bm{\Psi}(c(v_{j}))$ is the embedding of structural context $c(v_{j})$.

To explore the contextual information, we consider two specific strategies to sample structural contexts $c(v_{i})$ and $c(v_{j})$ for nodes $v_{i}$ and $v_{j}$. The details are provided below:

• •

  Context/neighboring nodes. Following DeepWalk [20], we consider nodes occurring in random walk context as structural context. To construct context node embeddings, a possible solution is to perform a transformation on context node attributes similarly using Eq. (1). However, the attribute proximity among context nodes would override the structural proximity between nodes $v_{i}$ and $v_{j}$ reflected by their common neighbors. This could be exacerbated in case of attribute noise. To retain context node identity, as a better way to construct context node embeddings, we perform a linear transformation on one-hot encodings of node IDs, which is equivalent to looking up a learnable embedding table $\bm{\Psi}\in\mathbb{R}^{|\mathcal{V}|\times d}$ with node IDs.

• •

  Context subgraphs. For each node, we also consider its surrounding subgraphs as structural context, as they provide a macroscopic view to measure the structural similarity between nodes. Following [56], we use a short random walk with fixed length starting from each node as a sample of its context subgraph. For a sampled context subgraph $\mathcal{G}_{c}$, its context embedding is constructed by summing the context embeddings of its nodes:

  $$\bm{\Psi}(\mathcal{G}_{c})=\sum_{v\in\mathcal{G}_{c}}\bm{\Psi}(v),$$

  where $\bm{\Psi}(v)$ is constructed via looking up a learnable embedding table $\bm{\Psi}\in\mathbb{R}^{|\mathcal{V}|\times d}$ with $v$’s node ID.

For node pair $\langle v_{i},v_{j}\rangle$, the contextualized self-supervised learning task aims to minimize the following loss:

where $z_{v_{i},c(v_{i})}$ and $z_{v_{j},c(v_{j})}\in\{0,1\}$ are the ground-truth labels that respectively indicate whether $c(v_{i})$ and ${c(v_{j})}$ are nodes $v_{i}$’s and $v_{j}$’s structural contexts, with 1 for true and 0 for false. To allow the contextualized self-supervised learning task to learn non-context patterns, for each sampled positive structual context, we sample $k$ negative contexts with label 0. Negative context nodes are sampled as nodes not occurring in random walk context and negative context subgraphs are sampled as the sets of negative context nodes with the same node number as positive context subgraphs, which are unnecessarily connected in the training graph.

To boost the link prediction task, the contextualized self-supervised learning task can be leveraged with two training strategies as follows:

• •

  Joint Training. We jointly train the supervised link prediction task and the self-supervised learning task by minimizing the joint loss:

  $$\mathcal{L}_{v_{i},v_{j}}=\mathcal{L}_{v_{i},v_{j}}^{link}+\mathcal{L}_{v_{i},v_{j}}^{context}.$$

  (8)

• •

  Pretraining. We can also first train the weight matrix $\bm{W}_{emb}$ used for constructing node embeddings by minimizing the contextualized self-supervised learning loss in Eq. (7), and then finetune the model paremeters $\bm{W}_{emb}$ and $\bm{W}_{link}$ by minimizing the link prediction loss in Eq. (4).

Algorithm 1 details the training procedure of CSSL with the joint training strategy. In Step 1, we sample a list of positive and negative structural contexts for each node $v_{i}\in\mathcal{V}$ with the following procedure: we first generate $\gamma$ random walks with length $l$ starting from node $v_{i}$; then nodes that occur following the starting node $v_{i}$ in random walks are collected as its positive context nodes with size $\gamma\cdot(l-1)$, or $\gamma$ positive context subgraphs are sampled for it by assembling its $l-1$ context nodes along each random walk; finally we sample $k$ negative context nodes/subgraphs for each positive context node/subgraph, and $v_{i}$’s context node/subgraph list is formed by combining all sampled positive and negative context nodes/subgraphs. In Step 2, the algorithm samples a set of negative edges with size $|\mathcal{E}_{tr}|$ that are not observed in $\mathcal{G}_{tr}$. Then, the model is trained with stochastic gradient descent by sampling a minibatch of edges iteratively in Steps 3-12, where at each iteration, before performing stochastic gradient descent, for each sampled edge $\langle v_{i},v_{j}\rangle$, we respectively sample a context node/subgraph $c(v_{i})$ and $c(c_{j})$ together with the corresponding context label $z_{v_{i},c(v_{i})}$ and $z_{v_{j},c(v_{j})}$ for nodes $v_{i}$ and $v_{j}$ from their context node/subgraph lists. Finally, the existence of the unobserved links are inferred with the learned model parameters in Step 13.

Algorithm 2 details the training procedure of CSSL with the pretraining strategy. The algorithm first samples a list of positive and negative structural contexts for each node in $\mathcal{V}$ (Step 1), as well as a set of negative edges with size $|\mathcal{E}_{tr}|$ that are not observed in $\mathcal{G}_{tr}$ (Step 2), as is done in Algorithm 1. Then, in Steps 3-11, the model is pretrained with stochastic gradient descent by sampling a minibatch of edges and descending their context prediction loss at each iteration. After that, in Steps 12-16, the model is finetuned with the supervised link prediction task, by sampling a minibatch of edges and descending their link classification loss iteratively. Finally, in Step 17, the existence of the unobserved links is inferred with the learned model parameters.

Seven real-world networks are used in the experiments. The detailed statistics are summarized in Table I.

• •

  Cora and Citeseer¹¹1https://linqs.soe.ucsc.edu/data are two citation networks where nodes represent papers and links represent the citations between papers. Each paper is described by a fixed-dimension binary bag-of-words attribute vector, with each dimension indicating the occurrence/absence of each word from a fixed-size dictionary; the value of 1 indicates the corresponding word occurs in the paper, and 0 otherwise.

• •

  WebKB and Wiki\footreffn:linqs are two webpage networks where nodes indicate webpages and links indicate hyperlinks among webpages. Node attributes are binary bag-of-words representations of webpages.

• •

  Facebook²²2https://snap.stanford.edu/data/ego-Facebook.html and Google+³³3https://snap.stanford.edu/data/ego-Gplus.html are two social networks respectively formed by the combination of a group of Facebook and Google+ ego networks. A set of user profile values are used to construct binary node attribute vectors according to their presence/absence.

• •

  DBLP is the subgraph of the DBLP bibliographic network⁴⁴4https://aminer.org/citation (version 3), formed by the papers in the area of Database and Data Mining and the citations between them. Paper titles are used to construct binary bag-of-words node attribute vectors. The publication years of papers are also collected as their time stamps.

On Cora, Citeseer, WebKB, Wiki, and Facebook, we perform transductive link prediction experiments. To ensure the training links are sparse (with roughly equal numbers of links and nodes) and prevent generating too many isolated nodes, we randomly split all links into training, validation and test sets, with the ratios as 45%/5%/50%, 67.5%/7.5%/25%, 54%/6%/40%, 9%/1%/90%, and 4.5%/0.5%/95%, respectively. The sizes of the three sets are given in Table I. Following the standard setting of learning-based link prediction, for each edge, we randomly sample one negative edge, which is not observed in the current network. We repeat the random split ten times to avoid bias in evaluating the link prediction performance. DBLP is used to evaluate the performance of predicting links for out-of-sample nodes in the inductive settings, by using paper publication years to partition in-sample and out-of-sample nodes. We use Google+ to evaluate the scalability of the proposed CSSL framework by extracting subnetworks with varying scales.

We compare the proposed CSSL link prediction framework with context/neighboring node prediction (CSSL_Neigh) and context subgraph prediction (CSSL_Subgraph) with three groups of competitive link prediction baselines:

We use the area under the ROC curve (AUC) to evaluate the link prediction performance of all methods. For GraphSAGE, Attri2Vec and CSSL variants, we start 10 random walks with length 5 from each node. For each starting node, the nodes that follow in random walks are collected as its context nodes/subgraphs. For CSSL variants, we sample $k=1$ negative context node/subgraph for each positive context node/subgraph. For all methods, embedding dimension is set to 128. GraphSAGE, VGAE and the attribute-oriented embedding component of DEAL use two hidden layers and the number of neurons at the first hidden layer is set to 256. Other parameters of baselines are set to their default values. To train the proposed CSSL joint training variants, we use 100 epochs and a batch size of 20. To train the proposed CSSL pretraining variants, we use 40 pretraining epochs, 100 finetuning epochs and a batch size of 20. CSSL uses the Weighted-L2 as the default aggregation operator to generate edge embeddings. For VGAE, SEAL, DEAL and CSSL variants, the validation edges are used to select the best epoch. For embedding based methods (GraphSAGE, Attri2Vec, DIG and GMI), the best aggregation operators are selected from Average, Hadamard, Weighted-L1 and Weighted-L2 defined in Eq. (2) with the validation edges.

We first compare the proposed CSSL framework with all baselines under the transductive link prediction setting. Table II reports the link prediction AUC scores of CSSL variants and baselines on Cora, Citeseer, WebKB, Wiki and Facebook. The best and the second best performers are respectively highlighted by bold and underline. On each network, we also conduct paired t-test between the best performer and its competitors. The methods that are significantly worse than the best performer at 0.05 level are marked with $\bullet$ ⁵⁵5Tables VI-X use the same notations..

Table II shows that the proposed CSSL framework achieves the best overall link prediction performance on the five networks. CSSL_Neigh performs best on Cora and Citeseer, while the best performers on WebKB and Wiki are respectively CSSL_Neigh_Pretr and CSSL_Subgraph_Pretr. CSSL_Neigh_Pretr performs slightly worse than the best performer VGAE on Facebook, while VGAE does not perform well on other four networks. The performance gain of CSSL_Neigh and CSSL_Subgraph over the counterpart without self-supervision (CSSL_Ablated) and the single task based link prediction baselines proves the usefulness of structural context prediction as a self-supervised learning task for link prediction, and also the ability of CSSL to mitigate the negative impact of link sparsity via self-supervision to reinforce link prediction.

The pretraining based CSSL variants perform comparably with joint training based counterparts. Among different self-supervised learning tasks, CSSL_Neigh and CSSL_Subgraph perform better than CSSL_Attr. This suggests that, as a self-supervised learning task, structural context prediction is more effective in boosting link prediction, as compared to the attribute prediction based task.

In reality, networks are often dynamically changing with new nodes constantly joining, which may have only very few or even no connections to the existing nodes. Therefore, we also compare the proposed CSSL framework with the baseline methods in terms of their abilities to infer links for out-of-sample nodes on DBLP. We split the DBLP network according to a threshold on the years when papers were published. Papers published before the threshold year together with citations between these papers form the in-sample network, and papers published after the threshold year are considered as out-of-sample nodes, with the threshold year varying from 2000 to 2009. The statistics of the splits are summarized in Table V. We train the link prediction model on in-sample networks and then predict the links of out-of-sample nodes. Two cases are considered for the inductive link prediction setting: 1) no links are observed connecting to out-of-sample nodes. Only Attri2Vec, DEAL and CSSL variants work in this case, and 2) sparse (10%) links of out-of-sample nodes are observed, including the links connecting to the in-sample training graph and the links among the out-of-sample nodes.

Tables V-V report the link prediction performance in the two cases, with the best and the second best performers highlighted by bold and underline, respectively. From Tables V-V, we can observe that the proposed CSSL variants (CSSL_Neigh_Joint and CSSL_Sugbgraph_Joint) consistently achieve the best overall performance in predicting the links for out-of-sample nodes in both two cases, while GraphSAGE, DIG, GMI and SEAL fail to achieve satisfactory performance when out-of-sample nodes have sparse neighborhood structure. By exploiting structural context prediction based self-supervised learning, our approach can learn the mapping from node attributes to link existence with stronger generalization ability than baseline methods, thus enabling its better inductive capacity.

To validate the robustness of CSSL against attribute noise, on Citeseer and WebKB, we randomly flip the binary attribute values with a ratio for each node. The rational behind is to simulate some real-world cases in social networks, where users might deliberately hide sensitive profile features or provide false profile metadata to protect their privacy, such as fabricating their gender or providing fake affiliations. We vary the ratio from 5% to 25% by an increment of 5%.

Tables VI-VII compare CSSL_Neigh and CSSL_Subgraph with baselines in terms of different attribute noise ratios on Citeseer and WebKB, where the best and the second best performers are respectively highlighted by bold and underline. As can be seen, CSSL_Neigh and CSSL_Subgraph consistently outperform other baselines with all attribute noise ratios on Citeseer and WebKB, except for the case with very high noise ratios of 20% and 25% on Citeseer. Still for the highly noise ratio, CSSL_Subgraph_Joint performs comparably with the best performer SEAL that mainly leverages network structure. SEAL learns a mapping from the local neighborhood subgraph of paired nodes to the link existence with a graph neural network [57], where the mapping is dominated by the convolution over node structural label encoding, so the link prediction performance of SEAL is not largely impacted by the increase in attribute noise levels. On the other hand, CSSL has the capability of extracting useful information from noisy node attributes to alleviate the negative impact of noise, leading to the overall better link prediction performance at different noise ratios.

We also investigate the convergence property of the proposed CSSL framework on Cora, as a case study. Fig. 4 plots the link prediction performance change of all the proposed CSSL variants (CSSL_Neigh_Joint, CSSL_Neigh_Pretr, CSSL_Subgraph_Joint and CSSL_Subgraph_Pretr), when the number of epochs used for joint training or finetuning after pretraining varies from 1 to 500. We can see that, the performance of all CSSL variants converges fast after no more than 25 epochs, and then retains at a stable level.

We also study the scalability of the proposed CSSL with an increase of network scale, by evaluating the CPU time consumed by running CSSL_Neigh_Joint with 10 epochs on the sampled Google+ subnetworks of different scales. Fig. 4 plots the change of the running time as network size (the number of edges) increases from 5000 to $10^{7}$, with both axes in logarithm scale. As we can see, the proposed CSSL scales almost linearly with network size $|\mathcal{E}_{tr}|$.

As reported, the proposed CSSL uses Weighted-L2 as the aggregation operator to form edge embeddings from paired node embeddings. Now, we compare its link prediction performance using Weighted-L2 against the other three operators, namely, Average, Hadamard and Weighted-L1. Table VIII compares the performance of CSSL variants using different aggregation operators on Cora. For each variant, the result of the best operator is highlighted by bold and paired t-test is used to compare it with other operators. In general, Weighted-L1 and Weighted-L2 deliver the best link prediction performance, while Hadamard performs better than Average. According to the homophily phenomenon, node attribute similarity/difference is the key evidence for predicting the link existence. As Weighted-L1 and Weighted-L2 are better at characterizing node attribute difference, this translates to their superior link prediction performance.

We also study the sensitivity of CSSL variants with respect to four important parameters, including the dimension $d$ of node/edge embeddings, the number of random walks $\gamma$ starting from each node, the random walk length $l$ for sampling structural context, and the number of sampled negative context nodes/subgraphs $k$ for each positive context. We take turns to fix any three parameters as default values and test the sensitivity of CSSL variants to the remaining one. Fig. 5 shows how CSSL variants perform on Cora by varying the four parameters. Compared to other variants, CSSL_Neigh_Joint is more sensitive to embedding dimension. Yet, all CSSL variants tend to have stable performance with respect to varying values of these parameters.

So far, we have validated the efficacy of the proposed CSSL framework on attributed networks. As a flexible learning framework, CSSL is also capable of handling non-attributed networks where node attributes are unavailable. In this case, CSSL takes one-hot node representations as node attributes to infer the link existence. The baseline methods DIG, GIM, VGAE, SEAL, and DEAL use the same one-hot representation scheme to construct node features. Furthermore, we compare CSSL variants with three competitive embedding based methods that use only network structure: Node2Vec [21], the first- and second-order LINE (LINE1 and LINE2) [22], and four heuristics-based methods with their definitions provided in Table IX.

Table X reports the link prediction results on five networks without node attributes, with the best and the second best performers highlighted by bold and underline, respectively. Clearly, as compared to Table II, significant performance drops can be observed when node attributes are not incorporated for all methods except for SEAL. As SEAL does not well leverage node attributes for link prediction, it even achieves better performance on Cora, Citeseer, WebKB and Wiki at the absence of node attributes.