GPT-GNN: A Generative Pre-training Framework for Large-Scale Heterogeneous Graphs

GPT-GNN is an open-source implementation of a top-conference paper from KDD 2020, proposing a new framework for initializing graph neural networks via generative pre-training. This method can be applied to large-scale heterogeneous graphs, and its effectiveness has been verified on the Open Academic Graph (OAG) and Reddit datasets, providing a new paradigm for the pre-training of graph neural networks.

Graph Neural Networks (GNNs) have achieved significant success in recent years in fields such as social network analysis, recommendation systems, knowledge graphs, and drug discovery. However, unlike in the fields of computer vision and natural language processing, research on pre-training for graph neural networks is relatively lagging. Traditional GNN training usually uses random initialization followed by supervised learning on specific downstream tasks, which faces two main challenges:

1. Scarcity of labeled data: It is difficult to obtain large amounts of high-quality labeled data in many graph application domains
2. Insufficient generalization ability: Models trained on specific tasks are hard to transfer to other related tasks

Inspired by the successful experience of the GPT (Generative Pre-Training) series models in the natural language processing field, researchers have begun to explore pre-training methods for graph neural networks. GPT-GNN was born in this context, proposing a new generative pre-training paradigm that enables the model to acquire general graph representation capabilities by learning to reconstruct the attributes and structure of the graph.

The core innovation of GPT-GNN lies in defining the pre-training task as a generative problem: given a partially masked graph, the model needs to predict the masked node attributes and edge connections. This self-supervised learning approach does not require manual annotation and can automatically mine supervision signals from the original graph structure.

The attribute generation task requires the model to predict the attribute features of masked nodes based on the neighborhood structure and neighbor attributes of the nodes. This task forces the model to learn how to effectively fuse structural information with attribute information to form meaningful node representations.

Specifically, the model needs to:

• Aggregate feature information from neighboring nodes
• Capture the position and role of nodes in the graph structure
• Reconstruct the original attributes of masked nodes

The edge generation task requires the model to predict whether a connection exists between two nodes. This task forces the model to learn the semantic similarity and structural correlation between nodes, thereby capturing the topological characteristics of the graph.

By optimizing both the attribute generation and edge generation tasks simultaneously, the node representations learned by GPT-GNN not only contain rich semantic information but also encode the topological structural features of the graph.

Graph data in the real world is often heterogeneous, containing multiple types of nodes and edges. For example, in an academic graph, there are multiple node types such as papers, authors, institutions, and fields, as well as multiple edge types such as writing, citation, and affiliation.

GPT-GNN is specifically designed for heterogeneous graphs with a pre-training framework that supports:

• Unified processing of multi-type nodes and edges
• Type-aware neighbor sampling strategy
• Heterogeneous message passing mechanism

To efficiently process large-scale graph data, GPT-GNN introduces an adaptive embedding queue mechanism. This mechanism maintains a fixed-size embedding cache that stores the representation vectors of historical nodes, thereby avoiding recalculating the embeddings of all nodes in each iteration and significantly improving training efficiency.

The hyperparameter for queue size can be configured via --queue_size, with a default value of 256. A larger queue can store more historical information but increases memory overhead; a smaller queue is more lightweight but may lose some historical context.

GPT-GNN supports two types of attribute decoders:

1. Vector Decoder (vec): Directly predicts the attribute vector of nodes, suitable for numerical or dense vector features
2. Text Decoder (text): Designed for text attributes, using pre-trained word vector models (such as Word2Vec) for text generation

This flexible design allows GPT-GNN to adapt to different types of graph data and application scenarios.