Enhancing Size Generalization in Graph Neural Networks through Disentangled Representation Learning

Zheng Huang, Qihui Yang, Dawei Zhou, Yujun Yan
Proceedings of the 41st International Conference on Machine Learning, PMLR 235:20365-20381, 2024.

Abstract

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v235-huang24ac, title = {Enhancing Size Generalization in Graph Neural Networks through Disentangled Representation Learning}, author = {Huang, Zheng and Yang, Qihui and Zhou, Dawei and Yan, Yujun}, booktitle = {Proceedings of the 41st International Conference on Machine Learning}, pages = {20365--20381}, year = {2024}, editor = {Salakhutdinov, Ruslan and Kolter, Zico and Heller, Katherine and Weller, Adrian and Oliver, Nuria and Scarlett, Jonathan and Berkenkamp, Felix}, volume = {235}, series = {Proceedings of Machine Learning Research}, month = {21--27 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v235/main/assets/huang24ac/huang24ac.pdf}, url = {https://proceedings.mlr.press/v235/huang24ac.html}, abstract = {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.} }
Endnote
%0 Conference Paper %T Enhancing Size Generalization in Graph Neural Networks through Disentangled Representation Learning %A Zheng Huang %A Qihui Yang %A Dawei Zhou %A Yujun Yan %B Proceedings of the 41st International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2024 %E Ruslan Salakhutdinov %E Zico Kolter %E Katherine Heller %E Adrian Weller %E Nuria Oliver %E Jonathan Scarlett %E Felix Berkenkamp %F pmlr-v235-huang24ac %I PMLR %P 20365--20381 %U https://proceedings.mlr.press/v235/huang24ac.html %V 235 %X 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.
APA
Huang, Z., Yang, Q., Zhou, D. & Yan, Y.. (2024). Enhancing Size Generalization in Graph Neural Networks through Disentangled Representation Learning. Proceedings of the 41st International Conference on Machine Learning, in Proceedings of Machine Learning Research 235:20365-20381 Available from https://proceedings.mlr.press/v235/huang24ac.html.

Related Material