Quantifying Distributional Invariance in Causal Subgraph for IRM-Free Graph Generalization

TL;DR

The paper tackles graph out-of-distribution generalization by proposing an IRM-free method, IDG, that discovers causal subgraphs under distributional shifts. It formalizes an Invariant Distribution Criterion, showing that the causal subgraph $G_c$ yields the smallest cross-environment shift and remains within training support, while shifts primarily affect the non-causal part $G_s$. A theoretical-empirical link between distribution shift and representation norms is established, motivating a norm-guided objective that maximizes subgraph representation norms to identify $G_c$ without environment labels. The IDG framework uses a GIN-based extractor/predictor, top-$r$ edge masking, and a two-stage training regimen, achieving state-of-the-art results on GraphOOD and DrugOOD while offering efficient inference. These findings advance robust graph learning by enabling IRM-free causal discovery and improving reliability under domain shifts.

Abstract

Out-of-distribution generalization under distributional shifts remains a critical challenge for graph neural networks. Existing methods generally adopt the Invariant Risk Minimization (IRM) framework, requiring costly environment annotations or heuristically generated synthetic splits. To circumvent these limitations, in this work, we aim to develop an IRM-free method for capturing causal subgraphs. We first identify that causal subgraphs exhibit substantially smaller distributional variations than non-causal components across diverse environments, which we formalize as the Invariant Distribution Criterion and theoretically prove in this paper. Building on this criterion, we systematically uncover the quantitative relationship between distributional shift and representation norm for identifying the causal subgraph, and investigate its underlying mechanisms in depth. Finally, we propose an IRM-free method by introducing a norm-guided invariant distribution objective for causal subgraph discovery and prediction. Extensive experiments on two widely used benchmarks demonstrate that our method consistently outperforms state-of-the-art methods in graph generalization.

Figures (11)

• Figure 1: Our intuition: causal subgraphs exhibit smaller distributional shifts across environments than non-causal ones. For instance, in Motif, the causal subgraph comprises three typical motifs across environments, while the base graph shows diverse. Similarly, molecular properties rely on a small, stable set of substructures while non-causal components vary widely across environments.
• Figure 2: (a) Average activation values and L2 norms in GNN on training (in-distribution) and test sets (out-of-distribution). (b) Accuracy and norms under varying perturbation ratios—higher ratios indicate greater distributional shifts. Motif-basis, Motif-size, and Graph-SST2 from left to right.
• Figure 3: SVD results of weight matrices in GIN (3 graph convolutional layers, each consists of 2 transforms) trained with ERM objective. Singular values are sorted in descending order.
• Figure 4: IDG framework. The extractor selects the subgraph $Z$ and feeds it to the predictor. Subgraph distributional shift are minimized by maximizing the representation norm of $Z$.
• Figure 5: From left to right: (1) Training loss vs. epochs in IDG and baseline (2) Testing loss vs. epochs in IDG and baseline (3) Testing Accuracy vs. epochs in IDG and baseline (4) t-SNE results of IDG representations (5) t-SNE results of baseline representations

Theorems & Definitions (10)

• Lemma 1
• Lemma 2
• Lemma 3
• Theorem 1
• Theorem 2
• proof
• proof
• proof
• proof
• proof