PISA: Prioritized Invariant Subgraph Aggregation

TL;DR

PISA introduces a dynamic aggregation method for invariant subgraphs in graph data, enhancing out-of-distribution generalization and outperforming previous approaches in diverse datasets.

Contribution

It proposes a novel MLP-based prioritized aggregation technique for multiple invariant subgraphs, improving robustness over existing methods.

Findings

Achieves up to 5% higher accuracy on 15 datasets.

Outperforms prior methods in OOD generalization.

Demonstrates effectiveness on DrugOOD benchmark.

Abstract

Recent work has extended the invariance principle for out-of-distribution (OOD) generalization from Euclidean to graph data, where challenges arise due to complex structures and diverse distribution shifts in node attributes and topology. To handle these, Chen et al. proposed CIGA (Chen et al., 2022b), which uses causal modeling and an information-theoretic objective to extract a single invariant subgraph capturing causal features. However, this single-subgraph focus can miss multiple causal patterns. Liu et al. (2025) addressed this with SuGAr, which learns and aggregates diverse invariant subgraphs via a sampler and diversity regularizer, improving robustness but still relying on simple uniform or greedy aggregation. To overcome this, the proposed PISA framework introduces a dynamic MLP-based aggregation that prioritizes and combines subgraph representations more effectively. Experiments on 15 datasets, including DrugOOD (Ji et al., 2023), show that PISA achieves up to 5% higher classification accuracy than prior methods.