Towards Effective Graph Rationalization via Boosting Environment Diversity

TL;DR

This work tackles the problem of graph generalization under distribution shifts by introducing GRBE, a two-module framework that performs environment augmentation in the original graph space. PRSE refines rationale subgraph extraction through mask-based learning and contrastive signals, while EDA generates diverse environment subgraphs via mixup-like merging across graphs followed by synthesis of augmented graphs. The combined objective integrates rationalization accuracy, augmentation reliability, and mutual information-based regularization, leading to substantial gains in rationalization and classification across multiple benchmarks, especially under larger shifts. By operating in the original graph space, GRBE achieves finer-grained and more diverse augmentations, offering practical improvements for robust graph learning in real-world OOD settings.

Abstract

Graph Neural Networks (GNNs) perform effectively when training and testing graphs are drawn from the same distribution, but struggle to generalize well in the face of distribution shifts. To address this issue, existing mainstreaming graph rationalization methods first identify rationale and environment subgraphs from input graphs, and then diversify training distributions by augmenting the environment subgraphs. However, these methods merely combine the learned rationale subgraphs with environment subgraphs in the representation space to produce augmentation samples, failing to produce sufficiently diverse distributions. Thus, in this paper, we propose to achieve an effective Graph Rationalization by Boosting Environmental diversity, a GRBE approach that generates the augmented samples in the original graph space to improve the diversity of the environment subgraph. Firstly, to ensure the effectiveness of augmentation samples, we propose a precise rationale subgraph extraction strategy in GRBE to refine the rationale subgraph learning process in the original graph space. Secondly, to ensure the diversity of augmented samples, we propose an environment diversity augmentation strategy in GRBE that mixes the environment subgraphs of different graphs in the original graph space and then combines the new environment subgraphs with rationale subgraphs to generate augmented graphs. The average improvements of 7.65% and 6.11% in rationalization and classification performance on benchmark datasets demonstrate the superiority of GRBE over state-of-the-art approaches.

Figures (8)

• Figure 1: An example to illustrate the difference between the existing general augmentation in the representation space and our proposed mixup augmentation in the original graph space. Specifically, (a) Input graphs are divided into rationale and environment subgraphs. (b) General augmentation in the representation space splices the representations of rationale and environment subgraphs to create a new sample. (c) Our proposed Mixup augmentation operates in the original graph space by sampling environment subgraphs from a mixed distribution and attaching them to the rationale subgraph to generate diverse graphs.
• Figure 2: The unsupervised clustering results of the environment subgraph representations on the Spmotif-0.9 dataset. (a) The representations are learned by the existing GREA method operated in the representation space. (b) The representations are obtained by our proposed GRBE method, which operates in the original graph space. Each node denotes an environment subgraph representation, and different colors represent different categories of environment subgraph representations.
• Figure 3: The overall framework of our proposed GRBE method. It performs environment diversity augmentation within the original graph space and can be divided into two main components: (a) PRSE guides the GNN model in learning a precise subgraph division. (b) EDA generates more diverse training graphs by exploring unknown environment subgraphs.
• Figure 4: Ablation studies. We report the mean - 0.5*std of the rationalization and classification performance of the different variants of GRBE.
• Figure 5: Hyper-parameter sensitive analysis. We report the mean - std of classification and rationalization performance under five critical hyper-parameters $\alpha, \beta, \gamma, r_{aug}, r_s$ on the Spmotif-0.7 dataset.

Theorems & Definitions (1)

• Definition 1