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

Apr 9, 2026

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

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.

Simon Zhang, Ryan P. DeMilt, Kun Jin +1

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

Ask PDF

May 29, 2024

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

Continuous-Time Dynamic Graphs (CTDGs) enable fine-grained modeling of evolving relational systems. However, most existing CTDG representation learning methods are tailored to in-distribution settings and exhibit limited robustness under out-of-distribution (OOD) shifts. Although recent causal approaches learn invariant representations via interventions, they are primarily designed for static or discrete-time graphs and become computationally prohibitive for CTDGs due to the combinatorial explosion of structural and temporal variations. To address these challenges, we propose CIR, a framework grounded in a novel structural causal model termed the ICCM. To avoid exhaustive interventions, we leverage the Normalized Weighted Geometric Mean (NWGM) to efficiently approximate interventional predictions. We further instantiate ICCM within a practical deep learning architecture that jointly captures invariant structural and temporal patterns through dedicated subgraph extractors, and maintains an environment memory bank to model distributional shifts across evolving contexts. Extensive experiments demonstrate that CIR consistently outperforms existing methods under diverse OOD scenarios.

Lanting Fang, Yulian Yang, Yawei Zhang +3

1Beijing Institue of Technology · Southeast University · Wuhan University

Causal Representation Learning Structural Causal Models OOD Generalization Dynamic Graph Learning Invariant Representation Learning

Ask PDF

Jun 11, 2026

Is Spurious Correlation Removal Always Learnable?

Invariant learning can fail even when the invariant structure is statistically identifiable. We show a conditional computational barrier: under a black-box samplable supervised sparse recovery primitive motivated by average-case sparse-recovery reductions, there exist samplable multi-environment instances with a one-dimensional predictive invariant subspace (

k=1

) that are learnable with polynomial samples by exhaustive search, while any polynomial-time constant-accuracy recovery algorithm would contradict the primitive. We further quantify environment diversity by a separation parameter

γ

, which controls identifiability and the curvature of invariance objectives. Under sufficient diversity and local Gaussian regularity, the minimax risk is

\mathbb{E}[\dist(\hat{V},V_{\mathrm{inv}})^2]=Θ(k(d-k)/(n|\mathcal{E}|))

, and under label-induced shifts a phase transition occurs at

n^*\propto k(d-k)/(|\mathcal{E}|γ^2)

with refined estimation error scaling proportional to

1/γ^2

. Synthetic and real datasets illustrate the predicted gaps and transitions and motivate simple diversity diagnostics.

Yibo Zhou, Bo Li, Hai-Miao Hu +3

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

Spurious Correlation Robustness Invariant Representation Learning Statistical Learning Theory

Ask PDF