PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yongqiang Chen, Yonggang Zhang, Yatao Bian, Han Yang, Kaili Ma, Binghui Xie, Tongliang Liu, Bo Han, James Cheng

arXiv:2202.05441v3cs.LG

TL;DR

Graph OOD generalization is difficult because shifts can affect attributes and structures, while environment partitions may be unavailable. CIGA models these shifts causally and extracts invariant, label-critical subgraphs; across 16 datasets, it outperforms existing methods by up to 10%, under stated theoretical assumptions.

  • Problem

    Graph OOD studies are limited, while graph shifts across attributes and structures and unavailable environment partitions complicate identifying invariant information.

  • Method

    CIGA uses structural causal models and a decomposed featurizer-classifier GNN to extract invariant subgraphs and predict labels from them.

  • Results

    Across 16 synthetic and real-world datasets with varied shifts, CIGA significantly outperforms existing methods by up to 10%.

  • Takeaways & Limitations

    The paper shows that focusing on invariant subgraphs can provide OOD generalization guarantees under the modeled graph-generation assumptions.

  • Takeaways & Limitations

    The theoretical objective can require an unknown invariant-subgraph size, and unconstrained optimization may select the full graph as a trivial solution.

Abstract

from arXiv · show

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.

1 Introduction

Graph distribution shifts can occur across attributes, structures, sizes, and other properties, challenging invariant identification and weakening OOD generalization. CIGA addresses this by modeling graph shifts causally and extracting invariant subgraphs, with experiments showing improved performance across diverse datasets.

  • Graph representation learning often assumes training and test graphs share a distribution, but real distribution shifts can seriously degrade deployed models.
  • Graph shifts can affect attributes and structures through graph sizes, subgraph densities, homophily, and different label-correlated modes, complicating invariant identification.
  • CIGA models graph distribution shifts with three structural causal models and seeks invariant, label-critical subgraphs containing most information about label causes.
  • CIGA decomposes graph learning into invariant subgraph identification and label prediction, using a contrastive objective that maximizes intra-class mutual information.
  • 10%: CIGA significantly outperforms existing methods by up to 10% across 16 synthetic and real-world datasets with varied distribution shifts.

2 OOD Generalization on Graphs through the Lens of Causality

The paper characterizes graph generation and distribution shifts with structural causal models separating invariant and varying latent factors. Analysis shows that ERM and existing OOD objectives can fail under entangled shifts and limited support or unavailable environment labels.

  • OOD graph generalization seeks a GNN trained on training environments that performs well across unseen environments, but requires assumptions about those environments.
  • 2.2 Graph Generation Process: The graph generation model separates latent variables into invariant C and environment-affected S, which jointly control observed graph generation.
  • 2.2 Graph Generation Process: FIIF and PIIF structural causal models represent distinct interactions between invariant features, spurious features, environments, and labels.
  • 2.3 Challenges of OOD Generalization on Graphs: ERM can overfit graph shortcuts, causing performance to drop dramatically as data bias strengthens and producing greater variance when graph-size shifts are mixed.
  • 2.3 Challenges of OOD Generalization on Graphs: IRM and related OOD objectives can also fail, including catastrophically when test environments lack sufficient support overlap with training environments.
  • 2.3 Challenges of OOD Generalization on Graphs: Existing methods face entangled graph-specific spurious correlations and usually unavailable environment labels, motivating an invariance principle adapted to graphs.

3 Invariance Principle for OOD Generalization on Graphs

