Pre-Trained Models for Heterogeneous Information Networks

Research on NRL traces back to dimensionality reduction techniques [18, 19, 20, 21], which utilize feature vectors of nodes to construct an affinity graph and then calculate eigenvectors. Graph factorization models [22] represent a graph as an adjacency matrix, and generate a low-dimensional representation via matrix factorization. Such models suffer from high computational costs and data sparsity, and cannot capture the global network structure [5].

Random walks or paths in a network are being used to help preserve the local and global structure of a network. DeepWalk [4] leverages random walks and applies the SkipGram word2vec model to learn network embeddings. node2vec [6] extends DeepWalk; it adopts a biased random walk strategy to explore the network structure. LINE [5] harnesses first- and second-order proximities to encode local and neighborhood structure information.

The aforementioned approaches are designed for homogeneous networks; other methods have been introduced for heterogeneous networks. PTE [23] defines the conditional probability of nodes of one type generated by nodes of another, and forces the conditional distribution to be close to its empirical distribution. metapath2vec [7] has a heterogeneous SkipGram with its context window restricted to one specific type. HINE [8] uses a metapath based proximity and minimizes the distance between nodes’ joint probability and empirical probabilities. HIN2Vec [9] uses Hadamard multiplication of nodes and metapaths to capture features.

What we contribute to NRL on top of the work listed above is an efficient and effective method for representation learning for HINs based on graph neural networks (GNNs).

GNN models have shown promising results for representing networks. Efforts have been devoted to generalizing convolutional operations from visual data to graph data. Bruna et al. [24] propose a spectral graph theory based graph convolution operation. \AcpGCN [10] adopt localized first-order approximations of spectral graph convolutions to improve scalability. There is a line of research to improve spectral GNN models [25, 26, 27, 28], but it processes the whole graph simultaneously, leading to efficiency bottlenecks. To address the problem, spatial GNN models have been proposed [29, 30, 31, 32]. GraphSAGE [32] leverages a sampling strategy to iteratively sample neighboring nodes instead of the whole graph. Gao et al. [31] utilize a sub-graph training method to reduce memory and computational cost.

GNN fuse neighboring nodes or walks in graphs so as to learn a new node representation [11, 33, 34]. The main difference with convolution based models is that graph attention networks introduce attention mechanisms to assign higher weights to more important nodes or walks. GAT [11] harnesses masked self-attention layers to apply different weights to different nodes in a neighborhood, to improve efficiency on graph-structured data. GIN [35] models injective multiset functions for neighbor aggregation by parameterizing universal multiset functions with neural networks.

The above GNN models have been devised for homogeneous networks as they aggregate neighboring nodes or walks regardless of their types. Targeting HINs, HetGNN [12] first samples a fixed number of neighboring nodes of a given node and then groups these based on their types. Then, it uses a neural network architecture with two modules to aggregate the feature information of the neighboring nodes. One module is used to encode features of each type of node, the other to aggregate features of different types. HGT [36] uses node- and edge-type dependent parameters to describe heterogeneous attention over each edge; it also uses a heterogeneous mini-batch graph sampling algorithm for training.

What we contribute to GNNs is that the traditional NRL and GNN models listed above need to be re-trained all over again for different datasets and tasks, while our proposal PF-HIN only needs fine-tuning using a small number of task-specific parameters for a specific task and dataset, after pre-training via self-supervision tasks.

There exist relatively few approaches for pre-training a GNN model for downstream tasks. InfoGraph [37] maximizes the mutual information between graph-level embeddings and sub-structure embeddings. Hu et al. [38] pre-train a GNN at the level of nodes and graphs to learn local and global features, showing performance improvements on various graph classification tasks. Our proposed model, PF-HIN, differs as we focus on node-level transfer learning and pre-train our model on a single (large-scale) graph.

Hu et al. [39] design three pre-training tasks: denoising link reconstruction, centrality score ranking, and cluster preserving. GPT-GNN [17] adopts HGT [36] as its base GNN and uses attribute generation and edge generation as pre-training tasks. Hu et al. [17] only conduct their downstream tasks on the same dataset that was used for pre-training. GCC [16] designs subgraph instance discrimination as a pre-training task and uses contrastive learning to train GNNs, with its base GNNs as GIN; then it transfers its pre-trained model to different datasets. However, it is designed for homogeneous networks, not apt to exploit heterogeneous networks. Hwang et al. [40] proposes to adopt meta-path prediction as a pre-training task, which is to predict if two nodes are connected via a meta-path. The authors have not assigned a name for this method, so use MPP to refer to it in this paper.

