GraphSHINE: Training Shift-Robust Graph Neural Networks with Environment Inference

Qitian Wu, Fan Nie, Chenxiao Yang, Tianyi Bao, Junchi Yan

Keywords: graph neural networks, distribution shifts

Abstract: Graph neural networks (GNNs) have achieved remarkable performance across predictive tasks on graph-structured data. However, a critical issue gaining increasing attention is their performance degradation when faced with out-of-distribution (OOD) testing nodes. This challenge is exacerbated by the fact that distribution shifts on graphs involve intricate interconnections between nodes, and the environment labels are often absent in data. In this paper, we adopt a bottom-up data-generative perspective and reveal a key observation that the crux of GNNs' failure in OOD generalization lies in the latent confounding bias from the environment. The latter misguides the model to leverage environment-sensitive correlations between ego-graph features and target nodes' labels, resulting in undesirable generalization on new unseen nodes. Building upon this analysis, we introduce a novel, provably generalizable approach for training robust GNNs under node-level distribution shifts, without prior knowledge of environment labels. Our method resorts to a new learning objective that coordinates two key components: 1) an environment estimator that infers pseudo environment labels, and 2) a mixture-of-expert GNN predictor with feature propagation units conditioned on the pseudo environments. We show that the new approach can counteract the confounding bias in training data and facilitate learning shift-robust predictive relations. Extensive experiment demonstrates that our model can effectively enhance generalization with various types of distribution shifts and yield up to 27.4% accuracy improvement over other graph OOD generalization methods.

Track: Graph Algorithms and Learning for the Web

Submission Guidelines Scope: Yes

Submission Guidelines Blind: Yes

Submission Guidelines Format: Yes

Submission Guidelines Limit: Yes

Submission Guidelines Authorship: Yes

Student Author: Yes

Submission Number: 1744