Joint Learning of Label and Environment Causal Independence for Graph Out-of-Distribution Generalization (2306.01103v3)

Abstract: We tackle the problem of graph out-of-distribution (OOD) generalization. Existing graph OOD algorithms either rely on restricted assumptions or fail to exploit environment information in training data. In this work, we propose to simultaneously incorporate label and environment causal independence (LECI) to fully make use of label and environment information, thereby addressing the challenges faced by prior methods on identifying causal and invariant subgraphs. We further develop an adversarial training strategy to jointly optimize these two properties for causal subgraph discovery with theoretical guarantees. Extensive experiments and analysis show that LECI significantly outperforms prior methods on both synthetic and real-world datasets, establishing LECI as a practical and effective solution for graph OOD generalization. Our code is available at https://github.com/divelab/LECI.

• The paper introduces LECI, a framework that exploits causal independence between labels and environments using a novel subgraph selector and staged adversarial training.
• Its subgraph selection mechanism isolates invariant features, improving predictions across graph datasets with significant structural and feature shifts.
• Empirical evaluations on GOOD benchmarks demonstrate that LECI outperforms multiple baseline methods, highlighting its potential for real-world applications.

Joint Learning of Label and Environment Causal Independence for Graph Out-of-Distribution Generalization

This paper presents an innovative method called LECI (Label and Environment Causal Independence), which is targeted at enhancing Out-of-Distribution (OOD) generalization for graph data. It addresses the challenge of achieving robustness against covariate shifts that often degrade the performance of Graph Neural Networks (GNNs) when deployed in unknown environments.

Methodology

The core of the LECI framework is designed around a subgraph selector that functions to isolate subgraphs, thereby targeting invariant causal features for robust prediction. This approach leverages the causal independence between labels and environment features to improve OOD generalization. A key component is the adversarial training mechanism where discriminators are trained to decouple environment-specific information from node representations.

• Subgraph Selector: The method employs a novel subgraph selector that adapts from the regular graph to a subgraph, focusing on correlating subnetworks essential for invariant prediction. This component uses edge probabilistic masks derived from node embeddings to enhance differentiation capabilities.
• Pure Feature Shift Consideration (PFSC): This technique involves a transformer-based model that eliminates environment-specific biases from node features, thereby aiding in the generation of generalizable node representations across different data distributions.
• Adversarial Training Strategy: The paper distinguishes itself by employing a staged adversarial training strategy. Initially, the discriminators for label and environment are trained independently until reaching stability, followed by integrating these with the subgraph selector to effectively enforce environment and label independence.

Experimental Evaluation

The empirical evaluations are extensive, involving a suite of benchmarks from the GOOD datasets as well as custom datasets like GOOD-Twitter and GOOD-Motif2, which were specifically constructed by following the original dataset split strategies. The results confirm that LECI consistently outperforms existing baselines, including IRM, VREx, Coral, DANN, and graph-specific methods like DIR, GSAT, and CIGA, especially on synthetic datasets where the spurious correlation filtering is crucial.

Importantly, on datasets with severe structure and feature shifts, LECI showcased its capability by achieving marked improvements over other models, even under stringent conditions of test data irrelevance for hyperparameter tuning.

Implications and Future Prospects

The paper's contributions hold substantial promise for advancing approaches in GNNs regarding their deployment in real-world scenarios characterized by domain shifts. The subgraph-centric model with capability for decomposing influential factors presents theoretical robustness as well as practical applicability. Additionally, its framework may be vital for fields like drug discovery or social network analysis where domain-specific biases are prevalent.

Future research could explore the integration of other invariant learning techniques in conjunction with LECI to further boost generalizability. Moreover, the method presents opportunities for further optimization in terms of computational efficiency given its intricate adversarial constructs.

Thus, the LECI framework is a significant step forward in crafting models resilient to domain shifts in graphical data, contributing to the foundational understanding of causality in machine learning applications within diverse environments.