CIGA aligns a featurizer and classifier with causal graph-generation mechanisms to extract invariant subgraphs for label prediction. Its objectives and theorem provide invariant-GNN guarantees under stated assumptions, while the framework has size and implementation limitations.

  • CIGA generalizes the invariance principle to graphs under structural causal models and provides theoretical guarantees for OOD generalization.
  • 3.2 Invariant Graph Learning Framework: The framework decomposes a GNN into a featurizer g that identifies a subgraph and a classifier f_c that predicts labels from it.
  • 3.2 Invariant Graph Learning Framework: CIGA aims to maximize label information in the extracted subgraph while making it independent of environment effects, despite unavailable environment labels.
  • 3.2 Invariant Graph Learning Framework: When invariant-subgraph size is unknown, unconstrained mutual-information optimization can include spurious features, with the full graph becoming a trivial solution.
  • 3.2 Invariant Graph Learning Framework: The CIGAv2 objective adds label information from the complementary subgraph to discourage spurious information from entering the estimated invariant subgraph.
  • 3.3 Theoretical Analysis and Practical Discussions: Under the theorem’s assumptions, every solution to Eq. 3 with fixed invariant-subgraph size and every solution to Eq. 4 induces an invariant GNN.
  • 3.3 Theoretical Analysis and Practical Discussions: CIGA’s architecture admits multiple implementations, but the paper uses interpretable GNN prototypes and notes that existing interpretable models do not generally provide the same OOD guarantees.

4 Empirical Studies

The empirical studies evaluate CIGA across synthetic and real-world graph datasets, distribution shifts, and comparison methods. Results show strong OOD performance and robustness to key hyperparameters.

  • Experimental setup: CIGA is compared with ERM, interpretable GNNs, OOD objectives, and methods that do not require environment labels.
  • Experimental setup: Evaluation reports accuracy, ROC-AUC for DrugOOD, and Matthews correlation coefficient for TU datasets using repeated validation-based model selection.
  • OOD performance: Up to 10%, CIGAv1 and CIGAv2 outperform all baselines with lower variance on structure and mixed distribution shifts.
  • Ablation comparisons: Neither DIR nor CNC achieves satisfactory OOD performance across experiments, supporting CIGA’s decomposition and objective.
  • Hyperparameter sensitivity: Both CIGAv1 and CIGAv2 remain robust across α and β values, although extreme coefficients can weaken invariance or optimization.

5 Conclusions

The paper studies graph OOD generalization through causality and introduces CIGA to model distribution shifts and learn invariant graph representations. The framework shows promising theoretical and empirical performance across varied shifts.

  • CIGA models potential graph distribution shifts with structural causal models and generalizes the invariance principle to graph classification.
  • The framework demonstrates promising theoretical and empirical OOD generalization ability under a variety of distribution shifts.

Checklist

The paper frames graph OOD generalization using causal and invariant-learning concepts, relates CIGA to prior graph methods, and documents assumptions, limitations, and reproducibility details.

  • Checklist: The checklist states that the paper describes its contributions, limitations, theoretical assumptions, proofs, compute resources, and reproducibility materials.
  • Related work: The paper distinguishes CIGA from explainability methods, node-classification formulations, and Euclidean OOD approaches by targeting graph classification under multiple shifts.
  • Causality and OOD generalization: The invariance principle relies on independent causal mechanisms and seeks representations of causal relationships that remain stable across environments.
  • Limitations: CIGA’s expressivity is constrained by the graph encoder’s ability to distinguish invariant subgraphs and non-isomorphic invariant graphs across classes.

B.4 Discussions on limitations of CIGA and future directions

The discussion identifies practical limitations of CIGA’s graph-generation modeling, contrastive sampling, architecture choices, training cost, and parameter tuning.

  • CIGA’s generic graph-generation SCM may benefit from practitioners’ inductive knowledge for more precise modeling and better OOD generalization.
  • The implementation uses supervised contrastive learning without sophisticated augmentation or sampling, leaving improved contrastive strategies as future work.
  • The interpretable architectures used for concept verification require more training time and may be unsuitable for resource-limited settings such as Edge-AI.
  • CIGA requires additional tuning effort, and its approximation may be limited by the chosen architectures.

C Full Structural Causal Models on Graph Generation

