[Paper Review] Handling Distribution Shifts on Graphs: An Invariance Perspective

Qitian Wu, Hengrui Zhang+2 more|arXiv (Cornell University)|Feb 5, 2022

Health, Environment, Cognitive Aging30 citations

TL;DR

The paper formulates out-of-distribution generalization for node-level graph predictions and introduces Explore-to-Extrapolate Risk Minimization (EERM), using adversarial graph editors to enforce invariant features and improve GNN robustness under distribution shifts.

ABSTRACT

There is increasing evidence suggesting neural networks' sensitivity to distribution shifts, so that research on out-of-distribution (OOD) generalization comes into the spotlight. Nonetheless, current endeavors mostly focus on Euclidean data, and its formulation for graph-structured data is not clear and remains under-explored, given two-fold fundamental challenges: 1) the inter-connection among nodes in one graph, which induces non-IID generation of data points even under the same environment, and 2) the structural information in the input graph, which is also informative for prediction. In this paper, we formulate the OOD problem on graphs and develop a new invariant learning approach, Explore-to-Extrapolate Risk Minimization (EERM), that facilitates graph neural networks to leverage invariance principles for prediction. EERM resorts to multiple context explorers (specified as graph structure editers in our case) that are adversarially trained to maximize the variance of risks from multiple virtual environments. Such a design enables the model to extrapolate from a single observed environment which is the common case for node-level prediction. We prove the validity of our method by theoretically showing its guarantee of a valid OOD solution and further demonstrate its power on various real-world datasets for handling distribution shifts from artificial spurious features, cross-domain transfers and dynamic graph evolution.

Motivation & Objective

• Formulate the OOD problem for node-level tasks on graphs by exploiting ego-graphs and local data generation.
• Develop an invariance-based learning paradigm to extract stable, invariant features for graph predictions.
• Propose EERM to simulate multiple environments via adversarial graph edits and train models to minimize mean and variance of risks.
• Provide theoretical guarantees linking invariance-driven objectives to valid OOD solutions and upper bounds on OOD error.
• Empirically validate EERM across diverse real-world graph datasets with various GNN backbones and distribution shift scenarios.

Proposed method

• Define ego-graphs G_v around target nodes to decompose graph data for node-level OOD, enabling a two-step sampling view.
• Propose Assumption 1 (Invariance Property) and a recursive, BFS-tree-based readout r_v to capture invariant predictive features.
• Introduce the EERM objective that minimizes a weighted combination of the mean and variance of risks across multiple environments (Eq. 4).
• Use K auxiliary context generators g_wk(G) to produce K graph views via adversarial graph edits; train generators to maximize risk variance while the predictor minimizes it (Eq. 5).
• Employ REINFORCE to optimize edge-editing actions when constructing graph views, enabling tractable learning despite non-differentiability.
• Provide theoretical results (Theorems 1–3) linking the invariant representation to optimal OOD solutions and bounding OOD generalization error.

Experimental results

Research questions

• RQ1 How can OOD generalization be formulated for node-level tasks on graphs given the inter-connected data and graph structure?
• RQ2 Can an invariance-based representation be learned for graphs that yields stable prediction across environments?
• RQ3 Does enforcing mean-variance risk minimization over adversarial graph edits improve extrapolation to unseen environments?
• RQ4 What theoretical guarantees connect the proposed invariant learning objective to valid OOD solutions and bounded OOD error?
• RQ5 Do experiments across diverse datasets and GNN backbones show robust improvements under distribution shifts?

Key findings

• The proposed EERM method consistently outperforms standard empirical risk minimization on several datasets under artificial, cross-domain, and dynamic graph shifts.
• Invariant readouts based on recursive BFS-tree aggregation capture stable predictive features across environments.
• Adversarial graph edits generated by K context generators, optimized with REINFORCE, effectively simulate multiple environments for robust learning.
• Theoretical analyses establish that minimizing the EERM objective enforces sufficiency and invariance in the learned representation, leading to valid OOD solutions and reduced OOD error bounds.
• Empirical results show EERM improves generalization across multiple datasets and GNN backbones (e.g., GCN, GAT, GraphSAGE, GCNII, GPRGNN) under various distribution shifts.