Learning Causally Invariant Representations for Out-of-Distribution Generalization on Graphs

TL;DR

This paper introduces CIGA, a framework for learning causally invariant subgraph representations to improve out-of-distribution generalization on graph data, addressing challenges posed by diverse distribution shifts.

Contribution

CIGA leverages causal models and an information-theoretic objective to identify invariant subgraphs, enhancing OOD robustness without requiring environment labels.

Findings

CIGA outperforms baselines on 16 datasets.

It achieves superior OOD generalization in drug discovery tasks.

The method effectively captures invariant causal subgraphs.

Abstract

Despite recent success in using the invariance principle for out-of-distribution (OOD) generalization on Euclidean data (e.g., images), studies on graph data are still limited. Different from images, the complex nature of graphs poses unique challenges to adopting the invariance principle. In particular, distribution shifts on graphs can appear in a variety of forms such as attributes and structures, making it difficult to identify the invariance. Moreover, domain or environment partitions, which are often required by OOD methods on Euclidean data, could be highly expensive to obtain for graphs. To bridge this gap, we propose a new framework, called Causality Inspired Invariant Graph LeArning (CIGA), to capture the invariance of graphs for guaranteed OOD generalization under various distribution shifts. Specifically, we characterize potential distribution shifts on graphs with causal models, concluding that OOD generalization on graphs is achievable when models focus only on subgraphs containing the most information about the causes of labels. Accordingly, we propose an information-theoretic objective to extract the desired subgraphs that maximally preserve the invariant intra-class information. Learning with these subgraphs is immune to distribution shifts. Extensive experiments on 16 synthetic or real-world datasets, including a challenging setting -- DrugOOD, from AI-aided drug discovery, validate the superior OOD performance of CIGA.