Enhanced Pre-training of Graph Neural Networks for Million-Scale Heterogeneous Graphs

TL;DR

This work tackles the scarcity of labeled data in large-scale heterogeneous graphs by introducing PHE, a pre-training framework with two contrastive tasks: a structure-aware task that uses a network-schema subspace and attention to capture fine-grained heterogeneity, and a semantic-aware task that uses a perturbation subspace of semantic neighbors to address semantic mismatch. A transformer-based heterogeneous encoder supports these tasks, and a scalable mini-batch sampling strategy enables pre-training on graphs with millions of nodes. Empirically, PHE outperforms strong baselines on the Open Academic Graph, demonstrates robust transferability across domains, and shows flexibility across architectures with acceptable training efficiency. The approach offers a practical pathway to learn transferable, high-quality representations for diverse downstream tasks in large heterogeneous graphs.

Abstract

In recent years, graph neural networks (GNNs) have facilitated the development of graph data mining. However, training GNNs requires sufficient labeled task-specific data, which is expensive and sometimes unavailable. To be less dependent on labeled data, recent studies propose to pre-train GNNs in a self-supervised manner and then apply the pre-trained GNNs to downstream tasks with limited labeled data. However, most existing methods are designed solely for homogeneous graphs (real-world graphs are mostly heterogeneous) and do not consider semantic mismatch (the semantic difference between the original data and the ideal data containing more transferable semantic information). In this paper, we propose an effective framework to pre-train GNNs on the large-scale heterogeneous graph. We first design a structure-aware pre-training task, which aims to capture structural properties in heterogeneous graphs. Then, we design a semantic-aware pre-training task to tackle the mismatch. Specifically, we construct a perturbation subspace composed of semantic neighbors to help deal with the semantic mismatch. Semantic neighbors make the model focus more on the general knowledge in the semantic space, which in turn assists the model in learning knowledge with better transferability. Finally, extensive experiments are conducted on real-world large-scale heterogeneous graphs to demonstrate the superiority of the proposed method over state-of-the-art baselines. Code available at https://github.com/sunshy-1/PHE.

Figures (10)

• Figure 1: Shortcomings of existing methods. (a) The fine-grained contributions are not considered. (b) The semantic mismatch is not considered.
• Figure 2: The demonstration of the proposed model. #1) The overall framework includes two components: the structure-aware pre-training task and the semantic-aware pre-training task. In the structure-aware pre-training task, enhanced query samples are used to capture high-order properties of the heterogeneous graph. In semantic-aware pre-training task, enhanced positive/negative samples are used to guide the GNN to learn more transferable knowledge. #2) and #3) describe the proposed modules. The module in #2) utilizes a structure-enhanced network schema for contrastive learning to capture fine-grained heterogeneous information. The module in #3) constructs positive/negative pairs through parameter perturbation to address the issue of semantic mismatch. #4) After obtaining the pre-trained GNN, we fine-tune it on the downstream tasks. There are two modes for fine-tuning, namely full fine-tuning mode and frozen fine-tuning mode. The full fine-tuning mode is to fine-tune both the pre-trained GNN and the downstream task adapter (e.g., the classification head tailored for the downstream node classification task) simultaneously, while the frozen fine-tuning mode is to freeze the parameters of pre-trained GNN and only fine-tune the parameters of the adapter.
• Figure 3: Display of the dynamic negative sample queue. Only the most recent batch of negative samples are kept in the queue for use in the current batch.
• Figure 4: Visualizing the distribution of embeddings.
• Figure 5: Experiments on the generalization of the proposed method on other types of graph data (i.e., Reddit)