Graph Data Augmentation with Contrastive Learning on Covariate Distribution Shift

TL;DR

Graph classification under covariate distribution shift is improved by MPAIACL, which leverages contrastive learning in latent space to disentangle stable from environmental features. Building on AIA, it strengthens both the stable feature generator and the adversarial augmenter with InfoNCE and triplet/Wasserstein objectives, respectively. The approach is theoretically motivated and empirically validated on diverse OOD datasets, showing robust generalization and competitive gains over baselines. This work highlights the practical value of latent-space information for out-of-distribution robustness in graph learning.

Abstract

Covariate distribution shift occurs when certain structural features present in the test set are absent from the training set. It is a common type of out-of-distribution (OOD) problem, frequently encountered in real-world graph data with complex structures. Existing research has revealed that most out-of-the-box graph neural networks (GNNs) fail to account for covariate shifts. Furthermore, we observe that existing methods aimed at addressing covariate shifts often fail to fully leverage the rich information contained within the latent space. Motivated by the potential of the latent space, we introduce a new method called MPAIACL for More Powerful Adversarial Invariant Augmentation using Contrastive Learning. MPAIACL leverages contrastive learning to unlock the full potential of vector representations by harnessing their intrinsic information. Through extensive experiments, MPAIACL demonstrates its robust generalization and effectiveness, as it performs well compared with other baselines across various public OOD datasets. The code is publicly available at https://github.com/flzeng1/MPAIACL.

Figures (8)

• Figure 1: The stable features and environment features. Stable features capture the underlying patterns of the entire graph, providing a robust representation of the graph's intrinsic structure. Environmental features are subject to variation for the label.
• Figure 2: Correlate distribution shift and covariate distribution shift.
• Figure 3: Traditional graph data augmentation. Traditional data augmentation strategies may destroy the structure of stable features.
• Figure 4: Result of the experiment in section \ref{['sec: motivation']}. (a) and (b) visualize the stable feature and environment feature distribution in the latent space. (c) Visualize the average Euclidean distance of the stable feature and environment feature between AIA and MPAIACL. (d) Visualize the accuracy comparison between different datasets and methods.
• Figure 5: The overview of MPAIACL. The explanation of the notation used in the figure is located at the bottom and the bottom left. The training process consists of two distinct phases: (1) Strengthen the Stable Feature Generator (SFG), and (2) Strengthen the Adversarial Augmenter (AA). Initially, the input graph is processed through a shared GNN encoder, which generates the graph embeddings. In phrase (1), SFG generates the stable features. Then, the stable feature $\mathcal{F}_s$ would make a data augmentation with the environment feature $\mathcal{F}_e$, which is generated by AA to obtain the augmentation feature $\mathcal{F}_a$. Subsequently, MPAIACL utilizes contrastive learning to pushing away $\mathcal{F}_s$ from $\mathcal{F}_e$. Finally, the results of the contrastive and $\mathcal{F}_a$ are used to strengthen the SFG. In phrase (2), $\mathcal{F}_e$ uses Wasserstein distance to push away from $\mathcal{F}_s$ and the original graph embedding. The result of the Wasserstein distance is used to strengthen AA.

Theorems & Definitions (2)

• Definition 1: Correlation shift
• Definition 2: Covariate shift