Disentangling Invariant Subgraph via Variance Contrastive Estimation under Distribution Shifts
Authors: Haoyang Li
Xin Wang
Xueling Zhu
Weigao Wen
Wenwu Zhu
Venue: ICML 2025 poster (2025)
Keywords: Disentanglement
Graph Neural Network
Distribution Shift
Primary area: general_machine_learning
TL;DR: We propose to learn disentangled invariant subgraph via self-supervised contrastive variant subgraph estimation for achieving satisfying OOD generalization.

Abstract: Graph neural networks (GNNs) have achieved remarkable success, yet most are developed under the in-distribution assumption and fail to generalize to out-of-distribution (OOD) environments. To tackle this problem, some graph invariant learning methods aim to learn invariant subgraph against distribution shifts, which heavily rely on predefined or automatically generated environment labels. However, directly annotating or estimating such environment labels from biased graph data is typically impractical or inaccurate for real-world graphs. Consequently, GNNs may become biased toward variant patterns, resulting in poor OOD generalization. In this paper, we propose to learn disentangled invariant subgraph via self-supervised contrastive variant subgraph estimation for achieving satisfactory OOD generalization. Specifically, we first propose a GNN-based invariant subgraph generator to disentangle the invariant and variant subgraphs. Then, we estimate the degree of the spurious correlations by conducting self-supervised contrastive learning on variant subgraphs. Thanks to the accurate identification and estimation of the variant subgraphs, we can capture invariant subgraphs effectively and further eliminate spurious correlations by inverse propensity score reweighting. We provide theoretical analyses to show that our model can disentangle the ground-truth invariant and variant subgraphs for OOD generalization. Extensive experiments demonstrate the superiority of our model over state-of-the-art baselines.

Disentangling Invariant Subgraph via Variance Contrastive Estimation under Distribution Shifts
Haoyang Li Xin Wang'
Xueling Zhu?
Weigao Wen Wenwu Zhu
Abstract
1. Introduction
Graph ubiquitous in our daily life, which has been widely used to represent the complex relationships between enti- ties in many fields, including social network (Qiu al., 2018; Li et al., 2019), knowledge representation (Wang et al., 2017), recommendation systems (Wu et al., 2022b; Li et al., 2021a), multimedia (Hudson & Manning, 2019), computer vision (Zellers et al., 2018), natural language pro- cessing (Marcheggiani & Titov, 2017), etc. Among these popular applications, graph-level prediction tasks constitute major branch. Generally speaking, the label for graph depends the information in its critical subgraph (i.e., the part that has invariant and truly predictive relations to the label in the task). rather than the whole graph (Luo et al., 2020; Yu et al., 2021; Fan et al., 2022). For example, the sol- ubility of one molecule can be determined by the predictive functional group, rather than the molecular scaffold (Duve- naud al., 2015; Hu al., 2020). The label of one graph from the MNIST superpixel dataset (Dwivedi al., 2023; Fan et al., 2022) has deterministic and invariant relations with the digit subgraph, rather than the background sub- graph (Figure 1). Several representative works (Ying et al., 2018; Luo et al., 2020) are developed to explore capturing such critical subgraphs for learning graph representations ef- fectively, which can largely improve graph-level prediction performance.
Graph neural networks (GNNs) have achieved remarkable success, yet most are developed un- der the in-distribution assumption and fail to generalize out-of-distributioo (OOD) environ- ments. To tackle this problem, some graph in- variant learning methods aim to learn invariant subgraph against distribution shifts, which heav- ily rely on predefined or automatically generated environment labels. However, directly annotating or estimating such environment labels from biased graph data typically impractical or inaccurate for real-world graphs. Consequently, GNNs may become biased toward variant patterns, resulting in poor OOD generalization. In this paper, we propose to learn disentangled invariant subgraph via self-supervised contrastive variant subgraph estimation for achieving satisfactory OOD gener- alization. Specifically, we first propose a GNN- based invariant subgraph generator to disentangle the invariant and variant subgraphs. Then, es- timate the degree of the spurious correlations by conducting self-supervised contrastive learning on variant subgraphs. Thanks to the accurate iden- tification and estimation of the variant subgraphs, we can capture invariant subgraphs effectively and further climinate spurious correlations by in- verse propensity score reweighting. We provide theoretical analyses to show that our model can disentangle the ground-truie truth invariant and vari- ant subgraphs for OOD generalization. Extensive experiments demonstrate the superiority of our model over state-of-the-art baselines.
Despite their remarkable progress, the existing approaches are generally built upon the in-distribution (L.D.) hypoth- esis, namely the testing and training graphs are sampled from an identical distribution. Yet the graph data genera- tion mechanism in real-worle scenarios uncontrollable and unobservable (Bengio al., 2019), that the distribu- tion shifts can widely exist, making the (OOD) generalization become one of the most crucial issues to be handled (Li et al., 2022b). The existing subgraph- based graph methods fail to capture critical subgraphs due to lacking OOD generalization ability, whose performance can degenerate significantly under distribution shifts. Al- though several pioneering works (Wu al., 2022c:a; Li et al., 2022d; Yang et al., 2022: Li et al., 2023b) have been proposed to handle graph distribution shifts via learning en- vironment invariant subgraphs (which are defined as critical
DDeparmeent of Computer Science and Technology, BN- Rist, Tsinghua University, Beijing. China 'Department of Ra- diology. Xiangya Hospital, Central South University, Chang- sha, Hunan, China Allibaba Group. Correspondence to: Xin Wang <xin. Sann_wang@@ttt waan@@ssinguaaccu.ccomm <wwzhu tsu.ghu..cd..cn>>.
Wenwu Zhu
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Disentangling Invariant Subgraph via Variance Contrastive Estimation under Distribution Shifts
subgraphs that have invariant relations to the labels in any environment) against distribution shifts, they heavily rely on the predefined or automatically generated environment labels, i.e., multiple training environments. However, the environment labels are unavailable in most scenarios and directly annotating or generating environment labels is also impractical or inaccurate, especially in the graph dataset with severe bias (Fan et al.. 2022; Qi al., 2022), which largely limits these methods to capture the truly invariant subgraphs for OOD generalization.
6
3
2
Figure 1. Example: the critical subgraph (i.e., digit subgraph. de- noted by the bold lines) has truly predictive relations with the label (i.e., the digit) in the MNIST superpixel dataset (Fan et al., 2022; Dwivedi al., 2023).
To address this issue, we study the problem of handling graph distribution shifts with severe bias by explicitly es- variant patterns and further discovering invariant subgraphs for OOD generalization, which remains largely unexplored in the literature due to the following challenges:
timating environment-relann
guarantee. We conduct extensive experiments to demon- strate the superiority of our method over state-of-the-an baselines.
The contributions of our work are summarized as follows:
It is highly non-trivial to disentangle invariant and vari- ant subgraph patterns from their complex interactions within input graphs.
We propose learning invariant subgraph via variant sub- graph contrastive estimation, which can handle graph distribution shifts with severe bias. To the best of our knowledge, this is the first work to study how to explic- itly utilize the variant subgraphs to facilitate accurate identification of invariant subgraphs under distribution shifts.
It is challenging to estimate the degree of spurious corre- lations between the variant subgraphs and labels, since the variant subgraphs reflect the environment-related information.
It is also challenging to learn invariant subgraphs based on the estimated degree of spurious correlations for OOD generalized predictions.
We propose three mutually promoted modules to dis- entangle invariant and variant subgraphs, estimate the degree of spurious correlations, and make predictions based on the invariant subgraphs.
In this paper, we propose a novel Varlant subgraph learn- ing based on VAriance Contrastive Estimation (VIVACE) method, which able capture invariant subgraphs for achieving satisfactory OOD generalization under distribu- tion shifts with severe bias. Specifically, we first propose invariant and variant subgraph identification module to dis- entangle potentially invariant and variant patterns for input graphs. Then, we find that the variance information behind graph data is also important and should not be directly ig- nored as in most existing literature. So we propose a variant subgraph contrastive estimation module to explicitly model the degree of the spurious correlations between variant sub- graphs and labels with self-supervised contrastive training, which is useful for predicting the graph labels under dis- tribution shifts. Finally, we propose an inverse propensity weighting based invariant subgraph prediction module to reweight the invariant subgraph predictions for eliminating the spurious correlations and achieving OOD generalization. In this way, our VIVACE method can capture and utilize invariant subgraphs with stable power to make predictions under distribution shifts with severe bias. We provide com- prehensive theoretical analyses to show that our proposed method can disentangle ground-truth invariant and variant subgraphs for achieving OOD generalization with strong
We theoretically prove that our method can disentangle the ground-truth invariant and variant subgraphs which is a significant step towards OOD generalized predictions.
Extensive experiments on several graph classification benchmark datasets demonstrate the superiority of our proposed method over state-of-the-art baselines.
The rest of the paper is organized as follows. We first formulate the problem in Section 2. In Section 3, we present the details of our proposed VIVACE method. We present the experimental results to show the effectiveness of the method in Section 4, including quantitative comparisons, ablation studies, hyper-parameter sensitivity, etc. We review the related works in Section 5. Finally, we conclude this work in Section 6.
2. Problem Formulation
The OOD generalization under distribution shifts on graphs can be formulated as: Problem The duitt--disttibution generalization on graphs is to find an optimal graph predictor (-()) that maps input
2

---

Disentangling Invariant Subgraph via Variance Contrastive Estimation under Distribution Shifts
graph G to the label y and performs well on all environments E:
the invariant and variant subgraphs can be obtained:
G1 = + Gv = GA,G,
(2)
(-)- = arg min sup R(fle), 1 c
(1)
where ((-) is instantiated as a learnable edge mask matrix M on G. Therefore,
where R(fle)- = EE [() y)] is the risk of the predic- tor on the environment e, and e denotes a loss function. We further decompase f(-) = w h, where i(-) is the encoder that maps input graph into the d-dimensional representation, and w--) is the classifier.
G = (X, A), Gv = (X, M) e A),
(3)
Here, we obtain the edge mask matrix M with a learnable GNN instead of directly learning the edge mask for each graph independently. It is because this design can identify the invariant patterss in global view shared across the entire dataset and also easily generalize to handle unseen test graphs without retraining (Li et al., 2022d; Fan et al., 2022; Wu et al., 2022c). Specifically, each entry Mi, of the edge mask matrix is calculated:
Following the literature (Li et al., 2022b), we make the assumption:
Assumption 1. Given input graph G, there exists an invari-
ant subgraph G7 satisfying:
Invariance assumption: e' E, (y)r))
pe (y]G?).
M.,, = Sigmoid (MIPP (Z! 2 1) z(m) GNNM
b. Sufficiency assumption: y wr(hr(G7)) te, E G, where h denotes graph encoder, wy is the classifier, 1 indicates statistical independence, and e is random noise.
(4)
where ((m) is the node embedding for calculating the edge mask and M.P is multilayer perceptron. Finally, we can disentangle the invariant subgraph Gr,; and variant sub- graph Gv,, for the i-th input graph G, of the whole dataset). We use soft edge weights to decompose input graph into in- variant and variant subgraphs instead of hard splits, which is common in the literature for better optimization (Fan et al., 2022; Wu et al., 2022c). The Sigmoid function is employed to project the mask values into the interval (0. 1), indicating the probability of the edge being classified into invariant subgraph.
The invariance assumption means that there exists a sub- graph inside the input graph that has invariant relations to the label across different environments. The sufficiency as- sumption means that the invariant subgraph has sufficient predictive capability for predicting the graph label. Overall, the graph OOD generalization problem can be solved by finding the invariant subgraph for the input graph.
3. Method
3.2. Variant Subgraph Contrastive Estimation
In this section, we introduce the details of our proposed VIVACE, which mainly consists of three key modules: in- variant & variant subgraph identification module, variant subgraph contrastive estimation module, and inverse propen- sity weighting based invariant subgraph prediction module. The framework of VIVACE is shown in Figure 2.
Since GNNs tend to exploit the spurious correlations be- tween the variant subgraphs and the labels to make predic- tions, which leads to poor OOD generalization, the existing literature aims to encourage GNNs only focusing on in- variant subgraphs but ignoring the variant subgraphs. For achieving such a goal, they rely on the ideal predefined or au- tomatically generated multiple environments for training. It is because only when we can explicilly observe environment- discriminative features among multiple environments, we can capture the variant subgraph that has variant correlations under different environments, so that we can disregard such variant subgraph but in turn capture the invariant subgraph that has invariant relationships between predictive graph structural information and the label for OOD generalization.
3.1. Invariant & Variant Subgraph Identification
The graph invariant learning literature (Li et al., 2022d; Fan al., 2022; Wu et al., 2022c) generally assumes that each graph G consists of an invariant subgraph G2 G (which is dominant to the label in the task and also has invariant relations to the label across different environments). and a variant subgraph (which could form spurious correlations and have variant relations with the label in different envi- ronments). Since variant subgraph is the complement of invariant subgraph in terms of the input graph, the accu- rate identification and estimation for variant subgraph can promote the modeling for invariant subgraph.
However, when the training graph data is with severe bias (c.g., imbalanced and less diverse), is nearly impossible to obtain ideal multiple environments for learning invariant subgraphs. The variant pattern can not be fully disregarded from the invariant pattern (Qi et al., 2022), so that the per- formance in unseen testing graphs with distribution shifts is
Given one input graph G = A), where X is the node feature matrix and A is the adjacency matrix of G, we first adopt the proposed invariant subgraph generator (--) to dis- entangle the invariant pattern in G under distribution shifts,
Note that for simplification, we omit the index i when there is no ambiguity.
3

