# Learning Invariant Representations of Graph Neural Networks via Cluster Generalization

arXiv: https://arxiv.org/abs/2403.03599 · HTML: https://mlsanity.com/paper/2403.03599/

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 40 references that are in the map.
- First posted on arXiv 6 Mar 2024.

## Builds on (most cited first, 30 of 30)

- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Open Graph Benchmark: Datasets for Machine Learning on Graphs](https://mlsanity.com/paper/2005.00687/) (2020)
- [Variational Graph Auto-Encoders](https://mlsanity.com/paper/1611.07308/) (2016)
- [Predict then Propagate: Graph Neural Networks meet Personalized PageRank](https://mlsanity.com/paper/1810.05997/) (2018)
- [Revisiting Semi-Supervised Learning with Graph Embeddings](https://mlsanity.com/paper/1603.08861/) (2016)
- [Simple and Deep Graph Convolutional Networks](https://mlsanity.com/paper/2007.02133/) (2020)
- [Link Prediction Based on Graph Neural Networks](https://mlsanity.com/paper/1802.09691/) (2018)
- [Generalizing to Unseen Domains: A Survey on Domain Generalization](https://mlsanity.com/paper/2103.03097/) (2021)
- [Heterogeneous Graph Attention Network](https://mlsanity.com/paper/1903.07293/) (2019)
- [A Fair Comparison of Graph Neural Networks for Graph Classification](https://mlsanity.com/paper/1912.09893/) (2019)
- [Improving Out-of-Distribution Robustness via Selective Augmentation](https://mlsanity.com/paper/2201.00299/) (2022)
- [Discovering Invariant Rationales for Graph Neural Networks](https://mlsanity.com/paper/2201.12872/) (2022)
- [Design Space for Graph Neural Networks](https://mlsanity.com/paper/2011.08843/) (2020)
- [Probabilistic Face Embeddings](https://mlsanity.com/paper/1904.09658/) (2019)
- [Interpretable and Generalizable Graph Learning via Stochastic Attention Mechanism](https://mlsanity.com/paper/2201.12987/) (2022)
- [Handling Distribution Shifts on Graphs: An Invariance Perspective](https://mlsanity.com/paper/2202.02466/) (2022)
- [Graph Clustering with Graph Neural Networks](https://mlsanity.com/paper/2006.16904/) (2020)
- [Attributed Graph Clustering: A Deep Attentional Embedding Approach](https://mlsanity.com/paper/1906.06532/) (2019)
- [GOOD: A Graph Out-of-Distribution Benchmark](https://mlsanity.com/paper/2206.08452/) (2022)
- [Interpreting and Unifying Graph Neural Networks with An Optimization Framework](https://mlsanity.com/paper/2101.11859/) (2021)
- [Uncertainty Modeling for Out-of-Distribution Generalization](https://mlsanity.com/paper/2202.03958/) (2022)
- [Graph-Structured Representations for Visual Question Answering](https://mlsanity.com/paper/1609.05600/) (2016)
- [Shift-Robust GNNs: Overcoming the Limitations of Localized Graph Training Data](https://mlsanity.com/paper/2108.01099/) (2021)
- [New Benchmarks for Learning on Non-Homophilous Graphs](https://mlsanity.com/paper/2104.01404/) (2021)
- [Debiased Graph Neural Networks with Agnostic Label Selection Bias](https://mlsanity.com/paper/2201.07708/) (2022)
- [Graph Invariant Learning with Subgraph Co-mixup for Out-Of-Distribution Generalization](https://mlsanity.com/paper/2312.10988/) (2023)
- [High-dimensional Linear Discriminant Analysis: Optimality, Adaptive Algorithm, and Missing Data](https://mlsanity.com/paper/1804.03018/) (2018)
- [A Convex Optimization Approach to High-Dimensional Sparse Quadratic Discriminant Analysis](https://mlsanity.com/paper/1912.02872/) (2019)
- [A Simple and Powerful Framework for Stable Dynamic Network Embedding](https://mlsanity.com/paper/2311.09251/) (2023)

## Built on by (most cited first, 10 of 10)

- [Towards Invariance to Node Identifiers in Graph Neural Networks](https://mlsanity.com/paper/2502.13660/) (2025)
- [Graphs Generalization under Distribution Shifts](https://mlsanity.com/paper/2403.16334/) (2024)
- [Topology-Aware Dynamic Reweighting for Distribution Shifts on Graph](https://mlsanity.com/paper/2406.01066/) (2024)
- [Causal invariant geographic network representations with feature and structural distribution shifts](https://mlsanity.com/paper/2503.19382/) (2025)
- [Quantifying Explanation Quality in Graph Neural Networks using Out-of-Distribution Generalization](https://mlsanity.com/paper/2602.07708/) (2026)
- [DeCaf: A Causal Decoupling Framework for OOD Generalization on Node Classification](https://mlsanity.com/paper/2410.20295/) (2024)
- [On the Utilization of Unique Node Identifiers in Graph Neural Networks](https://mlsanity.com/paper/2411.02271/) (2024)
- [Scale-aware Message Passing For Graph Node Classification](https://mlsanity.com/paper/2411.19392/) (2024)
- [Causal Neighbourhood Learning for Invariant Graph Representations](https://mlsanity.com/paper/2602.17934/) (2026)
- [CGRL: Causal-Guided Representation Learning for Node-Level Out-of-Distribution Generalization](https://mlsanity.com/paper/2603.24304/) (2026)

## Abstract (arXiv)

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.