PISA: Prioritized Invariant Subgraph Aggregation for Out-of-Distribution Generalization on Graphs

Ali Ghasemi; Farooq Ahmad Wani; Maria Sofia Bucarelli; Fabrizio Silvestri

cs.LGNov 27, 2025

PISA: Prioritized Invariant Subgraph Aggregation for Out-of-Distribution Generalization on Graphs

Authors: Ali Ghasemi, Farooq Ahmad Wani, Maria Sofia Bucarelli, Fabrizio Silvestri

Organizations: Sapienza University of Rome · Université Côte d’Azur, CNRS, Inria, I3S

Abstract

Invariant learning on graphs aims to build predictors that rely on causal substructures rather than on environment-specific shortcuts. Current methods extract either a single invariant subgraph (CIGA) or, more recently, a set of diverse invariant subgraphs (SuGAr). We observe that the second half of the multi-subgraph pipeline, "how the extracted subgraphs are combined" has gone essentially unexamined; existing methods average branch predictions uniformly or select one greedily. We show that this is precisely where accuracy is lost. A minimum-variance combination argument establishes that uniform averaging is optimal only when every branch carries the same total error covariance with the ensemble---a knife edge balance that second order exchangeability guarantees and that the diversity regularizer used to obtain multiple subgraphs is explicitly designed to break. We then show that invariance is closed under post composition; because the branches are frozen and read the graph only through their invariant subgraphs, any measurable aggregator of their outputs remains invariant, and the aggregate's environment gap is bounded by the branches' residual invariance violation uniformly over aggregators. Learning the combination rule by plain empirical risk minimization is therefore safe, and aggregator capacity trades off against estimation error rather than robustness. PISA instantiates this with a lightweight instance-conditional MLP over stacked branch logits, trained on frozen branches at negligible cost. Across 15 synthetic and real benchmarks, PISA improves the average over the strongest prior method by 2.96 points on synthetic shifts and 1.74 points on realistic shifts, with per-dataset gains up to 5.90 points, and nearly doubles SuGAr's improvement over ERM on the six DrugOOD tasks (from +2.76 to +5.07 ROC-AUC).

Figures & tables

Panel of Figure 3 — Figure 3: (a) A graph sampler stochastically drops edges to create perturbed views. Each view is passed to a featurizer $g^{i}$ predicting edge-importance weights; a diversity regularizer drives the featurizers towards distinct invariant subgraphs $\widehat{G}_{c}^{i}$ , and classifier $f^{i}_{c}$ predicts from each. (b) Branch logits are stacked into $\mathbf{P}$ and fed to a lightweight MLP that assigns instance-specific importance to each branch. All $g^{i},f^{i}_{c}$ are frozen in Phase II; only the aggregator is trained, with a standard empirical loss.

Method	SPMotif			SUMotif			AVG
Method	$\text{bias}=0.33$	$\text{bias}=0.6$	$\text{bias}=0.9$	$\text{bias}=0.33$	$\text{bias}=0.6$	$\text{bias}=0.9$	AVG
IRM [ 15 ]	$63.98_{\pm 8.51}$	$61.58_{\pm 12.85}$	$47.14_{\pm 12.13}$	$61.39_{\pm 13.10}$	$58.48_{\pm 15.47}$	$48.35_{\pm 14.91}$	$56.82$
V-REx [ 27 ]	$69.18_{\pm 7.34}$	$58.76_{\pm 11.51}$	$43.81_{\pm 13.21}$	$63.24_{\pm 15.63}$	$65.23_{\pm 14.18}$	$44.03_{\pm 12.40}$	$57.38$
IB-IRM [ 16 ]	$62.30_{\pm 11.27}$	$59.78_{\pm 12.77}$	$46.19_{\pm 12.10}$	$71.13_{\pm 11.04}$	$56.55_{\pm 13.60}$	$47.27_{\pm 13.31}$	$57.20$
EIIL [ 28 ]	$63.82_{\pm 12.43}$	$59.42_{\pm 13.16}$	$42.51_{\pm 11.89}$	$67.66_{\pm 10.27}$	$55.64_{\pm 13.11}$	$41.22_{\pm 8.56}$	$55.04$
GREA [ 29 ]	$74.53_{\pm 8.89}$	$68.26_{\pm 10.53}$	$48.45_{\pm 13.96}$	$68.17_{\pm 17.64}$	$54.39_{\pm 12.37}$	$51.83_{\pm 15.03}$	$60.94$

