Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Handling Distribution Shifts on Graphs: An Invariance Perspective

Qitian Wu, Hengrui Zhang, Junchi Yan, David Wipf

TL;DR

This work tackles out-of-distribution generalization for node-level prediction on graphs, where inter-node dependencies and structural information complicate standard OOD solutions. It introduces Explore-to-Extrapolate Risk Minimization (EERM), which employs multiple adversarial graph editors to synthesize diverse virtual environments from a single observed graph and optimizes a bilevel objective that minimizes risk variance across environments and the mean risk with respect to the data. The authors provide invariance- and information-theoretic analysis showing that enforcing invariance and sufficiency of a learned representation yields a valid OOD solution and bounds the OOD error. Extensive experiments across datasets with artificial shifts, cross-domain transfers, and dynamic graph evolution demonstrate that EERM consistently outperforms ERM and robustly generalizes to unseen or evolving graphs, with backbone-agnostic applicability to GCN, GAT, GraphSAGE, and related architectures.

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.

Handling Distribution Shifts on Graphs: An Invariance Perspective

TL;DR

This work tackles out-of-distribution generalization for node-level prediction on graphs, where inter-node dependencies and structural information complicate standard OOD solutions. It introduces Explore-to-Extrapolate Risk Minimization (EERM), which employs multiple adversarial graph editors to synthesize diverse virtual environments from a single observed graph and optimizes a bilevel objective that minimizes risk variance across environments and the mean risk with respect to the data. The authors provide invariance- and information-theoretic analysis showing that enforcing invariance and sufficiency of a learned representation yields a valid OOD solution and bounds the OOD error. Extensive experiments across datasets with artificial shifts, cross-domain transfers, and dynamic graph evolution demonstrate that EERM consistently outperforms ERM and robustly generalizes to unseen or evolving graphs, with backbone-agnostic applicability to GCN, GAT, GraphSAGE, and related architectures.

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.
Paper Structure (40 sections, 9 theorems, 34 equations, 17 figures, 4 tables, 1 algorithm)

This paper contains 40 sections, 9 theorems, 34 equations, 17 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Problem Formulation
  3. Out-of-distribution Problem for Graph-Structured Data
  4. Invariant Features for Node-Level Prediction on Graphs
  5. Methodology
  6. Motivating Example
  7. Stable Learning with Explore-to-Extrapolate Risk Minimization
  8. Theoretical Discussions
  9. Relationship between Invariance Principle and OOD Problem
  10. Information-Theoretic Error for OOD Generalization
  11. Experiments
  12. Handling Distribution Shifts with Artificial Transformation
  13. Generalizing to Unseen Domains
  14. Extrapolating over Dynamic Data
  15. Discussions with Existing Works
  16. ...and 25 more sections

Key Result

Proposition 1

Let the risk under environment $e$ be $R(e) = \frac{1}{|V|} \sum_{v\in V} \mathbb E_{\mathbf y| \mathbf G_{\mathbf v} = G_v}[\| \hat{y}_v - y_v \|_2^2]$. The unique optimal solution for objective $\min_{\theta} \mathbb E_{\mathbf e}[R(e)]$ would be $[\theta_1, \theta_2] = [\frac{1+\sigma_e^2}{2+\sig

Figures (17)

  • Figure 1: (a) The proposed approach Explore-to-Extrapolate Risk Minimization which entails $K$ context generators that generate graph data of different (virtual) environments based on input data from a single (real) environment. The GNN model is updated via gradient descent to minimize a weighted combination of mean and variance of risks from different environments, while the context generators are updated via REINFORCE to maximize the variance loss. (b) Illustration for our Assumption \ref{['assump-1']}. (c) The dependence among variables in the motivating example in Section \ref{['sec-method-1']}.
  • Figure 2: Results on Cora with artificial distribution shifts. We run each experiment with 20 trials. (a) The (distribution of) test accuracy of vanilla GCN using our approach for training and using Erm. (b) The (averaged) accuracy on the training set (achieved by the epoch where the highest validation accuracy is achieved) when using all the input node features and removing the spurious ones for inference. (c) The (averaged) test accuracy with different GNNs for data generation.
  • Figure 3: Experiment results on Amazon-Photo with artificial distribution shifts.
  • Figure 4: Test ROC-AUC on Twitch where we compare different GNN backbones.
  • Figure 5: Test F1 score on Elliptic where we group graph snapshots into 9 test sets (T1$\sim$ T9).
  • ...and 12 more figures

Theorems & Definitions (13)

  • Proposition 1
  • Proposition 2
  • Theorem 1
  • Theorem 2
  • Theorem 3
  • Lemma 1
  • proof
  • Lemma 2
  • proof
  • Lemma 3
  • ...and 3 more