Disentangling Invariant Subgraph via Variance Contrastive Estimation under Distribution Shifts

• Q1. Clarification on the reliance on accurate variant subgraph identification.

We would like to clarify that our method can provably learn accurate variant subgraphs with theoretical guarantee.

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^*$.

Proofs.

Denote the first contrastive loss term of Eq. (5) as $L_{ssl}$ and the second variance term as $L_{var}$.

$\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,

$$G_V^\prime \cap G_I^* = G_V^\prime \cap G\backslash G_V^* \neq \emptyset.$$

Therefore, 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.,

$$\exists\, k_1, k_2 = 1, \dots, |Y|, k_1 \neq k_2, \text{such that } \ell_{\text{ssl}}(\Phi^\prime, h_V; k_1) \neq \ell_{\text{ssl}}(\Phi^\prime, h_V; k_2).$$

Thus

$$L_{var}(\Phi^\prime) = \mathrm{Var}\left(\ell_{ssl}\left(\Phi\prime, h_V; k\right)\right) > 0.$$

However, $G_V^*$ excludes all the ground-truth invariant information that is sufficiently predictive to the graph label. We have

$$L_{var}(\Phi^*) = 0.$$

Thus,

$$L_{var}(\Phi^\prime) > L_{var}(\Phi^*).$$

This 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}$.

$\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).

We will add more detailed proofs in the revised paper.

• Q2. Additional discussions on ablation studies.

We 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.

• Q3. Why the baselines perform worse.

We 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.

• Q4. Explanations of training objective.

We 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.

• Q5. Would the approach be extended to node or link predictions?

In 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.

We thank the reviewer for the valuable feedback. We addressed all the comments. Please kindly find the detailed responses to the comments below.

• Q1. Figure 1 is not clear to show the method’s training procedure.

Thank 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.

• Q2. Differences among the listed related works (line 392-400).

Thank 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.

• Q3. There are some typos.

Thank you for this comment. We have carefully proofread the paper and revised the following typos:

1. Line 412: we have revised the section name "5.1. Disentangled Graph Neural Network" into "Disentangled Graph Neural Network".
2. Line 322, we have revised the expression "well handling distribution shifts" into "well handle distribution shifts".

We will update the revised paper.

• Q4. What are the computational trade-offs of using inverse propensity weighting?

Thank 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.

References:

[1] Discovering Invariant Rationales for Graph Neural Networks

[2] Causal attention for interpretable and generalizable graph classification

[3] Handling Distribution Shifts on Graphs: An Invariance Perspective.

[4] Learning Invariant Graph Representations Under Distribution Shifts

[5] Empowering Graph Invariance Learning with Deep Spurious Infomax

[6] Does invariant graph learning via environment augmentation learn invariance?

[7] Debiasing Graph Neural Networks via Learning Disentangled Causal Substructure

[8] Interpretable and generalizable graph learning via stochastic attention mechanism

[9] Size-invariant graph representations for graph classification extrapolations

[10] Sizeshiftreg: a regularization method for improving size-generalization in graph neural networks

[11] Graph out-of-distribution generalization with controllable data augmentation

[12] Graphmetro: Mitigating complex distribution shifts in gnns via mixture of aligned experts

• Q1. Theoretical analysis on the rationale of the approach.

We 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.

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^*$.

Proofs.

Denote the first contrastive loss term of Eq. (5) as $L_{ssl}$ and the second variance term as $L_{var}$.

$\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,

$$G_V^\prime \cap G_I^* = G_V^\prime \cap G\backslash G_V^* \neq \emptyset.$$

Therefore, 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.,

$$\exists\, k_1, k_2 = 1, \dots, |Y|, k_1 \neq k_2, \text{such that } \ell_{\text{ssl}}(\Phi^\prime, h_V; k_1) \neq \ell_{\text{ssl}}(\Phi^\prime, h_V; k_2).$$

Thus

$$L_{var}(\Phi^\prime) = \mathrm{Var}\left(\ell_{ssl}\left(\Phi^\prime, h_V; k\right)\right) > 0.$$

However, $G_V^*$ excludes all the ground-truth invariant information that is sufficiently predictive to the graph label. We have

$$L_{var}(\Phi^*) = 0.$$

Thus,

$$L_{var}(\Phi^\prime) > L_{var}(\Phi^*).$$

This 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}$.

$\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).

• Q2. Detailed analysis on hyper-parameter sensitivity.

We 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.

• Q3. Computation cost of the modules.

Denote 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:

1. For the invariant and variant subgraph identification module, the time complexity is $O(|E|d+|V|d^2)$ from the GCN component.
2. 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.
3. For the inverse propensity weighting based invariant prediction module, the time complexity is $O(B)$.

Finally, 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.

• Q4. Reason to use GCE loss.

We adopt GCE loss to fit the spurious correlations as it has been shown that GCE loss can emphasize spurious correlations.

• Q5. Impact of $q$.

As 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 performance, but our method is not very sensitive to it by outperforming the best baselines within a wide range of hyper-parameters choices.

• Q1. Theoretical analyses.

We have added Theorem 1 to show our method can accurately identify the invariant and variant subgraphs for OOD generalization.

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^*$.

Proofs.

Denote the first contrastive loss term of Eq. (5) as $L_{ssl}$ and the second variance term as $L_{var}$.

$\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,

$$G_V^\prime \cap G_I^* = G_V^\prime \cap G\backslash G_V^* \neq \emptyset.$$

Therefore, 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.,

$$\exists\, k_1, k_2 = 1, \dots, |Y|, k_1 \neq k_2, \text{such that } \ell_{\text{ssl}}(\Phi^\prime, h_V; k_1) \neq \ell_{\text{ssl}}(\Phi^\prime, h_V; k_2).$$

Thus

$$L_{var}(\Phi^\prime) = \mathrm{Var}\left(\ell_{ssl}\left(\Phi^\prime, h_V; k\right)\right) > 0.$$

However, $G_V^*$ excludes all the ground-truth invariant information that is sufficiently predictive to the graph label. We have

$$L_{var}(\Phi^*) = 0.$$

Thus,

$$L_{var}(\Phi^\prime) > L_{var}(\Phi^*).$$

This 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}$.

$\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).

We will add more detailed proofs in the revised paper.

• Q2.1. Hyperparameter analyses.

We 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.

• Q2.2. 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.

• Q3. Time complexity of each module.

Denote 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:

1. For the invariant and variant subgraph identification module, the time complexity is $O(|E|d+|V|d^2)$ from the GCN component.
2. 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.
3. For the inverse propensity weighting based invariant prediction module, the time complexity is $O(B)$.

Finally, the overall time complexity of our method is $O(|E|d+|V|d^2)$.