GraphAlign: Pretraining One Graph Neural Network on Multiple Graphs via Feature Alignment

TL;DR

GraphAlign tackles the problem of training one graph neural network across heterogeneous graphs with misaligned node features. It introduces a modular feature-alignment pipeline—a language-model feature encoder, per-graph normalization, and a mixture-of-feature-experts—that pretrains a unified GNN within existing graph SSL frameworks and transfers effectively to unseen graphs. Empirical results on OGB and knowledge-graph tasks show improvements over naive joint training and strong cross-domain transfer, including few-shot scenarios, with an in-context inference strategy that requires no task-specific prompts. This approach broadens the practicality of graph pretraining by enabling cross-graph generalization without relying on task labels or extensive prompt engineering.

Abstract

Graph self-supervised learning (SSL) holds considerable promise for mining and learning with graph-structured data. Yet, a significant challenge in graph SSL lies in the feature discrepancy among graphs across different domains. In this work, we aim to pretrain one graph neural network (GNN) on a varied collection of graphs endowed with rich node features and subsequently apply the pretrained GNN to unseen graphs. We present a general GraphAlign method that can be seamlessly integrated into the existing graph SSL framework. To align feature distributions across disparate graphs, GraphAlign designs alignment strategies of feature encoding, normalization, alongside a mixture-of-feature-expert module. Extensive experiments show that GraphAlign empowers existing graph SSL frameworks to pretrain a unified and powerful GNN across multiple graphs, showcasing performance superiority on both in-domain and out-of-domain graphs.

Figures (4)

• Figure 1: Linear probing results on OGB node classification datasets. individually-pretrain denotes that we train an individual GNN for each dataset. vanilla jointly-pretrain represents training one GNN using all datasets without the incorporation of any designs. GraphAlign empowers us to train one GNN that can achieve superior performance across various datasets and shows a clear advantage over vanilla jointly-pretraining.
• Figure 2: Overview of GraphAlign. Given graphs from different domains, we first utilize a language model functioning as an encoder to project their node attributes into a semantic dimension. Subsequently, we apply feature normalization to each graph individually. To further capture subtle differences, a mixture-of-feature-experts module is designed and implemented to allow each node to adaptively select feature transformations. Finally, we are enabled to pretrain a unified GNN on this aligned feature distribution with any self-supervised learning methods. This GNN can be applied to downstream graphs and tasks unseen during training.
• Figure 3: Feature distribution comparison between w/ and w/o normalization.
• Figure 4: Ablation on the number of MoF layers and MoF projectors. Linear represents using linear transformation as projectors and MLP represents using multi-layer perception as projectors. Using Linear and employing MoF at the input layer is the best.