The full SCMs model graph generation by separating invariant and varying latent factors, their interactions, and how they produce observed graph structure and attributes. The framework covers multiple interaction modes and can specialize to graph families when additional generative knowledge is available.

  • Graph generation SCM: The latent graph-generation model partitions Z into invariant C and environment-affected S, with E denoting environments.C can control functional groups, while S captures varying factors analogous to style in image models.
  • Graph generation SCM: The graph-generation SCM maps latent variables to invariant and spurious subgraphs, then combines them into the observed graph G.Separate generation mappings control adjacency and features for Gc and Gs before a final mapping produces G.
  • Distribution shifts: Structure-level properties such as degrees, sizes, and subgraph densities, and attribute-level properties such as homophily, can vary across environments and spuriously correlate with labels.These correlations motivate modeling multiple interaction modes rather than assuming a single form of graph shift.
  • Interaction modes: FIIF, PIIF, and MIIF distinguish whether spurious features are directly or indirectly controlled by invariant factors, or contain both interaction types.FIIF satisfies (S, E) ⊥⊥Y |C, while MIIF partitions S into directly and indirectly controlled components.
  • Specialized SCMs: The SCMs are designed to cover many graph-generation processes and can generalize to specialized cases when additional knowledge about the graph family is incorporated.The graph-generation and FIIF SCMs together generalize to the graphon SCM, while prior graph models resemble FIIF or PIIF settings.

D.2 More discussions about failure case study in Sec. 2.3

The failure-case study examines whether deeper message passing or more expressive GNN architectures resolve graph OOD failures. The reported discussions indicate that stronger aggregation and expressivity can remain sensitive to spurious structure and attributes, especially under mixed shifts.

  • Failure-case study: Experiments test whether ERM, IRM, deeper message passing, and more expressive GNNs improve OOD generalization on synthetic graph classification tasks.The study evaluates models classifying motif-containing graphs under controlled distribution shifts.
  • Adding more message passing turns: More message passing can denoise some shifts, but OOD performance drops more sharply as bias increases when spurious features spread through the graph.Mixed structure- and attribute-level shifts make these spurious correlations harder to disentangle.
  • Adding more message passing turns: Deeper ERM-trained message-passing GNNs can mix invariant and spurious node information, making OOD generalization more difficult or even impossible.Topology features encoded by hidden layers may also become spuriously correlated with labels through varying latent factors.
  • Using more expressive GNNs: Higher GNN expressivity does not by itself remove OOD vulnerability when labels are spuriously correlated with subgraphs.The authors report similar failures for k-GNNs, despite their greater expressivity than basic GNN variants.

D.3 More empirical results about failure case study in Sec. 2.3

The case study shows that existing GNNs fail under increasingly complex graph distribution shifts, while ERM and IRM can fail theoretically when test support is not covered by training support. These findings motivate learning invariant subgraphs, for which CIGA’s objective yields an invariant predictor under the stated graph-generation assumptions.

  • Empirical failure cases: All existing GNN methods are sensitive to increasingly strong structure-level and attribute-level spurious correlations.As spurious correlation grows, models increasingly overfit shortcuts in either graph structure or node attributes.
  • Empirical failure cases: Entangling attribute and graph-size shifts produces worse and more unstable GNN performance than single shifts.Higher variance signals unstable OOD behavior as multiple shift types interact.
  • Empirical failure cases: More message-passing layers and more expressive k-GNN architectures do not make OOD performance immune to distribution shifts.Deeper or more expressive models can overfit shortcuts more sharply as spurious-correlation strength increases.
  • Theoretical failure cases: ERM and IRM fail to generalize OOD almost surely for the stated graph settings when test-environment support is not contained in training-environment support.The failure applies both to graph features generated by the FIIF model and to graphs with globally identical node features and corresponding labels.
  • Objective behavior: The information-theoretic baseline can have unstable convergence because estimating I(Ĝc, G) is difficult for complex graphs, whereas CIGA’s contrastive objective is reported to converge more stably.The contrastive objective is presented as a stability advantage relative to the information-bottleneck estimation.
  • Invariant-subgraph motivation: A predictor based on invariant subgraphs can minimize worst-case risk across environments because P(Y | Gc) remains invariant under the graph-generation assumptions.This motivates CIGA variants that extract invariant subgraphs without requiring environment information.

E.2 Proof for theorem 3.1 (i)

Under the theorem’s assumptions, the ground-truth invariant subgraph Gc uniquely optimizes the objective, yielding an invariant GNN predictor.

  • The theorem assumes invertible graph-generation mappings and equally distributed samples across training environments.
  • The proof interprets the first term as preserving label information and the second as enforcing within-class stability across environments.
  • Maximizing I(bGc; Y) identifies a subgraph retaining label information, with bGc = Gc achieving I(C; Y) = H(Y).
  • The conditional mutual-information objective I(bGc; eGc|Y) eliminates spurious subsets from Gs because their inclusion lowers the objective.
  • The ground-truth Gc is the only maximizer of Eq. 15, so solving the objective elicits an invariant GNN.