---

Disentangling Invariant Subgraph via Variance Contrastive Estimation under Distribution Shifts
Variant Subgraph Contrastive Estimation
! Invariant & Variant Subgraph Identification
Var(f Minimize Varlance
Maximize Agreement
Agreemett
Rast
Maximize/Agreement
fest
Gv
A
Variant Subgraphs
Estimated degree of spurious correlations
Variant Subgraphs
1 Gv
Weighted Predictions
Loss
R
Invariant Subgraph Generator ())
Predictions
Labels
[0.3, 0.7)
[0.4, 0.6)
G1
(0.2, 0.5)
(0.1, 0.9)
(0. 1)
Graph Dataset
(0.2, 0.8)
(0.1, 0.9)
[0,
:'
Invariant
Inverse Propensity
(0.7, 0.3)
(0.6,
(1. 0)
Labels: Y
Predictor O
Weighting
(0.8, 0.2)
(0.9, 0.1)
0)
Invariant Subgraph Gj
[0.8, 0.2)
(0.9, 0.1]
(1. 0)
Variant Subgraph Gy
Invariant Subgraphs
3 Invariant Subgraph Prediction via Inverse
Propensity Weighting
Figure 2. The framework of the proposed method VIVACE. Training Stage: Given the input graph dataset, the invariant subgraph generator (-)) first disentangles the invariant subgraph G, and variant subgraph Gy for each graph G. Then, the variant subgraphs are utilized for estimating the degree of the spurious correlation in self-supervised contrastive manner. Finally, the invariant subgraph predictions are reweighted based on the accurate estimation the spurious correlations between the variant subgraphs and labels, so as eliminate the spurious correlations for OOD generalization. Testing Stage: We directly adopt the optimized invariant subgraph generator d" and predictor fi to make OOD generalized predictions y = fi ( (G)).
not stable. Therefore, in this module, we propose novel strategy to explicitly estimate the degree of the spurious correlations between the variant subgraph Gy and its label y for each input graph G. i.e., p(y]Gv)) Gv).
contrastive loss for the graphs whose label is k, i.e.
( (0,
exp @(hv (Gv.), T wink E, exp 6(1v ((hv by
Σ log
(6)
Let fv wy hy denote the variant subgraph predictor mapping the variant subgraph into the corresponding pre- dicted label, which consists of the variant subgraph encoder fy to learn representations and variant subgraph classifier wy that makes predictions based on the learned represen- tations. Since the invariant subgraph is truly predictive to the label, namely fully including the label information, the variant subgraph captures the environment-rclated informa- tion that is not available as the supervisions for training hv. Therefore, we propose the following training objective with self-supervisions:
where Tj is the number of training graph whose label is k and N is the total number of training graph in the dataset. Note that EM Ta N. Gyy, the augmented graph of Gy.,, where adopt common graph augmentation strate- gies (You et al., 2020). This objective can maximize the agreement between Gv and Gry- & is the cosine similarity with temperature т, i.c., ó(a. b) = COSINE(a, b)/r and COSINE(a, b) = a bTb b/([aa]/ [bll2)-
However, only utilizing the self-supervised contrastive loss in Eq. (6) train the variant subgraph encoder not enough, since the invariant patterns will also be captured by the en- coder during training. Therefore, inspired by the invariant learning literature (Krueger et al., 2021; Qi et al., 2022), we adopt the variation of the contrastive loss functions on all groups of graphs divided by the labels as the regular- izer to encourage the encoder by only capturing the variant patterns. Intuitively, this regularizer can help the encoder
YY Σ
Evisal (D.hy)
la,( hv; ko) (Peat (b, k):
Y
(5) Here Y denotes the label set of the graph dataset, fest is the contrastive loss for training the identification and pre- diction the variant subgraph, and Var(-) denotes the vari- ation. More specifically, lsst ④b, hy; k) is defined as the

---

Disentangling Invariant Subgraph via Variance Contrastive Estimation under Distribution Shifts
ignore the divergence among different groups but focus on their consistency. Since the training graphs are divided into different groups according to their labels, the invariant patterns that fully reflect the label information will be disre- garded and only the variant patterns will be captured by the variant subgraph encoder hy.
invariant subgraph predictor mapping the invariant subgraph into the corresponding predicted label. We adopt the fol- lowing inverse propensity weighted (IPW) loss (Jung et al., 2020; Seaman & Vansteelandt, 2018; Qi et al., 2022) to optimize the invariant subgraph predictor:
(x S1)
Σ P(y]Gvv CI
(8)
For better clarifications, we further explain how to prevent Gv from including invariant information. Rather than only optimizing the contrastive loss Eq. (6), optimize the contrastive learning objective Eq. (5) with the second vari- ance term to achieve this goal (where the contrastive loss in Eq. (6) is only the first term), as shown in Algorithm 1. The second variance term in Eq. (5) is proposed to mini- mize the difference (variance) among the contrastive loss from all Gy with different labels 2,2..... | where | is the number of unique labels. If Gv contains the invariant information with label by mistake, the variance term will in- crease since there exists a significant difference among label k and the other labels 1,...,k k - 1, k + I,.... [Y]. There- fore, when the second variance term is properly optimized, it can encourage Gv to exclude invariant information to the label for achieving a minimum variance term. Finally, will lead the variant subgraph encoder capture the subgraph that only contains the variant subgraph.
where CE is the standard cross-entropy loss, and the propen- sity score function p(y(Gy) (Fan et al., 2022; Qi et al., 2022)
is calculated by the following equation:
CEly, fv(v))) P(y|Gv) + CEly, (vVV))
(9)
Intuitively, for the input G. we expect the invariant subgraph generator Φ(-) and predictor fr(-) can capture the invariant
ground-truth relation p(y (G)) for OOD generalization. Therefore, we need to explicitly estimate the degree of the spurious correlations between the variant subgraph and the label:
(10) If p(y[Gv) is large, the Eq. (9) can under-weight the pre- diction loss to reduce the over-large spurious correlations. And a small p(y] ((y) means the variant subgraph has very few impacts on the prediction under distribution shifts, so the prediction loss should be up-weighted to encourage the invariant subgraph generator and predictor for capturing invariant patterns and improving OOD generalization.
p(ulGv) plu (G)()))
Based on the trained variant subgraph encoder hy, we fur- ther train the variant classifier wy via the generalized cross entropy (GCE) loss (Lee et al., 2021; Fan et al., 2022) to fit the spurious correlations of the variant classifier:
(hv (Gv))) 1 (f
For the optimization, we jointly optimize the self-superrised contrastive objective in Eq. (5), the supervised objectives in Eq. (8) and Eq. (7) as follows:
(v( = GCT(y, wy
9
(7) where w (fv())) means the predicted probability on the label y, i.e., the softmax output of the y-th dimensional- ity in practice. And q E (0, 1] is a hyperparameter that controls the degree of fitting the spurious correlations. Fi- nally, we can rely on the trained variant subgraph predictor fv = wv o by to accurately estimate the degree of the spurious correlations between the variant subgraph and its label p(y(Gy), ), whose implementation will be introduced in the next module.
4' f'. argmin ( (v (Av hv).
(11)
3.4. Optimization Procedure
We present the pseudocode of VIVACE in Algorithm 1 to show the training procedure.
At the testing stage, we directly adopt the optimized $* to capture the invariant subgraph for each input testing graph Gte and further feed it into optimized fî to make prediction as follows:
3.3. Invariant Subgraph Prediction via Inverse Propensity Weighting
Ú = (7 (Gte))
(12)
Although we do not have any annotations to specify the variant subgraph in the dataset, we are able to estimate the relations between the variant subgraph and the label, which can largely help to learn accurate relations between the invariant subgraph and the label for OOD generalization. Specifically, we adopt propensity score techniques in causal- ity (Jung et al., 2020; Seaman & Vansteelandt, 2018; Qi et al., 2022) to reweight the invariant subgraph predictions for eliminating the spurious correlations. Let fi denote the
Since the predictions only depend on the subgraphs that have invariant relations with the labels, the OOD generalization performance can be largely improved.
3.5. Theoretical Analyses
We theoretically demonstrate that the proposed VIVACE model can achieve OOD generalization by disentangling the ground-truth invariant and variant subgraphs.
5

---

Disentangling Invariant Subgraph via Variance Contrastive Estimation under Distribution Shifts
Algorithm 1 The training procedure of VIVACE.
4. Experiments
Input: The graph dataset
Output: An optimized invariant subgraph generator ((-) and predictor f-) mapping each graph to its label
In this section, we conduct extensive experiments to verify that our VIVACE method can effectively handle distribu- tion shifts even on the severely biased graph datasets by capturing the invariant subgraphs, including the experimen- tal setup, quantitative comparisons, ablation studies, the impact of the papper-praaeters, etc.
1: while not converge do
2: for sampled minibatch B of the graph dataset do
3: for each graph G and the corresponding label y in B do
4: Generate the edge mask matrix M by Eq. (4),
5:
Generate the invariant subgraph Gy and the variant subgraph Cv by Eq. (3). 6: end for
4.1. Experiment Setup
7: Calculate the contrastive learning objective by Eq. (5).
Baselines. We consider several representative methods that are widely used in the literature (Li et al., 2022a; Fan et al., 2022) as the haselines:
8:
Calculate the objective of variant subgraph predictor by
Eq. (7).
9:
Obtain the propensity score function for reweighting by
Eq. (9).
10:
Calculate the objective of invariant subgraph predictor
GCN (Kipf & Welling, 2017): It follows the recursive neighborhood aggregation scheme and is considered one of the most popular GNNs.
by Eq. (8).
11:
Obtain the overall objective by Eq. (1I).
12:
Update model parameters by backpropagation.
13:
end for
14:
end while
GIN (Xu et al., 2019): It is also one famous GNN and has shown to be one of the most expressive GNNs in the representation learning of graphs.
Theorem 1. Denote the optimal invariant subgraph gen-
erator 50 that disentangles the ground-truth invariant sub- graph G and variant subgruph Gi given the inpul graph G. where G7 satisfies Assumption 1 and denote the comple- ment as Giy = G. Assume the second variance term of Eq. (5) is minimized, we have that the first contrastive loss term is minimized iff the invariant subgraph generator +b equals x²
FactorGCN (Yang et al., 2020): It is representative graph disentangling model for graph classification.
DiffPool (Ying et al., 2018): It graph pooling method that learns the cluster assignment for each node and outputs the coarsened graph.
LDD (Lee al., 2021): It can learn debiased representa- tions via disentangled feature augmentation, which is general debiasing method.
The proof is also shown in Appendix. The theorem shows that our proposed method can identify the ground-truth in-
variant and variant subgraphs by the variance contrastive es- timation for accurately modeling the degree of the spurious correlations and further eliminating the spurious correlations to achieve OOD generalization.
DIR (Wu et al., 2022c): It generates multiple environ- ments from biased graph data by conducting interven- tions on graphs and further capturing invariant explain- able subgraphs for predictions.
DisC (Fan al., 2022): It is notable debiasing method for GNNs, which aims to explicitly distinguish causal or bias patterns from the input graphs.
3.6. Complexity Analysis
We provide the detailed time complexity analysis of the proposed VIVACE method as follows. The time complex- ity of VIVACE is O(FF|d + [V]d"). where [V| and |F] indicate the number of nodes and edges of the input graph, respectively, and d is the dimensionality of the representa- tions. Specifically, we use GCN (Kipf & Welling, 2017), a message-passing GNN, to instantiate the GNN components in our method, which has complexity of E[E d [V| d²). We disentangle the invariant and variant subgraphs by gen- erating mask for the existing edges of the input graph, so the time complexity is O(E d + [V] dd). The variant sub- graph contrastive estimation module and inverse propensity weighting based invariant prediction module do not intro- duce a higher time complexity. Therefore, the time complex- ity of our proposed VIVACE method is comparable with the baselines, demonstrating its promising efficiency.
Datasets. We consider five widely adopted datasets in the lit- erature for comprehensive evaluations. First, following (Fan et al., 2022; Dwivedi et al., 2023), we adopt three datasets, i.c., CMNIST, CFashion and CKuzushiji, which are con- verted from image datasets using superpixels (Knyazev et al., 2019), since they have controllable bias degrees and clear unnaan-dddrsttddaalle ground -tunt invariant subgraphs for evaluations. Specifically, for CMNIST, each graph is con- verted from an image in MNIST (LeCun et al., 1998). The task is to classify each graph into the corresponding hand- written digit. And the spurious correlations are introduced via colorizing the background based on the correlations with the label. The degree of such correlation is controlled by r. We consider datasets with biased r = (0.8, 0.9, 0.95)
6

---

Disentangling Invariant Subgraph via Variance Contrastive Estimation under Distribution Shifts
Table 1. Exxerimental results (%) of our method and baselines. The evaluation metric is accuracy for CMNIST, CFashion, and CKuzushiji, and ROC-AUC for MOLSIDER and MOLHIV. + denotes the standard deviation. The best results are in bold for each row. Our VIVACE outperforms the baselines in all comparisons, indicating its superiority against graph distribution shifts.
Bias
Methods
Dataset
r
GCN
GIN
FactorGCN
DifiPool
DIR
LDD
DisC
VIVACE
0.8
50.43=4.13
57.75=0.78
72.30::1.18
73.79=0.02
9.9810.33
64.95:1.222
82.60-10.93
0..71=0.74
CMNIST
0.9
28.97:14.40
36.78+5.55
62.35.+5.07
66.45=00778
9.96:::.23
56.65+22118
78.14:2..14
79.46 1.87
0.95
13.50
16.04:1.14
42.50++.91
47.12=11004
10.03++0.27
46.83-+2.88
63.47:15.65
6..72:::461
0.8
63,60-0.553
64.25+0.46
61.23+1.11
6..82=0.53
13.02:+:.92
63.85+1117
66.85+1.11
1.20
CFashion
0.9
57.22-10.93
58. 08.03:1:.40
13.50:1..29
57.50-0.39
12,80
64.30:10.9
65,33:14.70
65.38::..18
0.95
47.6911.42
49.74:1060
45.78112.0
50.86=0.20
11.98
62.28:10.48
63.93:11.50
1.27
0.8
38.45+1.100
41.83+0.78
42
45.46-0065
10.35+0.32
42.38+0.33
55.53++2.29
55,58 1.87
CKuzushiji
0.9
28.351.0.79
30.09
32.3512.79
36.18-0.19
10.72:10.27
38.75:0.49
48.13:1..59
18.1511.91
0.95
20.70-1..88
21.18+1.63
23.87:10.12
27.45=0.26
10.5910.46
33.08:10.59
36.63111.3
37.01=11677
MOLSIDER
59.62+ 1.82
57.61+1.48
53.32+1.75
10.21+1.55
17.74+1.63
58.83+1.62
59.31+1.87
62.15+1.110
MOLHIV
7..13+1.01
75.63+1441
57.18411.54
76.32=1.48
77.05=10.57
76.91+1.81
76.97-11.03
78.11-10.22
for training and unbiased r for testing. We also adopt similar strategies to construct CFashion and CKuzushiji from Fashion-MNIST (Xiao et al., 2017) and Kuzushiji- MNIST (Clanuwat et al., 2018) datasets. Also, we consider two datasets, MOLSIDER and MOLHIV from Open Graph Benchmark (Hu et al., 2020). The default split separates structuralyy different molecules with different scaffolds into different subsets, i.e., training/vliddatoonnesttgg sets. We report the accuracy for CMNIST, CFashion, and CKuzushiji, ROC-AUC for MOLSIDER and MOLHIV.
based invariant subgraph prediction module, showing the remarkable OOD generalization ability in practice.
As the degree of spurious correlations increases, namely grows larger, the performance of all the methods tends to decrease since there exists a larger degree of distribution shifts between testing and training graph data. Nevertheless, our proposed VIVACE is able to keep the most relatively stable performances, and demonstrates the effectiveness in handling graph distribution shifts. When increases from 0.8 to 0.95 on CFashion dataset, GCN drops by about 15% accuracy, while our method with GCN as the backbone can drop no more than 4%, which verifies that our method can better capture the invariant and variant subgraphs, and thus well handle distribution shifts.
4.2. Experiment Results
The experimental results reported in Table 1. We have the following observations.
Our proposed VIVACE method consistently and signifi- cantly outperforms the GNN backbones (i.e., GCN and GIN) on the CMNIST, CFashion, and CKuzushiji datasets. When compared with the graph disentangling method Fac- torGCN or representative graph pooling method, our pro- posed VIVACE also achieves substantial performance gains. Besides, the representative graph invariant learning method DIR achieves unsatisfactory prediction performances. It is because DIR relies on datasets without severe bias to create interventional distributions, so that the invariant patterns captured by DIR are not accurate when this assumption is not valid. The general debiasing method LDD and graph debiasing method DisC show promising performance gains upon the other baselines, but our VIVACE still performs better than them in all comparisons. One plausible reason is that they can not fully leverage the informative variant patterns for capturing invariant subgraphs to achieve OOD generalized predictions. Our method can capture accurate invariant and variant subgraphs simultaneously by the mu- tually promoted two modules respectively, i.e., variant sub- graph contrastive estimation module and propensity score
We also conduct comparisons on more challenging and large-scale OGB datasets. Since the molecules with dif- ferent scaffolds are naturally split into different train- ing/validdtioon/tssinng sets, the generalization between molecules with different scaffolds is difficult. Neverthe- less, our proposed VIVACE also achieves improvements upon the best results of the baselines. For example, VI- VACE increases the ROC-AUC by 1.94% on MOLSIDER and I. 1.06% on MOLHIV against the strongest baselines re- spectively. Although the baseline DisC shows competitive performances on CMNIST, CFashion, CKuzushiji, it fails to achieve promising OOD generalization results on these two more challenging OGB datasets. In contrast, VIVACE achieves significant and consistent performance gains in all comparisons, demonstrating its effectiveness against distri- bution shifts.
Training Dynamics. In Figure 5 in Appendix, we plot the loss and accuracy in the training process on CMNIST (r = 0.9), while the results on the other datasets show similar patterns. We can empirically observe the convergence of
7

---

Disentangling Invariant Subgraph via Variance Contrastive Estimation under Distribution Shifts
80
79.0
77,28
GCNI
78.0
Performance (%)
75
w/ GIN
w/o Var.
77.0
10-1
10a4
10a
w/o IPW
70
(a) The sensitivity of a.
65
80.0
Acouracy
60
78.0
77.0
55
0.8
0.8
CMNIST
MOLSIDER
(b) The sensitivity of
Figure 3. Ablation studies. "w/ GCNII and "w/ GIN' denote adopt- ing different backbones. "w/o Var. and "w/o IPW" denote remov- ing the variant subgraph contrastive module and inverse propensity weighting module.
Figure 4. Sensitivity analysis for some important hyper-parameter in our model. The and grey lines denote results of VIVACE and the best result of all baselines, respectively. VIVACE achieves better results with a wide range of hyper-paramett choices.
our proposed method. The loss and accuracy will converge before reaching the maximal training epoch.
4.4. Hyper-parameter Sensitivity
We analyze the hyper-parameter sensitivity of our proposed method in terms of the invariance regularizer coefficient a in the variant subgraph contrastive estimation module and the coefficient y that controls the degree of fitting the spurious correlations in Eq. (7).
4.3. Ablation Study
We perform comprehensive ablation studies to further vali- date the effectiveness of the key modules in our method.
Variants 'w/GCNII' "w/ GIN' means that place the instantiation of the GNN components from GCN (Kipf& & Welling, 2017) to the other GNN back- bone GCNII (Chen et al., 2020) or GIN (Xu et al., 2019).
For simplicity, only report the results CMNIST (r= 0.9), while the results on other datasets show similar pat- terns. In Figure 4, the red and grey lines denote the results of VIVACE and the best result of all baselines, respectively, The arper-paramtter a has an influence on the OOD gener- alization performance, indicating that we need properly balance the self-supervised contrastive learning loss and the invariance regularizer term. A small value may be sufficient to encourage the invariance among different en- vironments effectively, while a very large value may affect the self-supervised contrastive training. We also observe the hyperparameter has a moderate impact on the model performance, verifying the significance to fit the spurious correlations. Overall, although an appropriate choice of the hyper-parameters can further achieve generalization perfor- mance gains, our method is not sensitive to them and is able to outperform the best baselines within wide range of hyper- parameters choices.
Variant 'w/o Var.' means that we remove the variant subgraph contrastive estimation module and Variant 'w/o IPW' further removes the inverse propensity weighting module.
The results are obtained from the CMNIST (r - 0.9) dataset. From Figure 3, we can first observe that our method is also compatible with the other popular GNNs even with slight performance improvements. By default, we use GCN (Kipf Welling, 2017) as the backbone GNN model. Although the backbone GNN is replaced with GCNII or GIN model, VIVACE keep relatively stable performances even achieve gains. Specifically, the generalization perfor- mance further increases by 0.65% when using GIN as the backbone that is a more expressive GNN.
Besides, the performance of the variants 'w/o Var.' and "w/o IPW' drop drastically, indicating that it is important explicitly identify variant subgraphs and further estimate their impact the labels for the final goal of learning invariant subgraphs for OOD generalization. If the spurious correlations of all variant subgraphs with the labels are treated equally, the accurate identification of the invariant subgraphs will also be affected, leading to the unsatisfactory OOD generalization performance.
5. Related Works
In this section, we review the related works of graph neural network, disentangled representation learning, and OOD generalization.
Graph Neural Network. Graph data has been widely used to model the complex relationships between different enti- ties. In the field graph machine learning, graph neural networks (GNNs) (Kipf Welling, 2017; Veličković et al.,
8

---

Disentangling Invariant Subgraph via Variance Contrastive Estimation under Distribution Shifts
2018: Li et al., 2023a; 2025; Chen et al., 2025), following message passing paradigm to iteratively update node rep- resentations by the neighbors, have drawn ever-increasing attention and achieved enormous success in various appli- cations across wide-ranging applications. Graph-level pre- diction task is one of the popular applications and can be abstracted as one main branch in addition to the node-level and link-level prediction tasks. The prediction of a graph is generally based on its critical subgraph rather than the whole input graph (Luo et al., 2020; Yu et al., 2021; Xuan et al., 2019), where the critical subgraph is regarded as the part that is truly predictive to the label in the task. There are some works ( (Ying et al., 2018; Luo et al., 2020; Sun et al., 2021; Monti et al., 2018) proposed to improve the graph-level prediction performances by capturing the criti- cal subgraph for learning graph representations effectively. However, these methods are built upon the in-distribution (L.D. hhD.ohhsssssss and fail to generalize under distribution shifts.
in-the-wild unseen data (Feng et al., 2023), OOD general- ization on graphs has drawn great attention (Li et al., 2022a; Gui et al., 2022; Cai et al., 2024). Several famous works are proposed to tackle this problem on graphs by learn- ing subgraph backed by different theories or assumptions, including causality (Wu et al., 2022c; Sui et al., 2022), in- variant learning (Wu et al., 2022a; Li et al., 2022d; Yao et al., 2024; Chen al., 2024c; Zhang et al.. 2023b), disen- tanglement (Fan et al., 2022; Li et al., 2024; Zhang et al., 2024a), information bottleneck (Miao et al., 2022). Differ- ent from these works that output explainable or invariant subgraphs under distribution shifts, some works directly learn generalizable graph representations for the problems where distribution shifts exist on graph size (Bevilacqua et al., 2021: Buffelli et al., 2022) or the other structural patterns (Wu al., 2024). And the learned representations are expected to remain invariant across different environ- ments. For capturing invariant subgraphs, most of them heavily rely on the predefined or automatically generated environment labels, i.e., multiple training environments to specify the variant information, and further learn invariant subgraphs. However, the environment labels are unavail- able and directly generating environment labels during the training process is also impractical, leading to inaccurate modeling for the variant patterns and the estimation of the degree of the spurious correlations. Some of them rely on the strong causality assumptions that are hardly guaranteed to hold true all the time in real-world scenarios. How to explicitly estimate the impact of environment- variant information and further eliminate the spurious correlations remains largely unexplored.
Disentangled Representation Learning. Disentangled rep- resentation learning (DRL. (Wang et al., 2024) aims to learn representations capable of identifying the underlying factors behind the observable data (Zhang et al., 2024b; Wang et al., 2022; Zhang et al., 2023a; Wang et al., 2025a). In addition to its success in image and video (Wang et al., 2025b; Chen et al., 2024b;a), DRL is also a promising direction in the domain of graph neural network (Li et al., 2021b; 2022c; Zhang et al., 2023c). For instance, DisenGCN (Ma et al., 2019) is the pioneering approach in this area, introducing method to disentangle node representations by dynami- cally extracting factors that contribute to the formation of edges between a node and its neighbors. Building upon this, IPGDN (Liu et al., 2020) not only separates the fac- tors connecting nodes to their neighbors but also ensures these factors to be independent. Both methods are guided by supervision from downstream node classification tasks. FactorGCN (Yang et al., 2020) further advances disentan- gled learning by applying factoring mechanism at the input graph level. This approach disentangles input graph to create distinct factors, which are subsequently treated as separate graphs. However, their generalization abilities under distribution shifts remain relatively underexplored.
6. Conclusion
In this paper, we study learning invariant subgraph via self- supervised contrastive variant subgraph estimation, which can well handle graph distribution shifts with severe bias for OOD generalization. We find that the variance infor- mation behind graph data is also important and should not be directly ignored but be explicitly captured. We propose a novel VIVACE model, which consists of three tailored modules, i.e., invariant and variant subgraph identification module, variant subgraph contrastive estimation module, and inverse propensity weighting based invariant subgraph prediction module. The main technical contribution lies in that design self-supervised contrastive learning on variant subgraphs to explicitly model the degree of the spu- rious correlations and further design the inverse propensity weighting strategy to reweight the invariant subgraph pre- dictions to climinate the spurious correlations for OOD generalization. We conduct extensive experiments on sev- eral graph classification benchmark datasets. The results demonstrate the superiority of our proposed method over state-of-the-art baselines against distribution shifts.
cittrf-Distribution (OOD) Generalization. Most ma- chine learning models rely on the assumption that testing and training data are identically distributed. However, this assumption can be easily violated due to the widely existing but uncontrollable distribution shifts in the real world (Shen et al., 2021; Du et al., 2022; Yang et al., 2023; Sui et al., 2023). The performance will degenerate significantly if the machine learning models do not have strong OOD general- ization ability. Graph neural networks, as the most popular models in the graph community recently, also face the same obstacle. Considering the increasing demand for handling

---

Disentangling Invariant Subgraph via Variance Contrastive Estimation under Distribution Shifts
Acknowledgements
Chen, M., Wei, Z., Huang, Z., Ding, B., and Li, Y. Simple and deep graph convolutional networks. In International conference machine learning, pp. 1725-1735. PMLR, 2020.
This work supported by National Natural Science Foun- dation of China N2.6222220,, 62473380, 62272481, 61972460, 62372470, 62172440), Beijing National Re- search Center for Information Science and Technology der Grant No.BNR2023TD00306 The Science and Technol- ogy Innovation Program of Hunan Province (2023RC1029).
Chen, Y., Bian, Y., Zhou, K., Xie, B., Han, B., and Cheng. J. Does invariant graph learning via environment augmenta- tion learn invariance? Advances in Neural Information Processing Systems, 36, 2024c.
Impact Statement
Clanuwat, T., Bober-Irizar, M., Kitamoto, A., Lamb, A., Yamamoto, K., and Ha, D. Deep learning for classical japanese literature. arXiv preprint arXiv:1812.01718, 2018.
This paper presents work whose goal advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Du, X., Wu, Z.., Feng, F., He, X., and Tang. J. Invariant representation learning for multimedia recommendation. In Proceedings of the 30th ACM International Conference on Multimedia, pp. 619-628, 2022.
References
Bengio, Y., Deleu, T., Rahaman, N., Ke, R., Lachapelle, S., Bilaniuk, O., Goyal, A., and Pal, C. A meta-transfer objective for learning to disentangle causal mechanisms. International Conference on Learning Representations, 2019.
Duvenaud, D. K., Maclaurin, D., Iparraguirre, J., Bom- barell, R., Hirzel, T., Aspuru-Guzii A., and Adams, P. Convolutional networks graphs for learning molecular fingerprints. Neural Information Processing Systems, 28, 2015.
Bevilacqua, B., Zhou, Y., and Ribeiro, B. Size-invariant graph representations for graph classification extrapola- tions. In International Conference on Machine Learning, pp. 837-851. PMLR, 2021.
Dwivedi, V. P., Joshi, C. K., K.,, Luu, A. T., Laurent, T., Ben- gio, Y., and Bresson, X. Benchmarking graph neural networks. Journal of Machine Learning Research, 24 (43):1-48, 2023.
Buffelli, D., Liò, P., and Vandin, F. Sizeshiftreg: a regu- larization method for improving size- sezereeneraliztiion in graph neural networks. Advances in Neural Information Processing Systems, 35:31871-31888 2022.
Fan, S., Wang. X., Mo, Y.. Shi, C., and Tang. J. Debi- asing graph neural networks via learning disentangled causal substructure. In Advances in Neural Information Processing Systems, 2022.
Cai, J., Wang, X., Li, H., Zhang, Z.,, and Zhu, W. Multi- modal graph neural architecture search under distribution shifts. In Proceedings of the AAAI Conference on Artifi- cial Intelligence, volume 38, pp. 8227-8233 2024.
Feng, W., Li, H., Wang. X., Duan, X., Qian, Z., Liu, W., and Zhu, W. Multimedia cognition and evaluation in open environments. In Proceedings of the 1st Interna- tional Workshop Multimedia Content Generation and Evaluation: New Methods and Practice, pp. 9-18, 2023.
Chen, H., Wang, X., Zhang, Y., Zhou, Y., Zhang, Z., Tang, S., and Zhu, W. Disenstudio: Customized multi-subjes subject text-to-video generation with disentangled spatial control. In Proceedings of the 32nd ACM International Confer- on Multimedia, pp. 3637-3646, 2024a.
Fey, M. and Lenssen, J. E. Fast graph representation learning with PyTorch Geometric. In ICLR Workshop Repre- sentation Learning on Graphs and Manifolds, 2019.
Gui, S., Li, X., Wang, L., and Ji, S. Good: A graph out-of- distribution benchmark. Advances in Neural Information Processing Systems, 35:2059-2073, 2022.
Chen, H., Zhang, Y., Wu, S., Wang, X., Duan, X., Zhou, Y., disentan- gled tuning for subject-driven text-to-image generation. International Conference on Learning Representations, 2024b.
and Zhu, W. Disenbooth:
Hu, W., Fey, M., Zitnik, M., Dong, Y., Ren, H., Liu, B., Catasta, M., and Leskovec, J. Open graph benchmark: Datasets for machine learning on graphs. Neural Infor- mation Processing Systems, 2020.
Chen, H., Wang, X., Zhang, Z., Li, H., Feng, L., and Zhu, W. Autogfm: Automated graph foundation model with adaptive architecture customization. In International Con- ference on Machine Learning, 2025.
Hudson, D. and Manning, D. Learning by abstraction:
The neural state machine. Neural Information Processing Systems, 2019.
10

---

Disentangling Invariant Subgraph via Variance Contrastive Estimation under Distribution Shifts
Jung, Y., Tian, J.. and Bareinboim, E. Learning causal effects via weighted empirical risk minimization. Neural Information Processing Systems, 33: 1697- 2020.
Li, H., Zhang, Z., Wang, X., and Zhu, W. Learning invari- ant graph representations for out-of-distribution general- ization. In Advances in Neural Information Processing Systems, 2022d.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Li, H., Wang, learning: A survey. International Joint Conference on Artificial Intelligence, 2023a.
X., and Zhu, W. Curriculum graph machine
Kipf, T. N. and Welling, M. Semi-supervised classifica- tion with graph convolutional networks. In International Conference Learning Representations, 2017.
Li, H., Zhang, Z., Wang, X., and Zhu, W. Invariant node representation learning under distribution shifts with mul- tiple latent environments. ACM Transactions on Informa- tion Systems, 42(1):1-30, 2023b.
Knyazev, B., Taylor, G. W., and Amer, M. Understanding attention and generalization in graph neural networks. Advances in neural information processing systems, 32, 2019.
Li, H., Wang, X., Zhang, Z.,, Chen, H., Zhang, Z., and Zhu, W. Disentangled graph self-supervised learning for diitt-riiitiittti generalization. In International Conference on Machine Learning, 2024.
Krueger, D., Caballero, E., Jacobsen, J.-H., Zhang, A., Bi- nas, J., Zhang, D., Le Priol, R., and Courville, A. Out- of-distribution generalization via risk extrapolation (rex). In International Conference on Machine Learning, pp- 5815-5826. PMLR, 2021.
Li, H., Wang, X., Zhang, Z.., Wu, Z., Xiao, L., and Zhu, W. Self-supervised masked graph autoencoder via structure- aware curriculum. In International Conference on Ma- chine Learning, 2025.
LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P. Gradient- based learning applied to document recognition. Proceed- ings of the IEEE, 1998.
Liu, Y., Wang, X., Wu, S., and Xiao, Z. Independence promoted graph disentangled networks. In Association for the Advancement of Artificial Intelligence, volume 34, pp.
Lee, J., Kim, E., Lee, J., Lee, J., and Choo, J. Learn- ing debiased representation via disentangled feature aug- mentation. Neural Information Processing Systems, 34: 2021.
2020.
Li, H., Cui, P., Zang, C., Zhang, T., Zhu, W., and Lin, Y. Fates of microscopic social ecosystems: Keep alive or dead? In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 668-676, 2019.
Luo, D., Cheng, W., Xu, D., Yu, W., Zong. B., Chen, H., and Zhang, X. Parameterized explainer for graph neural network. Neural Information Processing Systems, 2020.
Ma, J., Cui, P., Kuang, K., Wang. X., and Zhu, W. Disen- tangled graph convolutional networks. In International conference on machine learning, pp. 2212-4221 PMLR, 2019.
Li, H., Wang, X., Zhang, Z., Ma, J., Cui, P., and Zhu, W. Intention- aware sequential recommendation with struc- tured intent transition. IEEE Transactions Knowledge and Data Engineering. 34(11):5403-5414 2021a.
Marcheggiani, D. and Titov, Encoding sentences with graph convolutional networks for semantic role labeling. arXiv preprint arXiv: 170.004820. 2017.
Li, H., Wang, X., Zhang, Z., Yuan, Z., Li, H., and Zhu, W. Disentangled contrastive learning on graphs. Advances in Neural Information Processing Systems, 34, 2021b.
Miao, S., Liu, M., and Li, Interpretable and generaliz- able graph learning via stochastic attention mechanism. In International Conference on Machine Learning, pp.
Li, H., Wang, X., Zhang. Z., and Zhu, W. Ood-gnn: Out-of-distribation generalized graph neural network. IEEE Transactions on Knowledge and Data Engineer- ing, 2022a.
15524-15543, 2022.
Monti, F., Otness, K., and Bronstein, M. M. Motifnet: motif-based graph convolutional network for directed graphs. In 2018 IEEE Data Science Workshop (DSW), pp. 225-228. IEEE, 2018.
Li, H., Wang X., Zhang, Z., and Zhu, W. Out-of-distributio generalization on graphs: survey. arXiv preprint arXiv:2202.07987, 2022b.
Li, H., Zhang, Z., Wang, X., and Zhu, W. Disentangled graph contrastive learning with independence promotion. IEEE Transactions on Knowledge and Data Engineering, 2022c.
Qi, J., Tang, K., Sun, Q., Hua, X.-S., and Zhang, H. Class is invariant to context and vice versa: on learning invari- ance for of disttiiibbtti generalization. In Computer Vision-ECCV 2022, pp. 92-109, 2022.

---

Disentangling Invariant Subgraph via Variance Contrastive Estimation under Distribution Shifts
Qiu, J., Tang, J., Ma, H., Dong, Y., Wang, K., and Tang, Deepinf: Social influence prediction with deep learning. In Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 1199-2119, 2018.
Wang, X., Pan, Z., Chen, H., and Zhu, W. Divico: Dis- entangled visual token compression for efficient large vision-language model. IEEE Transactions Circuits and Systems for Video Technology. 2025b.
Wu, Q. Zhang, H., Yan, J., and Wipf, D. Handling distri- bution shifts on graphs: An invariance perspective. In International Conference on Learning Representations, 2022a.
Seaman, S. R. and Vansteelandt, S. Introduction to double robust methods for incomplete data. Statistical science: a review journal of the Institute of Mathematical Statistics, 33(2):84,, 2018.
Wu, S., Sun, F., Zhang, W., Xie, X., and Cui, B. Graph neural networks in recommender systems: a survey, ACM Computing Surveys, 55(5):1337, 2022b.
Shen, Z., Liu, J., He, Y., Zhang, X., Xu, R., Yu, H., and Cui, P. Towards -fit---iisttitttionn generalization: survey. arXiv preprint arXiv: / 2021.
Wu, S., Cao, K., Ribeiro, B., Zou, J., and Leskovec, J. Graphmetro: Mitigating complex distribution shifts in gnns via mixture of aligned experts. Neural Information Processing Systems, 2024.
Sui, Y., Wang. X., Wu, J., Lin, M., He, X., and Chua, T.-S.
Causal attention for interpretable and generalizable graph classification. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 1696-1705, 2022.
Wu, Y.-X., Wang, X., Zhang, A., He, X., and seng Chua, T. Discovering invariant rationales for graph neural net- works. In International Conference on Learning Repre- sentations, 2022c.
Sui, Y., Wu, Q., Wu, J., Cui, L., Zhou, J., Wang, X., and He, X. Unleashing the power of graph data augmentation on covariate distribution shift. In Thirty- seventh Conference on Neural Information Processing Systems, 2023.
Xiao, H.,, Rasul, K., and Vollgraf, R. Fashion-m novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv: 1ee.1177.7.777777 2017.
Sun, Q., Li, J., Peng, H., Wu, J., Ning, Y., Yu, S., and He, L. Sugar: Subgraph neural network with reinforcement pooling and supervidd mutual information mecha- nism. In Proceedings of the Web Conference 2021, pp. 2081-2091 2021.
Xu, K., Hu, W., Leskovec, J., and Jegelka, S. How powerful are graph neural networks? In International Conference on Learning Representations, 2019.
Xuan, Q., Wang, J., Zhao, M., Yuan, J., Fu, C., Ruan, Z., and Chen, G. Subgraph networks with application to structural feature space expansion. IEEE Transactions on Knowledge and Data Engineering, 33(6):277--278,, 2019.
Veličković, P.. Cucurull, G., Casanova, A., Romero, A., Lio, P., and Bengio, Y. Graph attention networks. In International Conference on Learning Representations, 2018.
Wang, Mao, Z., Wang, B., and Guo, L. Knowledge graph embedding: A survey of approaches and applications. IEEE Transactions on Knowledge and Data Engineering, 2017.
Yang, N., Zeng, K., Wu, Q., Jia, X., and Yan, J. Learning substructure invariance for out-of-distriuttion molecular representations. Advances in Neural Information Process- ing Systems, 35:12964-1222 12978, 2022.
Wang, X., Chen, H.. Zhou, Y., Ma. J., and Zhu, W. Dis- entangled representation learning for recommendation. IEEE Transactions Pattern Analysis and Machine In- telligence, 45(1) 45(1):408 424, 2022.
Yang, Y., Feng, 7., Song, M., and Wang, X. Factorizable graph convolutional networks. Advances in Neural Infor- mation Processing Systems, 33:20286-20296 2020.
Yang, Z., He, X., Zhang. J., Wu, J., Xin, X., Chen, J., and Wang, X. A generic learning framework for sequential recommendation with distribution shifts. In Proceedings of the 46th International ACM SIGIR Conference on Re- search and Development in Information Retrieval, 2023.
Wang, X., Chen, H., Wu, Z., Zhu, W., et al. Disentangled representation learning. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2024.
Wang, X., Chen, H., Pan, Z., Zhou, Y., Guan, C., Sun, L., and Zhu, W. Automated disentangled sequential recom- mendation with large language models. ACM Transac- tions on Information Systems, 43(2):1-29, 2025a.
Yao, T., Chen, Y., Chen, Z., Hu, K., Shen, Z., and Zhang, K. Empowering graph invariance learning with deep spuri- ous infomax. arXiv preprint arXiv:2407. 1083, 2024.
12

---

Disentangling Invariant Subgraph via Variance Contrastive Estimation under Distribution Shifts
Ying, Z., You, J., Morris, C., Ren, X., Hamilton, W., and Leskovec, J. Hierarchical graph representation learning with differentiable pooling, Neural Information Process- ing Systems, 31, 2018.
You, Y., Chen, T., Sui, Y., Chen, T., Wang, Z., and Shen, Y. Graph contrastive learning with augmentations. Ad- vances in Neural Information Processing Systems, 33: 5812-5823, 2020.
Yu, J., Xu, T., Rong, Y., Bian, Y., Huang, J., and He, R. Graph information bottleneck for subgraph recognition. In International Conference on Learning Representations, 2021.
Zellers, R.. Yatskar, M., Thomson, S., and Choi, Y. Neural motifs: Scene graph parsing with global context. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 5831-5840, 2018.
Zhang, Y., Wang, X., Chen, H., and Zhu, W. Adaptive disentangled transformer for sequential recommendation. In Proceedings of the 29th ACM SIGKDD conference knowledge discovery and data mining, pp. 3434-3445, 2023a.
Zhang, Z., Wang. X., Zhang. Z., Qin, Z., Wen, W., Xue, H., Li, H., and Zhu, W. Spectral invariant learning for dynamic graphs under distribution shifts. Advances in Neural Information Processing Systems, 36:6619-6633 2023b.
Zhang, Z., Wang, X., Zhang, Z., Shen, G., Shen, S., and Zhu, W. Unsupervised graph neural architecture search with disentangled self-supervision. Advances in Neu- ral Information Processing Systems, 36:73175-73190, 2023c.
Zhang, Z., Wang, X., Chen, Hi, H., and Zhu, W. Dis- entangled dynamic graph attention network for out-of- distribution sequential recommendation. ACM Transac- tions on Information Systems, 13(1):1-42, 2024a.
Zhang, Z., Wang, X., Qin, Y., Chen, H., Zhang, Z., Chu, X., and Zhu, W. Disentangled continual graph neural architecture search with invariant modular supernet. In International Conference on Machine Learning, 2024b.
13

---

Disentangling Invariant Subgraph via Variance Contrastive Estimation under Distribution Shifts
A. Notations
For better clarifications, we first summarize the key notations in the proposed method and the corresponding descriptions in Table 2.
Table
The summary the key notations and correspouding descriptions.
Notation
Description
The number of graphs
The input graph
X,A
The node feature and adjacency matrix
The invariant subgraph generator
G1 ((G)
The invariant subgraph of G
Gy = G\GI
The variant subgraph of G
M
The learnable edge mask matrix
fr/fv
The invariant/variant subgraph predictor
hy
The variant subgraph encoder
uv
The variant subgraph classifier
l
The loss function
Y
The number of different labels
B. Proof of Theorem 1
Proof. Denote the first contrastive loss term of Eq. (5) as:
|
(ov
Lgs!
1 |Y|
Σ Ca kv)
log
(13)
Y
Eny= exp (o (hr(Gvs),
wi
Denote the second variance term of Eq. (5)
L Var (I, ky;k))
k = L,....., |Y].
(14)
:: To prove the optimal invariant subgraph generator d* can minimize the contrastive loss Lasts i.c.,
Th arg min Lasts
(15)
assume there exists another क' + 46, where the input graph G is disentangled into G (G', Gly) and
Lust c
(+))
(16)
This implics that Gy includes the ground-truth invariant subgraph information, i.c.,
Gyy n G = Gyy (GVG)) + 0.
(17)
Therefore, the contrastive loss ( lv; k) among the environments partitioned by the graph label A 1,..., |Y| is dependent on the graph label itself, i.e.,
k1, K2 E {l..... |.]. K1 + kiz such that (o. lly k2) o Lsal (2, lyv kz).
(18)
Thus,
L ( Var(esst (v; O, = YY)
(19)
However, Gi excludes all the ground-truth invariant information that sufficiently predictive to the graph label. We have
(ssi("" lhv & = 1) = lsst( hv; k = 2) Exsi( "/v; k = (V),
(20)
14

---

Disentangling Invariant Subgraph via Variance Contrastive Estimation under Distribution Shifts
i..,
Lyyar ( Var (Csat( ,hv;k)) 0, k= 1, 1....... Y Laar Larr (
(21)
Thus,
(22)
which contradicts the assumption that the variance term Lvar is minimized. Therefore, we prove that the optimal invariant subgraph generator hh minimizes the contrastive loss Leel-
=: To prove that minimizing the contrastive loss Last implies + = = i.e., the uniqueness of Φ*, assume there exists another invariant subgraph generator 4' + d- derived by minimizing Lest, where G = (G',, GV-) and
np = arg min Lyst-
(23)
Since Lest is minimized, Gyy preserves all the intrinsic features of the ground-truth variant subgraph, i.e.,
GV C C.v.
(24)
Meanwhile, because the second variance term L minimized across graph label partitions 1,..., |Y), only ground-truth variant patterns be included in Gy, i.e.
Gv C C
(25)
Therefore,
Gyy = GV + d' = -
(26)
Thus, we conclude that there exists a unique 4* that minimizes the first contrastive loss term of Eq. (5).
C. Additional Experimental Details
C.1. Datasets
The datasets are publicly available as follows:
CMNIST: http:://ttpt////e// C/NNST/rtrrr//aaa...eeeccc/////d with license unspecified
CFashion: http://ttttsss...toocoo///////et -nniiiinoo-n-nntt with MIT License
CKuzushiji: C http:////gtttt.oc//////oooa/olaa with CC BY-SA 4.0 License
/Mt.D//tttttr///occccccccc///eer with MIT License
wtt.tn/..tt////oh.....errrrrrrrr with MIT License
C.2. Implementations
We use GCN (Kipf & Welling, 2017) as the backbone GNN model. The hyper-parameter q in Eq. (7) is 0.7. We adopt the Adam optimizer (Kingma Ba, 2014). Note that we adopt the default hyperparameter settings following (Fan et al., 2022) for fair comparison. We report the mean values with standard deviations of four repeated experiments.
All the experiments conducted with:
Operating System: Ubuntu 18.04.1 LTS
CPU: Intel(R) Xeon(R) CPU ES-2699 v4@2.20GHz
GPU: NVIDIA GeForce GTX TITAN Xp with 12GB of Memory
Software: Python 3.6.5; NumPy 1.19.2; PyTorch 1.10.1; PyTorch Geometric 2.0.3 (Fey Lenssen, 2019)
15

---

Disentangling Invariant Subgraph via Variance Contrastive Estimation under Distribution Shifts
7.5
100
6.0
80
Performance (%)
4.5
60
Loss
Loss
Performance
3.0
40
1.5
20
0.0
0
50
100
150
200
250
300
Epoch
Figure 5. The training dynamics on CMNIST (.9)). We can observe the convergence of our proposed method in practice.
D. Training Dynamics
We present the loss and accuracy during the training process the CMNIST dataset 0.9) in Figure 5. Similar trends observed across other datasets. The results empirically demonstrate the convergence of our proposed method, with both the loss and accuracy stabilizing well before the maximum training epoch is reached.
E. Limitations
potential limitation is that we only focus on OOD generalized prediction tasks on static and homogenous graphs, which are most common in the graph learning community. is worth exploring extend this work into dynamic and heterogeneous graphs in future.
16


Reviews:
[{"review_id": "zOOrjLQonV", "reviewer": "Reviewer_sHjp", "summary": "This manuscript studies out-of-distribution generalization issue in graph neural networks. The authors propose learning invariant subgraphs via variant subgraph contrastive estimation, which can handle graph distribution shifts with severe bias. The key innovation is leveraging contrastive learning on variant subgraphs to estimate spurious correlations, which is then mitigated using inverse propensity weighting. This method explicitly addresses the scenario where environment labels are either unavailable or unreliable, significantly enhancing the robustness of GNNs to severe biases in datasets.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "- **Q1. Theoretical analyses.**\n\nWe have added **Theorem 1** to show our method can accurately identify the invariant and variant subgraphs for OOD generalization. \n\n**Theorem 1.** Denote the optimal invariant subgraph generator $\\Phi^*$ that disentangles the ground-truth invariant subgraph $G_I^*$ and variant subgraph $G_V^*$ given the input graph $G$, where $G_I^*$ satisfies Assumption 2.1 and denote the complement as $G_V^* = G\\backslash G_I^*$. Assume the second variance term of Eq. (5) is minimized, the first contrastive loss term is minimized *iff* the invariant subgraph generator $\\Phi$ equals $\\Phi^*$.\n\n*Proofs.*\n\nDenote the first contrastive loss term of Eq. (5) as $L_{ssl}$ and the second variance term as $L_{var}$.\n\n$\\Leftarrow$: To prove $\\Phi^* = \\arg\\min_{\\Phi} L_{ssl}$, assume there exists another $\\Phi^\\prime \\neq \\Phi^*$ such that $L_{ssl}(\\Phi^\\prime) < L_{ssl}(\\Phi^*)$. It implies that $G_V^\\prime$ includes the ground-truth invariant subgraph information,\n$$\nG_V^\\prime \\cap G_I^* = G_V^\\prime \\cap G\\backslash G_V^* \\neq \\emptyset.\n$$\nTherefore, the contrastive loss $\\ell_{ssl}\\left(\\Phi, h_V; k\\right)$ among the environments partitioned by the graph label $k = 1,\\dots,|Y|$ is dependent on the label itself, i.e.,\n$$\n\\exists\\, k_1, k_2 = 1, \\dots, |Y|, k_1 \\neq k_2, \\text{such that } \n\\ell_{\\text{ssl}}(\\Phi^\\prime, h_V; k_1) \\neq \\ell_{\\text{ssl}}(\\Phi^\\prime, h_V; k_2).\n$$\nThus\n$$\nL_{var}(\\Phi^\\prime) = \\mathrm{Var}\\left(\\ell_{ssl}\\left(\\Phi^\\prime, h_V; k\\right)\\right) > 0.\n$$\nHowever, $G_V^*$ excludes all the ground-truth invariant information that is sufficiently predictive to the graph label. We have\n$$\nL_{var}(\\Phi^*) = 0.\n$$\nThus, \n$$\nL_{var}(\\Phi^\\prime) > L_{var}(\\Phi^*).\n$$\nThis contradicts the assumption that the variance term $L_{var}$ is minimized. Finally we prove that the optimal invariant subgraph generator $\\Phi^*$ can minimize the contrastive loss $L_{ssl}$.\n\n$\\Rightarrow$: To prove minimizing the contrastive loss $L_{ssl}$ can derive  $\\Phi$ equals $\\Phi^*$, assume there exists another invariant subgraph generator $\\Phi^\\prime$ derived by minimizing $L_{ssl}$, i.e., $\\Phi^\\prime \\neq \\Phi^*$, where $G = (G_I^\\prime, G_V^\\prime)$ and $\\Phi^\\prime = \\arg\\min_{\\Phi} L_{ssl}$.  Because $L_{ssl}$ is minimized, $G_V^\\prime$ preserves all the intrinsic features of the ground-truth variant subgraph, i.e., $G_V^* \\subseteq G_V^\\prime$. Since the second variance term $L_{var}$ is minimized, only ground-truth variant patterns can be included in $G_V^\\prime$, i.e., $G_V^\\prime \\subseteq G_V^*$. Therefore, we have $G_V^\\prime = G_V^*$, so that $\\Phi^\\prime = \\Phi^*$. Finally, we prove that there exists a unique $\\Phi^*$ that can minimize the first contrastive loss term of Eq. (5).\n\nWe will add more detailed proofs in the revised paper.\n\n\n\n- **Q2.1. Hyperparameter analyses.**\n\nWe want to clarify that $\\alpha$ is the coefficient to control the balance between the contrastive loss and the invariance regularizer. A large $\\alpha$ encourages the invariance among different training environments and a small $\\alpha$ can learn informative representations but may not be sufficient to encourage the invariance. $q$ is a hyperparameter in GCE loss, which controls the degree of fitting spurious correlations. A small $q$ pays more attention to correctly classified samples and suffers more from noisy samples, while a large $q$ means the model becomes less sensitive and prevents overfitting to the spurious correlations. Fig. 4 shows our method outperforms the best baselines within a wide range of hyper-parameters choices.\n\n\n\n- **Q2.2. Ablation studies.**\n\nThe first two ablated versions \"w/ GCNII\" and \"w/ GIN\" denote replacing the backbone with GCNII and GIN. The next two ablated versions \"w/o Var.\" and \"w/o IPW\" denote removing the variant subgraph contrastive module and further removing inverse propensity weighting module. Fig. 3 shows that the performance nearly keeps unchanged for the first two ablated versions and drops significantly for the next two ablated versions, indicating that (1) our method is compatible with the other popular GNNs, and (2) it is important to explicitly identify variant subgraphs for estimating its degree of spurious correlations and remove it by reweighting. \n\n\n\n- **Q3. Time complexity of each module.**\n\nDenote the number of nodes and edges as  $|V|$ and $|E|$ of the input graph, the representation dimensionality as $d$, and the batch size is $B$. Our method mainly consists of three modules:\n\n1. For the invariant and variant subgraph identification module, the time complexity is $O(|E|d+|V|d^2)$ from the GCN component. \n2. For the variant subgraph contrastive module, the time complexity is $O(B^2d)$, which is mainly from the pair-wise similarity calculation within a batch of graphs. \n3. For the inverse propensity weighting based invariant prediction module, the time complexity is $O(B)$.\n\nFinally, the overall time complexity of our method is $O(|E|d+|V|d^2)$.", "claims_and_evidence": "The claims made by the authors regarding the benefits of leveraging contrastive learning to estimate and mitigate spurious correlations for better generalization are convincing. The provided empirical evidence supports these claims effectively.", "methods_and_evaluation_criteria": "The proposed approach is reasonable. The inverse propensity weighting, coupled with self-supervised contrastive learning, effectively addresses the identified issues. The evaluation criteria using established datasets and metrics are appropriate and effectively highlight the strengths of the method for the problem.", "theoretical_claims": "The assumptions made for invariant subgraph learning are acceptable given the problem setting. However, the detailed theoretical discussions and proofs on why this method can solve this OOD problem are missing.", "experimental_designs_or_analyses": "The experiments provide strong empirical validations for the proposed method. Extensive experiments on several graph classification benchmark datasets demonstrate the superiority of the proposed method over baselines. However, the analyses on the important hyperparameters are weak.", "relation_to_broader_scientific_literature": "The paper is well-situated within the literature on graph learning and OOD generalization.", "essential_references_not_discussed": "The important references are fully discussed from my point of view.", "other_strengths_and_weaknesses": "### strengths\n(1) The research problem is interesting and important to the community. As the real-world application of GNNs continues to expand, improving robustness against real-world out-of-distribution scenarios becomes increasingly crucial I think.  \n(2) The proposed method is technically sound. The technical details are easy to understand. The use of variance contrastive estimation and inverse propensity weighting to mitigate spurious correlations is particularly novel.  \n(3) The comparative results against baselines validate the effectiveness of the approach. The presented experimental results indicate consistent improvements over several state-of-the-art baselines, validating the effectiveness and robustness of the proposed approach under varying degrees of bias.  \n\n### weaknesses:\n(1) While the paper is methodologically solid and clearly explained, the authors should provide more detailed theoretical analyses of the methods.  \n(2) The discussions on the experimental results in the experiment section can also be more detailed. Deeper insights into hyperparameter analyses and ablation studies would improve clarity.  \n(3) The discussions on the time complexity are limited. More comprehensive analysis on each module is missing.", "other_comments_or_suggestions": "I would suggest that the authors incorporate more rigorous theoretical analyses into the proposed method. And I also strongly encourage the authors to incorporate more details on the time complexity of each module.", "questions_for_authors": "Can you provide more theoretical analyses on the method to explain why it works?\nCan you incorporate the time complexity of each module? Do these variant subgraph contrastive estimation and inverse propensity weighting modules introduce unacceptable time complexity?\nCan you explain why the hyperparameter $\\alpha$ and $q$ can influence the performance?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Vv4SswYX63", "reviewer": "Reviewer_cgCh", "summary": "This study addresses a critical problem in GNNs regarding their limited generalization capabilities under distribution shifts. Current approaches mainly use correlations in graph patterns rather than discovering fundamental causal substructures for predictions. To overcome this limitation, the paper jointly considers the identification of both invariant and variant subgraphs. Specifically, the method estimates the impact of spurious correlations induced by variant subgraphs and leverages this estimation to enhance the learning of invariant subgraphs. The proposed model demonstrates substantial performance gains over representative baselines, and comprehensive ablation studies confirm the effectiveness of each designed module.\n\n## update after rebuttal\n\nI will keep my positive opinion towards the paper after rebuttal.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedback. We addressed all the comments. Please kindly find the detailed responses to the comments below.\n\n- **Q1. Figure 1 is not clear to show the method’s training procedure.**\n\nThank you for this comment. We would like to follow your suggestion to improve Figure 1 by incorporating additional details to better illustrate our method. Specifically, we have made the following three improvements: (1) We have included the key equations used in the method directly into the figure for the readers to connect the pipeline in the figure with the corresponding details in the text. (2) We have refined the pipeline by adding more detailed step-by-step flows indicated by arrows, making the process clearer and easier to follow. (3) We have highlighted more technical details with concrete examples and included a small legend that explains the meaning of specific symbols, colors, or arrow styles. We will update Figure 1 following your suggestion in the revised paper.\n\n\n\n- **Q2. Differences among the listed related works (line 392-400).**\n\nThank you for this question. We have revised the discussions on the related works (line 392-400) as follows: \"Several famous works are proposed to tackle this problem on graphs by learning subgraph backed by different theories or assumptions, including causality [1-2], invariant learning [3-6], disentanglement [7], information bottleneck [8]. Different from these works that output explainable or invariant subgraphs under distribution shifts, some works directly learn generalizable graph representations for the problems where distribution shifts exist on graph size [9-10] or the other structural patterns [11-12]. And the learned representations are expected to remain invariant across different environments. We will also add more discussions on related works in the revised paper.\n\n\n\n- **Q3. There are some typos.**\n\nThank you for this comment. We have carefully proofread the paper and revised the following typos: \n\n1. Line 412: we have revised the section name \"5.1. Disentangled Graph Neural Network\" into \"Disentangled Graph Neural Network\".\n2. Line 322, we have revised the expression \"well handling distribution shifts\" into \"well handle distribution shifts\".\n\nWe will update the revised paper.\n\n\n\n- **Q4. What are the computational trade-offs of using inverse propensity weighting?**\n\nThank you for this question. We would like to clarify that it will not have additional computation cost to use inverse propensity weighting. Specifically, in the variant subgraph contrastive estimation module, the time complexity of estimating the degree of the spurious correlations is $O(B^2d)$, which is mainly from the pair-wise similarity calculation within a batch of graphs, where $d$ is the representation dimensionality and $B$ is the batch size. After that, we calculate the inverse propensity weights, whose time complexity is  $O(B)$. So the overall time complexity of inverse propensity weighting is $O(B^2d)$, which is significantly lower than the complexity of message-passing GNNs used in our invariant and variant subgraph identification module as well as the other baselines, whose time complexity is $O(|E|d+|V|d^2)$ where $|E|$ and $|V|$ denote the number of edges and nodes. We will add these detailed analyses in the revised paper to clarify the efficiency of our method.\n\n\n\n**References:**\n\n[1] Discovering Invariant Rationales for Graph Neural Networks\n\n[2] Causal attention for interpretable and generalizable graph classification\n\n[3] Handling Distribution Shifts on Graphs: An Invariance Perspective.\n\n[4] Learning Invariant Graph Representations Under Distribution Shifts\n\n[5] Empowering Graph Invariance Learning with Deep Spurious Infomax\n\n[6] Does invariant graph learning via environment augmentation learn invariance?\n\n[7] Debiasing Graph Neural Networks via Learning Disentangled Causal Substructure\n\n[8] Interpretable and generalizable graph learning via stochastic attention mechanism\n\n[9] Size-invariant graph representations for graph classification extrapolations\n\n[10] Sizeshiftreg: a regularization method for improving size-generalization in graph neural networks\n\n[11] Graph out-of-distribution generalization with controllable data augmentation\n\n[12] Graphmetro: Mitigating complex distribution shifts in gnns via mixture of aligned experts", "claims_and_evidence": "The paper's claims about improved generalization are clear and convincing. The paper claims that the method improves out-of-distribution generalization in graph classification tasks by explicitly modeling spurious correlations through contrastive learning and mitigating their impact via inverse propensity weighting. The empirical results support the claims by outperforming existing graph OOD generalization baselines across five datasets, including CMNIST, CFashion, CKuzushiji, MOLSIDER, and MOLHIV.", "methods_and_evaluation_criteria": "The proposed method and evaluations make sense for the graph OOD problem. The use of contrastive learning for estimating variant subgraph effects is novel and well-motivated. Traditional methods only focus on learning invariant subgraph. But the proposed method focuses on learning both the variant subgraph and invariant subgraph. The key idea is inspiring. As said above, the evaluations included existing graph OOD generalization baselines and five common datasets (CMNIST, CFashion, CKuzushiji, MOLSIDER, and MOLHIV).", "theoretical_claims": "The application of inverse propensity weighting is grounded in causal inference literature and is employed in a principled manner to correct for spurious correlations. The theoretical foundations of the method align with prior works on invariant learning and causal representation learning.", "experimental_designs_or_analyses": "I have checked the soundness/validity of the experimental designs or analyses. I think the experimental designs are good but the analyses are a little limited.", "relation_to_broader_scientific_literature": "The paper belongs to the literature on OOD generalization.", "essential_references_not_discussed": "There are no additional references that need to be discussed.", "other_strengths_and_weaknesses": "The other strengths are summarized as follows:\n\n-\tThe work made strong methodological contributions with an interesting idea.\n-\tThe method effectively disentangles invariant and variant subgraphs, which is a novel approach to handling distribution shifts.\n-\tThe empirical results show consistent improvements across diverse datasets.\n\nThe other weaknesses are summarized as follows:\n\n-\tFigure 1 is not clear to show the method’s training procedure.\n-\tThe discussions on the related works are arbitrary. For some contents, the authors just simply listed the references without necessary discussions.\n-\tTypos: line 412, it should remove ‘5.1’; line 322, “well handling distribution shifts” should be “well handle distribution shifts”.", "other_comments_or_suggestions": "The authors should address the weaknesses above. The model framework in Figure 1 can be clearer. And the discussions on the related works are arbitrary. For the related work of OOD Generalization part, they just listed the relevant works (even in several lines), which should be revised. The differences among these works should be introduced.", "questions_for_authors": "What are the computational trade-offs of using inverse propensity weighting?\n\nWhat are the differences among the listed OOD generalization methods from line 392 to line 400.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4aoe1HXdBO", "reviewer": "Reviewer_S2qE", "summary": "The submission explores the challenge of out-of-distribution generalization in GNNs. The authors propose a novel model that enhances out-of-distribution generalization by explicitly identifying invariant subgraphs and leveraging contrastive learning on variant subgraphs. Their approach consists of three key components: (1) distinguishing invariant and variant subgraphs, (2) applying contrastive learning to variant subgraphs to estimate the degree of spurious correlations, and (3) predicting invariant subgraphs with inverse propensity weighting to mitigate these spurious correlations. The model is evaluated on multiple benchmark datasets under varying degrees of distribution shifts, demonstrating its superiority over existing methods.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "- **Q1. Clarification on the reliance on accurate variant subgraph identification.**\n\nWe would like to clarify that our method can provably learn accurate variant subgraphs with theoretical guarantee. \n\n**Theorem 1.** Denote the optimal invariant subgraph generator $\\Phi^*$ that disentangles the ground-truth invariant subgraph $G_I^*$ and variant subgraph $G_V^*$ given the input graph $G$, where $G_I^*$ satisfies Assumption 2.1 and denote the complement as $G_V^* = G\\backslash G_I^*$. Assume the second variance term of Eq. (5) is minimized, the first contrastive loss term is minimized *iff* the invariant subgraph generator $\\Phi$ equals $\\Phi^*$.\n\n*Proofs.*\n\nDenote the first contrastive loss term of Eq. (5) as $L_{ssl}$ and the second variance term as $L_{var}$.\n\n$\\Leftarrow$: To prove $\\Phi^* = \\arg\\min_{\\Phi} L_{ssl}$, assume there exists another $\\Phi^\\prime \\neq \\Phi^*$ such that $L_{ssl}(\\Phi^\\prime) < L_{ssl}(\\Phi^*)$. It implies that $G_V^\\prime$ includes the ground-truth invariant subgraph information,\n$$\nG_V^\\prime \\cap G_I^* = G_V^\\prime \\cap G\\backslash G_V^* \\neq \\emptyset.\n$$\nTherefore, the contrastive loss $\\ell_{ssl}\\left(\\Phi, h_V; k\\right)$ among the environments partitioned by the graph label $k = 1,\\dots,|Y|$ is dependent on the label itself, i.e.,\n$$\n\\exists\\, k_1, k_2 = 1, \\dots, |Y|, k_1 \\neq k_2, \\text{such that } \n\\ell_{\\text{ssl}}(\\Phi^\\prime, h_V; k_1) \\neq \\ell_{\\text{ssl}}(\\Phi^\\prime, h_V; k_2).\n$$\nThus\n$$\nL_{var}(\\Phi^\\prime) = \\mathrm{Var}\\left(\\ell_{ssl}\\left(\\Phi\\prime, h_V; k\\right)\\right) > 0.\n$$\nHowever, $G_V^*$ excludes all the ground-truth invariant information that is sufficiently predictive to the graph label. We have\n$$\nL_{var}(\\Phi^*) = 0.\n$$\nThus, \n$$\nL_{var}(\\Phi^\\prime) > L_{var}(\\Phi^*).\n$$\nThis contradicts the assumption that the variance term $L_{var}$ is minimized. Finally we prove that the optimal invariant subgraph generator $\\Phi^*$ can minimize the contrastive loss $L_{ssl}$.\n\n$\\Rightarrow$: To prove minimizing the contrastive loss $L_{ssl}$ can derive  $\\Phi$ equals $\\Phi^*$, assume there exists another invariant subgraph generator $\\Phi^\\prime$ derived by minimizing $L_{ssl}$, i.e., $\\Phi^\\prime \\neq \\Phi^*$, where $G = (G_I^\\prime, G_V^\\prime)$ and $\\Phi^\\prime = \\arg\\min_{\\Phi} L_{ssl}$.  Because $L_{ssl}$ is minimized, $G_V^\\prime$ preserves all the intrinsic features of the ground-truth variant subgraph, i.e., $G_V^* \\subseteq G_V^\\prime$. Since the second variance term $L_{var}$ is minimized, only ground-truth variant patterns can be included in $G_V^\\prime$, i.e., $G_V^\\prime \\subseteq G_V^*$. Therefore, we have $G_V^\\prime = G_V^*$, so that $\\Phi^\\prime = \\Phi^*$. Finally, we prove that there exists a unique $\\Phi^*$ that can minimize the first contrastive loss term of Eq. (5).\n\nWe will add more detailed proofs in the revised paper.\n\n\n\n- **Q2. Additional discussions on ablation studies.**\n\nWe have added additional discussions on ablation studies. The first two ablated versions \"w/ GCNII\" and \"w/ GIN\" denote replacing the backbone with GCNII and GIN. The next two ablated versions \"w/o Var.\" and \"w/o IPW\" denote removing the variant subgraph contrastive module and further removing inverse propensity weighting module. Fig. 3 shows that the performance nearly keeps unchanged for the first two ablated versions and drops significantly for the next two ablated versions, indicating that (1) our method is compatible with the other popular GNNs, and (2) it is important to explicitly identify variant subgraphs for estimating its degree of spurious correlations and remove it by reweighting. More details will be added in the revised paper.\n\n\n\n- **Q3. Why the baselines perform worse.**\n\nWe have revised the discussions into sec. 4.2 as follows. Some baselines (e.g., GCN and GIN) did not consider the specific design for generalization under distribution shifts. And some baselines (e.g., FactorGCN and DIR) also did not show good performance since their assumptions might be invalid under severe bias. Also, some debiased methods (e.g., LDD and DisC) did not explicitly capture the spurious correlations for each input graph. Therefore, the baselines perform worse than our method. We will add the discussions above to the revised paper.\n\n\n\n- **Q4. Explanations of training objective.**\n\nWe would like to clarify that our method is one joint framework to optimize Eq. (11), including (1) the self-supervised contrastive objective in Eq. (5) that can ensure the accurate identification of invariant and variant subgraphs, (2) the objective in Eq. (7) for accurately estimating the degree of the spurious correlations, and (3) the objective in Eq. (8) to learn the predictions on the invariant subgraph after reweighting.\n\n\n\n- **Q5. Would the approach be extended to node or link predictions?**\n\nIn this paper, we mainly focus on the graph-level prediction task. But our method can be easily extended to the node or link prediction tasks which we leave for future work.", "claims_and_evidence": "The authors claim that this is the first work to explicitly utilize variant subgraphs to help capture invariant subgraphs under distribution shifts and that the three mutually promoted modules significantly enhance performance over state-of-the-art baselines. This claim is supported by their experimental results, which consistently show the superiority of their method across various datasets with varying bias levels.", "methods_and_evaluation_criteria": "The method aligns well with existing literature on invariant learning and causality based OOD generalization. The three proposed modules are motivated, particularly the introduction of variance contrastive estimation for variant subgraph learning, which differentiates this work from previous approaches that only focus on learning invariant representations.", "theoretical_claims": "The problem formulation, which defines the OOD generalization objective and the role of invariant subgraphs, is clearly stated and follows a well-defined causality framework. The use of inverse propensity weighting is a reasonable and theoretically sound approach to mitigating spurious correlations.", "experimental_designs_or_analyses": "The experimental design is well-structured and includes strong baseline comparisons. The authors compare against graph learning methods, including standard GNN architectures (GCN, GIN), pooling-based methods (DiffPool), and recent invariant learning approaches (DIR, LDD, DisC).", "relation_to_broader_scientific_literature": "The work builds upon prior research (invariant learning) effectively.", "essential_references_not_discussed": "The authors should cite or discuss more recent graph OOD papers.", "other_strengths_and_weaknesses": "This work addresses an important problem in graph machine learning and presents a novel solution with strong theoretical and empirical backing. The proposed approach is innovative in its explicit estimation of spurious correlations and its use of contrastive learning to refine invariant subgraph identification. The extensive experimental validation strengthens its contribution. The method is relatively efficient, as shown by the complexity analysis, which indicates that VIVACE maintains comparable computational cost to existing baselines. The scalability of the method makes it suitable for real-world applications beyond the datasets tested.\n\nOne potential limitation is the reliance on the accuracy of the variant subgraph identification step. If the model fails to accurately disentangle variant subgraphs, the effectiveness of the entire approach could be questioned. While the authors’ ablation studies suggest that their method is robust, additional discussions on how this challenge would be useful. Another limitation is that they could further elaborated on the underperformance of the baselines in the experiments. A more detailed discussion on this issue would provide valuable insights into the limitations of existing methods.", "other_comments_or_suggestions": "In addition to addressing these limitations, the authors should provide more explicit explanations of the training objective in the main text, particularly for the contrastive learning module.", "questions_for_authors": "Could you clarify the mechanism to ensure the accuracy of the variant subgraph identification step? Would the approach be extended to node or link predictions?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kwtFSZWjQL", "reviewer": "Reviewer_GXRP", "summary": "This paper presents VIVACE for learning invariant subgraphs under distribution shifts using variance contrastive estimation. The authors propose a three-module framework to disentangle invariant and variant subgraphs, estimate the impact of spurious correlations, and employ inverse propensity weighting for predictions. The framework's effectiveness is validated by the experiments. The experimental results across multiple benchmarks demonstrate the method's superiority against existing approaches. The results indicate that VIVACE achieves better robustness to distribution shifts and effectively captures invariant subgraphs while mitigating the influence of spurious correlations.\n\n## update after rebuttal Thanks for the response. I will keep my score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "- **Q1. Theoretical analysis on the rationale of the approach.**\n\nWe would like to clarify that the rationale of our method is to achieve OOD generalization by **accurately disentangling invariant and variant subgraphs**. We have added the following theorem.\n\n**Theorem 1.** Denote the optimal invariant subgraph generator $\\Phi^*$ that disentangles the ground-truth invariant subgraph $G_I^*$ and variant subgraph $G_V^*$ given the input graph $G$, where $G_I^*$ satisfies Assumption 2.1 and denote the complement as $G_V^* = G\\backslash G_I^*$. Assume the second variance term of Eq. (5) is minimized, the first contrastive loss term is minimized *iff* the invariant subgraph generator $\\Phi$ equals $\\Phi^*$.\n\n*Proofs.*\n\nDenote the first contrastive loss term of Eq. (5) as $L_{ssl}$ and the second variance term as $L_{var}$.\n\n$\\Leftarrow$: To prove $\\Phi^* = \\arg\\min_{\\Phi} L_{ssl}$, assume there exists another $\\Phi^\\prime \\neq \\Phi^*$ such that $L_{ssl}(\\Phi^\\prime) < L_{ssl}(\\Phi^*)$. It implies that $G_V^\\prime$ includes the ground-truth invariant subgraph information,\n$$\nG_V^\\prime \\cap G_I^* = G_V^\\prime \\cap G\\backslash G_V^* \\neq \\emptyset.\n$$\nTherefore, the contrastive loss $\\ell_{ssl}\\left(\\Phi, h_V; k\\right)$ among the environments partitioned by the graph label $k = 1,\\dots,|Y|$ is dependent on the label itself, i.e.,\n$$\n\\exists\\, k_1, k_2 = 1, \\dots, |Y|, k_1 \\neq k_2, \\text{such that } \n\\ell_{\\text{ssl}}(\\Phi^\\prime, h_V; k_1) \\neq \\ell_{\\text{ssl}}(\\Phi^\\prime, h_V; k_2).\n$$\nThus\n$$\nL_{var}(\\Phi^\\prime) = \\mathrm{Var}\\left(\\ell_{ssl}\\left(\\Phi^\\prime, h_V; k\\right)\\right) > 0.\n$$\nHowever, $G_V^*$ excludes all the ground-truth invariant information that is sufficiently predictive to the graph label. We have\n$$\nL_{var}(\\Phi^*) = 0.\n$$\nThus, \n$$\nL_{var}(\\Phi^\\prime) > L_{var}(\\Phi^*).\n$$\nThis contradicts the assumption that the variance term $L_{var}$ is minimized. Finally we prove that the optimal invariant subgraph generator $\\Phi^*$ can minimize the contrastive loss $L_{ssl}$.\n\n$\\Rightarrow$: To prove minimizing the contrastive loss $L_{ssl}$ can derive  $\\Phi$ equals $\\Phi^*$, assume there exists another invariant subgraph generator $\\Phi^\\prime$ derived by minimizing $L_{ssl}$, i.e., $\\Phi^\\prime \\neq \\Phi^*$, where $G = (G_I^\\prime, G_V^\\prime)$ and $\\Phi^\\prime = \\arg\\min_{\\Phi} L_{ssl}$.  Because $L_{ssl}$ is minimized, $G_V^\\prime$ preserves all the intrinsic features of the ground-truth variant subgraph, i.e., $G_V^* \\subseteq G_V^\\prime$. Since the second variance term $L_{var}$ is minimized, only ground-truth variant patterns can be included in $G_V^\\prime$, i.e., $G_V^\\prime \\subseteq G_V^*$. Therefore, we have $G_V^\\prime = G_V^*$, so that $\\Phi^\\prime = \\Phi^*$. Finally, we prove that there exists a unique $\\Phi^*$ that can minimize the first contrastive loss term of Eq. (5).\n\n\n\n- **Q2. Detailed analysis on hyper-parameter sensitivity.**\n\nWe clarify that $\\alpha$ is the coefficient to control the balance between the contrastive loss and the invariance regularizer. A large $\\alpha$ encourages the invariance among different training environments and a small $\\alpha$ can learn informative representations but may not be sufficient to encourage the invariance. $q$ is a hyperparameter in GCE loss, which controls the degree of fitting spurious correlations. A small $q$ pays more attention to correctly classified samples and suffers more from noisy samples, while a large $q$ means the model becomes less sensitive and prevents overfitting to the spurious correlations. Fig. 4 shows our method outperforms the best baselines within a wide range of hyper-parameters choices.\n\n\n\n- **Q3. Computation cost of the modules.**\n\nDenote the number of nodes and edges as  $|V|$ and $|E|$ of the input graph, the representation dimensionality as $d$, and the batch size is $B$. Our method mainly consists of three modules:\n\n1. For the invariant and variant subgraph identification module, the time complexity is $O(|E|d+|V|d^2)$ from the GCN component. \n2. For the variant subgraph contrastive module, the time complexity is $O(B^2d)$, which is mainly from the pair-wise similarity calculation within a batch of graphs. \n3. For the inverse propensity weighting based invariant prediction module, the time complexity is $O(B)$.\n\nFinally, the overall time complexity of our method is mainly induced by the message-passing GNN. The variant subgraph contrastive estimation and inverse propensity weighting modules will not have extra higher time complexity.\n\n\n\n- **Q4. Reason to use GCE loss.**\n\nWe adopt GCE loss to fit the spurious correlations as it has been shown that GCE loss can emphasize spurious correlations.\n\n\n\n- **Q5. Impact of $q$.**\n\nAs shown in Figure 4, as a hyperparameter in GCE loss to control the degree of fitting the spurious correlations, $q$ has a moderate impact on the model\nperformance, but our method is not very sensitive to it by outperforming the best baselines within a wide range of hyper-parameters choices.", "claims_and_evidence": "The claims are supported with extensive experimental compare. The authors compare VIVACE against baseline GNNs, showing consistent improvements in OOD generalization. The ablation studies provide further evidence for the effectiveness of the proposed modules. To be specific, removing the variant subgraph contrastive module or the inverse propensity weighting module leads to a substantial drop in performance, indicating the important role of these components. The hyperparameter sensitivity analysis further demonstrates that VIVACE is robust across a wide range of settings.", "methods_and_evaluation_criteria": "The framework combining self-supervised contrastive learning and causal inference techniques is sound.  The evaluation criteria are appropriate. The authors report classification accuracy for synthetic datasets (e.g., CMNIST, CFashion, CKuzushiji) and ROC-AUC for real-world datasets (e.g., MOLSIDER, MOLHIV), ensuring fair comparisons with prior works. The inclusion of multiple runs and standard deviation reporting strengthens the statistical significance of the results.", "theoretical_claims": "The theoretical foundations of the method are ground by the causality. The authors formalize the problem using causal invariance assumptions, ensuring that the learned subgraphs have stable predictive power across environments. But a deeper theoretical analysis on the rationale of the approach should be introduced.", "experimental_designs_or_analyses": "The experiments are comprehensive, covering both controlled and real-world scenarios. The authors ensure significant performance gains by including multiple runs. The method shows consistent improvements in those datasets.", "relation_to_broader_scientific_literature": "It builds upon prior work on graph neural networks and causality. The discussion of related work is comprehensive and contextualizes the contributions well.", "essential_references_not_discussed": "No critical references appear to be missing.", "other_strengths_and_weaknesses": "Pros:  \n- The paper is novel overall, which introduces variance contrastive estimation as a self-supervised learning technique to explicitly model and quantify spurious correlations in graph data.    \n- Unlike prior work that either assumes predefined environments or relies on heuristic-based disentanglement, VIVACE directly estimates the variant subgraphs, making the process adaptive to various real-world datasets.    \n-  This methodology advances existing invariant learning approaches by integrating contrastive learning with causal inference, effectively modeling and mitigating the impact of spurious correlations.  \n\nCons:  \n- Detailed analysis for some experiment results are weak (e.g., hyper-parameter sensitivity).\n- Since VIVACE introduces additional modules (e.g., variant subgraph contrastive estimation, inverse propensity weighting), it is unclear how much extra computation is required.\n- Some designs in the method lacked the detailed explanations (e.g., GCE loss in equ. 7).", "other_comments_or_suggestions": "- Clarifying the computational cost compared to baselines would be helpful.   \n- Equ. 7 introduces the Generalized Cross-Entropy (GCE) loss, but the reason for using GCE instead of standard cross-entropy is not fully discussed.", "questions_for_authors": "- Could you discuss more on the computational cost compared to baselines? \n-  Why was the Generalized Cross-Entropy (GCE) loss chosen rather than Cross-Entropy (CE) loss? \n-  What impact does the hyperparameter q in GCE have on model performance?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper proposes a novel method to learn disentangled invariant subgraph via self-supervised contrastive variant subgraph estimation for achieving satisfying OOD generalization. All reviewers recommended the acceptance. Therefore, I followed the reviewers' recommendations and gave the acceptance.
Decision: Accept (poster)