Self-Supervised Learning of Contextual Embeddings for Link Prediction in Heterogeneous Networks

In this work, we generate the pre-training subgraphs $G_{C}$ for each node in the graph using a random walk strategy. Masking and predicting nodes in the random walks helps SLiCE learn the global connectivity patterns in $G$ during the pre-training phase. In fine tuning phase, the context subgraphs for link prediction are generated using following approaches: (1) Shortest Path strategy considers the shortest path between two nodes as the context. (2) Random strategy, on the other hand, generates contexts following the random walks between two nodes, limited to a pre-defined maximum number of hops. Note that the context generation strategies are generic and can be applied for generating contexts in many downstream tasks such as link prediction (Zhang and Chen, 2018), knowledge base completion (Socher et al., 2013) or multi-hop reasoning (Das et al., 2017; Hamilton et al., 2018).

In our experiments, context subgraphs are generated for each node $v$ in the graph during pre-training and for each node-pair during fine-tuning. Each generated subgraph $g_{c}\in G_{c}$ is encoded as a set of nodes denoted by $g_{c}=(v_{1},v_{2},\cdots,v_{|V_{c}|})$, where $|V_{c}|$ represents the number of nodes in $g_{c}$. Different from the sequential orders enforced on graph sampled using pre-defined metapaths, the order of nodes in this set is not important. Therefore, our context subgraphs are not limited to paths, and can handle tree or star-shaped subgraphs.

We first represent each node $v_{i}$ in the context subgraph as a low-dimensional vector representation by $h_{i}=\bm{W}_{e}f_{attr}(v_{i})$, where $f_{attr}(.)$ is a function that returns a stacked vector containing the structure-based embedding of $v_{i}$ and the embedding of its attributes. $\bm{W}_{e}$ is the learnable embedding matrix. We represent the input node embeddings in $g_{c}$ as $\bm{H}_{c}=(h_{1},h_{2},\cdots,h_{|V_{c}|})$. It is flexible to incorporate the node and relation attributes (if available) for attributed networks (Cen et al., 2019) in the low-dimensional representations or initialize them with the output embeddings learnt from other global feature generation approaches that capture the multi-relational graph structure (Grover and Leskovec, 2016; Dong et al., 2017; Wang et al., 2019a; Yun et al., 2019; Vashishth et al., 2020).

There are multiple approaches for generating global node features in heterogeneous networks (see “related work”). Our experiments show that the node embeddings obtained from random walk based skip-gram methods (RW-SG) produces competitive performance for link prediction tasks. Therefore, in the proposed SLiCE model, we mainly consider the pre-trained node representation vectors from node2vec for initialization of the node features.

Our model pre-training is performed by training the self-supervised contextualized node prediction task. More specifically, for each node in $G$, we generate the node context $g_{c}$ with diameter (defined as the largest shortest pair between any pair of nodes) using the aforementioned context generation methods and randomly mask a node for prediction based on the context subgraph. The graph structure is left unperturbed by the masking procedure. Therefore, the pre-training is learnt by maximizing the probability of observing this masked node $v_{m}$ based on the context $g_{c}$ in the following form.

where $\theta$ represents the set of model parameters. The procedure for pre-training is given in Algorithm 1. In this algorithm, lines 2-8 generate context subgraphs for nodes in the graph and further applies a random masking strategy to process the data for pre-training. Lines 9-10 learn the pre-trained global node features and initialize them as the node embeddings in SLiCE. In lines 13-15, we apply the contextual translation layers on the context subgraphs, aggregate the output of different layers as the contextual node embeddings in line 16 and update the model parameters with contextual node prediction task. In a departure from traditional skip-gram methods (Park et al., 2020) that predicts a node from the path prefix that precedes it in a random walk, our random masking strategy forces the model to learn higher-order relationships between nodes that are arbitrarily connected by variable length paths with diverse relational patterns.