E.3 Proof for theorem 3.1 (ii)

For CIGAv2, the combined information objective identifies the underlying invariant subgraph and thereby produces an invariant GNN under the stated assumptions.

  • The theorem assumes invertible generation mappings and equally distributed samples across training environments.
  • The invariant subgraph carries at least as much label information as the spurious subgraph, I(Gc; Y) ≥ I(Gs; Y).
  • The proof considers estimated subgraphs containing only invariant portions or invariant portions mixed with spurious portions.
  • Maximizing conditional mutual information prevents adding spurious subsets when those subsets provide no additional benefit to the objective.
  • Only the underlying Gc solves Eq. 28, and its minimizer elicits an invariant GNN predictor.

F Details of Prototypical CIGA Implementation

The prototypical implementation decomposes CIGA into a subgraph featurizer and classifier, using edge-weight sampling and contrastive information objectives.

  • CIGA uses interpretable architectures experimentally, while more sophisticated GNN architectures can be incorporated.
  • CIGA decomposes graph classification into a featurizer g: G → Gc and a classifier fc: Gc → Y.
  • The featurizer predicts soft edge weights, and the largest-weight edges form the estimated invariant subgraph when a sampling ratio is selected.
  • CIGAv2 combines invariant-subgraph prediction, contrastive mutual-information estimation, and a spurious-subgraph objective, whereas CIGAv1 omits the third term.
  • The contrastive objective uses same-label extracted subgraphs as positive samples and different-label subgraphs as negative samples.

G.1 Details about the datasets

The experiments evaluate CIGA across synthetic and real-world graph distribution shifts, with standardized architectures, optimization settings, baselines, and repeated validation-based model selection.

  • Datasets: The experiments cover synthetic motif datasets, TU graph-size shifts, sentiment graphs, ColoredMnist graph shifts, and DrugOOD binding-affinity datasets.
  • Datasets: SPMotif uses three graph motifs—House, Cycle, and Crane—and injects distribution shifts into training data while keeping validation and testing distributions unchanged.
  • Datasets: TU experiments use Matthews correlation coefficient because the generated splits have class imbalance.
  • Implementation: All methods use the same graph-encoder architecture, with default three-layer GNNs, batch normalization, and jumping-knowledge residual connections.
  • Implementation: CIGA searches selected loss weights by validation performance without exhaustive hyperparameter tuning, while baseline implementations use corresponding validation-based choices.
  • Evaluation: Experiments repeat runs across datasets, select models by validation performance, and report the corresponding mean and standard deviation.

G.4 Additional Analysis

Additional analyses find CIGA robust across hyperparameter settings, effective in a single training environment, and computationally efficient while providing interpretable predictions.

  • Hyperparameter sensitivity: CIGAv1 and CIGAv2 remain robust to different α and β values across datasets and distribution shifts.Very small α may fail to guarantee invariant subgraphs, while very small β can retain spurious-subgraph parts; overly large β may eliminate invariant-subgraph parts.
  • Running time analysis: CIGA is the only OOD method that outperforms ERM by a non-trivial margin with relatively low additional computational overhead.Training-time comparisons were conducted on DrugOOD-Scaffold with the GNN backbone, batch size, devices, and hardware fixed.
  • Single environment OOD generalization: Both CIGAv1 and CIGAv2 maintain state-of-the-art performance in the single training environment setting.The authors hypothesize that mutual-information enforcement between estimated invariant subgraphs helps retain invariance in this setting.
  • Interpretation visualization: CIGA automatically provides interpretations for predictions through its interpretable GNN architecture, which may facilitate human understanding in practice.Visualizations color ground-truth motif nodes pink and use edge-color intensity to represent relative attention strength.
  • Interpretation visualization: DrugOOD visualizations show attention patterns for molecules sharing labels, potentially providing insights to domain experts developing new drugs.The visualizations cover activate and inactivate examples across Assay, Scaffold, and Size settings.
Loading 2202.05441v3…