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

TL;DR

This paper introduces an IRM-free approach for identifying causal subgraphs in graphs, leveraging distributional invariance to improve out-of-distribution generalization without environment annotations.

Contribution

It formalizes the Invariant Distribution Criterion, establishes the relationship between distributional shift and representation norm, and proposes a norm-guided invariant distribution method.

Findings

Outperforms state-of-the-art methods on benchmark datasets

Effectively captures causal subgraphs with smaller distributional variations

Provides theoretical proof of the Invariant Distribution Criterion

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.