The SLiCE model is further fine-tuned on the contextualized link prediction task by generating multiple fine-grained contexts for each specific node-pair that is under consideration for link prediction. Based on the predicted scores, this stage is trained by maximizing the probability of observing a positive edge ($e_{p}$) given context ($g_{cp}$), while also learning to assign low probability to negatively sampled edges ($e_{n}$) and their associated contexts ($g_{cn}$). The overall objective is obtained by summing over the training data subsets with positive edges ($D_{p}$) and negative edges ($D_{n}$). Algorithm 2 shows the process of fine-tuning step. In this algorithm, lines 2-7 generate context subgraphs of the node-pairs for link prediction task and process the data for fine-tuning in the same manner described in pre-training. Lines 8-18 perform the fine-tuning with link prediction task.

We compute the probability of the edge between two nodes $e=(v_{i},v_{j})$ as the similarity score $S(v_{i},v_{j})=\sigma(\tilde{h}_{i}^{T}\cdot\tilde{h}_{j})$ (Abu-El-Haija et al., 2018), where $\tilde{h}_{i}$ and $\tilde{h}_{j}$ are contextual embeddings of $v_{i}$ and $v_{j}$ learnt based on a context subgraph, respectively. $\sigma(\cdot)$ represents sigmoid function.

We use four public benchmark datasets covering multiple applications: e-commerce (Amazon), academic graph (DBLP), knowledge graphs (Freebase), and social networks (Twitter). We use the same data split for training, development, and testing as described in previous works (Cen et al., 2019; Abu-El-Haija et al., 2018; Vashishth et al., 2020). In addition, we also introduce a new knowledge graph from the publicly available real-world Medical Information Mart for Intensive Care (MIMIC) III dataset²²2https://mimic.physionet.org/ in the healthcare domain. We generated equal number of positive and negative edges for the link prediction task. Table 2 provides the basic statistics of all datasets. The details of each dataset is provided below.

• •

  Amazon³³3https://github.com/THUDM/GATNE/tree/master/data : includes the co-viewing and co-purchasing links between products. The edge types, also_bought and also_viewed, represent that products are co-bought or co-viewed by the same user, respectively.

• •

  DBLP⁴⁴4https://github.com/Jhy1993/HAN/tree/master/data: includes the relationships between papers, authors, venues, and terms. The edge types include paper_has_term, published_at and has_author.

• •

  Freebase⁵⁵5https://github.com/malllabiisc/CompGCN/tree/master/data_compressed: is a pruned version of FB15K with inverse relations removed. It includes links between people and their nationality, gender, profession, institution, place of birth/death, and other demographic features.

• •

  Twitter³: includes links between tweets users. The edge types included in the network are re-tweet, reply, mention, and friendship/follower.

• •

  Healthcare: includes relations between patients and their diagnosed medical conditions during each hospital admission along with relations to procedures and medications received. To ensure data quality, we use a 5-core setting, i.e., retaining nodes with at least five neighbors in the knowledge graph. The codes for generating this healthcare knowledge graph from MIMIC III dataset are also available at⁶⁶6https://github.com/pnnl/SLICE.

We compare our model against the following state-of-the-art network embedding learning methods. The first four methods learn static embeddings and the remaining methods learn contextual embeddings.

• •

  TransE (Bordes et al., 2013) treats the relations between nodes as the translation operations in a low-dimensional embedding space.

• •

  RefE (Chami et al., 2020) incorporates hyperbolic space and attention-based geometric transformations to learn the logical patterns of networks.

• •

  node2vec (Grover and Leskovec, 2016) is a random-walk based method that was developed for homogeneous networks.

• •

  metapath2vec (Dong et al., 2017) is an extension of node2vec that constrains random walks to specified metapaths in heterogeneous network.

• •

  GAN (Graph Attention Networks) (Abu-El-Haija et al., 2018) is a graph attention network for learning node embeddings based on the attention distribution over the graph walk context.

• •

  GATNE-T (General Attributed Multiplex HeTerogeneous Network Embedding) (Cen et al., 2019) is a metapath-constrained random-walk based method that learns relation-specific embeddings.