What we contribute to graph pre-training on top of the work listed above is that we are able to not only fine-tune the proposed model across different tasks and different datasets, but can also deal with heterogeneous networks.

A heterogeneous information network (HIN) is a graph $G=(V,E,T)$, where $V$ denotes the set of nodes and $E$ denotes the set of edges between nodes. Each node and edge is associated with a type mapping function, $\phi:V\rightarrow T_{V}$ and $\varphi:E\rightarrow T_{E}$, respectively. $T_{V}$ and $T_{E}$ denote the sets of node and edge types. A HIN is a network where $|T_{V}|>1$ and/or $|T_{E}|>1$.

A visual presentation of the proposed model, pre-training and fin-tuning heterogeneous information network (PF-HIN), is given in Figure 1. Below, we describe the node sequence generation procedure, the input representation, followed by the pre-training and fine-tuning stages of PF-HIN.

We first transform the structure of a node’s neighborhood to a sequence of length $k$. To measure the importance of nodes based on the structural roles in the graph, node centrality is proposed in [41]. This framework makes use of three centrality metrics²²2This framework is extensible in the sense that additional metric that is of particular interest to the user can be explicitly supplemented after the three metrics, as long as the combination exhibits better performance in downstream tasks. In our implementation, we stick to the basic version consisted of three metrics to demonstrate the effectiveness of the framework., i.e., 1. betweenness centrality, 2. eigencentrality, and 3. closeness centrality. Betweenness centrality is calculated as the fraction of shortest paths that pass through a given node. Eigencentrality measures the influence of a node on its neighbors. Closeness centrality computes the total length of the shortest paths between the given node and others. We assign learnable weights to these metrics.

To capture heterogeneous features of a node’s neighbor, we adopt a so-called rank-guided heterogeneous walk to form a sequence. The walk is with restart; it will iteratively travel from a node to its neighbors. It starts from node $v$, and it first reaches out to a node with a higher rank, which is what makes the walk rank-guided. This walk will not stop until it collects a pre-determined number of nodes. In order to assign the model with a sense of heterogeneity, we constrain the number of different types to be collected in the sequence so that every type of node can be included. We group nodes into mini-sequences, where a mini-sequence is a sequence of nodes having the same type [12]. In each mini-sequence, the nodes are sorted based on each node’s rank, which serve as a kind of position information.

Importantly, unlike traditional sampling strategies like random walks, breadth first search or depth first search, our sampling strategy is able to extract important and influential neighboring nodes for each node by selecting nodes with a higher rank; this allows us to capture more representative structural information of a neighborhood. Our sampling strategy collects all types of node for each node while traditional strategies ignore the nodes’ types. Nodes of the same type are grouped in mini-sequences so that further type-based analysis can be conducted to capture the heterogeneous features of a HIN. More empirical results with analysis are provided in Section 5.3.

After generating the sequences, we learn the input embeddings of each node in the sequence using a Bi-LSTM layer. Given the input sequence $\{x_{1},x_{2},\ldots,x_{n}\}$, in which $x_{i}\in\mathbb{R}^{d\times 1}$, a Bi-LSTM is used to capture the interaction relationships between nodes. The Bi-LSTM is composed of a forward and a backward LSTM layer. The LSTM layer is defined as follows:

where $h_{i}\in\mathbb{R}^{d/2}\times 1$ is the output hidden state of node $i$, $\odot$ represents the element-wise product, $\mathbf{W}\in\mathbb{R}^{(d/2)\times(d/2)}$ and $b\in\mathbb{R}^{d/2\times 1}$ are learnable parameters, which denote weight and bias, respectively; $\mathbf{j}_{i}$, $\mathbf{f}_{i}$, $\mathbf{o}_{i}$ are the input gate vector, forget gate vector and output vector, respectively. We concatenate the hidden state of the forward and backward LSTM layers to form the final hidden state of the Bi-LSTM layer. For each type of node, we adopt different Bi-LSTMs so as to extract type-specific features.

