# Enhancing Size Generalization in Graph Neural Networks through Disentangled Representation Learning

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

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 55 references that are in the map.
- First posted on arXiv 7 Jun 2024.

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

- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [How Powerful are Graph Neural Networks?](https://mlsanity.com/paper/1810.00826/) (2018)
- [A Comprehensive Survey on Graph Neural Networks](https://mlsanity.com/paper/1901.00596/) (2019)
- [Open Graph Benchmark: Datasets for Machine Learning on Graphs](https://mlsanity.com/paper/2005.00687/) (2020)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Graph Contrastive Learning with Augmentations](https://mlsanity.com/paper/2010.13902/) (2020)
- [Deep Graph Infomax](https://mlsanity.com/paper/1809.10341/) (2018)
- [MoleculeNet: A Benchmark for Molecular Machine Learning](https://mlsanity.com/paper/1703.00564/) (2017)
- [GNNExplainer: Generating Explanations for Graph Neural Networks](https://mlsanity.com/paper/1903.03894/) (2019)
- [TUDataset: A collection of benchmark datasets for learning with graphs](https://mlsanity.com/paper/2007.08663/) (2020)
- [GraphCodeBERT: Pre-training Code Representations with Data Flow](https://mlsanity.com/paper/2009.08366/) (2020)
- [Graph Contrastive Learning with Adaptive Augmentation](https://mlsanity.com/paper/2010.14945/) (2020)
- [Masked Label Prediction: Unified Message Passing Model for Semi-Supervised Classification](https://mlsanity.com/paper/2009.03509/) (2020)
- [Parameterized Explainer for Graph Neural Network](https://mlsanity.com/paper/2011.04573/) (2020)
- [Explainability in Graph Neural Networks: A Taxonomic Survey](https://mlsanity.com/paper/2012.15445/) (2020)
- [AllenNLP: A Deep Semantic Natural Language Processing Platform](https://mlsanity.com/paper/1803.07640/) (2018)
- [Invariant and Equivariant Graph Networks](https://mlsanity.com/paper/1812.09902/) (2018)
- [On Explainability of Graph Neural Networks via Subgraph Explorations](https://mlsanity.com/paper/2102.05152/) (2021)
- [Molecular Graph Convolutions: Moving Beyond Fingerprints](https://mlsanity.com/paper/1603.00856/) (2016)
- [Flexibly Fair Representation Learning by Disentanglement](https://mlsanity.com/paper/1906.02589/) (2019)
- [Training Graph Neural Networks with 1000 Layers](https://mlsanity.com/paper/2106.07476/) (2021)
- [Machine Learning on Graphs: A Model and Comprehensive Taxonomy](https://mlsanity.com/paper/2005.03675/) (2020)
- [Relational Pooling for Graph Representations](https://mlsanity.com/paper/1903.02541/) (2019)
- [Learning Causally Invariant Representations for Out-of-Distribution Generalization on Graphs](https://mlsanity.com/paper/2202.05441/) (2022)
- [Data Augmentation for Deep Graph Learning: A Survey](https://mlsanity.com/paper/2202.08235/) (2022)
- [Self-supervised Graph-level Representation Learning with Local and Global Structure](https://mlsanity.com/paper/2106.04113/) (2021)
- [From Local Structures to Size Generalization in Graph Neural Networks](https://mlsanity.com/paper/2010.08853/) (2020)
- [Size-Invariant Graph Representations for Graph Classification Extrapolations](https://mlsanity.com/paper/2103.05045/) (2021)
- [Fairness by Learning Orthogonal Disentangled Representations](https://mlsanity.com/paper/2003.05707/) (2020)
- [Convergence and Stability of Graph Convolutional Networks on Large Random Graphs](https://mlsanity.com/paper/2006.01868/) (2020)

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

- [A Survey of Deep Graph Learning under Distribution Shifts: from Graph Out-of-Distribution Generalization to Adaptation](https://mlsanity.com/paper/2410.19265/) (2024)
- [Graph Neural Network Generalization with Gaussian Mixture Model Based Augmentation](https://mlsanity.com/paper/2411.08638/) (2024)
- [Exploring Consistency in Graph Representations:from Graph Kernels to Graph Neural Networks](https://mlsanity.com/paper/2410.23748/) (2024)
- [From Small to Large: Generalization Bounds for Transformers on Variable-Size Inputs](https://mlsanity.com/paper/2512.12805/) (2025)
- [Score-based Conditional Out-of-Distribution Augmentation for Graph Covariate Shift](https://mlsanity.com/paper/2410.17506/) (2024)
- [Learning Repetition-Invariant Representations for Polymer Informatics](https://mlsanity.com/paper/2505.10726/) (2025)
- [Dimension-Free Scaling Laws for Size Generalization: An Exponential Separation](https://mlsanity.com/paper/2610.19279/) (2026)

## Abstract (arXiv)

Although most graph neural networks (GNNs) can operate on graphs of any size, their classification performance often declines on graphs larger than those encountered during training. Existing methods insufficiently address the removal of size information from graph representations, resulting in sub-optimal performance and reliance on backbone models. In response, we propose DISGEN, a novel and model-agnostic framework designed to disentangle size factors from graph representations. DISGEN employs size- and task-invariant augmentations and introduces a decoupling loss that minimizes shared information in hidden representations, with theoretical guarantees for its effectiveness. Our empirical results show that DISGEN outperforms the state-of-the-art models by up to 6% on real-world datasets, underscoring its effectiveness in enhancing the size generalizability of GNNs. Our codes are available at: https://github.com/GraphmindDartmouth/DISGEN.