• •

  RGCN (Relational Graph Convolutional Networks)  (Schlichtkrull et al., 2018) learns multi-relational data characteristics by assigning a different weight for each relation.

• •

  CompGCN (Composition-based Graph Convolutional Networks) (Vashishth et al., 2020) jointly learns the embedding of nodes and relations for heterogeneous graph and updates a node representation with multiple composition functions.

• •

  HGT (Heterogeneous Graph Transformer)  (Hu et al., 2020) models the heterogeneity of graph by analyzing heterogeneous attention over each edge and learning dedicated embeddings for different types of edges and nodes. We adapt the released implementation of node classification task to perform link prediction task.

• •

  asp2vec (Multi-aspect network embedding)  (Park et al., 2020) captures the interactions of the pre-defined multiple aspects with aspect regularization and dynamically assigns a single aspect for each node based on the specific local context.

All evaluations were performed using NVIDIA Tesla P100 GPUs. The results of SLiCE are evaluated under the parameter settings described in Section 3.6. The results of all baselines are obtained with their original implementations. Note that for all baseline methods, the parameters not specially specified here are under the default settings. We use the implementation provided in KGEmb⁷⁷7https://github.com/HazyResearch/KGEmb for both TransE and RefE. node2vec⁸⁸8https://github.com/aditya-grover/node2vec is implemented by sampling 10 random walks with a length of 80. The original implementation of metapath2vec⁹⁹90https://ericdongyx.github.io/metapath2vec/m2v.html is used by generating 10 walks for each node as well. We set the learning rate to be {0.1, 0.01, 0.001} and reported the best performance for GAN¹⁰¹⁰10https://github.com/google-research/google-research/tree/master/
graph_embedding/watch_your_step. GATNE-T is implemented by generating 20 walks with a length of 10 for each node. The results of CompGCN¹¹¹¹11https://github.com/malllabiisc/CompGCN are obtained using the multiplicative composition of node and relation embeddings. We adapt the Deep Graph Library (DGL) based implementation¹²¹²12https://github.com/dmlc/dgl/tree/master/examples/pytorch/hgt of HGT and RGCN to perform the link prediction task. We use the original implementation of asp2vec¹³¹³13https://github.com/pcy1302/asp2vec for the evaluation. For all baselines, the dimension of embedding is set to 128.

Figure  3 shows an example subgraph from DBLP dataset between “Jiawei Han” and “$P_{28406}$”, a paper on frequent itemset mining. Paths 1 to 5 (shown in different legend) show different contexts subgraphs present between the two nodes. We observe that the link prediction score between Jiawei Han and paper $P_{28406}$ varies, depending on the context subgraph provided to the model.

We also observe that SLiCE assigns high link prediction scores for all context subgraphs (paths 1-4) where all the nodes on the path share strong semantic association with other nodes in the context subgraph. The lowest score is achieved for path-5 which contains a conference node $C_{19054}$ (SIAM Data Mining). As a conference publishes papers in multiple topics, we hypothesize that this breaks the semantic association across nodes, and consequently lowers the probability of the link compared to other context subgraphs where all nodes are closely related.

It is important to note, that a node can share a semantic association with another node separated by multiple hops on the path, and thus would be associated via a higher-order relation. We explore this concept further in the following subsections.

We provide the visualization of the semantic association matrix $\bar{A}^{k}_{ij}$ as defined in Eq. (1) to investigate how the node dependencies evolve through different layers in SLiCE. Given a node pair ($v_{i}$, $v_{j}$) in the context subgraph $g_{c}$, a high value of $\bar{A}^{0}_{ij}$, indicates a strong global dependency of node $v_{i}$ on $v_{j}$. While a high value of $\bar{A}^{k}_{ij}(k\geq 1)$ (the association after applying more translation layers) indicates a prominent high-order relation in the subgraph context.

