Pruning Spurious Subgraphs for Graph Out-of-Distribution Generalization

Le Song (GenBio AI MBZUAI) · Eric Xing (CMU/MBZUAI/GenBio) · Haoxuan Li (Peking University) · Tongliang Liu (Mohamed bin Zayed University of AI) · Tianjun Yao (Mohamed bin Zayed University of Artificial Intelligence) · Yongqiang Chen (MBZUAI/CMU) · Zhiqiang Shen (Mohamed bin Zayed University of Artificial Intelligence)
distribution shiftsepsilon-probability alignmentexperimental validationgraph neural networksgraph size constraintinvariant subgraphood generalizabilityout-of-distribution generalizationperformance degradationpredictive labelpruning-based methodregularization termsspurious edgesstate-of-the-art methodstheoretical analysis

Graph Neural Networks (GNNs) often encounter significant performance degradation under distribution shifts between training and test data, hindering their applicability in real-world scenarios. Recent studies have proposed various methods to address the out-of-distribution (OOD) generalization challenge, with many methods in the graph domain focusing on directly identifying an invariant subgraph that is predictive of the target label. However, we argue that identifying the edges from the invariant subgraph directly is challenging and error-prone, especially when some spurious edges exhibit strong correlations with the targets. In this paper, we propose $\texttt{PrunE}$, the first pruning-based graph OOD method that eliminates spurious edges to improve OOD generalizability. By pruning spurious edges, $\texttt{PrunE}$ retains the invariant subgraph more comprehensively, which is critical for OOD generalization. Specifically, $\texttt{PrunE}$ employs two regularization terms to prune spurious edges: 1) _graph size constraint_ to exclude uninformative spurious edges, and 2) _$\epsilon$-probability alignment_ to further suppress the occurrence of spurious edges. Through theoretical analysis and extensive experiments, we show that $\texttt{PrunE}$ achieves superior OOD performance and outperforms previous state-of-the-art methods significantly.