After generating the input embeddings via a Bi-LSTM, we adopt masked node modeling (MNM) as our pre-training task. We randomly mask a percentage of the input nodes and then predict those masked nodes. We conduct this task on the type-based mini-sequences generated by the aforementioned rank guided heterogeneous walk. For each group of nodes with the same type, we randomly mask nodes in the mini-sequence. Given the mini-sequence of type $t$, denoted as $\{x_{1}^{t},x_{2}^{t},\ldots,x_{n}^{t}\}$, we randomly choose 15% of the nodes to be replaced. And for a chosen node $x_{i}^{t}$, we replace its token with the actual [MASK] token with 80% probability, another random node token with 10% probability and the unchanged $x_{i}^{t}$ with 10% probability. The masked sequence is fed into the bi-directional transformer encoders. The embeddings generated via the Bi-LSTM are used as token embeddings, while the rank information is transferred as position embeddings. After the transformer module, the final hidden state $h_{i}^{t_{L}}$ corresponding to the [MASK] token is fed to a feedforward layer. The output is used to predict the target node via a softmax classification layer:

where $z_{i}^{t}$ is the output of the feedforward layer, $\mathbf{W}^{\mathrm{MNM}}\in V^{t}\times d$ is the classification weight shared with the input node embedding matrix, $V^{t}$ is the number of nodes in the $t$-type mini-sequence, $d$ is the dimension of the hidden state size, $\mathbf{p}_{i}^{t}$ is the predicted distribution of $x_{i}^{t}$ over all nodes.

For training, we use the cross-entropy between the one-hot label $\mathbf{y}_{i}^{t}$ and the prediction $\mathbf{p}_{i}^{t}$:

where $y_{m}^{t}$ and $p_{m}^{t}$ are the $m$-th components of $\mathbf{y}_{i}^{t}$ and $\mathbf{p}_{i}^{t}$, respectively. We adopt a smoothing strategy by setting $y_{m}^{t}=\epsilon$ for the target node and $y_{m}^{t}=\frac{1-\epsilon}{V^{t}-1}$ for each of the other nodes. By doing so, we loosen the restriction that a one-hot label corresponds to only one answer.

After the masked node modeling module, we design another pre-training task, i.e., adjacent node prediction (ANP), to further capture the relationship between nodes. Unlike the MNM task, which operates on type-based mini-sequences, we perform the ANP task on full sequences, and we compare two full sequences to see whether their starting nodes are adjacent or not. The reason that we do not perform the ANP task on type-based mini-sequences is that given $k$ types of nodes, there will be $k(k-1)/2$ mini-sequence pairs to be analyzed, which is very time-consuming.

In our setting, for node $v$ with sequence $X_{v}$ and node $u$ with sequence $X_{u}$, 50% of the time we choose $u$ to be the actual adjacent node of $v$ (labeled as IsAdjacent), and 50% of the time we randomly choose $u$ from the corpus (labeled as NotAdjacent). Given the classification layer weights $\mathrm{W}^{\mathrm{ANP}}$, the scoring function $s_{\tau}$ of whether the node pair is adjacent is shown as follows:

where $s_{\tau}\in\mathbb{R}^{2}$ is a binary vector with $s_{\tau 0},s_{\tau 1}\in[0,1]$ and $s_{\tau 0}+s_{\tau 1}=1$, $C\in\mathbb{R}^{H}$ denotes the hidden vector of classification label used in a transformer architecture [42]. Considering the positive adjacent node pair $\mathbb{S}^{+}$ and a negative adjacent node pair $\mathbb{S}^{-}$, we calculate a cross-entropy loss as follows:

where $y_{\tau}$ is the label (positive or negative) of that node pair.

Our two pre-training tasks share the same transformer architecture. To increase the training speed of our model, we adopt two parameter reduction techniques to lower memory requirements, inspired by the ALBERT architecture [43]. Instead of setting the node embedding size $Q$ to be equal to the hidden layer size $H$ like BERT does, we make more efficient use of the total number of model parameters, dictating that $H\gg Q$. We adopt factorized embedding parameterization, which decomposes the parameters into two smaller matrices. Rather than mapping the one-hot vectors directly to a hidden space with size $H$, we first map them to a low-dimensional embedding space with size $Q$, and then map it to the hidden space. Additionally, we adopt cross-layer parameter sharing to further boost efficiency. Traditional sharing mechanisms either only share the feed forward network parameters across layers or only the attention parameters. We share all parameters across layers.