Table 1: OOD generalization on synthetic graphs (accuracy %). Best in bold , second best underlined . AVG is the mean over the six columns.

Method	EC-Assay	EC-Scaffold	EC-Size	SST5	Twitter	CMNIST	Ki-Assay	Ki-Scaffold	Ki-Size	AVG
ERM [ 35 ]	$75.57_{\pm 1.23}$	$64.21_{\pm 0.89}$	$63.30_{\pm 1.19}$	$\underline{44.21_{\pm 0.91}}$	$63.84_{\pm 1.61}$	$10.26_{\pm 0.62}$	$73.30_{\pm 1.67}$	$70.45_{\pm 0.30}$	$74.00_{\pm 1.55}$	$59.90$
IRM [ 15 ]	$77.10_{\pm 2.55}$	$64.32_{\pm 0.42}$	$62.33_{\pm 0.86}$	$42.77_{\pm 1.26}$	$60.42_{\pm 1.06}$	$15.15_{\pm 3.66}$	$75.10_{\pm 3.38}$	$69.32_{\pm 1.84}$	$76.25_{\pm 0.73}$	$60.31$
V-REx [ 27 ]	$75.57_{\pm 2.17}$	$64.73_{\pm 0.53}$	$62.80_{\pm 0.89}$	$42.48_{\pm 1.67}$	$60.50_{\pm 2.05}$	$17.12_{\pm 5.68}$	$74.16_{\pm 1.46}$	$71.40_{\pm 2.77}$	$76.68_{\pm 1.35}$	$60.60$
IB-IRM [ 16 ]	$64.70_{\pm 2.50}$	$62.62_{\pm 2.05}$	$58.28_{\pm 0.99}$	$43.02_{\pm 1.94}$	$60.80_{\pm 2.50}$	$13.06_{\pm 1.97}$	$71.98_{\pm 3.26}$	$69.55_{\pm 1.66}$	$70.71_{\pm 1.95}$	$57.19$
EIIL [ 28 ]	$64.20_{\pm 5.40}$	$62.88_{\pm 2.75}$	$59.58_{\pm 0.96}$	$43.79_{\pm 1.19}$	$60.15_{\pm 1.44}$	$11.80_{\pm 0.42}$	$74.24_{\pm 2.48}$	$69.63_{\pm 1.46}$	$76.56_{\pm 1.37}$	$58.09$
GREA [ 29 ]	$66.87_{\pm 7.53}$	$63.14_{\pm 2.19}$	$59.20_{\pm 1.42}$	$43.29_{\pm 0.85}$	$59.92_{\pm 1.48}$	$13.92_{\pm 3.43}$	$73.17_{\pm 1.80}$	$67.82_{\pm 4.67}$	$73.52_{\pm 2.75}$	$57.87$

Table 2: OOD generalization under realistic shifts. ROC-AUC for the six DrugOOD tasks (EC-

\ast

, Ki-

\ast

), accuracy elsewhere. Best in bold , second best underlined . AVG is the mean over the nine columns.

Panel of Figure 4 — Figure 4: Ablation studies on SPMotif-0.6. Left: aggregator depth. Middle: effect of the number of branches $n$ . Right: parameter sharing.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 5 — Figure 5: Full SCMs for graph distribution shifts. Each panel depicts a distinct interaction pattern between invariant ( $C$ ) and spurious ( $S$ ) latent factors and the environment $E$ . In all panels $E$ has no directed path into $C$ , which is the content of Assumption 1 .

Dataset	Backbone	LR	Batch	r	Contrast	Hidden dim.	Dropout	Agg.
SPMotif	3-layer GCN	$10^{-3}$	32	.25	2	32/32/32	0	2-layer
SUMotif	3-layer GCN	$10^{-3}$	32	.25	2	32/32/32	0	2-layer
EC-Assay	4-layer GIN	$10^{-3}$	128	.80	8	32/32/32	.1	3-layer
EC-Scaffold	4-layer GIN	$10^{-3}$	128	.80	32	32/32/32	.1	3-layer
EC-Size	4-layer GIN	$10^{-3}$	128	.80	16	32/32/32	.1	3-layer
Ki-Assay	4-layer GIN	$10^{-3}$	128	.80	8	32/32/32	.1	3-layer

