Diffusion-based Graph Masked Autoencoders for Out-of-Distribution Generalization
Keywords: Machine Learning; Deep learning; Graph learning; Self-supervised learning
Abstract: Graph Out-of-Distribution (GraphOOD) problems have become increasingly significant in the field of graph neural networks. Graph Neural Networks (GNNs) are particularly vulnerable to performance degradation when facing distribution shifts. This is due to the intricate interconnections between nodes in graph data and the lack of environmental labels, making it difficult to ensure model reliability. Recent advances in computer vision have shown that Diffusion Models(DMs) have strong generalization capabilities, providing a natural advantage in mitigating the effects of distribution shifts. Specifically, DMs can effectively capture and generate details of data distributions through a stepwise denoising process, thereby enhancing model robustness. However, applying diffusion to GraphOOD problems presents challenges, such as learning invariant knowledge that remains unaffected by distribution shifts. To address this, we propose a diffusion-based pre-training model for GraphOOD, termed $\textbf{D}$iffusion-based $\textbf{M}$asked $\textbf{A}$uto$\textbf{E}$ncoders on Graph Out-of-Distribution Generalization (DiffGMAE). Firstly, we propose a novel empirical risk minimization (ERM) approach that enhances the data by progressively adding noise, called the NoisedERM module, which aims to learn invariant features and avoid corrupting the discrete information of the original graph. Then, we design a self-supervised learning module called DiGMAE, which replaces the traditional MAE decoder with a diffuse-based denoising process. The aim is to use the invariant features obtained by NoisedERM for conditional diffusion and improve the robustness of the model in a self-supervised way to cope with the distribution shift of GraphOOD problem. We demonstrate significant improvements in DiffGMAE on OOD benchmarks. In addition, our ablation experiments show that the diffusion process is superior to traditional graph generation methods in solving OOD problems. The implementation code is available in \textbf{Supplementary material} for reproducibility.
Primary Area: transfer learning, meta learning, and lifelong learning
Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Reciprocal Reviewing: I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.
Submission Number: 14135
Loading