We denote the number of transformer layers as $L$, and the number of self-attention heads as $A$. For our parameter settings we follow the configuration of ALBERT [43], where $L$ is set to 12, $H$ to 768, $A$ to 12, $Q$ to 128, and the total number of parameters is equal to 12M. For the procedure of our pre-training task, see Algorithm 1.

The self-attention mechanism in the transformer allows PF-HIN to model many downstream tasks. Fine-tuning can be realized by simply swapping out the proper inputs and outputs, regardless of the single node sequence or sequence pairs used. For each downstream task, the task-specific inputs and outputs are simply plugged into PF-HIN and all parameters are fine-tuned end-to-end. Here, we introduce four tasks: 1. link prediction, 2. similarity search, 3. node classification, and 4. node clustering as downstream tasks.

Specifically, in link prediction, we predict whether there is a link between two nodes, and the inputs are the node sequence pairs. To generate output, we feed the classification label into the sigmoid layer, so as to predict the existence of a link between two nodes. The only new parameters are the classification layer weights $W\in\mathbb{R}^{2\times H}$, where $H$ is the size of hidden state.

In similarity search, in order to measure the similarity between two nodes, we use the node sequence pairs as input. We leverage the token-level output representations to compute the similarity score of two nodes.

In node classification, we only use a single node sequence as input and generate the classification label via a softmax layer. To calculate the classification loss, we only need to add classification layer weights $W\in\mathbb{R}^{K\times H}$ as new parameters, where $K$ is the number of classification labels and $H$ is the size of hidden state.

In node clustering, we also use a one node sequence as input and then put the token-level output embeddings to a clustering model, so as to cluster the data.

Experimental details for these downstream tasks are introduced in Section 4 below.

We adopt the open academic graph OAG³³3https://www.openacademic.ai/oag/ as our pre-training dataset, which is a heterogeneous academic dataset. It contains over 178 million nodes and 2.223 billion edges with five types of node: 1. author, 2. paper, 3. venue, 4. field, and 5. institute.

For downstream tasks, we transfer our pre-trained model to four datasets: 1. OAG-mini, 2. DBLP, 3. YELP, and 4. YAGO. OAG-mini is a subset extracted from OAG; the authors are split into four areas: machine learning, data mining, database, and information retrieval. DBLP⁴⁴4http://dblp.uni-trier.de is also an academic dataset with four types of node: 1. author, 2. paper, 3. venue, and 4. topic; the authors are split into the same areas as those in OAG-mini. YELP⁵⁵5https://www.yelp.com/dataset\_challenge is a social media dataset, with restaurants reviews and four types of node: 1. review, 2. customer, 3. restaurant, and 4. food-related keywords. The restaurants are separated into 1. Chinese food, 2. fast food, and 3. sushi bar. YAGO⁶⁶6https://old.datahub.io/dataset/yago is a knowledge base and we extracted a subset of it containing movie information, having five types of node: 1. movie, 2. actor, 3. director, 4. composer, and 5. producer. The movies are split into five types: 1. action, 2. adventure, 3. sci-fi, 4. crime, and 5. horror. The dataset statistics are shown in Table I.

We first choose network embedding methods to directly train the downstream datasets for specific tasks as baselines: DeepWalk [4], LINE [5] and node2vec [6]; they were originally applied to homogeneous information networks. DeepWalk and node2vec leverage random walks, while node2vec uses a biased walk strategy to capture the network structure. LINE uses the local and neighborhood structural information via first-order and second-order proximities.

We include three state-of-the-art algorithms devised for HINs: metapath2vec [7], HINE [8], HIN2Vec [9]. They are all based on metapaths, but differ in the way they use metapath features: metapath2vec adopts heterogeneous SkipGrams, HINE proposes a metapath-based notion of proximity, and HIN2Vec utilizes the Hadamard multiplication of nodes and metapaths.

We also include other GNN models, i.e., GCN [10], GAT [11], GraphSAGE [32] and GIN [35], which were originally devised for homogeneous information networks. GCN and GraphSAGE are based on convolutional operations, while GCN requires the Laplacian of the full graph, and GraphSAGE only needs a node’s local neighborhood. GAT employs an attention mechanism to capture the correlation between central node and neighboring nodes. GIN uses parameterizing universal multiset functions with neural networks to model injective multiset functions for neighbor aggregation.

