Pre-Trained Models for Heterogeneous Information Networks

TL;DR

This paper introduces PF-HIN, a self-supervised pre-training framework for heterogeneous information networks that enhances transferability and efficiency across multiple downstream tasks using transformer encoders.

Contribution

The paper proposes a novel pre-training and fine-tuning framework for heterogeneous information networks that reduces training costs and improves performance on various tasks.

Findings

PF-HIN outperforms state-of-the-art methods on four benchmark tasks.

Pre-training with self-supervised tasks improves transferability.

Transformer-based architecture effectively captures network features.

Abstract

In network representation learning we learn how to represent heterogeneous information networks in a low-dimensional space so as to facilitate effective search, classification, and prediction solutions. Previous network representation learning methods typically require sufficient task-specific labeled data to address domain-specific problems. The trained model usually cannot be transferred to out-of-domain datasets. We propose a self-supervised pre-training and fine-tuning framework, PF-HIN, to capture the features of a heterogeneous information network. Unlike traditional network representation learning models that have to train the entire model all over again for every downstream task and dataset, PF-HIN only needs to fine-tune the model and a small number of extra task-specific parameters, thus improving model efficiency and effectiveness. During pre-training, we first transform the neighborhood of a given node into a sequence. PF-HIN is pre-trained based on two self-supervised tasks, masked node modeling and adjacent node prediction. We adopt deep bi-directional transformer encoders to train the model, and leverage factorized embedding parameterization and cross-layer parameter sharing to reduce the parameters. In the fine-tuning stage, we choose four benchmark downstream tasks, i.e., link prediction, similarity search, node classification, and node clustering. PF-HIN consistently and significantly outperforms state-of-the-art alternatives on each of these tasks, on four datasets.