Subgraph Aggregation for Out-of-Distribution Generalization on Graphs

TL;DR

This paper introduces SuGAr, a novel framework that learns multiple diverse invariant subgraphs to improve out-of-distribution generalization in graph neural networks, outperforming existing methods.

Contribution

SuGAr is the first approach to learn multiple invariant subgraphs for enhanced OOD generalization in graphs, using a tailored sampler and diversity regularizer.

Findings

Achieves up to 24% improvement in OOD generalization.

Outperforms state-of-the-art methods on synthetic and real datasets.

First to study multiple invariant subgraphs for graph OOD generalization.

Abstract

Out-of-distribution (OOD) generalization in Graph Neural Networks (GNNs) has gained significant attention due to its critical importance in graph-based predictions in real-world scenarios. Existing methods primarily focus on extracting a single causal subgraph from the input graph to achieve generalizable predictions. However, relying on a single subgraph can lead to susceptibility to spurious correlations and is insufficient for learning invariant patterns behind graph data. Moreover, in many real-world applications, such as molecular property prediction, multiple critical subgraphs may influence the target label property. To address these challenges, we propose a novel framework, SubGraph Aggregation (SuGAr), designed to learn a diverse set of subgraphs that are crucial for OOD generalization on graphs. Specifically, SuGAr employs a tailored subgraph sampler and diversity regularizer to extract a diverse set of invariant subgraphs. These invariant subgraphs are then aggregated by averaging their representations, which enriches the subgraph signals and enhances coverage of the underlying causal structures, thereby improving OOD generalization. Extensive experiments on both synthetic and real-world datasets demonstrate that \ours outperforms state-of-the-art methods, achieving up to a 24% improvement in OOD generalization on graphs. To the best of our knowledge, this is the first work to study graph OOD generalization by learning multiple invariant subgraphs. code: https://github.com/Nanolbw/SuGAr