We also select HetGNN [12], HGT [36] as models for comparison; both have been devised for HIN embeddings. HetGNN samples heterogeneous neighbors, grouping them based on their node types, and then aggregates feature information of the sampled neighboring nodes. HGT has node- and edge-type dependent parameters to characterize heterogeneous attention over each edge.

Aside from those network embedding methods, for a fair comparison, we also choose GPT-GNN [17], GCC [16] and MPP [40] to run the entire pre-training and fine-tuning pipeline. GPT-GNN utilizes attribute generation and edge generation tasks to pre-train GNN, with HGT as its base GNN. GCC adopts subgraph instance discrimination as a pre-training task, taking GIN as its base GNN. MPP proposes to adopt meta-path prediction as the pre-training task.

For pre-training, we set the generated sequence length $k$ to 20. The dimension of the node embedding is set to 128 and the size of hidden state is set to 768. On transformer layers, we use 0.1 as the dropout probability. The Adam learning rate is initiated as 0.001 with a linear decay. We use 256 sequences to form a batch and the training epoch is set to 20. The training loss is the sum of the mean masked node modeling likelihood and the mean adjacent node prediction likelihood.

In fine-tuning, most parameters remain the same as in pre-training, except the learning rate, batch size and number of epochs. We use grid search to determine the best configuration. The learning rate is chosen from $\{0.01,0.02,0.025,0.05\}$. The training epoch is chosen from $\{2,3,4,5\}$. The batch size is chosen from $\{16,32,64\}$. The optimal parameters are task-specific. For the other models, we adopt the best configurations reported in the source publications.

We report on statistical significance with a paired two-tailed t-test and we mark a significant improvement of PF-HIN over GPT-GNN for $p<0.05$ with ^▲.

We evaluate PF-HIN over four downstream tasks, i.e., link prediction, similarity search, node classification and node clustering.

To evaluate the efficiency of our fine-tuning framework compared to other models, we conduct an analysis of the computational costs. Specifically, we analyze the running time of each model on each task, using the early stopping mechanism. Due to space limitations we only report the results on the DBLP dataset; the results for the remaining datasets are qualitatively similar. See Table VI. We use standard hardware (Intel (R) Core (TM) i7-10700K CPU + GTX-2080 GPU); the time reported is wall-clock time, averaged over 10 runs.

PF-HIN’s running time is longer than of the three traditional homogeneous models DeepWalk, LINE and node2vec, which are based on random walks; it is not as high as any of the other models. GNN based models like GCN, GraphSAGE, GAT, GIN and HetGNN have a higher running time than all other models, since the complexity of traditional deep neural networks is much higher than other algorithms. GPT-GNN, GCC and PF-HIN have relatively short running times; this is because pre-trained parameters help the loss function converge much faster. Among the pre-train and fine-tune models, PF-HIN is the most efficient one since the complexity of the transformer encoder we use is lower than that of HGT in GPT-GNN and that of GIN in GCC.

We analyze the effect of the pre-training tasks, the bi-directional transformer encoders, the components of the input representation, and the rank-guided heterogeneous walk sampling strategy.

We conduct a sensitivity analysis of the hyper-parameters of PF-HIN. We choose two parameters for analysis: the maximum length of the input sequence and the dimension of the node embedding. For each downstream task, we only choose one metric for evaluation: AUC for link prediction, AUC for similarity search, micro-F1 value for node classification, and NMI value for node clustering. Figures 2 and 3 show the results of our parameter analysis.

Figure 2 shows the results for the maximum length of the input sequence. The performance improves rapidly when the length increases from 0 to 20. A short node sequence is not able to fully express neighborhood information. When the length reaches 20 or longer, the performance stabilizes, and longer sequence lengths may even hurt the performance. Given a node, its neighboring information can be well represented by its direct neighborhood, however, including far-away nodes may introduce noise. According to this analysis, we choose the length of the input sequence to be 20 so as to balance effectiveness and efficiency.

As to the dimension of node embeddings, Figure 3 shows that the performance improves as the dimension increases, for all tasks and datasets. The higher dimensions are able to capture more features. PF-HIN is not very sensitive to the dimension we choose, especially once it is at least 128. The performance gap is not very large between dimension 128 and 256. Thus, we choose 128 as our setting for the dimension of node embeddings for efficiency considerations.