Table of Contents
Fetching ...

Learning Invariant Representations of Graph Neural Networks via Cluster Generalization

Donglin Xia, Xiao Wang, Nian Liu, Chuan Shi

TL;DR

This work addresses the vulnerability of Graph Neural Networks to structure distribution shifts by introducing Cluster Information Transfer (CIT), a plug-in that learns invariant node representations through cluster-based information transfer in the embedding space. CIT alternates between clustering representations to form cluster centers, transferring nodes across clusters with Gaussian perturbations, and training with a mixed objective that blends cross-entropy with cluster-aware losses. Theoretical analysis demonstrates that transferring cluster information mitigates the adverse effects of cluster changes during structure shift, and empirical results across perturbation, multiplex, and multigraph benchmarks show consistent generalization improvements over strong baselines. The method is lightweight, compatible with common GNN architectures, and offers a practical path toward robust graph learning in the presence of unknown structure shifts.

Abstract

Graph neural networks (GNNs) have become increasingly popular in modeling graph-structured data due to their ability to learn node representations by aggregating local structure information. However, it is widely acknowledged that the test graph structure may differ from the training graph structure, resulting in a structure shift. In this paper, we experimentally find that the performance of GNNs drops significantly when the structure shift happens, suggesting that the learned models may be biased towards specific structure patterns. To address this challenge, we propose the Cluster Information Transfer (CIT) mechanism (Code available at https://github.com/BUPT-GAMMA/CITGNN), which can learn invariant representations for GNNs, thereby improving their generalization ability to various and unknown test graphs with structure shift. The CIT mechanism achieves this by combining different cluster information with the nodes while preserving their cluster-independent information. By generating nodes across different clusters, the mechanism significantly enhances the diversity of the nodes and helps GNNs learn the invariant representations. We provide a theoretical analysis of the CIT mechanism, showing that the impact of changing clusters during structure shift can be mitigated after transfer. Additionally, the proposed mechanism is a plug-in that can be easily used to improve existing GNNs. We comprehensively evaluate our proposed method on three typical structure shift scenarios, demonstrating its effectiveness in enhancing GNNs' performance.

Learning Invariant Representations of Graph Neural Networks via Cluster Generalization

TL;DR

This work addresses the vulnerability of Graph Neural Networks to structure distribution shifts by introducing Cluster Information Transfer (CIT), a plug-in that learns invariant node representations through cluster-based information transfer in the embedding space. CIT alternates between clustering representations to form cluster centers, transferring nodes across clusters with Gaussian perturbations, and training with a mixed objective that blends cross-entropy with cluster-aware losses. Theoretical analysis demonstrates that transferring cluster information mitigates the adverse effects of cluster changes during structure shift, and empirical results across perturbation, multiplex, and multigraph benchmarks show consistent generalization improvements over strong baselines. The method is lightweight, compatible with common GNN architectures, and offers a practical path toward robust graph learning in the presence of unknown structure shifts.

Abstract

Graph neural networks (GNNs) have become increasingly popular in modeling graph-structured data due to their ability to learn node representations by aggregating local structure information. However, it is widely acknowledged that the test graph structure may differ from the training graph structure, resulting in a structure shift. In this paper, we experimentally find that the performance of GNNs drops significantly when the structure shift happens, suggesting that the learned models may be biased towards specific structure patterns. To address this challenge, we propose the Cluster Information Transfer (CIT) mechanism (Code available at https://github.com/BUPT-GAMMA/CITGNN), which can learn invariant representations for GNNs, thereby improving their generalization ability to various and unknown test graphs with structure shift. The CIT mechanism achieves this by combining different cluster information with the nodes while preserving their cluster-independent information. By generating nodes across different clusters, the mechanism significantly enhances the diversity of the nodes and helps GNNs learn the invariant representations. We provide a theoretical analysis of the CIT mechanism, showing that the impact of changing clusters during structure shift can be mitigated after transfer. Additionally, the proposed mechanism is a plug-in that can be easily used to improve existing GNNs. We comprehensively evaluate our proposed method on three typical structure shift scenarios, demonstrating its effectiveness in enhancing GNNs' performance.
Paper Structure (30 sections, 2 theorems, 17 equations, 7 figures, 4 tables, 1 algorithm)

This paper contains 30 sections, 2 theorems, 17 equations, 7 figures, 4 tables, 1 algorithm.

Key Result

Theorem 1

The decision boundary of fisher classifier is affected by the cluster information.

Figures (7)

  • Figure 1: The node classification accuracy of GCN, GAT, APPNP and GCNII on generated data with structure shift. The x-axis is probability of edges (%).
  • Figure 2: The overall framework of our proposed CIT mechanism on GNNs consists two parts: the traditional GNNs and Cluster Information Transfer (CIT) mechanism. After getting node representations from GNNs, we conduct CIT mechanism on node representations before the last layer of GNNs, which transfers the node to another cluster to generate new representations for training.
  • Figure 3: We show two parts of one graph. The orange and green points represent two clusters in one graph. The circle is aggregating scope of cluster. And the red point represents the target node we transfer from orange cluster to green cluster.
  • Figure 4: ROC-AUC on Twitch where we compare different GNN backbones.
  • Figure 5: Analysis of the probability of transfer.
  • ...and 2 more figures

Theorems & Definitions (4)

  • Theorem 1
  • Theorem 2
  • proof
  • proof