IENE: Identifying and Extrapolating the Node Environment for Out-of-Distribution Generalization on Graphs

TL;DR

This paper introduces IENE, a novel method for improving out-of-distribution generalization in graph neural networks by identifying and extrapolating node environments at feature and structural levels, validated through extensive experiments.

Contribution

IENE uniquely combines node-level environmental identification with extrapolation techniques to enhance GNNs' robustness under distribution shifts.

Findings

Outperforms existing OOD generalization methods on multiple datasets

Effectively identifies invariant features and structures

Provides theoretical analysis and proofs of the approach

Abstract

Due to the performance degradation of graph neural networks (GNNs) under distribution shifts, the work on out-of-distribution (OOD) generalization on graphs has received widespread attention. A novel perspective involves distinguishing potential confounding biases from different environments through environmental identification, enabling the model to escape environmentally-sensitive correlations and maintain stable performance under distribution shifts. However, in graph data, confounding factors not only affect the generation process of node features but also influence the complex interaction between nodes. We observe that neglecting either aspect of them will lead to a decrease in performance. In this paper, we propose IENE, an OOD generalization method on graphs based on node-level environmental identification and extrapolation techniques. It strengthens the model's ability to extract invariance from two granularities simultaneously, leading to improved generalization. Specifically, to identify invariance in features, we utilize the disentangled information bottleneck framework to achieve mutual promotion between node-level environmental estimation and invariant feature learning. Furthermore, we extrapolate topological environments through graph augmentation techniques to identify structural invariance. We implement the conceptual method with specific algorithms and provide theoretical analysis and proofs for our approach. Extensive experimental evaluations on two synthetic and four real-world OOD datasets validate the superiority of IENE, which outperforms existing techniques and provides a flexible framework for enhancing the generalization of GNNs.