PISA: Prioritized Invariant Subgraph Aggregation

TL;DR

The paper addresses the challenge of out-of-distribution generalization for graph data, where distribution shifts affect both node features and topology. It extends invariant subgraph learning by proposing PISA, which trains multiple invariant GNN branches with diversity injections and then uses a dynamic MLP to weight subgraph evidence per instance. Key contributions include (i) learning multiple invariant subgraphs via sampling and a diversity penalty, (ii) a cross-branch contrastive objective to align same-label subgraphs, and (iii) a second-stage adaptive aggregator that captures nonlinear interactions among subgraphs. Across 15 datasets, including DrugOOD, PISA achieves state-of-the-art robustness, with up to 5% improvements over prior methods and strong performance in multi-subgraph scenarios, demonstrating practical gains in real-world graph tasks with distribution shifts.

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.

Figures (6)

• Figure 1: CIGA ciga: the featurizer extracts an invariant subgraph used by the classifier to make predictions.
• Figure 2: Aspirin contains multiple causal functional groups. SuGAr sugar captures multiple subgraphs.
• Figure 3: Phase I of Prioritized Invariant Subgraph Aggregation (PISA). The task is to classify graphs by their discriminative motifs (e.g., House, Cycle). A graph sampler stochastically drops edges from the input graph to produce a set of perturbed graph views. Each view is processed by its corresponding featurizer $g^{i}$, which predicts edge-importance weights. A diversity regularizer encourages the featurizers to extract distinct invariant subgraphs $\widehat{G}^{i}_{c}$. The associated classifier $f^{i}_{c}$ then produces a prediction from each $\widehat{G}^{i}_{c}$.
• Figure 4: Phase II of PISA (Dynamic Aggregation). Branch-wise prediction scores are stacked and fed to a lightweight MLP, which learns to aggregate them and to assign instance-specific importance weights to each branch. During this phase, all featurizers $g^{i}$ and classifiers $f^{i}_{c}$ are frozen; only the MLP is trained using a standard empirical loss (e.g., cross-entropy) on the final aggregated scores.
• Figure 5: Ablation Studies