Appendix

Table 3: Dataset-specific hyperparameters used for PISA. Hidden dimensions are reported as branch encoder / branch classifier / aggregator.

	$n=2$	$n=5$	$n=10$	$n=15$
Phase-I params.	42,412	106,030	212,060	318,090
Phase-II params.	129	225	385	545
Phase-I train (s)	$1083.21_{\pm 9.18}$	$2440.18_{\pm 33.78}$	$4571.75_{\pm 25.31}$	$6829.79_{\pm 53.71}$
Phase-II train (s)	$420.92_{\pm 146.51}$	$664.68_{\pm 69.87}$	$1117.40_{\pm 7.31}$	$1774.37_{\pm 398.55}$
Inference (s)	$1.78_{\pm 0.01}$	$3.21_{\pm 0.04}$	$5.60_{\pm 0.03}$	$7.93_{\pm 0.16}$

Appendix

Table 4: Computational scaling with the number of branches

n

on SPMotif-0.6. Times are reported in seconds (mean

\pm

std).

Explore similar work

Adversarial Label Invariant Graph Data Augmentations for Out-of-Distribution Generalization
Apr 9, 2026Simon Zhang, Ryan P. DeMilt, Kun Jin +1OOD Generalization Covariate Shift
Department of Computer Science, Purdue University, West Lafayette, USA · Department of Computer Science and Engineering, The Ohio State University, Columbus Ohio, USA · Department of Industrial and Systems Engineering, The Ohio State University, Columbus Ohio, USA
Invariant Graph Representations for Continuous-Time Dynamic Graphs Under Distribution Shifts
May 29, 2024Lanting Fang, Yulian Yang, Yawei Zhang +3Causal Representation Learning Structural Causal Models
1Beijing Institue of Technology · Southeast University · Wuhan University
Is Spurious Correlation Removal Always Learnable?
Jun 11, 2026Yibo Zhou, Bo Li, Hai-Miao Hu +3Spurious Correlation Robustness Invariant Representation Learning
Beijing Key Laboratory of Digital Media, School of Computer Science and Engineering, Beihang University, Beijing 100191, China · State Key Laboratory of Virtual Reality Technology and Systems, Beihang University, Beijing 100191, China · Hangzhou Innovation Institute of Beihang University, Hangzhou 310051, China +1

Apr 9, 2026J/K move · Enter open · S save

Adversarial Label Invariant Graph Data Augmentations for Out-of-Distribution Generalization

Simon Zhang, Ryan P. DeMilt, Kun Jin, Cathy H. Xia

Department of Computer Science, Purdue University, West Lafayette, USA · Department of Computer Science and Engineering, The Ohio State University, Columbus Ohio, USA · Department of Industrial and Systems Engineering, The Ohio State University, Columbus Ohio, USA

OOD Generalization Covariate Shift Data Augmentation Counterfactual Data Augmentation Adversarial Training

Open paper AskPDF arXiv

Out-of-distribution (OoD) generalization occurs when representation learning encounters a distribution shift. This occurs frequently in practice when training and testing data come from different environments. Covariate shift is a type of distribution shift that occurs only in the input data, while the concept distribution stays invariant. We propose RIA - Regularization for Invariance with Adversarial training, a new method for OoD generalization under convariate shift. Motivated by an analogy to

Q

-learning, it performs an adversarial exploration for counterfactual data environments. These new environments are induced by adversarial label invariant data augmentations that prevent a collapse to an in-distribution trained learner. It works with many existing OoD generalization methods for covariate shift that can be formulated as constrained optimization problems. We develop an alternating gradient descent-ascent algorithm to solve the problem in the context of causally generated graph data, and perform extensive experiments on OoD graph classification for various kinds of synthetic and natural distribution shifts. We demonstrate that our method can achieve high accuracy compared with OoD baselines.

Abstract

Explore similar work

Adversarial Label Invariant Graph Data Augmentations for Out-of-Distribution Generalization

Invariant Graph Representations for Continuous-Time Dynamic Graphs Under Distribution Shifts

Is Spurious Correlation Removal Always Learnable?