Unifying Invariance and Spuriousity for Graph Out-of-Distribution via Probability of Necessity and Sufficiency

TL;DR

The paper addresses graph out-of-distribution generalization by proposing PNSIS, a framework that unifies invariant subgraph extraction with Probability of Necessity and Sufficiency (PNS) and an ensemble approach that leverages spurious subgraphs. It introduces dual invariant subgraph extractors (sufficient and necessary) whose outputs are regularized via a PNS-based upper bound that incorporates a Graph Structure Distance across environments, enabling robust invariant discovery. The method is extended with a spurious-subgraph classifier and an ensemble prediction strategy to further boost generalization, with practical implementations using GCN-based extractors and Monte Carlo estimation. Empirical results on synthetic and real-world graph benchmarks show substantial improvements over state-of-the-art methods, highlighting the practical impact of integrating causality-inspired invariants with spurious information for robust graph learning.

Abstract

Graph Out-of-Distribution (OOD), requiring that models trained on biased data generalize to the unseen test data, has a massive of real-world applications. One of the most mainstream methods is to extract the invariant subgraph by aligning the original and augmented data with the help of environment augmentation. However, these solutions might lead to the loss or redundancy of semantic subgraph and further result in suboptimal generalization. To address this challenge, we propose a unified framework to exploit the Probability of Necessity and Sufficiency to extract the Invariant Substructure (PNSIS). Beyond that, this framework further leverages the spurious subgraph to boost the generalization performance in an ensemble manner to enhance the robustness on the noise data. Specificially, we first consider the data generation process for graph data. Under mild conditions, we show that the invariant subgraph can be extracted by minimizing an upper bound, which is built on the theoretical advance of probability of necessity and sufficiency. To further bridge the theory and algorithm, we devise the PNSIS model, which involves an invariant subgraph extractor for invariant graph learning as well invariant and spurious subgraph classifiers for generalization enhancement. Experimental results demonstrate that our \textbf{PNSIS} model outperforms the state-of-the-art techniques on graph OOD on several benchmarks, highlighting the effectiveness in real-world scenarios.

Figures (5)

• Figure 1: Illustration of graph OOD methods with invariant subgraph learning, (a) In the training phase, the existing invariant extraction methods might lead to sufficient but not necessary subgraphs (pink block) and the necessary but not sufficient subgraphs (green block) according to two extreme optimization goals. The PNS invariant extraction method can extract the correct invariant subgraph. (b) In the test phase, the conventional invariant methods, that extract sufficient but not necessary latent subgraphs, might generate wrong subgraphs and further lead to suboptimal results (pink block). The methods that extract the necessary but not sufficient subgraphs might lead to the loss of semantic information (green block). When the noise-label data exists, the model might generate ambiguous predictions even if the correct subgraphs have been extracted. (yellow block). Ideal performance can be achieved by combining necessary and sufficient invariant subgraphs as well as spurious subgraphs.
• Figure 2: PIIF SCMarjovsky2019invariant. Inside this graph, the noises are omitted for brevity, where the grey and white nodes denote the latent and observed variables, respectively.
• Figure 3: The illustration of the PNSIS framework. (a) The left side of the figure denotes the subgraph extractors that are used to extract the sufficient and necessary invariant subgraphs by optimizing the upper bound. (b) The right side of the figure denotes the ensemble inference phase, which includes an invariant subgraph classifier and a spurious subgraph classifier
• Figure 4: Ablation study on the Molsider, Molbace, Moltox21, and Molbbbp datasets in the OGBG benchmark. We explore the impact of different components in the PNSIS method.
• Figure 5: Visualization of molecule examples in the OGBG benchmark. Modes with different colors denote different atoms, and edges denote different chemical bonds.

Theorems & Definitions (8)

• definition 1
• definition 2
• definition 3
• definition 4
• definition 5
• theorem 1
• theorem 2
• theorem 3