Figure 4 shows weights of semantic association matrix for the context generated for node pair (N0: Summarizing itemset patterns: a profile-based approach (Paper), N1: Jiawei Han (Author)). Nodes in the context consist of N2: Approach (Topic), N3: Knowledge-base reduction (Paper), N4: Redundancy (Topic) and N5: Handling Redundancy in the Processing of Recursive Database Queries (Paper). We observe that at layer 1 (Figure 4(a)), the association between source node N0 and target node N1 is relatively low. Instead, they both assign high weights on N4. However, the dependencies between nodes are dynamically updated when applying more learning layers, consequently enabling us to identify higher-order relations. For example, the dependency of N1 on N0 becomes higher from layer 3 (Figure 4(c)) and N0 primarily depends on itself without highly influenced by other nodes in layer 4 (Figure 4(d)). This visualization of semantic association helps to understand how the global embedding is translated into the localized embedding for contextual learning.

Metapaths provide a symbolic interpretation of the higher-order relations in a heterogeneous graph. We analyze the ability of SLiCE to learn relevant metapaths that characterize positive semantic associations in the graph. We observe from Table 4 that SLiCE is able to match existing metapaths and also identify new metapath patterns for prediction of each relation type. For example, to predict the paper-author relationship, SLiCE learns three shortest metapaths, including “TPA” (authors publish with the same topic), “APA” (co-authors) and “CPA”(authors published in the same conference).

Interestingly, our learning suggests that longer metapath “APCPA”, which is commonly used to sample academic graphs for co-author relationship, is not as highly predictive of a positive relationship, i.e., “all authors who publish in the same conference do not necessarily publish together”. Overall, the metapaths reported in Table 4 are consistent with the top ranked paths in Figure 3. These metapaths demonstrate SLiCE’s ability to discover higher order semantic associations and perform interpretable link prediction in heterogeneous networks.

Figure 5 provides the distribution of similarity scores for both positive and negative edges obtained by SLiCE. We compare against embeddings produced by node2vec (Grover and Leskovec, 2016) which is one of the best performing static embedding methods (see Table 3) and CompGCN (Vashishth et al., 2020) a relation based contextualization method. We observe that for node2vec and CompGCN, the distribution of similarity scores across positive and negative edges overlaps significantly for all datasets. This indicates that the embeddings learnt from global methods or relation specific contextualization cannot efficiently differentiate the positive and negative edges in link prediction task.

On the contrary, SLiCE increases the margin between the distributions of positive and negative edges significantly. It brings node embeddings in positive edges closer and shifts nodes in negative edges farther away in the low-dimensional space. This indicates that the generated subgraphs provide informative contexts during link prediction and enhance embeddings such that it improves the discriminative capability of both positive and negative node-pairs.

We investigate the performance of SLiCE and the closest performing contextual learning method, asp2vec (Park et al., 2020), as a function of query node pair properties. Given each method, we select all query node pairs drawn from both positive and negative samples that are associated with an incorrect prediction. Next, we compute the average degree of the nodes in each such pair. We opt for degree and ignore any type constraint for it’s simplicity and ease of interpretation. Fig. 6 shows the distribution of these incorrect predictions as a function of average degree of query node pair. It can be seen that, for Amazon and Healthcare datasets, most of the incorrect predictions are concentrated around the query pairs with low and medium values of average degree. However, SLiCE has fewer errors than asp2vec for such node pairs. This can be attributed to the aspect-oriented nature of asp2vec, which maps each node to a fixed number of aspects. Since nodes in a graph may demonstrate varying degree of aspect-diversity, mapping a node with low diversity to more aspects (than it belongs to) reduces asp2vec’s performance. SLiCE adopts a complementary approach, where it considers the subgraph context that connects the query nodes, leading to better contextual representations.

In Figure 7, we provide the link prediction performance with micro-F1 score on four datasets by varying four parameters used in SLiCE model, including number of heads, number of contextual translation layers, number of nodes in contexts (i.e., walk length), and the number of (context) walks generated for each node in pre-training. The performance shown in these plots are the averaged performance by fixing one parameter and varying other three parameters.

