[Paper Review] Masked Label Prediction: Unified Message Passing Model for Semi-Supervised Classification

Yunsheng Shi, Zhengjie Huang+4 more|arXiv (Cornell University)|Sep 8, 2020

Advanced Graph Neural Networks30 references102 citations

TL;DR

UniMP unifies feature and label propagation via a Graph Transformer and introduces masked label prediction to prevent leakage, achieving state-of-the-art semi-supervised classification on Open Graph Benchmark datasets.

ABSTRACT

Graph neural network (GNN) and label propagation algorithm (LPA) are both message passing algorithms, which have achieved superior performance in semi-supervised classification. GNN performs feature propagation by a neural network to make predictions, while LPA uses label propagation across graph adjacency matrix to get results. However, there is still no effective way to directly combine these two kinds of algorithms. To address this issue, we propose a novel Unified Message Passaging Model (UniMP) that can incorporate feature and label propagation at both training and inference time. First, UniMP adopts a Graph Transformer network, taking feature embedding and label embedding as input information for propagation. Second, to train the network without overfitting in self-loop input label information, UniMP introduces a masked label prediction strategy, in which some percentage of input label information are masked at random, and then predicted. UniMP conceptually unifies feature propagation and label propagation and is empirically powerful. It obtains new state-of-the-art semi-supervised classification results in Open Graph Benchmark (OGB).

Motivation & Objective

• Motivate combining GNN-based feature propagation with label propagation in a single framework.
• Propose UniMP, a Graph Transformer-based model that ingests both node features and label embeddings for propagation.
• Introduce masked label prediction to prevent label leakage during training while leveraging labels during inference.
• Demonstrate state-of-the-art semi-supervised classification performance on large-scale OGBN datasets.

Proposed method

• Represent partially observed labels in the same embedding space as node features via label embedding.
• Use a multi-head Graph Transformer to propagate and attend over feature and label information.
• Incorporate a gated residual mechanism to mitigate over-smoothing across layers.
• Unify feature and label propagation into a single propagation rule, approximating decoupled feature and label dynamics.
• Apply a masked label prediction training objective to prevent label leakage, inspired by BERT-style masking.
• Evaluate on ogbn-products, ogbn-proteins, and ogbn-arxiv with ablation studies.

Experimental results

Research questions

• RQ1 Can feature and label propagation be effectively unified within a single message-passing framework?
• RQ2 Does masking input labels during training improve generalization and avoid leakage in a label-inclusive GNN?
• RQ3 How much do label embeddings and propagation contribute to performance across large-scale, diverse OGBN datasets?
• RQ4 What is the impact of label coverage and neighbor information on model accuracy across datasets?

Key findings

• Achieves new state-of-the-art results on three OGBN datasets: ogbn-products, ogbn-proteins, and ogbn-arxiv.
• UniMP with Graph Transformer and label embeddings outperforms prior methods that only use feature propagation or only regularize with LPA.
• Masked label prediction effectively prevents self-loop label leakage and supports reliable training.
• On ogbn-products, UniMP attains 0.8256 test accuracy (vs. baselines).
• On ogbn-proteins, UniMP achieves 0.8642 test ROC-AUC (best among reported methods).
• On ogbn-arxiv, UniMP attains 0.7311 test accuracy (competitive with state-of-the-art).