[Paper Review] Learning Causally Invariant Representations for Out-of-Distribution Generalization on Graphs

Yongqiang Chen, Yonggang Zhang+7 more|arXiv (Cornell University)|Feb 11, 2022

Health, Environment, Cognitive Aging49 citations

TL;DR

CIGA proposes a causality-inspired framework to extract invariant subgraphs that preserve label-causing information, enabling provable OOD generalization for graph classification without environment labels. It introduces FIIF/PIIF-based causal models and two practical objectives to identify invariant subgraphs and improve OOD robustness, with strong empirical results on 16 datasets including DrugOOD.

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.

Motivation & Objective

• Motivate OOD generalization on graphs and the challenges from structure- and attribute-level shifts.
• Propose CIGA to identify invariant subgraphs that encode the causes of labels under FIIF/PIIF SCMs.
• Provide information-theoretic objectives and practical implementations to extract invariant subgraphs.
• Theoretically justify and empirically validate OOD guarantees across diverse graph datasets.

Proposed method

• Decompose GNN into a featurizer g that extracts an invariant subgraph G_c and a classifier f_c that predicts Y from G_c.
• Formulate an information-theoretic objective maximizing I(G_c_hat; Y) with a constraint to keep G_c_hat independent of the environment E.
• Introduce CIGA v2 by adding a mutual-information-based constraint to minimize leakage of G_s information into G_c_hat while preserving predictive power.
• Approximate mutual information terms with contrastive learning objectives (Eq. 5) for practical training.
• Provide theoretical guarantees (Theorem 3.1) that solutions to the proposed objectives yield invariant GNNs under FIIF/PIIF SCMs.
• Discuss practical implementation choices using interpretable GNN architectures and hinge-like losses to enforce information bounds.

Experimental results

Research questions

• RQ1 How can invariance principles be extended to graph data to guarantee OOD generalization?
• RQ2 Can we identify invariant subgraphs G_c that capture the causes of labels and are independent of environment shifts?
• RQ3 How can information-theoretic objectives be operationalized without environment labels to extract invariant structure?
• RQ4 Do the proposed CIGA objectives yield provable OOD guarantees across diverse graph shifts and real-world datasets?

Key findings

• CIGA significantly improves OOD performance over existing methods across 16 synthetic and real-world datasets.
• CIGA v2 achieves notable gains, e.g., avg improvements over competing methods in reported experiments.
• The framework provides theoretical guarantees that invariant GNNs emerge under FIIF/PIIF SCMs (Theorem 3.1).
• Contrastive-like approximations enable practical optimization of mutual-information objectives without environment labels.
• CIGA outperforms strong baselines on challenging DrugOOD settings.