On the other hand, when varying the parameter number of layers, we observe that applying four layers of contextual translation provides the best performance on all the datasets; the performance dropped significantly when stacking more layers. This indicates that four contextual translation layers are sufficient to capture the complex higher-order relations over various knowledge graphs. Based on these analysis, we set the default values for both number of heads and the number of layers to be 4, and generate one walk for each node with a length of 6 in the pre-training step.

To explore the impact of pre-trained global node features on the performance of SLiCE, we performed an ablation study by analyzing a variant of SLiCE with four contextual translation layers. More specifically, the pre-trained global embeddings are disabled in SLiCE, termed $\textsc{SLiCE}_{w/o\ GF}$. The results are provided in Table 3. We observe that without initialization using the pre-trained global embeddings, the model performance of SLiCE decreased on all five datasets in both metrics. The reason being that, compared to the random initialization, the global node features are able to represent the role of each node in the global structure of the knowledge graph. By further applying the proposed contextual translation layers, they can collaboratively and efficiently provide contextualized node embeddings for downstream tasks like link prediction.

To investigate the effect of fine-tuning stage on learning the contextual node embeddings, we disable the fine tuning layer for supervised link prediction task, termed $\textsc{SLiCE}_{w/o\ FT}$ and show the results in Table 3. Compared to the baseline methods, it still achieves competitive performance. We attribute this to the effectiveness of capturing higher-order relations through the contextual translation layers. While compared to the full SLiCE model, the performance of $\textsc{SLiCE}_{w/o\ FT}$ degrades slightly on Amazon, DBLP, and Freebase datasets, but significantly decreases on both Twitter and MIMIC datasets. This can be attributed to the fact that supervised training with the link prediction task is able to learn the fine-grained contextual node embedding for link prediction task.

In this subsection, we mainly investigate the impact of the following three parameters on the overall time complexity of the SLiCE model: (1) number of contextual translation layers, (2) number of context subgraphs, and (3) length of context subgraph.

• •

  We study the scalability of the SLiCE model when the number of context translation layers are varied and the corresponding plots are provided in Figure 8(a). The $x$-axis and $y$-axis represent the number of layers and running time in seconds, respectively. The plots indicate that increasing the number of layers does not significantly increase the training time of the SLiCE model.

• •

  In Figure 8(b), we demonstrate the impact of the number of context subgraphs on the time complexity. Increasing the number of context subgraphs generated for each node in pre-training and each node-pair for fine-tuning raises the number of training edges which further increases the training time of the model.

These two plots empirically verify the analysis about the model complexity, discussed in Section 3.5, that the proposed SLiCE model is approximately linear to the number of edges in the graph and does not depend on other parameters such as the number of contextual translation layers and the number of nodes in the graph. In addition, we also vary the length (number of nodes) of the context subgraph in both plots. The plot shows that even doubling the context length will not significantly increase the running time. This time complexity analysis, combined with the performance results in Table 3 and the parameter sensitivity analysis in Figure 7 can jointly provide the guidelines for parameter selection.

In the complexity analysis discussed in Section 3.5, we approximated the total number of training edges in the entire graph as $N_{E}\approx|V|*N_{cpn}$, where $|V|$ represents the number of nodes in graph $G$ and $N_{cpn}$ denotes the number of context subgraphs generated for each node. This estimation also provides us guidelines for determining the number of context subgraphs for each node $N_{cpn}$. By incorporating $N_{E}=\alpha_{T}|E|$ into the approximation (where $|E|$ is the total number of edges in graph $G$ and $\alpha_{T}$ is the ratio of training split), we can estimate the number of context subgraphs per node as $N_{cpn}=\alpha_{T}|E|/|V|$. Table 5 shows the estimated numbers (with $\alpha_{T}=0.6$) for the five datasets used in this work. These estimations provide us an approximate range for the value of $N_{cpn}$ during the context generation step. Based on this analysis, in our experiments, we generally consider 1, 5, and 10 for the value of $N_{cpn}$ on all the five datasets in both pre-training and fine-tuning stages.