Subgraph Aggregation for Out-of-Distribution Generalization on Graphs

TL;DR

This work proposes a novel framework, SubGraph Aggregation (SuGAr), designed to learn a diverse set of subgraphs that are crucial for OOD generalization on graphs, and is the first work to study graph OOD generalization by learning multiple invariant subgraphs.

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

Figures (6)

• Figure 1: The molecule graph, Aspirin, contains two functional groups: -COOH (denoted with blue lines) and -COO- (denoted with red lines). If model can only capture one functional group as the invariant subgraph, either -COOH or -COO-, which will lead to suboptimal OOD generalization ability. However, SuGAr can capture all such functional groups in Aspirin molecule graph for promising performance.
• Figure 2: The illustration of SubGraph Aggregation(SuGAr). Our proposed method needs to classify graphs based on their motifs("House" and "Cycle") in two steps:(a) Training: A graph sampler randomly drops edges from the input graph to generate a set of different graphs. These graphs are then processed by their corresponding featurizers, $g^i$, to obtain predicted edge weights for each graph. The diversity regularizer encourages each $g^i$ to extract different subgraphs $\widehat{G}\xspace^i_c$ , and the corresponding classifier $f^i_c$ makes predictions based on $\widehat{G}\xspace^i_c$ .(b) Inference: During inference, a graph is fed into the well-trained $g^i$s, which extract a diverse set of subgraphs. The aggregator then merges these subgraphs to form $\widehat{G}\xspace^{m}_c$, which serves as the input for the trained classifiers $f_c$. Finally, the aggregator combines the predictions from $f_c$s to make the final decision.
• Figure 3: Ablation Studies for SuGAr (ENS).
• Figure 4: Ablation Studies for SuGAr(WA).
• Figure 5: Comparison of SuGAr and CIGA.