Does Invariant Graph Learning via Environment Augmentation Learn Invariance?
Authors: Yongqiang Chen
Yatao Bian
Kaiwen Zhou
Binghui Xie
Bo Han
James Cheng
Venue: NeurIPS 2023 poster (2023)
Keywords: Graph Neural Networks
Out-of-Distribution Generalization
Invariant Learning
TL;DR: We found impossibility results of learning invariant graph representations without environment partitions, established the minimal assumptions for the feasibility of the problem and propose a novel solution.

Abstract: Invariant graph representation learning aims to learn the invariance among data from different environments for out-of-distribution generalization on graphs. As the graph environment partitions are usually expensive to obtain, augmenting the environment information has become the de facto approach. However, the usefulness of the augmented environment information has never been verified. In this work, we find that it is fundamentally impossible to learn invariant graph representations via environment augmentation without additional assumptions. Therefore, we develop a set of minimal assumptions, including variation sufficiency and variation consistency, for feasible invariant graph learning. We then propose a new framework Graph invAriant Learning Assistant (GALA). GALA incorporates an assistant model that needs to be sensitive to graph environment changes or distribution shifts. The correctness of the proxy predictions by the assistant model hence can differentiate the variations in spurious subgraphs. We show that extracting the maximally invariant subgraph to the proxy predictions provably identifies the underlying invariant subgraph for successful OOD generalization under the established minimal assumptions. Extensive experiments on datasets including DrugOOD with various graph distribution shifts confirm the effectiveness of GALA.

Does Invariant Graph Learning via Environment Augmentation Learn Invariance?
Yongqiang Chen't Yatao Bian?, Kaiwen Zhou' 'The Chinese University of Hong Kong Tenceet AI Lab
(ygchen,kwzzouu www.hoooo@sseccc.eeeeeeeeeee
patao.bian@gmai.com
Binghui Xie', Bo Han', James Cheng
HHong Kong Baptist University
Charm1@@comphhhbueddu.. {bhxie21, jcheng (ccenng)ocee.cuukeeeee
Abstract
Invariant graph representation learning aims to learn the invariance among data from different environments for out-of-distributio generalization graphs. As the graph environment partitions are usually expensive to obtain, augmenting the envi- ronment information has become the de facto approach. However, the usefulness of the augmented environment information has never been verified. In this work, we find that it is fundamentally impossible to learn invariant graph representations via environment augmentation without additional assumptions. Therefore, we develop set of minimal assumptions, including variation sufficiency and variation sistency, for feasible invariant graph learning. We then propose a new framework Graph invAriant Learning Assistant (GALA). GALA incorporates an assistant model that needs to be sensitive to graph environment changes or distribution shifts. The correctness of the proxy predictions by the assistant model hence can differentiate the variations in spurious subgraphs. We show that extracting the maximally invari- ant subgraph to the proxy predictions provably identifies the underlying invariant subgraph for successful OOD generalization under the established minimal as- sumptions. Extensive experiments on datasets including DrugOOD with various graph distribution shifts confirm the effectiveness of GALA!.
Introduction
Graph representation learning with graph neural networks (GNNs) has proven to be highly successful in tasks involving relational information [24, 35, 68, 74, 75]. However, it assumes that the training and test graphs are independently drawn from the identical distribution (id.), which can hardly hold for many graph applications such as in Social Network, and Drug Discovery [25, 26, 37, 70, 86, 94]. The performance of GNNs could be seriously degenerated by graph distribution shifts, i.e., mismatches between the training and test graph distributions caused by some underlying environmental factors during the graph data collection process [16, 21, 28, 70, 87, 94]. To overcome the Out-of -fti--fbbtituttion (OOD) generalization failure, recently there has been a growing surge of interest in incorporating the invariance principle from causality [60] into GNNs [8, 17, 41, 45, 55, 72, 73, 77, 83). The rationale of the invariant graph learning approaches is identify the underlying invariant subgruph of the input graph, which shares an invariant correlation with the target labels across multiple graph
"ork done during internship at Tencent AI Lab. Code is available at https: httpp:///gtaia //github/b
com SathaaseGAALA.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Training Data
Proxy Predictions
"House"
"Cycle"
"House"
Graph Labels
(GP)
[G")
House
"Cyele"
Cycle
Figure I: An illustration of GALA with the task of classifying graphs according to whether there exists "House" or "Cycle" motif. Given the training data where the "House" subgraph often co-occurs with a "Grid" and the "Cycle" subgraph often co-occurs with a "Hexagon". An ERM trained environment assistant model will fit the spurious subgraph and therefore yield proxy predictions "House" or "Cycle" for any graphs containing "Grid" (left half) or "Hexagon" (right half), respectively. GALA first separates the samples according to the correctness of the proxy predictions into the sets of positive graphs {GP) (correct, in blue) and negative graphs {G") (incorrect, in green). Then, GALA extracts the maximally invariant subgraph among {GP) and {G") i.e., pulling graphs with the same graph label but from { and {G") closer in the latent space, hence identifies the invariant subgraph.
distributions from different environments [8. 72]. Thus, the predictions made merely based on the invariant subgraphs can be generalized to OOD graphs that come from a new environment [60].
As the environment labels or partitions on graphs are often expensive to obtain [8], augmenting the environment information, such as generating new environments [45, 72, 73] and inferring the environment labels [41, 77]. has become the de facto approach for invariant graph learning. However, little attention has been paid to verifying the fidelity (or faithfulness') of the augmented environment information. For example, if the generated environments or inferred environment labels induce higher bias or noise, would make the learning of graph invariance even harder. Although it looks appealing learn both the environment information and the graph invariance, the existing approaches could easily run into the "no free lunch" dilemma [71]. In fact, Lin et al. [43] found that there exist negative cases in the Euclidean regime where it is impossible to identify the invariant features without environment partitions. When it comes to the graph regime where the OOD generalization is fundamentally more difficult [8] than the Euclidean regime, it raises challenging research question:
When and how could one learn graph invariance without the environment labels?
In this work, we present a theoretical investigation of the problem and seek set of minimal assumptions the underlying environments for feasible invariant graph learning. Based family of simple graph examples (Def. 3.1), we show that existing environment generation approaches can fail to generate faithful environments, when the underlying environments are not sufficient to uncover all the variations of the spurious subgraphs (Prop. 3.2). On the contrary, incorporating even lead worse OOD performance. The failure of faithful
the generated environments environment generation implies the necessity of variation sufficiency (Assumption 3.3). Moreover, with sufficient environments, inferring faithful environment labels remains impossible. Since invariant and spurious subgraphs can have an arbitrary degree of correlation with labels, there exist multiple sets of training environments that have the same joint distribution of P(G, Y) but different invariant subgraphs. Any invariant graph learning algorithms will inevitably fail to identify the invariant subgraph in least one set of training environments (Prop. 3.4). Therefore, we need to additionally ensure the variation consistency (Assumption. 3.5), that is, the invariant and spurious subgraphs should have consistent relationship in the correlation strengths with the labels.
To resolve the OOD generalization challenge under the established assumptions, we propose a new framework Graph Ariant Learning Assistant (GALA). GALA incorporates an additional assistant model that needs to be prone to distribution shifts, to generate proxy predictions of the training samples. Different from previous environment inferring approaches [41, 77]. GALA does require explicit environment labels but merely proxy predictions to differentiate the variations in the spurious
'The fidelity or faithfulness refers to whether the augmented environment information can actually improve the OOD generalization graphs.

---

subgraphs. As shown in Fig. 1, we first fit an environment assistant model to the training distribution and then divide the training graphs into a positive set { and negative {G'), according whether the proxy predictions are correct or not, respectively. As spurious correlations tend to vary more easily than invariant correlations, the variations in spurious subgraphs are further differentiated and increased between {G*) and (G*), ). Then, only the invariant subgraph holds an invariant correlation with the label among {GP) and }. and hence can be identified by extracting the subgraphs that maximize the intra-class subgraph mutual information among {GP) j and {G') (Theorem 4.1).
We conduct extensive experiments to validate the effectiveness of GALA using 12 datasets with various graph distribution shifts. Notably, GALA brings improvements up to 30% in multiple graph datasets.
Our contributions can be summarized as follows:
We identify failure cases of existing invariant graph learning approaches and establish the minimal assumptions for feasible invariant graph learning;
We develop a novel framework GALA with provable identifiability of the invariant subgraph for OOD generalization on graphs under the assumptions;
We conduct extensive experiments to verify both our theoretical results and the superiority of GALA;
Notably, both our theory and solution differ from Lin et al. [43] fundamentally, as we do not rely on the auxiliary information and are compatible with the existing interpretable and generalizable GNN architecture for OOD generalization on graphs. Meanwhile, we provide a new theoretical framework that resolves the counterexample in Lin et al. [43] while enjoying provable identifiability.
2 Background and Preliminaries
We begin by introducing the key concepts and backgrounds of invariant graph learning, and leave more details in Appendix C. The notations used in the paper are given in Appendix A.
OOD generalization graphs. This work focuses on graph classification, while the re- sults generalize node classification as well using the same setting as in Wu et al. [72]. Specifically, we are given a set of graph datasets D= (D heehe collected from multiple environ- ments Eall Samples (G Y?") D' from the environment e are drawn independently from an identical distribution PP The goal of OOD generalization on graphs is to find a GNN f that minimizes the maximal loss among all environments, i.e., to minimize max RS, where R is the risk of f under environment e. We consider the same graph generation process proposed by Chen et al. [8] which is inspired by real-world drug discovery task [57] and covers a broad case of graph distribution shifts. As shown in Fig. 2, the generation of the observed graphs G and labels Y are controlled by a latent causal variable C and a spurious variable S. and S control Y and G by controlling the generation of the underlying invariant subgraph Ge and spurious subgraph G.. respectively. Since S can be affected by the environment E, the correlation between and G. can change arbitrarily when the environment changes. Besides, the interaction among C, S and Y at the latent space can be further categorized into Full Informative Invariant Features (FIIF) when Y L SIC. and Partially Informative Invariant Features (PIIF) when Y I. SIC.
E
(a) G-Gen. SCM (b) FIIF SCM
(c) PII SCM
Figure 2: SCMs graph distribution shifts [8].
To tackle the OOD generalization challenge on graphs from Fig. 2, the existing invariant graph learning approaches are generically designed to identify the underlying invariant subgraph G. to predict the label Y [8, 72]. Specifically, the goal of OOD generalization on graphs is to learn an invariant GNN f = fo g. which is composed of: invariant subgraph Ga: b) classifier fa Ge - that predicts the label Y based on the extracted G., where Ge refers to the space of subgraphs of G. The learning objectives of fa and g are formulated as max 9 I(G Y), s.t. G. JL E, G. g(G).
featurizer g
G. that estimates the
(1) Since E is not observed, many strategies are proposed to impose the independence of G. and E. prevalent approach is to augment the environment information. Based on the estimated invariant
3

---

subgraphs Ge and spurious subgraphs G.. Liu et al. [45], Wu et al. [72, 73] propose to generate new environments, while Li al. [41], Yang al. [77] propose infer the underlying environment labels. However, we show that they all fail to augment faithful environment information in Sec. 3.
Besides, Miao al. [55, 56], Yu et al. [81, 82, 83] adopt graph information bottleneck tackle FIIF graph shifts, but they cannot generalize to PIIF shifts, while Our work focuses on PIIF shifts as it more challenging when without environment labels [43]. Fan et al. [17] generalize [40] tackle severe graph biases, i.e., when H (S)Y) < H (C|Y). Chen et al. [8] propose a contrastive framework to tackle both FIIF and PIFF graph shifts, but is limited to (S) > HC))). In practice, as is usually unknown which correlation is stronger, we need a unified solution to tackle both cases.
Invariant learning without environment labels. In the Euclidean regime, there plentiful studies in invariant learning without environment labels. Creager et al. [12] propose a minmax formulation to infer the environment labels. Liu et al. [46] propose self- off-boosting framework based the estimated invariant and variant features. Deng et al. [14], Liu et al. [44]. Pezeshki et al. [61], Zhang et al. [85] propose infer labels based on the failures of an ERM model. However, Lin et al. [43] find failure cases of the aforementioned approaches that it is impossible to identify the invariant features without given environment labels in Euclidean data, and propose solution that leverages auxiliary environment information for invariant learning. As the OOD generalization on graphs poses more challenges [8]. whether it is feasible to learn invariant graph representations without any auxiliary environment information remains elusive.
3 Pitfalls of Environment Augmentation
Given only the mixed training data without environment partitions, is it possible to learn generate faithful environments infer the underlying environment labels that facilitate OOD generalization on graphs? In the discussion below, we adopt the two-piece graphs to instantiate the problem, which
is the simplistic version of the PIIF distribution shifts in Fig. 2(c), motivated by Kamath et al. [30].
Definition 3.1 (Two-piece graphs). Each environment is defined with two parameters, Oe. Be E [0, 1), and the datasel (G ye) E DD generated as follows:
(a) Sample Y E {-1,1] uniformly;
(b) Generate Ge and Gg via G. == - Rad(axe)), Ge G Sgen(Ye Rad(33)), where fga fG, map the input 11 to corresponding graph selected from given set, and Rad(a) is a random variable taking value with probability a and +1 with a
(c) Synthesize G by randomly assembling Ge and G.: == (Ges Gz).
We denote an environment e with (x) for simplicity. Different environments will have a different 3e thus P(Y|Ga) will change across different environments, while P(Y|G.) remains invariant.
3.1 Pitfalls of environment generation
We begin by discussing the cases where there are few environments, and generating new environ- ments is necessary [45, 72, 73]. Environment generation aims to provide some additional "virtual" environments E, such that the invariant subgraph can be identified via applying OOD risk to the joint dataset with the augmented data D = (Dele E E U
The generation of "virtual" environments is primarily based on the intermediate estimation of the invariant and spurious subgraphs, denoted as G. and G., respectively. Liu et al. [45], Wu et al. [73] propose DIR and GREA construct new graphs by assembling G. and G, from different graphs Specifically, given n samples {G', y y J==1. the new graph samples in E, is generated as follows:
G G ( Gg), Vi,j E (1....), =
which generates new environment e, with n² samples. Although both DIR and GREA gain some empirical success, the faithfulness of & remains questionable, as the generation is merely based on
We slightly abuse the superscript and subscript when denoting the ith sample to avoid confusion of double superscripts subscripts.

---

inaccurate estimations of the invariant and spurious subgraphs. Specifically, when G. contains parts of Gx assigning the same labels to the generated graph is more likely strengthen the spurious correlation between G, and Y. For example, when the model yields reversed estimation, i.e.,
G. G2 and G, Ge, the generated environment will destroy the invariant correlations.
Proposition 3.2. Consider the two-piece graph dataset & {(a, 81). (a, 32)} with a 2 31, 82 (e.g,, & = {00.25, 0.1), (0.25, 0.2) and its corresponding mixed environment Emix = {(ar (31 + 32) (e.g., Ei = {00.25, 0.15))). When G. = G, and G2 Ge it holds that the augmented environment e is also a two-piece graph dataset with
E, {00.5, (31 + 22) (e.g., E. (00.5, 0.15)))
The proof is given in Appendix E. This also extends to the adversarial augmentation [72, 83], which will destroy the actual G... As both DIR and GREA adopt the same environment generation procedure, we verify the failures of environment generation with GREA in Table of Sec. 5, where GREA can perform comparably with ERM. In fact, when the underlying environments are insufficient to differentiate the variations of the spurious features, is fundamentally impossible to identify the underlying invariant graph from the spurious subgraph. More formally, if 3G,. such that pet (Y|G.) = poa (Y for any C1, E2 E Eg. where Pe (Y|G,) is the conditional distribution P(Y under environment E Eall, it is impossible for any graph learning algorithm to identify Ge. We provide a formal discussion in Appendix E.2. The failure implies a fundamental requirement that E should uncover all the potential variations in the spurious subgraph.
Assumption 3.3. (Variation sufficiency) For graphs generated followingg Fig. 2, for any Go. Bel, e2 En, such that ped (Y. + eed G)), and ped (YIG.) = p
Assumption 3.3 aligns with the definition of invariance [8, 30] that the invariant subgraph Ge is expected to satisfy pe (Y Gz) = pr ( (Y|G,) ) for e1, C2 E Eall If there exists Gg satisfying the invariance condition as well, then it is impossible to tell G, from G, even with environment labels.
3.2 Pitfalls of environment inferring
Although environment sufficiency (Assumption 3.3) relieves the need for generating new environ- ments, is it possible infer the underlying environment labels via approaches such as Mole00D [77] and GIL [41], to facilitate invariant graph learning? Unfortunately, we find a negative answer.
Considering the two-piece graph examples & 2 0.1), (0.2, 0. when given the underlying environment labels, is easy to identify the invariant subgraphs from spurious subgraphs. However, when the environment labels are not available, we have the mixed data as En {00.2, 0.2)), where P(Y G.) = G.)) The identifiability of G2 is ill- ill-posed, as it does not affect the &r even if we swap G. and G.. More formally, considering the environment mixed from two two-piece graph environments {(a, } and { (a.32)) then we have & = {(a, (31 + /2) For each Eur, we can also find a corresponding E.¹ = [((8) 3 a2)))(()) with {3) a²)} and ((3g, a')). Then, let
(8) + /12 a' = (31 + B2)
We now obtain Enr and E which share the same joint distribution P(Y, G) while the underlying Ge completely different. More generally, we have the following proposition.
Proposition 3.4. There exist 2 two-piece graph training environments Enr and Eny 'that share the same joint distribution P(Y,C) G). Any learning algorithm will fail in either e or Ex¹
The proof is given in Appendix E.3. The experiments in Sec. validate that both Mole00D and GIL fail to infer faithful environment labels and even underperform ERM. It implies that whenever it allows the existence of an identical training distribution by mixing the environments, invariant graph learning is impossible. Therefore, we need an additional assumption that excludes the unidentifiable case. We propose to constrain the relationship between a (i.e., HY()G.) and &n (i.e., (Y))).))
Assumption 3.5. (Variation consistency) For all environments in En, (CYY) + H(S|Y).
Intuitively, Assumption 3.5 imposes the consistency requirement on the correlation strengths between invariant and spurious subgraphs with labels. For two-piece graphs with consistent variations, mixing up the environments will yield new environment with the same variation strength relationships.
5

---

Thus, Assumption 3.5 gets rid of the previous unidentifiable cases. Moreover, Assumption 3.5 also aligns with many realistic cases. For example, the relation of specific functional group (e.g.. -OH) with a molecule can hardly be reversed to that held upon the scaffold of the molecule, due to the data collection process. Therefore, Assumption 3.5 also resolves the counterexample proposed by Lin et al. [43]. Different from our work, Lin et al. [43] propose to incorporate additional auxiliary information that satisfies certain requirements to mitigate the unidentifialle case. However, such auxiliary information is often unavailable and expensive to obtain graphs. More importantly, the requirements are also unverifiable without more assumptions, which motivates us to consider the relaxed case implied by Assumption 3.5.
3.3 Challenges of environment augmentation
To summarize, the two assumptions constitute the minimal assumptions for feasible invariant graph learning. Failing to satisfy either of them while lacking additional inductive biases will result in the "no free lunch" dilemma [71] and suffer from the unidentifiabillty issue.
Corollary 3.6. (No Free Graph OOD Lunch) Without Assumption 3.3 or Assumption 3.5, there does not exist learning algorithm that captures the invariance of the two-piece graph environments.
Corollary 3.6 is a natural conclusion from the previous discussion. The proof is straightfor- ward and given in Appendix E.4. Assump- tion 3.3 and Assumption 3.5 establish the min- imal premises for identifying the underlying in-
Table Remaining challenges of invariant graph learning: existing works handle both cases.
HSSY H(CYY H(SYY) H(CYY)
DisC
CIGA
X
variant subgraphs. However, it also raises new challenges, as shown in Table. 1. Chen et al. [8]
GALA (Ours)
propose CIGA to maximize the intra-class mutual information of the estimated invariant subgraphs to tackle the case when H (ClY) < (S|Y). While for the case when (SSY) < H(CYY). Fan et al. [17] propose DisC that adopts GCE loss [40] to extract the spurious subgraph with a larger learning step size such that the left subgraph is invariant. However, both of them can fail when there is no prior knowledge about relations between (C)Y) and (S)Y). We verify the failures DisC and CIGA in Table. 2. The failure thus raises a challenging question:
Given the established minimal assumptions, is there framework that tackles both cases when H(CIY) < H(S)Y) and H(CYY) H(S(Y)?
4 Learning Invariant Graph Representations with Environment Assistant
We give affirmative answer by proposing new framework, GALA: Graph invariant Learning Assistant, which adopts an assistant model to provide proxy information about the environments.
4.1 Learning with An Environment Assistant
Intuitively, straightforward approach tackle the aforementioned challenge is to extend the framework of either DisC [17] or CIGA [8] to resolve the other case. As DisC always destroys the first learned features and tends to be more difficult extend (which is empirically verified in Sec. 5), we are motivated to extend the framework of CIGA to resolve the case when H(SIY) < H(CVV).
Understanding the success and failure of CIGA. The principle of CIGA lies in maximizing the intra-class mutual information of the estimated invariant subgraphs, i.e.,
max ((GG))) s.l. G E arg feig
1(G GgI))
G
where G: g(*)) and G ~ P(G]Y), i.e., G sampled from training graphs that share the same label Y as G. The key reason for the success of Eq. 3 is that, given the data generation process as in Fig. 2 and the same C, the underlying invariant subgraph G. maximizes the mutual information of subgraphs from any two environments, i.e., Vel, e2 E Eall.
G arg max G (C),
(4)

---

where G and G are the estimated invariant subgraphs corresponding to the same latent causal variable C' C=aa under the environments e,, C2, respectively. Since C' is not observable, CIGA adopts Y as proxy for C, as when H (S)Y) > H(C Ge maximizes 1 [( G (Y) and thus / (Ge: G2lY), However, when (S)Y) < H(C) the proxy no longer holds. Given the absence of E, simply maximizing intra-class mutual information favors the spurious subgraph Gg instead, i.e.,
G E arg m I(Gr; Y))
(5)
Invalidating spuriousness dominance. To mitigate the issue, we motivated to find proxy that samples Ge for Eq. 5, while preserving only the Ge the solution under both cases.
To begin with, we consider the case of H(S)) < H(C)Y), Although the correlation between G. and Y dominates the intra-class mutual information, Assumption 3.3 implies that there exists a subset of training data where P(Y|G,) varies, while P(Y]G.) remains invariant. Therefore, the dominance of spurious correlations no longer holds for samples from the subset. Incorporating samples from the subset into Eq. 3 as G invalidates the dominance of Gg. Denote the subset {G") ), then
Ge E arg max ( (Y),
(6)
where To E ( is sampled from the subset ( dominated by spurious correlations, while T E ( is sampled from the subset (Gr) where spurious correlation no long dominates, or is dominated by invariant correlations. We prove the effectiveness Eq. Theorem 4.1.
Environment assistant model A. To find the desired subsets {G and (G²). inspired by the success in tackling spuriousness-domnnate OOD generalization via learning from biased predictors [40, 44, 58, 853, we propose to incorporate an assistant model A that is prone to spurious correlations. Simply training A with ERM using the pprriossess-dommnatee data enables A to learn spurious correlations, and hence identifies the subsets where the spurious correlations hold or shift, according to whether the predictions of A are correct or not, respectively. Let A = arg max A (AA()); Y), we have (ह (o(G)))((G)) = Yi). (Õ.) = (o(CT)) * Y.).
Reducing to invariance dominance case. After showing that Eq. 6 resolves the spuriousness dominance case, still need show that Eq. 6 preserves Ge as the only solution when (S)Y) > HCC)Y. Considering training A with ERM using the anvariaccedoomnnae
data, A will learn both invariant correlations and spurious correlations [10, 17]. Therefore, (G") switches to the subset that dominated by spurious correlations, while (Čg) switches to the subset dominated by invariant correlations. Then, Eq. 6 establishes a lower bound for the intra-class mutual information, i.e.,
I [ G YY) S GHI))
(8)
where Gg (Gg), G2 (Gr) J. and Ges G are the same as Eq. 3. The inequality in Eq. 8 holds any sub- graph maximizes the left hand side can also be incorporated in right hand side, while the sampling space of G. and G in the right hand side (i.e., both G. and G: are sampled from the whole train set) is larger than that of the left hand side. The equality is achieved by taking the ground truth Ge as the solution for the featurizer We verify the correctness of Eq. 6 and Eq. 8 in Fig. 3(a).
Algorithm 1 GALA: Graph Ariant Learning Assistant
1: Input: Training data Drr: environment assistant A; fea- turizer GNN g; classifier GNN fe; length of maximum training epochs e; batch size
2: Initialize environment assistant A;
3: for p E .e. do Sample batch of data (G,, Yi from Duri Obtain Environment Assistant predictions ( for each sample G Vs (G,, Y.) do
6:
7: Find positive graphs with same y/ and different f;: Find negative graphs with different yı but same environment assistant prediction of:
8:
9: Calculate GALA risk via Eq. 9;
10: Update fe, via gradients from GALA risk;
11: end for
4.2 Practical implementations.
12: end for
13: return final model fa 9:
The detailed algorithm description of GALA is shown as Algorithm In practice, the environment assistant can have multiple imple- mentation choices so long as it is prone to distribution shifts. As discussed in Sec. 4.1, ERM trained
7

---

model can serve as reliable environment assistant, since ERM tends to learn the dominant features no matter whether the features are invariant or spurious. For example, when ((S))) HC ERM will first learn to use spurious subgraphs Gs to make predictions. Therefore, we can obtain {GP) by finding samples where ERM correctly predicts the labels, and {G") for samples where ERM predicts incorrect labels. In addition to label predictions, the clustering predictions of the hidden representations yielded by environment assistant models can also be used for sampling {G²) and } [85]. Besides, we can also incorporate models that are easier to overfit to the first dominant features to better differentiate {GP] } from {G"). When the number of positive or negative samples is imbalanced, we can upsample the minor group to avoid trivial solutions. In addition, the final GALA objective is given in Eq. 9 and implemented as in Eq. 18. We provide more discussions about the implementation options in Appendix F.
4.3 Theoretical analysis
In the following theorem, we show that the GALA objective derived in Sec. 4.1 can identify the underlying invariant subgraph and yields an invariant GNN defined in Sec. 2.
Theorem 4.1. Given i) the same data generation process as in Fig, 2; ii) Dir that satisfies variation sufficiency (Assumption 3.3) and variation consistency (Assumption 3.5); iii) {GP) and { are distinct subsets of Dur such that I [GP;; GY = O, YG = arg m I(G;; Y) under {GP), and VG. = arg max Gz T[;;; Y) under {Gm); suppose [G.l = Se, YGe, resolving the following GALA objective elicits an invariant GNN defined via Eq. 1,
1(Go; s.t. g E arg max I ( GNI)) f..g
3.C..ls..
where Go E (θ = gGG))} and T. (G = g(G")) are the estimated invariant subgraphh via from (G*} and {Gm), respectively.
The proof is given in Appendix E.5, Essentially, assumption iii) in Theorem 1 is an implication of the variation sufficiency (Assumption 3.3). When given the distinct subsets {GP} and {G") } with different relations of (ClY) and (SJY), since HCC))) remains invariant across different subsets, the variation happens mostly to the spurious correlations between S and Y. By differentiating spurious correlations into distinct subsets, maximizing the intra-class mutual information helps identify the true invariance. The fundamental rationale for why GALA resolves two seemingly conversed cases essentially relies on the commutative law of mutual information.
5 Experiments
We evaluated GALA with both synthetic and realistic graph distribution shifts. Specifically, we are interested in the following two questions: (a) Can GALA improve over the state-of-the-ar invariant graph learning methods when the spurious subgraph has a stronger correlation with the labels? (b) Will GALA affect the performance when the invariant correlations are stronger?
5.1 Datasets and experiment setup
We prepare both synthetic and realistic graph datasets containing various distribution shifts to evaluate GALA. We will briefly introduce each dataset and leave more details in Appendix G.
Two-piece graph datasets. We adopt BA-2motif; [50] to implement 4 variants of 3-class two-piece graph (Def. 3.1) datasets. The datasets contain different relationships of H (CIY) and (S)Y) by controlling the a and 8 in the mixed environment, respectively. We consider 4 cases of a a- 3, ranging from {+0.2, +0.1, 0.11, 0.22.. to verify our discussion in Sec. 4.3.
Realistic datasets. We also adopt datasets containing various realistic graph distribution shifts to comprehensively evaluate the OOD performance of GALA. We adopt 6 datasets from DrugOOD benchmark [28], which focuses on the challenging real-world task of Al-aided drug affinity prediction. The DrugOOD datasets include splits using Assay, Scaffold, and Size from the EC50 category (denoted as FCC50**) and the Ki category (denoted as Ki-*). We also adopt graphs converted from the ColoredMNIST dataset [3] using the algorithm from Knyazev et al. [36], which contains distribution
8

---

shifts in node attributes (denoted as CMN In addition, we adopt Graph-SST2 [84], where we split graphs with a larger average degree in the training set while smaller in the test set.
Experiment setup. We adopt the state-of-the-art OOD methods from the Fuclidean regime, including IRMv1 [3], VREx [38], EIII [13]] and IB-IRM [2], and from the graph regime, including GREA [45],
GSAT [55], CAL [64], MoleOOD [77]. GIL [41], DisC [17] and CIGA [8]. We exclude DIR [73] and GIB [81] as GREA and GSAT are their sophisticated variants. In addition to the ERM baseline that trained vanilla GNN with ERM objective, in two-piece motif datasets, we also include XGNN to demonstrate the failures of previous approaches, which is an interpretable GNN trained with ERM. We also exclude CIGAv2 [8] as GALA focuses on improving the contrastive sampling via environment assistant for the objective in C All methods use the same GIN backbone [75], the same interpretable GNN architecture as in [55], and optimization protocol for fair comparisons. We tune the hyperparmeters following the common practice. Details are given in Appendix G.2.
5.2 Experimental results and analysis
Proof-of-concepp study. The re- sults in two-piece graph datasets are reported in Table 2. It can be found that the previous environ- ment augmentation approaches fail either in datasets where the invariant correlations dominate or where the spurious correla- tions dominate, aligned with our discussions in Sec. 3. In par- ticular. GREA, CIGA and GIL achieve high performance when the invariant correlation domi- nates, but suffer great perfor- mance decrease when the spuri- ous correlations are stronger. Al- though DisC is expected to suc- ceed when spurious correlations dominate, DisC fails to outper- form others because of its exces- sive destruction of the learned in- formation. Mole00D also yields degraded performance, which could be caused by the failures to infer reliable environment labels. In contrast, GALA achieves
Table 2: OOD generalization performance under various invariant and spurious correlation degrees in the two-piece graph datasets. Each dataset is generated from variation of two-piece graph model, denoted as {a. b), where a refers to the invariant correla- tion strength and b refers to the spurious correlation strength. The blue shadowed entries are the results with the mean-1 sstd std
larger than the mean of the corresponding second best results.
Datasets
{0.8.
{0.8,
{0.8, 0.9)
(0.7, 0.9)
Avg.
ERM
17.3310.47
15.65
11.37
42.7313.82
61.77
IRM
78.32
1..13
50.7612.56 56
41.32+2.50
61.38
V-Rex
77.6920.38
74.9621.40
49.4711.58
41.65
60.94
IB-IRM
78.00=0.6s
70.
50.93+1.87 87
42.05+0.79
61.23
EIL
76.98.1..4
14.25
51.45
39.7112.64
60.60
XGNN
83.84:10.59
23.05 20
13.37:11.22
38.28
64.63
GREA
82.86 82.86=0.50
82.72+0.50
50.34+1.74
39.01
63.72
GSAT
80.5410.88
78.1111.20
48.6312.18
36.62.0.87
63.32
CAL
76.98
62.95
51.57+6.33
46.2343.98
59.43
MoleOOD
49,93
49.85
38.4914.25
34.81
43.27
GIL
83.51
82.67=1.18
51.76+1.32
40,07+2.61
64.50
DisC
60.47117.9
54.29
45.0617.82
39.4218.89
50.81
CIGA
84.03:20.53
83.21
57.87+3.38
43.62+t3.20
67.18
GALA
84.27=0034
83.65=00.44
+3.52
72.50+1.06
79.21
Oracle
84.7310038
85.42.0.25
84,2840.15
78.3810.19
consistently high performance under both cases and improves CIGA up to 30% under {0.7, 0.9} and 13% in average, which validates our theoretical results in Sec. 4.3.
Table 3: OOD generalization performance under realistic graph distribution shifts. The blue shadowed entries are the results with the mean- 1 st larger than the mean of the respective second best results.
Datasets
EC50-Assay
BC50-Sca
EC50-Size
Ki-Assay
Ki-Sca
Ki-Size
CMNIST-sp
Graph-SST2
Avg. (Rank)"
ERM
76.42.150 sm
64.56.1.55 95
61.61 se
74.61_a.9
69. 19.38211.:
76.63211 31
21 21.56a5.33
81 1..54g:.m 13
65.79 (6.50)
77.14
64.32.10.2
62.33 a2.33
75.10_3..
19.32.1
76.251o73
20.2511112
82.52.0.79
65.91 (6.13)
V-Rex
75.57=2ar
64.73=00.1
62 a0
74.16=1.0
71 71.40=2.
76.68
30.71+114
$1.11a:77
67.15 (5.25)
IB-IRM
64.70_250
58.2810 62.62.2.s
71.98_3.
69. 69.5511.m
70.71
23.5817.6
81.561002
62.87 (10.6)
EHL
64.20=5-m
62.88:2.75
59.5820ms
74.24=2.4
69.63-1144
76.56
23.55
82.4651.45
64.14 (8.00)
XGNN
72.99
63.62 as
62 5..55
72.40=3..0
21.0
73. 15
20 96
82. 82.55ac.65
65 65.03 (7.13)
GREA
66.87.753
14.
59.20 59.2011
73.17_1.00
07.82.14487
73.52,
12.77
82.40.1:.0
62.36
GSAT
76.07=195
63.58=1.35
61.12200 do
72.26=1.70
70. 1626 89
75. 25.78=2.00
15 15.24
80.57
64.35 (8.63)
CAL
75.10.971
64.7911.5
63.38
71.081483 75.22__1.71
72.9311.71
23.68114.8
82.38.11.0
66.07 (5.38)
DisC
61.9427776
10
57.64
$4.12
50.8319.30 55.35
30.26110.00
76.51
$6.59 (12.4)
MoleOOD
61.49
62. 12:11m
58. (..7411
15.10_0.73
3
73.69
21.0413.3
81. 16
61.76(11.00
GIL
70.56
61.59 13.16
60.4611.1
75.25 1.34
14.33
75.76.1223
12.55
83.31
CIGA
75.03=2.47
65.41=1.16
64.10=1.0 0s
73.95=2.0
71.874:.32
74.46=2:12
15.83++256
82 82.93
63.45 (5.88)
GALA 77.56
66.28
64.25
77.92_
73.17
22.4
68.9410.56
13.60.0.s
73.64 (1.00)
Orccle
84.77.0.s
82.6611111
84
91.08_1.41
88.58 1o.35
92.501..53
67.76.6..
91. 91.40....s
Averaged rank also reported in the parentheses because of datasct heterogencity. A lower rank is better.

---

OOD generalization in realistic graphs. The results in realistic datasets are reported in Table 3. Aligned with our previous discussion, existing environment augmentation approaches sometimes yield better performance than ERM, such as CAL in EC50-Size, Mole00D in Ki-Assay, GIL in Graph-SST2, or CIGA in EC50-Size, however, inevitably fail bring consistent improvements than ERM, due to the existence of failure cases. DisC is suspected to work only for graph distribution shifts on node features and bring impressive improvements in CMNIST-sp, but can destroy the learned information under more challenging settings. In contrast, GALA consistently outperform ERM by non-trivial margin in all datasets. Notably, GALA achieves near oracle performance in CMNIST-sp and improves CIGA by 53%. The consistent improvements of GALA confirm the effectiveness of GALA.
graphs
Negative graphs
CIGA
(0.80.0)
80
70 10mmo
0,00 performance
GALA-CIGAV2
spu
60
6-1obb lemally troglt
(0.7.0.9)
20-
50
75
100
4psamoningg
Cornelation
Ki-Sca
CHNBST-P
(a) Correlation strengths
(h) CIGAv2 compatibility
(c) Hyperparameter sensitivity
Figure 3: Ablation studies.
Correlation strengths of {GP} and { We conduct experiments with the two-piece graph datasets evaluated in Table 2 to verify the correctness of Eq. 6 and Eq. 8. Eq. 6 and Eq. 8 imply that the underlying invariant subgraph will be the subgraph that maximizes the mutual information among subgraphs from {G*} and {G*). matter whether the dominant correlation is spurious not. We measure the invariant and spurious correlation strengths in terms co-occur probability of the invariant and spurious subgraphs with the labels. The results shown in Fig. 3(a). It can be found that, under both cases, the underlying invariant subgraph maintains the predictivity with the label in an invariant manner. Hence. maximizing the intra-class subgraph mutual information between {GP) and {G"} } in GALA succeeds in identifying the underlying invariant subgraph.
CIGAv2 compatibility. Although GALA focuses the contrastive term in CIGA, both GALA and CIGA are compatible with the additional CIGAv2 term that facilitates constraining the graph sizes. To verify, compare the OOD performances of CIGA, CIGAv2, GALA, and GALA +CIGAv2 using two challenging datasets, Ki-Scaffold and CMNIST-sp. The results are given in Fig. 3(b). It can be found that, despite incorporating the additional CIGAv2 constraint, CIGA can not outperform GALA, while GALA can bring more improvements with the additional CIGAv2 constraint. In CMNIST-sp. since GALA already achieve the upper bound, incorporating CIGAv2 only achieve a similar result.
Hyperparameter sensitivity. We also test the hyperparameter sensitivity of GALA to the contrastive penalty weights well as the upsampling times that are introduced to mitigate the imbalance of positive and negative graphs. We conduct the experiments with two-piece graph dataset {0.7, 0.9). As shown in Fig. 3(c), it can be found that GALA is generically robust to different hyperparameter choices. In addition, when the penalty weight or the upsampling times turn to O, the performance will decrease a lot, which serves strong evidence for the effectiveness of GALA.
Computational analysis. We also conduct computational analysis of GALA and other methods, and defer the results to Table. 6 in Appendix G.4, due to space constraints. The results show that GALA costs only a competitive training time as environment generation based methods, while achieving much better OOD generalization performance.
6 Conclusions
We conducted retrospective study on the faithfulness of the augmented environment information for OOD generalization on graphs. By showing hardness cases and impossibility results the existing approaches, we developed a set of minimal assumptions for feasible invariant graph learning. Built upon the assumptions, we proposed GALA to learn the invariant graph representations guided by an environment assistant model. Extensive experiments with 12 datasets verified the superiority of GALA.
10

---

Acknowledgements
We thank the reviewers for their valuable comments. This work was supported by CUHK direct grant 4055146. BH was supported by the NSFC Young Scientists Fund No. 62006202, NSFC General Program No. 62376235, Guangdong Basic and Applied Basic Research Foundation No. 2022A151501 1652, HKBU Faculty Niche Research Areas No. RC-INRA-IG/22-22SC//0 and Tencent AI Lab Rhino-Biin Bird Gift Fund.
References
[1] I. Ahmad and P.-B. Lin. A nonparametric estimation of the entropy for absolutely continuous distributions (corresp.). IEEE Transactions on Information Theory, 22(3):372-375, 1976. (Cited on page 27)
[2] K. Ahuja, Caballero, Zhang, J.-C. Gagnon-Audet, Bengio. Mitliagkas, and Rish. Invariance principle meets information bottleneck for out-of-distribution generalization. In Advances Neural Information Processing Systems, 2021. (Cited on pages and 32)
[3] M. Arjovsky. L. Bottou, I. Gulrajani, and D. Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv: 1907v 028993 2019. (Cited on pages 8, 9, 30 and 32)
[4] M. I. Belghazi, A. Baratin, S. Rajeshwar. S. Ozair, Y. Bengio, A. Courville, and D. Hjelm. Mutual information neural estimation. In International Conference on Machine Learning. volume 80, pages 531-540, 10-15 Jul 2018. (Cited on page 27)
[5] B. Bevilacqua. Zhou, and B. Ribeiro. Size-invariant graph representations for graph classifi- cation extrapolations. In International Conference on Machine Learning, pages 837-851, 2021. (Cited on page 21)
[6] G. Chen, Y. Wang. F. Guo, Q. Guo, J. Shao, H. Shen, and X. Cheng. Causality and independence enhancement for biased node classification. In ACM International Conference on Information and Knowledge Management, pages 203-212 2023. (Cited on page 21)
[7] Y. Chen, Yang, Y. Zhang, K. Ma, Liu, Han, and J. Cheng. Understanding and improving graph injection attack by promoting unnoticeability. In International Conference on Learning Representations, 2022. (Cited on page 21)
[8] Y. Chen, Y. Zhang, Y. Bian, H. Yang, K. Ma, B. Xie, T. Liu, B. Han, and J. Cheng. Learning causally invariant representations for in Neural Information P Processing Systems, 2022. (Cited on pages 1, 2, 3, 4, 5, 6, 9, 19, 20, 21, 22, 25, 27, 31, 32 and 33)
generalization on graphs. In Advances
[9] Y. Chen, K. Zhou, Y. Bian, B. Xie, K. Ma, Y. Zhang, H. Yang, B. Han, and J. Cheng. Pareto invariant risk minimization. arXiv preprint, arXiv:2206.07766 2022. (Cited on page 30)
[10] Y. Chen, W. Huang, K. Zhou, Y. Bian, B. Han, and J. Cheng. Towards understanding feature learning in df-titiibbttion generalization. arXiv preprint xrXiv:2304.1132? 2023. (Cited on page 7)
[11] S. Chopra, R. Hadsell, and LeCun. Learning similarity metric discriminatively, with application to face verification. In 2005 IEEE Computer Society Conference Computer Vision and Pattern Recognition (CVPR 2005), 20-26 June 2005, San Diego, CA, USA, pages 539-546, 2005. (Cited on page 27)
[12] E. Creager, J. Jacobsen, and R. S. Zemel. Environment inference for invariant learning. In International Conference Machine Learning, pages 2189-2200, 2021. (Cited on pages 4, 21 and 32)
[13] E. Creager, J. Jacobsen, and R. S. Zemel. Environment inference for invariant learning. In International Conference on Machine Learning, volume 139, pages 2189-2200, 2021. (Cited on pages 9 and 32)

---

[14] Y. Deng, Y. Yang, B. Mirzasoleiman, and Q. Gu. Robust learning with progressive data expansion against spurious correlation. arXiv preprint, 2Xiv:2306.04949 2023. (Cited on pages 4 and 21)
[15] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. BERT:
Pre-training of deep bidirectional transformers for language understanding. In Conference of the North American Chapter the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pages 4171 4166, 2019. (Cited on page 30)
[16] M. Ding, K. Kong. J. Chen, J. Kirchenbauer, M. Goldblum, D. Wipf, F. Huang, and T. Goldstein. closer look distribution shifts and dittr-iiistiuutionn generalization on graphs. In NeurlPS 2021 Workshop on Distribution Shifts: Connecting Methods and Applications, 2021. (Cited on page 1)
[17] S. Fan, X. Wang, Y. Mo, C. Shi, and Tang. Debiasing graph neural networks via learning disentangled causal substructure. In Advances in Neural Information Processing Systems, 2022. (Cited on pages 4, 6, 7, 9, 21, 22 and 32)
[18] M. Fey and J. E. Lenssen. Fast graph representation learning with PyTorch Geometric. In ICLR Workshop on Representation Learning on Graphs and Manifolds, 2019. (Cited on page 33)
[19] J. Gao, Y. Zhou, J. Zhou, and B. Ribeiro. Double equivariance for inductive link prediction for both new nodes and new relation types. volume arXiv:2302.01313. 2023. (Cited on page 21)
[20] M. Gardner, J. Grus, M. Neumann, O. Tafjord, P. Dasigi, N. F. Liu, M. E. Peters, M. Schmitz, and I.. Zettlemoyer. Allennlp: A deep semantic natural language processing platform. arXiv preprint, arXiv:18033 wfXiv:100.007600, 2018. (Cited on page 30)
[21] S. Gui, X. Li, L. Wang, and S. Ji. GOOD: A graph out-of- benchmark. In Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2022. (Cited on page 1)
[22] S. Gui, M. Liu, X. Li, Y. Luo, and S. Ji. Joint learning of label and environment causal independence for graph -attof-diitribution generalization. arXiv preprint, arXivv2306.01 103, 2023. (Cited on page 21)
[23] Gulrajani and D. Lopez-Paz. In search of lost domain generalization. Conference on Learning Representations, 2021. (Cited on pages 30 and 32)
International
[24] W. L. Hamilton, Z. Ying, and J. Leskovec. Inductive representation learning on large graphs. In Advances in Neural Information Processing Systems, pages 1024-1034, 2017. (Cited on page 1)
[25] W. Hu, M. Fey, M. Zitnik, Y. Dong, H. Ren, B. Liu, M. Catasta, and J. Leskovee. Open graph benchmark: Datasets for machine learning on graphs. In Advances in Neural Information Processing Systems, 2020. (Cited on page
[26] K. Huang, T. Fu, W. Gao, Y. Zhao, Y. H. Roohani, J. Leskovec, C. W. Coley, C. Xiao, J. Sun, and M. Zitnik. Therapeutics data commons: Machine learning datasets and tasks for drug discovery and development. In Advances in Neural Information Processing Systems Datasets and Benchmarks Track (Round 1), 2021. (Cited on page 1)
[27] S. loffe and C. Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning, volume 37, pages 448-456. 2015. (Cited on page 31)
[28] Y. Ji, L. Zhang, J. Wu, B. Wu, L.-K. Huang, T. Xu, Y. Rong, L. Li, Ren, D. Xue, H. Lai, S. Xu, Feng, Liu, P. Luo, S. Zhou, J. Huang, P. Zhao, and Y. Bian. DrugOOD: Dutti--Distiibution (OOD) Dataset Curator and Benchmark for Al-aided Drug Discovery A Focus on Affinity Prediction Problems with Noise Annotations. arXiv preprint, arXiv:2201.096 2022. (Cited on pages 1, 8 and 31)
12

---

[29] W. Jin, T. Zhao, J. Ding, Y. Liu, J. Tang, and N. Shah. Empowering graph representation learning with test-time graph transformation. arXiv preprint, arXiv:2210.03561, 2022. (Cited on page 21)
[30] P. Kamath, A. Tangella, D. Sutherland, and N. Srebro. Does invariant risk minimization capture invariance? In International Conference on Artificial Intelligence and Statistics, pages 4069-4077, 2021. (Cited on pages 4 and 5)
[31] B. Kamhoua, L. Zhang, Y. Chen, H. Yang, M. KAILI, B. Han, B. Li, and Cheng. Exact shape correspondence via 2d graph convolution. In Advances in Neurul Information Processing Systems, 2022. (Cited on page 21)
[32] K. Kandasamy, A. Krishnamurthy, B. Poczos, L. Wasserman, and j. m. robins. Nonparametric von mises estimators for entropies, divergences and mutual informations. In Advances in Neural Information Processing Systems, volume 28, 2015. (Cited on page 27)
[33] P. Khosla, P. Teterwak, C. Wang, A. Sarna, Y. Tian, P. Isola, A. Maschinot, C. Liu, and D. Krishnan. Supervised contrastive learning. In Advances in Neural Information Processing Systems, volume 33, pages 18661-18673, 2020. (Cited on pages 27 and 29)
[34] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015. (Cited on pages 22 and 31)
[35] T. N. Kipf and M. Welling. Semi-supervised classification with graph convolutional networks. In International Conference on Learning Representations, 2017. (Cited on page 1)
[36] B. Knyazev, G. W. Taylor, and M. R. Amer. Understanding attention and generalization in graph neural networks. In Advances in Neural Information Processing Systems, pages 4204-4214, 2019. (Cited on pages 8 and 30)
[37] P. W. Koh, S. Sagawa, H. Marklund, S. M. Xie, M. Zhang, A. Balsubramani, W. Hu, M. Ya- sunaga, R. L. Phillips, I. Gao, T. Lee, E. David, Stavness, W. Guo, B. Earnshaw, Haque, S. M. Beery, J. Leskovee, A. Kundaje, E. Pierson, S. Levine, C. Finn, and P. Liang. WILDS: A benchmark of fi-the-wild distribution shifts. International Conference on Machine Learning. pages 5637-5664, 2021. (Cited page 1)
[38] D. Krueger, E. Caballero, Jacobsen, A. Zhang, Binas, D. Zhang, L. Priol, and A. C. Courville. Out-of-distributic generalization via risk extrapolation (rex). In International Conference on Machine Learning, pages 5826, 2021. (Cited on pages and 32)
[39] H. Lee, H. Park, and K. Yoon. Towards better generalization with flexible representation of multi-module graph neural networks. arXiv preprint, arXiv:2209.06589, 2022. (Cited on page 21)
[40] J. Lee, E. Kim, J. Lee, J. Lee, and J. Choo. Learning debiased representation via disentangled feature augmentation. Advances Neural Information Processing Systems, 2021. (Cited on pages 4, 6, 7 and 21)
[41] H. Li, 7. Zhang, X. Wang, and W. Zhu. Learning invariant graph representations for out- of-distribution generalization. In Advances in Neural Information Processing Systems, 2022. (Cited on pages 1, 2, 4, 5, 9, 21, 22, 32 an)
[42] X. Li, S. Gui, Luo, and S. Ji. Graph structure and feature extrapolation for out-of-distribution generalization. arXiv preprint, 07X1v:2306.08076, 2023. (Cited on page 21)
[43] Y. Lin, S. Zhu, L. Tan, and P. Cui. ZIN: When and how to learn invariance without environment partition? In Advances in Neural Information Processing Systems, 2022. (Cited on pages 2, 3, 4, 6 and 21)
[44] E. Z. Liu, B. Haghgoo, A. S. Chen, A. Raghunathan, P. W. Koh, S. Sagawa, P. Liang, and C. Finn. Just train twice: Improving group robustness without training group information. In International Conference on Machine Learning, pages 6781-6792, 2021. (Cited on pages 4, and 21)
13

---

[45] G. Liu, T. Zhao, J. Xu, T. Luo, and M. Jiang. Graph rationalization with environment- based augmentations. arXiv preprint arXiv:2206.028880 02886. 2022. (Cited on pages 4, 9, 21, 22 and 32)
[46] J. Liu, Z. Hu, P. Cui, B. Li, and Z. Shen. Heterogeneous risk minimization. In International Conference on Machine Learning, volume 139, pages 6804 6814, 2021. (Cited on pages and 21)
[47] S. Liu, T. Li, Y. Feng, N. Tran, H. Zhao, Q. Qiu, and P. Li. Structural re-weighting improves graph domain adaptation. In International Conference on Machine Learning, Proceedings of Machine Learning Research, pages 11778-217981 2023. (Cited on page 21)
[48] Y. Liu, X. Ao, F. Feng. Y. Ma, K. Li, T. Chua, and Q. He. FLOOD: A flexible invariant learning framework for out-of-distribttoon generalization on graphs. Proceedings of the 29th ACM SIGKDD Conference Knowledge Discovery and Data Mining, pages
2023.
(Cited on page 21)
[49] A. Lucie, M. A. Ter Hoeve, G. Tolomei, M. De Rijke, and F. Silvestri. Cf-gnnexplainer: Counterfactual explanations for graph neural networks. In International Conference on Artificial Intelligence and Statistics, pages 4499 4519, 2022. (Cited on page 21)
[50] D. Luo, W. Cheng, D. Xu, W. Yu, B. Zong, H. Chen, and X. Zhang. Parameterized explainer for graph neural network. In Advances in Neural Information Processing Systems, pages 19620-19631, 2020. (Cited on pages 8, 22 and 30)
[51] K. Ma, G. Yang, H. Yang, Y. Chen, and J. Cheng. Calibrating and improving graph contrastive learning. Transactions Machine Learning Research, 2023. ISSN 2835-8856. (Cited on page 19)
[52] S. Mahdavi, K. Swersky, Kipf, M. Hashemi, C. Thrampoulidis, and R. Liao. Towards better out-of-distribution generalization of neural algorithmic reasoning tasks. arXiv preprint arXiv:2211. 000692, 2022. (Cited on page 21)
[53] L. McInnes, J. Healy, N. Saul, and L. Grossberger. Umap: Uniform manifold approximation and projection. The Journal of Open Source Software, 3(29):861, 2018. (Cited on page 27)
[54] D. Mendez, A. Gaulton, A. Bento, J. Chambers, M. Veij. E. Felix, M. P. Magariños, F. M. Nowotka, M.
F. M. I. Hunter, L. Junco, G. Mugumbate, M. Rodríguez-López, F. Atkinson, N. Bosc, C. Radoux, A. Segura-Cabrest A. Hersey, and A. R. Leach. Chembl: towards direct deposition of bioassay data. Nucleic Acids 2019. (Cited on page 31)
Mosquera, P.
Research, 47 47Dattabaeelssuue:11933-00
[55] S. Miao, M. Liu, and P. Li. Interpretable and generalizable graph learning via stochastic attention mechanism. arXiv preprint arXiv:2201.12985 2022. (Cited on pages 4, 9, 21, 31 and 32)
[56] S. Miao, Y. Luo, M. Liu, and Li. Interpretable geometric deep learning via learnable randomness injection. In International Conference on Learning Representations, 2023, (Cited on pages and 21)
[57] C. Murray and D. Rees. The rise of fragment-based drug discovery. Nature chemistry, 1:187-92, 06 2009. (Cited on page 3)
[58] J. Nam, H. Cha, S. Ahn, J. Lee, and J. Shin. Learning from failure: Training debiased classifier from biased classifier. In Advances in Neural Information Processing Systems, 2020. (Cited on page
[59] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. Vetioo, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems, pages 8035, 2019. (Cited on page 33)
14

---

[60] J. Peters, P. Bühlmann, and N. Meinshausen. Causal inference by using invariant prediction: identification and confidence intervals. Journal of the Royal Statistical Society: Series (Statistical Methodology), 78(5):947-1012, 2016. (Cited on pages I and 2)
[61] M. Pezeshki, D. Bouchacourt, M. Ibrahim, N. Ballas, P. Vincent, and Lopez-Paz. Discovering environments with XRM. arXiv preprint, arXiv:2309.16744 2023. (Cited on pages 4 and 21)
[62] R. Salakhutdinov and G. E. Hinton. Learning a nonlinear embedding by preserving class neighbourhood structure. In International Conference on Artificial Intelligence and Statistics, pages 412-419, 2007. (Cited on page 27)
[63] R. Socher, A. Perelygin, J. Wu, J. Chuang, C. D. Manning, A. Y. Ng, and C. Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Conference on Empirical Methods in Natural Language Processing, pages 1631-1642, 2013. (Cited on page 30)
[64] Y. Sui, X. Wang. J. Wu, M. Lin, X. He, and T.-S. Chua. Causal attention for interpretable and generalizable graph classification. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, page 1696 1705, 2022. (Cited on pages 9 and 32)
[65] Tao, Q. Cao, H. Shen, Y. Wu, B. Xu, and X. Cheng. Idea: Invariant causal defense for graph adversarial robustness. arXiv preprint, arXiv:2305.15792, 2023. (Cited page 21)
[66] D. Ulyanov, A. Vedaldi, and V. S. Lempitsky, Instance normalization: The missing ingredient for fast stylization. arXiv preprint, arXiv: 160.0002,, 2016. (Cited page 31)
[67] A. van den Oord, Y. Li, and O. Vinyals. Representation learning with contrastive predictive coding. arXiv preprint, arXiv:18011 00778, 2018. (Cited on page 27)
[68] P. Veličković, G. Cucurull, A. Casanova, A. Romero, P. Liò, and Y. Bengio. Graph attention networks. In International Conference Learning Representations, 2018. (Cited page 1)
[69] T. Wang and P. Isola. Understanding contrastive representation learning through alignment and uniformity on the hypersphere. In International Conference on Machine Learning, pages 9929-9939, 2020. (Cited on page 27)
[70] Z. Wang, Y. Chen, Y. Duan, W. Li, B. Han, J. Cheng, and H. Tong. Towards out-of-distribution generalizable predictions of chemical kinetics properties. arXiv preprint, 1r.0v:2310003152, 2023. (Cited on page 1)
[71] D. Wolpert and W. Macready, No free lunch theorems for optimization. IEEE Transactions on Evolutionary Computation, 1(1):67-82, 1997. (Cited on pages 2 and 6)
[72] Q Wu, H. Zhang, J. Yan, and D. Wipf. Handling distribution shiftts on graphs: An invariance perspective. In International Conference on Learning Representations, 2022. (Cited on pages
1, 2. 3, 4, 5, 20 and 21)
[73] Y. Wu, Wang. Zhang, X. He, and T.-S. Chua. Discovering invariant rationales for graph neural networks. In International Conference on Learning Representations, 2022. (Cited on pages 2, 4, 9, 21 and 22)
[74] K. Xu, C. Li, Y. Tian, T. Sonobe, K. Kawarabayashi, and S. Jegelka. Representation learning on graphs with jumping knowledge networks. In International Conference on Machine Learning, pages 5449-5453 5458,,558440 2018. 8. (Cited on pages and 31)
[75] K. Xu, W. Hu, J. Leskovec, and S. Jegelka. How powerful are graph neural networks? In International Conference on Learning Representutions, 2019. (Cited on pages 1, 1, 20, 22 and 31)
[76] K. XU, M. Zhang. J. Li, S. S. Du, K. Kawarabayashi, and S. Jegelka. How neural networks extrapolate: From feedforward to graph neural networks. In International Conference on Learning Representations, 2021. (Cited on page 21)
15

---

[77] N. Yang. K. Zeng, Q. Wu, X. Jia, and J. Yan. Learning substructure invariance for out-of- distribution molecular representations. In Advances in Neural Information Processing Systems, 2022. (Cited on pages 1, 2, 4, 5, 9, 21, 22, 32 and 33)
[78] G. Yehudai, E. Fetaya, E. Meirom, G. Chechik, and H. Maron. From local structures to size generalization in graph neural networks. In International Conference on Machine Learning. pages 11975-11986, 2021. (Cited on page 21)
[79] R. Yeung. Information Theory and Network Coding. Springer New York, NY. 01 2008. ISBN 978-0-387-7923344 (Cited on page 25)
[80] Y. You, T. Chen, Z. Wang, and Y. Shen. Graph domain adaptation via theory-grounded spectral regularization. In The Eleventh International Conference on Learning Representations, 2023. (Cited on page 21)
[81] J. Yu, T. Xu, Y. Rong, Y. Bian, J. Huang, and R. He. Graph information bottleneck for subgraph recognition. In International Conference on Learning Representations, 2021. (Cited on pages and 21)
[82] J. Yu, T. Xu, Y. Rong, Y. Bian, J. Huang, and R. He. Recognizing predictive substructures with subgraph information bottleneck. IEEE transactions on pattern analysis and machine intelligence, 2021. (Cited on pages 4 and 21)
[83] J. Yu, J. Liang, and R. He. Mind the label shift of augmentation--assd graph OOD generalization. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2023. (Cited on pages
1, 4, 5 and 21)
[84] H. Yuan, H. Yu, S. Gui, and S. Ji. Explainability in graph neural networks: A taxonomic survey. arXiv preprint, arXiv:2012. 15445, 2020. (Cited on pages 9 and 30)
[85] M. Zhang, N. S. Sohoni, H. R. Zhang, C. Finn, and C. Ré. Correct- A contrastive approach for improving robustness to spurious correlations. arXiv preprins, arXiv:2203.01511, 2022. (Cited on pages 4, 7. 8, 21 and 27)
[86] X. Zhang. L. Wang, J. Helwig, Y. Luo, C. Fu, Y. Xie, M. Liu, Y. Lin, Z. Xu, K. Yan, K. Adams, M. Weiler, X. Li, T. Fu, Y. Wang. H. Yu, Y. Xie, X. Fu, A. Strasser, S. Xu, Y. Liu, Y. Du, A. Sax- ton, H. Ling, H. Lawrence, H. Stärk, S. Gui, C. Edwards, N. Gao, A. Ladera, T. Wu, E. F. Hof- gard, A. M. Tchrani, R. Wang. A. Daigavane, M. Bohde, J. Kurtin, Q. Huang. T. Phung, M. Xu, C. K. Joshi, S. V. Mathis, K. Azizzadenesheli, A. Fang, A. Aspuru-Guzik, E. Bekkers, M. M. Bronstein, M. Zitnik, A. Anandkumar, S. Ermon, P. Lið, R. Yu, S. Günnemann, J. Leskovec, H. Ji, J. Sun, R. Barzilay, T. S. Jaakkola, C. W. Coley, X. Qian, X. Qian, T. E. Smidt, and S. Ji. Artificial intelligence for science in quantum, atomistic, and continuum systems. arXiv preprint, arXiv:2307. 08423,, 2023. (Cited on page 1)
[87] Y. Zhang, M. Gong, T. Liu, G. Niu, X. Tian, B. Han, B. Schölkopf, and K. Zhang. Adversarial ro- bustness through the lens of causality. In International Conference Learning Representations, 2022. (Cited on page 1)
[88] J. Zhou, B. Bevilacqua, and B. Ribeiro. An ood multi-task perspective for link prediction with new relation types and nodes. arXiv preprint, (xii:2307.06644 2023. (Cited on page 21)
[89] Y. Zhou, G. Kutyniok, and B. Ribeiro. OOD link prediction generalization capabilities of message- message-passing GNNs in larger test graphs. In Advances in Neural Information Processing Systems, 2022. (Cited on page 21)
[90] Y. Zhou, G. Kutyniok, and B. Ribeiro. OOD link prediction generalization capabilities of message-passing GNNs in larger test graphs. In Advances in Neural Information Processing Systems, 2022. (Cited on page 21)
[91] Z. Zhou, J. Yao, J. Liu, X. Guo, Q. Yao, L. He, L. Wang, B. Zheng, and B. Han. Combating bilateral edge noise for robust link prediction. In Advances in Neural Information Processing Systems, 2023. (Cited on page 21)
16

---

[92] Z. Zhou, C. Zhou, X. Li, J. Yao, Q. Yao, and B. Han. On strengthening and defending graph reconstruction attack with markov chain approximation. In International Conference on Machine Learning, 2023. (Cited on page 21)
[93] Q Zhu, Y. Jiao, N. Ponomareva, J. Han, and B. Perozzi. Explaining and adapting graph conditional shift. arXiv preprint, arXiv:2306.0222 03256, 2023. (Cited on page 21)
[94] D. Zou, S. Liu, S. Miao, V. Fung, S. Chang, and P. Li. GDL-DS: benchmark for geometric deep learning under distribution shifts. arXiv preprint, abs/2310.08677, 2023. (Cited on page 1)
17

---

Appendix of GALA
Contents
A Notations
19
Limitations and Future Directions
19
C Full Details of the Background
19
More Details about the Failure Cases
21
E Proof's for Theorems and Propositions
23
E,1 Proof of Proposition 3.2
23
E,2 Complementary discussion for Sec, 3.1
23
E,3 Proof of Proposition 3.4
24
E,4 Proof of Corollary 3.6
24
E.5 Proof of Theorem 4.1
24
F More Discussions on Practical Implementations of GALA
27
More Details about the Experiments
30
G.1 Datasets
30
G.2 Baselines and Evaluation Setup
31
G.3 Software and Hardware
33
G.4 Computational analysis
18

---

A Notations
Typically, for graphs that appeared in the discussion, we will use the superscript to denote the sampling process (e.g.. G is the positive graph), and the subscript to denote the specific invariant (i.e., G.) or spurious subgraph (i.e., Gs). Graph symbols with G are the predicted graphs of model (i.e., the estimated invariant subgraph G.. Below, we list some examples of graphs involved in this paper.
Table 4: Notations for graphs involved in this paper
Symbols
Definitions
G
the graph space the space of subgraphs with respect to the graphs from the lahel space
Ge
y
Ge
g
graph
G (A,X)
graph with the adjacency matrix A e (0, (0,1) and node feature matrix X e Raxd set of graphs
[G}
graph sampled as positive samples
G"
graph sampled as negative samples graph sampled according to CIGA [8] the invariant subgraph with respect to G the spurious subgraph with respect to G the invariant subgraph of positive graph G" the spurious subgraph of positive graph GP
G
G.
G.
GP
C
G
the estimated invariant subgraph the estimated spurious subgraph
G.
G
the estimated invariant subgraph of a positive graph CP
C
estimated spurious subgraph a positive graph GP the part of the underlying invariant subgraph Ge appeared in G. the complementary part of LG. with respect to the invariant subgraph Ge
Δ C G.
ZC. Ge -GG
B Limitations and Future Directions
Although our work establishes a set of minimal assumptions for feasible invariant graph learning when the environment partitions and auxiliary information about the environment are both not available, our work is built upon the minimal availability of the environment knowledge. Nevertheless, there could exist some additional information that may be helpful for environment augmentation. Therefore, it remains interesting to explore more theoretically grounded strategies to discover and leverage more environment information for identifying the graph invariance. When the direct environment augmentation is not feasible, GALA provides suitable framework that could easily manipulate the environment assistant model or the partitioning of the positive and negative graphs, to select the spurious features via the additional information and better identify the graph invariance.
addition to the correlation strengths discussed in this work, there exist other factors, such as the size of spurious and invariant subgraphs, that affect the fitting of spurious and invariant patterns, another promising future direction is to discuss the influence of these factors to the design of environment assistant model and OOD generalization graphs.
Besides, better data partitioning strategy can be developed with uncertainty measures [51].
C Full Details of the Background
We give a more detailed background introduction about GNNs and Invariant Learning in this section.
Graph Neural Networks. Let = (A, X) denote a graph with n nodes and m edges, where A E (0, 1 is the adjacency matrix, and X Raxdd the node feature matrix with node feature dimension of d. In graph classification, given set of N graphs (G.) C G and their labels
19

---

[X))=1 =y=-R" from classes. Then, we train a GNN h with an encoder h : 9 - R² that learns meaningful representation hıg for each graph G to help predict their labels yo p(hc) with
downstream classifier p: Aº + y. The representation hGg is typically obtained by performing pooling with a READOUT function on the learned node representations:
hg = (E 2 E V}),
(10)
where the READOUT is permutation invariant function (e.g., SUM, MEAN) [75], and h. for the node representation of E K-th layer that is obtained by neighbor aggregation:
(K) stands
[()) <Wk a ((()) "H (u)) U (u))),
(11)
where (u) is the set of neighbors of node u, (--) is an activation function, c.g., ReLU, and a(-) is an aggregation function over neighbors, e.g., MEAN.
E
S
S1
(a) Graph generation SCM
(b) FIIF SCM
(c) PIIF SCM
(d) MIIF SCM
Figure 4: Full SCMs on Graph Distribution Shifts [8].
Graph generation process. This work focuses on graph classification, while the results generalize to node classification well using the same setting in Wu et al. [72]. Specifically, we given set of graph datasets D = {De}e collected from multiple environments Eall Samples (G;, Yt²) D from the same environment are considered as drawn independently from an identical distribution IP We consider the graph generation process proposed by Chen et al. [8] that covers a broad case of graph distribution shifts. Fig. 4 shows the full graph generation process considered in Chen al. [8]. The generation of the observed graph G and labels Y are controlled by a set latent causal variable C' and spurious variable S, i.e.,
G = ggen (C,
C' and S control the generation of G by controlling the underlying invariant subgraph G, and spurious subgraph G., respectively. Since S can be affected by the environment E, the correlation between Y, S and G2 can change arbitrarity when the environment changes. C and S control the generation of the underlying invariant subgraph Ge and spurious subgraph G.. respectively. Since S can be affected by the environment E, the correlation between Y. and G, can change arbitrarily when the environment changes. Besides, the latent interaction among C, S and Y can be further categorized into Full Informative Invariant Features (FIIF) when Y S and Partially Informative Invariant Features (PIIF) when Y y S/C. Furthermore, PIIF and FIIF shifts can be mixed together and yield Mixed Informative Invariant Features (MII)). as shown in Fig. We refer interested readers Chen et al. [8] for detailed introduction of the graph generation process.
Invariant graph representation learning. To tackle the OOD generalization challenge on graphs from Fig. 4, the existing invariant graph learning approaches generically aim to identify the underlying invariant subgraph Ge to predict the label Y [8, 72]. Specifically, the goal of OOD generalization on graphs is to learn an invariant GNN f = fe g. which is composed of two modules: a) featurizer g: g + Ge that extracts the invariant subgraph G.; b) classifier fc Ge > y that predicts the label Y based on the extracted G., where G. refers to the space of subgraphs of G. The learning objectives of fe and y are formulated
max I(G Y), s.t. G. 1 E, G. = G(G)
(12)
Since is not observed, many strategies are proposed to impose the independence of G. and E. A common approach is augment the environment information. For example, based on the
20

---

estimated invariant subgraphs G. and spurious subgraphs G, Liu et al. [45]. Wu et al. [72. 73] proposed generate new environments, while Li et al. [41], Yang et al. [77] proposed to infer the underlying environment labels. However, we show that it is fundamentally impossible to augment faithful environment information in Sec. 3. Miao et [55, 56], Yu et al. [81, 82, 83] adopt graph information bottleneck to tackle FIIF graph shifts, and they cannot generalize to PIIF shifts. Our work focuses on PIIF shifts, as it is more challenging when without environment labels [43]. Fan et al. [17] generalized [40] to tackle severe graph biases, i.e. when (S|Y) < H (C)Y). Chen et al. [8] proposed a contrastive framework to tackle both FIF and PIFF graph shifts, but limited to H(S > H(C However, in practice it is usually unknown whether (SIY) (C or H(S > HY) without environment information.
More OOD generalization on graphs. In addition to the aforementioned invariant learning ap- proaches, Bevilacqua et al. [5], Li et al. [42], Yehudai et al. [78], Zhou et al. [89] study the OOD generalization an extrapolation from small graphs to larger graphs in the task of graph classification and link prediction. In contrast, we study OOD generalization against various graph distribution shifts formulated in Fig. 4. In addition to the standard OOD generalization tasks studied in this paper, Mahdavi et al. [52], Xu et al. [76] study the OOD generalization in tasks of algorithmic reasoning on graphs. Jin et al. [29] study the test-time adaption in the graph regime. Kamhoua et al. [31] study the 3D shape matching under the presence of noises. Gui et al. [22] propose an independence constraint onto the target label and environment label to improve the OOD generalization when environment labels are available. Liu et al. [48] adopt a flexible framework to tackle shifting graph distributions. Chen et al. [7]. Tao et al. [65], Zhou al. [91, 92] study the OOD generalization on graphs from the adversarial robustness perspective.
In addition to graph classification, Chen et al. [6]. Wu et al. [72] study node classification. Liu et al. [47] propose a structural reweighting strategy to improve the OOD generalization of node classification. Lee et al. [39] propose to incorporate multiple modules to handle different degree modes in OOD node classification. You et al. [80], Zhu et al. [93] study unsupervised graph domain adaption.Gat et al. [19]. Zhou et al. [88, 90] study the OOD link prediction.
Besides, Lucic et al. [49] aims to find counterfactual subgraphs for explaining GNNs, which focuses on post-hoc explainability while this work focuses on intrinsic interpretability.
Invariant learning without environment labels. There are also plentiful studies in invariant learning without environment labels. Creager et al. [12] proposed a minmax formulation to infer the environment labels. Liu al. [46] proposed a self- belstoosting framework based on the estimated invariant and variant features. Liu et al. [44], Zhang et al. [85] proposed to infer labels based the predictions of an ERM trained model. Deng et al. [14]. Pezeshki et al. [61] improve the inference of group labels based on feature learning and prediction correctness. However, Lin et al. [43] found failure cases in Euclidean data where it impossible to identify the invariant features without given environment labels. Moreover, as the OOD generalization on graphs is fundamentally more difficult than Euclidean data [8], the question about the feasibility of learning invariant subgraphs without environment labels remains unanswered.
D More Details about the Failure Cases
We provide more empirical results and details about the failure case verification experiments in complementary to Sec. 3. The results are shown in Fig. 5. We compared different environment aug- mentation approaches the vanilla GNN model trained with ERM (termed ERM), and an interpretable GNN model trained with ERM (termed XGNN).
The failure cases are constructed according the two-piece graph generation models. The specific description is given as the following.
Definition D.I (3-class two-piece graphs). Each environment is defined with two parameters, Oley Se [0, 1}, and the dataset D. is generated as follows:
(a) Sample yf {0, 1,2) uniformly:
(b) Generate G. and G2 via:
Ge == fG Rad Re)( Go == (o Rad(B.))
21

---

CTM GALA
DWM DALA
000 2e0
000
spurious correlation strength
spurious correlation strength
spurious correlation strength
(a) Failures of env. generation
(b) Failures env, inferring
(c) Failures of resolving env, consis-
tency
00
GAEA
OGA
000 lest RCC
78
000 acc
DMA
65
80
000 toat
60
spurious correlation strength 64
07
13
spurious correlaiio strength
spurious correlation strength
(d) Failures of env. generation
(c) Failures of env. inferring
(f) Failures of resolving env. consis- tency
Figure 5: Failures of finding faithful environment information. Results shown in the figure are based on the 3 class two-piece graphs (Def. D.I), where the invariant correlation strength is fixed 0.7 while the spurious correlation strength is varied from 0.5 0.7. We can find that both environment augmentation and inferring approaches suffer from severe performance decreases or even underperform ERM and XGNN when the dominated correlation not suitable for the method. In contrast, GALA maintains strong OOD performance for both cases.
where Iga I respectively map input (0, to specific graph selected from given set, and Rad (a) is random variable with probability a taking uniformly random value from {0, 1, 2}, and a probability of - a taking the value of +1;
(c) Sythesize G by randomly concatenating Ge and G::
G = gen (Ge, Gs).
In experiments, we implement the 3-class two-piece graphs with the BA-motif's [50] model.
In experiments, we adopt a 3-layer GIN [75] with a hidden dimension of 32 and dropout rate of 0.0 as the GNN encoder. The XGNN architecture is implemented via two GNNs following the original implementation as CIGA. The optimization is proceeded with Adam [34] using learning rate of le 3. All experiments are repeated with 5 different random seeds of 21, 3, 4, 5). The mean and standard deviation are reported from the 5 runs.
We implement DIR [73], GREA [45], MoleOOD [77], GIL [41], DisC [17], and CIGA [8], according to the author provided codes (if available). As for the hyperparameters in each method, we use a penalty weight of 1e 2 for DIR following its original experiment in spurious motif datasets generated similarly using BA-motifs [73]. We use penalty weight of for GREA as we empirically it does not affect the performance by changing to different weights. For MoleOOD and GIL, we set the number of environments as We tune the penalty weights of MoleOOD with values from {le- 2, le 1, 1, 10} but did not observe much performance differences. We tune the penalty weights of GIL with values from {le- - 5, le 3, le - 1} recommended by the authors. For DisC, we tune only the y weight from (0.9, 0.7, 0.5) in the GCE loss as we did not observe performance differences by changing the weight of the other term. We tune the penalty weight of CIGA with values from (0.5, 1, 2, 4, 8, 16, 32) recommended by the authors.
22

---

E Proofs for Theorems and Propositions
E.1 Proof of Proposition 3.2
Proposition E.1. (Restatement of Proposition 3.2) Consider the two-piece graph dataset Err {aa, 31), (ar, }2)} with a 2 31, 32 (e.g., E, {00.25, 0.1), (0.25, 0.2) and its corresponding mixed environment ei = {(a, (31 + 32) (e.g., ei = {00.25, 0.15))). When G. = G. and
G. Go. it holds that the augmented environment Ey is also two-piece graph dataset with
& (0.55 (3, + 22) (e.g, & =
Proof. From Definition 3.1, we known that for each graph G; Exix {(a, (81 + B2) G, is the concatenation of G. and G² defined
G² = Jan (Y; Rad(a))) where Rad(.); denotes the ith sample of the random variable Rad(-).
= (Y; Rad (81 + (2)/2))),
Denote
GA fge (+1), GB = f (-1),
and
Gc f (+1), GD - (-1),
Considering applying the augmentation to 2n samples randomly sampled from egix, since the featurizer g separates each G Eir emix into G. G, and G. Ge. and the augmented graph G is obtained by
- فe G C2), Vi {1...n).
Then, the Oe, By in & be obtained by summing up the overall numbers of GA, GB, GD concatenated into 2n² samples in &
Specifically, can inspect the changes of the distributions of motif's and labels. Let 3 (31 + /2) without loss of generality, we focus on inspecting the changes given Y = +1, since the changes given Y = -1 is symmetric as Y +1. The original distribution is shown follows:
y=+1
GA (1 -1-
GB
Gc
a1 8yn
a(1
f)
GD
(1 a) 33
a.3n
Then, new distributions of the motifs and labels are determined by the number of original motifs identified as G. and Ge. respectively. When G. = G2 and G = Ge. in the new environment E,. given Y +1, Ge contributes (1 - 3)n 2n samples the "invariant" subgraph. More specifically, Gc will be concatenated with G and Gg by n times, respectively. Then we have the new distribution tables shown as follows:
Y=+1
GA
GB (1
Gc
1-33mm
8)m²
GD
3n²
Bn
Since given the same Y, the spurious subgraph Ge and Gp will still have the same chance being flipped, we have 3, = B. While as GA and Gg appear the same times given the same Y, it suffices to know that a, = 0.5.
E.2 Complementary discussion for Sec. 3.1
Proposition E.2. Given the same gruph generation process as Fig. 2, when there exists spurious subgraph G. such that pe: (Y|G.) = ped (Y,,) for any two environments ex, e2 E Em, where pe (Y is the conditional distribution P(Y G.) under environment Eall, il is impossible for any learning algorithm applied to fc differentiate Ge from G.-
23

---

Proof. Let G be the spurious subgraph such that p GY] pM (Y,)) for any two environ- ments e 62 E En. and G. be the invariant subgraph which t (Y = p (Y)) Ve,, e2 & by definition. Consider a learning algorithm applied to Se o that accepts the input of & and extracts subgraph G. - g(Y) as an estimation of the invariant subgraph for any G predict Y via fe(Ge) deterministic manner. the algorithm succeed extract Ge from E then there always exists Emix' with the desired spurious subgraph G" and underlying invariant subgraph G² such that G² Ge and Go G Due the deterministic nature, the algorithm fails to identify Ge in gmix'
E.3 Proof of Proposition 3.4
Proposition E.3. (Restatement of Proposition 3.4) There exist 2 two-piece graph training environ- ments err and Eyr that share the same joint distribution P(Y, G). Any learning algorithm will fail in either Ey Ey'
Proof. Let the mixed training environment of & and E be E = {(a,{1100000 B)}. Based on the definition of two-piece graphs (Definition 3.1), the joint distribution of the mixed training dataset (G Concat Ge, G.], Y) can be computed as
Y +1, -1, Ci (Ge) = Bit Y,
with probability 0.5,
with probability 0.5,
with probability (1 - -a) - 3),
(Go) & =Bir. (G.) Y,
with probability a(1 - B),
(G.) & Bit (G.) =Y,
with probability (1 -a)3,
(i.))(G.) Bit YY,
with probability a,3.
Here we use Bot (G.) to obtain the input bit of a subgraph Ge (or (fg(a) and Bit (C. for G.. respectively.
Any learning algorithm that tries to identify the invariant subgraph from this training dataset will compute a model that uses subgraph Ge, or subgraph G.. or both Ge and G, to predict Y determin- istically. Thus, as long as the joint distribution does not change, the resulting model will always identify the same invariant subgraph. Without loss of generality, let us assume that the model correctly identifies G. as the invariant subgraph for & {(a, 31), (a, (2)) with 3 = (31 + 82)
Now let the other training environment be E {(a1, 9), (a2, B)} with clear that since the mixed training environment of & is still {a, 3)}. the model keeps regarding Ge the invariant subgraph. However, for & the model fails to identify the invariance since now the invariant subgraph is G..
(al + a2) It is
E.4 Proof of Corollary 3.6
Corollary E.4. (Restatement of Corollary 3.6) Without Assumption 3.3 or Assumption 3.5, there does not exist learning algorithm that captures the invariance of the two-piece graph environments.
Proof. The proof for lacking Assumption 3.3 is identical to the proof for Proposition E.2. Consider a learning algorithm applied to fo o g that accepts the input of Emix and extracts a subgraph G. = gYY) estimation of the invariant subgraph for any G' to predict Y via fG. in deterministic manner. Without the holding Assumption 3.5, due Proposition 3.4, there exists gmix for each Emix that have the identical joint distribution but different underlying invariant subgraph. Thus, any learning algorithm that succeeds in either Eir or Eyr
guixx will fail in the other.
E.5 Proof of Theorem 4.1
Theorem E.5. (Restatement of Theorem 4.1) Given, i) the same data generation process as in
Fig. 2; ti) Dr, that satisfies variation sufficiency (Assumption 3.3) and variation consistency
24

---

(Assumption 3.5); tii) {GD} and {G"} are distinct subsets of Dur such that I[G G - 0,
VG arg max IGG Y) under (G"), and VG, arg max ;(G under {G"); suppose [Gel = Se: VG,, resolving the following GALA objective elicits an invariant GNN defined via Eq. 12,
max I(Gr; Y), s.t. E arg max T(C g.1GH1S..
(13)
Sc.9
where Go E (G = g(G") and Gh { = g(G")} are the estimated invariant subgraphs via from {G") and },m), respectively.
Proof. Without loss of generality, we assume that { has the same spurious dominance situation as Eg.. In other words, when (S))) <H(C)Y), the data distribution in {G*) also follows Y)) < H(C)Y), while H(SY) > H(ClY) in {G"). To proceed, we will use the language of Chen et al. [8].
We begin by discussing the case of H(S < H(C)Y). Given H(S)Y) < H(C)Y), we have H(S < H(CYY) in {GP) and (SSY) > H(C)Y) in {G"). Then, claim that
Ge E arg max
I( G.
(14)
Gg.
Otherwise, consider there exists subgraph of the spurious subgraph AGG C G Gg which takes up the space of A. S Ge from Gg. Then, let AGg = G - AG% we can inspect the changes to I[ G% (Y) led by AGg:
= -AH(GHY) AHI
=[H(TGg AQIIV) H(AO) A - [( AC,IC: Y) H(TG, AGgIO r) [HA o Y) H(A Tn
= =(H(AG:IAAGY))) - m)
(15)
where the last equality is obtained via expanding the conditional entropy. Then, considering the contents in Cg without loss of generality, divide all of the possible cases into two:
(i) G contains only the corresponding invariant subgraph G
iii contains subgraph from the corresponding spurious subgraph G., denoted as A S
For case (i), is easy to write Eq. 15 as:
AIG Y
YY HH - H( - C Y) H(A Gn r)
-HHLL AG H G. Y),
(16)
since HA Y) HHAG GN AदG, Y) H(AōgYY given C' 1 S for PIIF shifts. Then, suffices know that (I, Gh Y) conditioning on new variables will not increase the entropy [79].
For (ii), have:
AI GYlY)
H(A &n x)), - AGB H v) + [HAOGAA Y) H( T Y) H
DILGG,AAYYY -H
(17)
25

---

where we claim that H AGG, Y) H(AŌT)A C) = 0, and similarly conclude that AT(G;: G. & O. More specifically, can rewrite the first term in Eq. 17 as
H(AG V) H(ARYAGG C% Y) = H(AG - H(AG: Y) I[AGY; AG, Y) 0,
using the variation condition (i.e., assumption iii)) for LG2 under {G"), and AG; under {G"). After showing the success of GALA in tackling H(S < H(C))), it also suffices to know that the aforementioned discussion also generalizes to the other case, i.e. when ((S)Y) > H(C in {GP) and H(S)) < H(C) in {G*).
26

---

F More Discussions on Practical Implementations of GALA
In this section, we provide more implementation discussions about GALA in complementary Sec.
Objective implementation. As the estimation of mutual information could be highly expensive [4, 67], inspired by Chen al. [8], we adopt the contrastive learning approximates the mutual information between subgraphs in Eq. 9 [4, 11. 33, 62, 67]:
I(G% (2 R (o (
(G))
olkar hap)
(18)
log olhoy.hey) + Em alha. may)
where (Gg, G are subgraphs extracted by from {G*), ), { j that share the same label, respectively, {G M are subgraphs extracted by g from G that has a different label. (G Y) is the push- forward distribution of P = Y) by featurizer 9. P(G given the label Y. P + Y) refers to the distribution of given the label that is different from Y, G = ( G = ( G = g(G') are the estimated subgraphs, hāg hag hag are the graph presentations of the extracted subgraphs. is similarity measure. As M oc, Eq. 18 approximates I(Ō;; G7 (Y) [1, 32, 69].
Y) refers to the distribution of G
Environment assistant implementation. Theorem 4.1 shows the effectiveness GALA when given proper subsets of {GP) and {G*}. J. In practice, multiple forms. As discussed in Sec. 4.1, ERM trained model can serve as a reliable proxy. Since ERM tends to learn the first dominant features, when (SS)) H(C)Y), ERM will firstly learn to extract spurious subgraphs G, to make predictions. Therefore, we can obtain {G*) } by finding samples where ERM correctly predicts the labels, while {G") } for samples that ERM predicts an incorrect label. In addition to direct label predictions, we can also adopt clustering [85] to yield environment assistant predictions for better contrastive sampling. We provide the detailed description of the clustering based variant of GALA in Algorithm 2.
implement the environment assistant into
Algorithm 2 GALA: Clustering based Graph inv Ariant Learning Assistant
1: Input: Training data Dar; environment assistant A; featurizer gi classifier fe; length of maximum training epochs e: batch size b;
2: Initialize environment assistant A;
3: for p ee] do
4: Sample batch of data (Gi, YiJi from Der;
5: Obtain Environment Assistant predictions { using k-means clustering on the graph representations yielded by A;
6: for each sample G3, Vi {G,, Y. do
7:
Find postive graphs with same and different ;;: Find negative graphs with different yı but same environment assistant prediction e;; Calculate GALA risk via Eq. 18;
8:
9:
Update feag via gradients from GALA risk;
11: end for
end for
13: return final model fc o gg
Empirically, we find clustering based variants can provide better performance when the spurious correlations are well learned by the environment assistant model. More concretely, we plot the umap visualizations [53] of ERM trained environment assistant model as in Fig. 6, where we can find that clustering predictions provide better approximations to the underlying group labels.
Besides, we can also incorporate models that are easier to overfit to first dominant features to better differentiate (GP) from ]G"). To demonstrate the influence of different environment assistant implementations, we conduct more studies with interpretable GNNs with an interpretable ratio of 30% trained with ERM and also with CIGAv1 penalty of 4.
27

---

colored by group labers
catered predictoon
colored by chaser pred
186
(a) Colored by environment labels. (b) Colored by label predictions.
(c) Colored by cluster predictions.
Figure 6: Umap visualizations learned graph representations in ERM trained environment assistant model based on the 3-class two-piece graph {0.7, 0.9).
colored by group
relored prediction
relored by clattr pred
125
14
74
124
35
02
35
(a) Colored by environment labels. (b) Colored by label predictions.
(c) Colored by cluster predictions.
Figure 7: Umap visualizations of learned graph representations in an interpretable GNN model (ratio=30%) trained with ERM based on the 3-class two-piece graph {0.7, 0.9).
colorid group
caterdd predicilos
coloned by cluster
(a) Colored by environment labels. (b) Colored by label predictions.
(c) Colored by cluster predictions.
Figure 8: Umap visualizations of learned graph representations in an interpretable GNN model (ratio=30%) trained with ERM based on the 3-class two-piece graph (0.7, 0.9).
group
by prod
-3
(a) Colored by enviroument labels. (b) Colored by label predictions. (c) Colored by cluster predictions.
Figure 9: Umap visualizations of learned graph representations of interpretable GNN trained by ERM on EC50-Assay.
28

---

colored by group laber
citered prediction
colored by chastr pred
(a) Colored by environment labels. (b) Colored by label predictions.
(c) Colored by cluster predictions.
Figure 10: Umap visualizations of learned graph representations of an interpretable GNN trained by ERM on EC50-Scaffold.
(a) Colored by environment labels. (b) Colored by label predictions.
(c) Colored by cluster predictions.
Figure 11: Umap visualizations of learned graph representations of interpretable GNN trained by ERM on EC50-Size.
In Fig. and Fig, 8, it can be found that the interpretable GNN learns hidden representations that are better clustered with group labels. The clustering based predictions yields a better approximation of the underlying environment labels. Furthermore, when implementing the environment assistant model using interpretable GNN as well as CIGAvI penalty, which facilitates the overfitting to the spurious correlations, then the vanilla label predictions also yield good approximation of the underlying environment labels.
Although using the clustering predictions seem to be promising, we also find negative cases. For example, in DrugOOD datasets, the number of curated environment labels are much larger that learning a well clusterd hidden representations for the environment labels appears to be difficult. Shown as in Fig.
Fig. 11, the learned representations have poor quality for approximating the underlying environment labels. Empirically, also find that direct using label predictions in DrugOOD datasets generically yield better performance.
One-side contrastive sampling. The original supervised contrastive implementation [33] takes positive and negative samples within the batch using two-side contrastive sampling. That is. all the samples will be considered anchor points. However, when it is used to contrast samples from Gg and Ct. there could be undesired behaviors. First, it often happen that there are few negative cases when the spurious correlations are too strong. The samples from (GP) in a batch may pull the representations samples from (G*) } even closer, which makes the model further overfitted to the spurious correlations. Second, the sampling over G and G can be seen as hard positive and negative samples, that may impose too strong regularizations that preventing the learning of any correlations. Therefore, we propose to use one-side sampling. That is, only using the incorrectly predicted samples anchor points. We empirically observe one-side sampling could yield better performance in two-piece graphs.
Upsampling minority group samples. It is possible that the number of positive and negative graphs is imbalanced, especially when adopting the label predictions to sample positive and negative graphs. For example, when the ERM trained assistant model overfits the training distribution under the spuriousness- domeessnssss--dimeeedddee case, the number of negative graphs will be extremely small. Given an

---

extremely small number of negative samples for contrastive learning, the resulting mutual information estimation will be collapsed trivial solutions. Therefore, we propose simple strategy to mitigate the issue. We directly upsample the minority group samples. The minority group of samples will be repeated l times within the training set.
G More Details about the Experiments
In this section, we provide more details about the experiments, including the dataset preparation, baseline implementations, models and hyperparameters selection well the evaluation protocols.
G.1 Datasets
We provide more details about the motivation and construction method of the datasets that are used in experiments. Statisties of the datasets are presented in Table 5.
Table 5: Information about the datasets used in experiments. The number of nodes and edges are respectively taking average among all graphs
Dutasets
Training
# Validation
# Testing
# Classes
Nodes
Edges
Metries
Two-piece graphs [0.8, 0.6)
000
3,000
3, 1.000
3
26.14
36.21
ACC
Two-piece graphs [0.8, 0.7)
9, 000
3, 0,000
3, 3,000
3
26.18
36.27
ACC
Two-piece graphs {0.8, 0.9)
9, 9,000
3, 3,000
3, 0.000
3
26.13
36.22
ACC
Two-piece graphs {0.7, 0.9}
000
3,000
3, 0.000
3
26.13
36.22
ACC
CMNIST-sp
40, 000
5,000
15,000
2
56.90
373.85
ACC
Graph-SST2
28,
7,004
12, 893
2
10.20
ACC
EC50-Assay
4,978
761
2, 725
2
40.89
87.18
ROC-AUC
EC50-Scaffold
743
7.723
2, 762
35.54
75.56
ROC-AUC:
EC50-Size
5, 189
2, 495
2,505
2
35.12
75.30
ROC-AUC
Ki-Assay
8,490
4,741
4,720
2
32.66
71.38
ROC-AUC
Ki-Scaffold
5, 389
4,805
4,463
2
29.96
65.11
ROC-AUC
Ki-Size
605
4,486
4,558
2
30.35
66.49
ROC-AUC
Two-piece graph datasets. We construct 3-class synthetic datasets based on BAMotif [50] following Def. D.I, where the model needs to tell which one of three motifs (House, Cycle, Crane) the graph contains. For each dataset, we generate 3000 graphs for each class at the training set, 1000 graphs for each class at the validation set and testing set, respectively. Each dataset is defined with two variables {a, b) referring to the strength of invariant and spurious correlations. Given {a, b), we generate the training data following the percise generation process as Def. D.I. While for the generation validation sets, we use a b, = max(1/3,6 b 0.2) that facilitates the model selection for OOD generalization [9, 23J. While for the generation of test datascts, merely use b 0.33 that contains no distribution shifts, to fully examine to what extent the model learns the invariant correlations. During the construction, merely inject the distribution shifts in the training data while keeping the testing data and validation data without the biases.
CMNIST-sp. To study the effects of PIIF shifts, we select the ColoredMNIST dataset created in IRM [3]. We convert the MoloredMMist into graphs using the superpixel algorithm introduced by Knyazev et al. [36]. Specifically, the original Mnist dataset is assigned to binary labels where images with digits 0 0-4 4 are assigned to = and those with digits - are assigned to y = 1. Then, y will be flipped with probability of 0.25. Thirdly, green and red colors will be respectively assigned to images with labels 0 and an averaged probability of 0.15 (since we do not have environment splits) for the training data. While for the validation and testing data, the probability is flipped to 0.9.
Graph-SST2. Inspired by the data splits generation for studying distribution shifts on graph sizes, we split the data curated from sentiment graph data [84], that converts sentiment sentence classification datasets Graph-SST2 [63] into graphs, where node features are generated using BERT [15] and the edges are parsed by Biaffine parser [20]. Our splits are created according to the averaged degrees of each graph. Specifically, we assign the graphs as follows: Those that have smaller or equal to 50-th percentile averaged degree are assigned to training, those that have averaged degree large than
30

---

50-th percentile while smaller than 80-th percentile are assigned to the validation set, and the left are assigned test set.
To evaluate the OOD performance in realistic scenarios with realistic distribution shifts, also include three datasets from DrugOOD benchmark [28]. Dru- gOOD is a systematic OOD benchmark for Al-aided drug discovery, focusing on the task of drug target binding affinity prediction for both macromolecule (protein target) and small- molecule (drug compound). The molecule data and the notations are curated from realistic ChEMBL database [54]. Complicated distribution shifts can happen on different assays, -or-0----ap--or----ebaae scaffolds and molecule sizes. In particular, we select -ug000-1baap-oore -ccreeee 1e00000-1aap-co----5555 -cafffssaffodd
DrugOOD datasets.
and from the task of Ligand Based Affinity Prediction which uses ic50 measurement type and contains core level annotation noises. We directly use the data files provided by the authors." For more details, we refer interested readers to Ji et [28].
-inrey,
Brugoo--bbap-- saaap--cre-ee wwwe-oi-reeeeeeeee siee-iisscaffold.
brug00D-1bap-core-B
G.2 Baselines and Evaluation Setup
During the experiments, we do not tune the hyperparameters exhaustively while following the common recipes for optimizing GNNs. Details are as follows.
GNN encoder. For a fair comparison, we use the same GNN architecture as graph encoders for all methods. By default, we 3-layer GIN [75] with Batch Normalization [27] between layers and JK residual connections at the last layer [74]. The hidden dimension set to 32 for Two-piece graphs, CMNIST-sp, and 128 for SST2, and DrugOOD datasets. The pooling is by default mean function over all nodes. The only exception is DrugOOD datasets, where we follow the backbone used in the paper [28], i.e., 4-layer GIN with sum readout.
Interpretable GNN backbone. As mentioned in Sec. 2 that most of the existing invariant graph learning approaches adopt the interpretable GNN the basic backbone model for the whole predictor f = fe o g. where : G - Ge a featurizer GNN and fe : Ge + y is classifier GNN. first calculates the sampling weights as in G. for each edge. More formally, given graph G containing n nodes, soft mask is predicted through the following equation:
Z GNN(G) g M a( a(Z R
where calculates the sampling weights for each edge using a MLP: M,, MLP((Z 2,)). Based on the continuous sampling score M, g could sample discrete edges according to the predicted scores [55]. For two-piece graph datasets and DrugOOD datasets, we will directly reweight the messaging passing process along the edge, as we empirically find it yields more stable performance. While for CMNIST-sp and Graph-SST2,
the score to
will sample a ratio r% of all edges for each graph. The ratios adopted are 80% and 60%, respectively, following previous works [8, 28). Meanwhile, to improve the stability of the subgraph extractor, we adopt layernorm [66] following the practice of [55].
Besides, we also have various implementation options for obtaining the features in Ge, for further obtaining h as well as for obtaining predictions based on G., Following previous works [55], we will adopt the same GNN encoder for the two GNNs in the interpretable GNN backbone, and feed the raw graph inputs to the classifier GNN. The contrastive loss is obtained via the graph representations of the sampled subgraph by the classifier GNN. For classifying based on G., we use a separate MLP downstream classifier in the classifier GNN fc-
Optimization and model selection. By default, we Adam optimizer [34] with learning rate of le - 3 and a batch size of 128 for all models at all datasets. Except for CMNIST-sp, we use a batch size of 256 to facilitate the evaluation following previous works [55]. To avoid underfitting, we pre-train models for 20 epochs for all datasets by default. While in two-piece graphs, we find pre-training by 100 epochs yields more stable performance. To avoid overfitting, we also employ an early stopping of 5 epochs according to the validation performance. Meanwhile, dropout is also adopted for some datasets. Specifically, we use a dropout rate of 0.5 for all of the realistic graph datasets, following previous works [8, 28).
'aatps:/drugooo.gituuu
31

---

The final model is selected according to the performance at the validation set. All experiments are repeated with 5 different random seeds of (1, 5). The mean and standard deviation are reported from the 5 runs.
Implementations of Euclidean OOD methods. When implementing IRM [3], V-Rex [38] and IB-IRM [2], we refer the implementations from DomainBed [23]. Since the environment information not available, we perform random partitions on the training data obtain two equally large environments for these objectives following previous works [8, 12J. Moreover, we select the weights for the corresponding regularization from (0.01, 0.1, 10, 100} for these objectives according to the validation performances of IRM and stick to it for others, since we empirically observe that they perform similarly with respect to the regularization weight choice. For EIIL [13], use the author-reeeased implementations about assigning different samples the weights for being put in each
environment and calculating the IRM loss.
Implementations of invariant graph learning methods. We implement GSAT [55], GREA [45], CAL [64], MolcOOD [77], GIL [41], DisC [17], and CIGA [8], according the author provided codes (if available).
GREA [45]: We use penalty weight of for GREA we empirically it does not affect the performance by changing to different weights.
Interpretable ratio: same as others;
Penalty weight: 1;
Number of environments: N/A;
GSAT [55]: We follow the recommendations of the released implementations by the authors.
Interpretable ratio: 70%;
Penalty weight: 1;
Decay ratio: 10%;
Decay interval: pretrain epoch//2;
Number of environments: N/A;
CAL [64]: We follow the recommendations of the released implemennations by the authors.
Interpretable ratio: same as others;
Penalty weight: {0.1, 0.5, 1.0)::
Number of enviromments: N/A;
MoleOOD [77]: We tune the penalty weights of MoleOOD with values from {le - 2, le-le- 10) but did not observe much performance differences. Hence stick the penalty weight as 1 for all datasets.
Interpretable ratio: N/A;
Penalty weight: 1;
Number of environments: same others;
GIL [41]: We follow the recommendations of the paper.
Interpretable ratio: same others;
Penalty weight: {le 5, le 3, le 1):
Number of environments: same as others;
DisC [17]: We tune only the weight from {0.9, 0.7, 0.5) in the GCE loss as we did not observe performance differences by changing the weight of the other terms.
Interpretable ratio: same others;
q weight: (0.9, 0.7, 0.5):
Number of environments: same as others;
CIGA [8]: We follow the recommendations the released implementations by the authors..
Interpretable ratio: same as others;
Penalty weight: {0.5, 2, 4, 8, 16, 32];
Number of environments: N/A;
32

---

GALA:
Interpretable ratio: same as others;
Penalty weight: {0.5, 1,2, 4, 16, 32);
Environment assistant: {vanilla GNN, XGNN);
Sampling proxy: (label predictions, cluster predictions):
Number of environments: same others;
All of the graph learning methods adopt an interpretable GNN the backbone by default. The only exception is MoleOOD, we follow the original implementation while using shared GNN encoder for the variational losses ensure the fairness of comparison. Besides, for DisC, we find the soft masking implementation in two-piece graphs will incur a severe performance degeneration hence we use ratio of 25% for the interpretable GNN backbone.
For environment inferring methods, we search the number environments
Two-piece graphs: fixed 3 (since there are 3 spurious graphs);
CMNIST-sp: 2 (since there are environments);
Graph-SST2: (2,3,4) 4) following previous practice [41];
DrugOOD datasets: {2, 3, 5, 10,20) 20} following previous practice [77].;
Implementations of GALA. For fair comparison, GALA uses the same GNN architecture for GNN encoders as the baseline methods. By default, we fix the temperature to be in the contrastive loss, and merely search the penalty weight of the contrastive loss from (0.5. 1, 2, 4, 8, 16, 32) according to the validation performances, following the CIGA implementations [8]. By default, we implement the environment assistant as a ERM model, and adopt directly the environment assistant predictions to sample possible and negative graph pairs. Nevertheless, as discussed in Sec. 4 that there could be multiple implementation choices for the environment assistant and the of its predictions. We hence also try with XGNN based environment assistant model and clustering based proxy predictions. By default, the selection of the environment assistant model performed via best training performance, as which encourages a better fit to the dominant subgraph patterns, while we also try the model selection with best validation performance in DrugOOD datasets and find empirically sometimes leads to better performance. All the options for the selection of the environment assistant models depend on the validation performance. For Two piece graphs, CC50-Scaffold, EC50-Size, Ki-Assay, Ki-Scaffold, CMNIST-sp and Graph-SST2, we find implementing the environment assistant ERM model already yield impressive improvements. While for the other DrugOOD datasets, we implement the environment assistant an interpretable GNN trained with ERM and cluster the learned graph representations of the model sample positive and negative pairs.
Since GALA imposes a strong regularization to the data that may hinder the learning of graph representations, pre-train the model by 10 epochs using ERM and then impose the GALA penalty implemented as one-side contrastive loss as discussed in Sec. F. When the numbers of positive and negative pairs are extremely imbalanced, we will upsample the minor groups by a factor of (2, 3, 4}. depending on the validation performance.
G.3 Software and Hardware
We implement our methods with Torch [59] and Py Torch Geometrie [18]. We ran our experiments on Linux Servers installed with V100 graphies cards and CUDA 10.2.
33

---

G.4 Computational analysis
Table 6: Averaged total training time of different methods.
Datasets
Two-picce graphs
EC50-Assay
CMNIST-sp
Graph-SST2
ERM
435.85#2.14
10.45=10.27
315.84
374.31
XGNN
673.82=0.81
126.65 17.307
#91.49
42
GREA
1128.28+34.37
210.30-21.22
902.06+8.499
979.15
GSAT
1205.67
228.88-25.04
791.55
949.57
DisC
1244.68:4.70
207.50-17.72
932.40116..00
1280.77:551.99
MoleOOD
714.0616.53
136.39-17.e7
439.4919.10
712.31
GIL.
533..6211.42
279.30=25.39
919.53:14.15
133.36=114.008
CIGA
873.49#16.21
167.63+1.10
650.94+5.01
792.10=59.
GALA-cluster
811.41
147.97=2.05
121.61
765.32220.66
GALA-pred
793.27=8.68
149.89+2.71
644.78253.58
764.69#30.99
We calculate the average total training time of different methods at various datasets seconds. As shown in Table. 6, the training of GALA (no matter with clustering based sampling or prediction based sampling) does not bring much additional overhead than its counterpart CIGA. When considering the additional training time of the assistant model with ERM, GALA costs only competitive training time as environment generation based methods such as GREA and DisC. Notably, some methods such DisC and GIL sometimes may be slow converge even with the same early stop setting, which will cost even more time than the time cost by GALA plus the ERM training. Besides, the ERM training time (for assistant model) is not much long and usually around Smins (or 300seconds in the table).
34


Reviews:
[{"review_id": "0mS8r6IbFZ", "reviewer": "Reviewer_W1GV", "summary": "The paper proposes Graph invAriant Learning Assistant (GALA), targeting at out-of-distribution (OOD) generalization on graphs leveraging invariant graph representation learning. Authors argue that environment augmentation cannot achieve this without additional assumptions, leading to the development of minimal assumptions—variation sufficiency and variation consistency. GALA utilizes an assistant model to distinguish variations in spurious subgraphs sensitive to environment changes or distribution shifts. Empirical experimental results support the efficacy of GALA.\n\n**Post Rebuttal**: I have read the rebuttal of the authors and will keep my rating.", "questions": "-  The assistant model that needs to be sensitive to graph environment changes, but how do you guarantee that?", "limitations": "The authors may not explicitly discuss the limitations of the proposed method.", "rating": 5, "confidence": 1, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- This paper provides a thorough theoretical explanation with step-by-step illustration. It builds a strong theoretical case for the proposed framework, establishing a set of minimal assumptions required for invariant graph learning. Practically GALA has a new design using an assistant model to differentiate variations in spurious subgraphs.\n\n- Empirically, authors present extensive experimental results on both synthetic and real-world data, showing the effectiveness of the proposed method.", "weaknesses": "- Considering that scalability has always been an issue for GNNs, it would be better if the authors can also discuss the overhead of further introducing a new assistant module. \n- The proposed method seems specialized in the OOD generalization of graph structures, whereas the feature OOD is also a major problem in the graph OOD setting. Considering this the applicability of this method may be restrained.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time in reviewing our paper and your positive feedback! We hope our response below would make you more confident about the evaluation of our work.\n\n> W1 Computational overhead.\n\nWe have revised our manuscript to include the analysis. The additional overhead of GALA compared to a vanilla GNN lies in the training of the environment assistant model and the sampling with the assistant model.\n\nIn the table below, we calculate the average total training time of different methods at various datasets in seconds. It can be found that:\n- The training of GALA (no matter with clustering based sampling or prediction based sampling) does not bring much additional overhead than its counterpart CIGA. \n- When considering the additional training time of the assistant model with ERM, GALA cost only a competitive training time as environment generation based methods such as GREA and DisC. Notably, some methods such as DisC and GIL sometimes may be slow to converge even with the same early stop setting, which will cost even more time than the time cost by GALA plus the ERM training.\n- Besides, the ERM training time (for a assistant model) is not much long and usually around **5mins** (or 300seconds in the table). \n\n| Training Time (seconds) | Two-piece graphs | EC50-Assay      | CMNIST-sp         | Graph-SST2       |\n|---------------|------------------|-----------------|-------------------|------------------|\n| ERM           | 435.85-+-2.14    | 80.45-+-10.27   | 315.84-+-5.55     | 374.31-+-1.28    |\n| GREA          | 1128.28-+-34.57  | 210.30-+-21.23  | 902.06-+-8.49     | 979.15-+-18.3114 |\n| GSAT          | 1205.67-+-62.54  | 228.88-+-25.04  | 791.55-+-15.67    | 949.57-+-97.68   |\n| DisC          | 1244.68-+-4.76   | 207.50-+-17.72  | 932.40-+-76.99    | 1280.77-+-551.97 |\n| MoleOOD       | 714.06-+-6.53    | 136.39-+-17.87  | 439.49-+-9.10     | 712.31-+-81.62   |\n| GIL           | 533.46-+-11.42   | 279.30-+-25.39  | 919.53-+-14.1487  | 733.36-+-147.08  |\n| CIGA          | 873.49-+-16.21   | 167.63-+-1.10   | 650.94-+-5.01     | 792.10-+-59.12   |\n| GALA-cluster  | 811.41-+-3.20    | 147.97-+-2.05   | 756.41-+-21.63    | 765.32-+-20.86   |\n| GALA-pred     | 793.27-+-8.58    | 149.89-+-2.71   | 644.78-+-53.58    | 764.69-+-30.98   |\n\n> W2 OOD generalization on node features.\n\nWe need to clarify that, despite our discussion of failure cases focusing on graph structures, GALA works for distribution shifts on both graph structures and features, as demonstrated in experiments. For example, the CMNIST-sp dataset is converted from colored MNIST dataset where the colors have a spurious correlation with labels, and hence CMNIST-sp contains distribution shifts in node features.\n\nWe are happy to answer if there are any outstanding questions, and we’d be grateful if you could jointly consider our responses above when making a final evaluation of our work.\n"}, {"review_id": "UeoBy2EkSP", "reviewer": "Reviewer_waKT", "summary": "Invariant graph learning methods aim to extract invariant patterns among different training environments for OOD generalization on graphs. A plethora of approaches focuses on augmenting environment information from training samples as the environment label is difficult to obtain. This work provides both theoretical and empirical analysis of the faithfulness of the existing environment augmentation methods for learning invariant graph representations. They demonstrate that it is impossible to learn invariant graph representations by environment augmentation without a set of minimal assumptions developed in the paper, including variation sufficiency and variation consistency. Based on these assumptions, they then propose GALA, Graph invAriant Learning Assistant, a new graph invariant learning framework via an assistant model providing proxy predictions to differentiate spurious variations. Experiments on 12 real-world graph OOD datasets demonstrate GALA's effectiveness.\n", "questions": "1. The 3 questions mentioned in the Weaknesses part.\n2. Since this work is proposed to tackle OOD generalization on graphs, I wonder if the authors made some attempts on OOD generalization on Euclidean data?\n", "limitations": "The authors discussed the broader impacts in the appendix.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. This paper is very well-written and organized. The problem studying (i.e., whether the existing environment augmentation methods are faithful for learning invariant graph representations) is significant, and the motivating synthetic experiments clearly demonstrate the pitfalls of existing environment augmentation.\n2. Although the variation sufficiency and variation consistency assumptions seem intuitive, this paper is the first to develop these minimal assumptions systematically and utilize them to derive the impossibility results for existing methods.\n3. GALA is technically sound and the authors provide theoretical analysis to demonstrate GALA is able to identify invariant subgraph patterns.\n4. Experiments are comprehensive and extensive. The results of proof-of-concept and ablation studies are consistent with the theoretical results derived in the paper, and the real-world experiments demonstrate a considerable improvement among all the baselines.\n5. The authors did a great job of citing related works (in the appendix).", "weaknesses": "1. The environment assistant model is trained with vanilla ERM on spuriousness-dominated data to identify spurious correlations. My major concern is the vanilla ERM model should learn both spurious correlations and invariant patterns, and then how to guarantee the environment assistant model only learns environment-dependent patterns? I wonder if the authors have tried other biased ERM models as backbone and it would be great if they could elaborate more on this question.\n2. In this paper, the authors divide the training sets into spuriousness-dominated and invariance-dominated sets. I wonder what if the whole training sets only contains spurious and invariance patterns distributed equally? \n3. The authors did not discuss the computational complexity of GALA. I wonder if they could discuss it in detail. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful comments and your support of our work! Please see our detailed responses to your comments and suggestions below where we use reference numbers in our manuscript due to the character limit.\n\n> W1 How to guarantee the environment assistant model only learns environment dependent patterns? What about other biased ERM models?\n\nWe need to clarify that **we do not intend to let the environment assistant model learn only the environment dependent patterns**. Instead, the assistant model is also allowed to learn invariant patterns. Since spurious patterns are more likely to change between environments than invariant patterns (variation sufficiency), as shown in Fig. 3a, once the assistant model learns the spurious patterns, these spurious patterns can be differentiated from invariant ones across $\\{G^p\\}$ and $\\{G^n\\}$ according to the correctness of the predictions yielded by the assistant model. \n\nIn experiments, we tried two types of ERM biased models: one is the vanilla GNN and the other is the interpretable GNN, where both of them bring improvements when incorporated into GALA. Besides, we also discussed more implementation choices in Appendix F. The results show that one could implement the assistant model to better fit the spurious patterns when additional knowledge is provided, which highlights a promising future direction as discussed in Appendix B.\n\n> W2 What about training sets with equally correlated spurious and invariant patterns?\n\nAlthough this case is gotten rid of by the variation consistency assumption for the impossibility of identifying an invariant subgraph, results in Fig. 5 show that GALA still brings certain improvements to previous methods. We hypothesize that there exist other factors, such as the size of spurious and invariant subgraphs, affect the fitting of spurious and invariant patterns, and consider discussing more factors other than correlation strength to the design of the assistant model as well as OOD generalization on graphs as a promising future direction.\n\n> W3 Computational complexity of GALA.\n\nWe have revised our manuscript to include the analysis. The additional overhead of GALA compared to a vanilla GNN lies in the training of the environment assistant model and the sampling with the assistant model.\n\nIn the table below, we calculate the average total training time of different methods at various datasets in seconds. It can be found that:\n- The training of GALA (no matter with clustering based sampling or prediction based sampling) does not bring much additional overhead than its counterpart CIGA. \n- When considering the additional training time of the assistant model with ERM, GALA cost only a competitive training time as environment generation based methods such as GREA and DisC. Notably, some methods such as DisC and GIL sometimes may be slow to converge even with the same early stop setting, which will cost even more time than the time cost by GALA plus the ERM training.\n- Besides, the ERM training time (for a assistant model) is not much long and usually around **5mins** (or 300seconds in the table). \n| Training Time (seconds) | Two-piece graphs | EC50-Assay      | CMNIST-sp         | Graph-SST2       |\n|---------------|------------------|-----------------|-------------------|------------------|\n| ERM           | 435.85-+-2.14    | 80.45-+-10.27   | 315.84-+-5.55     | 374.31-+-1.28    |\n| GREA          | 1128.28-+-34.57  | 210.30-+-21.23  | 902.06-+-8.49     | 979.15-+-18.3114 |\n| GSAT          | 1205.67-+-62.54  | 228.88-+-25.04  | 791.55-+-15.67    | 949.57-+-97.68   |\n| DisC          | 1244.68-+-4.76   | 207.50-+-17.72  | 932.40-+-76.99    | 1280.77-+-551.97 |\n| MoleOOD       | 714.06-+-6.53    | 136.39-+-17.87  | 439.49-+-9.10     | 712.31-+-81.62   |\n| GIL           | 533.46-+-11.42   | 279.30-+-25.39  | 919.53-+-14.1487  | 733.36-+-147.08  |\n| CIGA          | 873.49-+-16.21   | 167.63-+-1.10   | 650.94-+-5.01     | 792.10-+-59.12   |\n| GALA-cluster  | 811.41-+-3.20    | 147.97-+-2.05   | 756.41-+-21.63    | 765.32-+-20.86   |\n| GALA-pred     | 793.27-+-8.58    | 149.89-+-2.71   | 644.78-+-53.58    | 764.69-+-30.98   |\n\n> Q2 What if applying GALA to Euclidean data?\n\nAlthough GALA is primarily proposed to tackle OOD generalization on graphs, it can also be applied to Euclidean data, as Euclidean data can be considered as special graphs (while methods developed in Euclidean data can hardly be generalized to graphs as shown in our experiments). For example, one may convert Euclidean data into graphs and apply GALA to it. The experiments on CMNIST-sp and Graph-SST2 show promising results of GALA in tackling the distribution shifts on Euclidean data-converted graphs. Therefore, we believe the success of GALA on graphs can be generalized to Euclidean data as well and consider it as a promising future direction.\n\nPlease let us know if you have any further questions. We’d be grateful if you could take the above responses into consideration when making the final evaluation of and supporting our work!\n"}, {"review_id": "r0UhqV3tr0", "reviewer": "Reviewer_m4Xg", "summary": "This paper discusses the challenges of learning invariance among data from different environments for out-of-distribution generalization on graphs. The authors propose a new framework called Graph invAriant Learning Assistant (GALA) that incorporates an assistant model to differentiate the variations in spurious subgraphs. The paper also includes experiments on 12 datasets, including DrugOOD, with various graph distribution shifts to verify the superiority of GALA.", "questions": "I am interested in knowing if the authors believe there exists a connection between this work and the explainability literature. In particular, the counterfactual explainability subfield (like the work from [A], just to provide a reference).\n\n_[A] Lucic, Ana, et al. \"Cf-gnnexplainer: Counterfactual explanations for graph neural networks.\" International Conference on Artificial Intelligence and Statistics. PMLR, 2022._\n", "limitations": "-", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper tackles an interesting problem. Distributions shifts between training and test sets hinder and limit the potential applicability of graph neural networks methods into real world scenarios.\n- To the best of my knowledge, the method proposed in this paper is novel.\n- The paper is, mostly, well written and makes an effort to be easily understandable.\n- The paper presents valid theoretical arguments to back its main claims.\n- Finally, the paper presents an extensive sets of experiments that back the validity of the proposed method.\n", "weaknesses": "- In the tables that contain results, the best values are bolded. However, the presentation of the results would be more fair if two different colours/markers would be used when the best result is within the confidence interval of another result.\n\n- The paper would be easier to follow if, along figure 2, a plausible real world example was used.\n\n- In the last part of the introduction section the authors convey their contributions. In my opinion, this part is one of the less clear in the paper. To improve the clarity, I suggest the authors make a bullet list of their contributions. In particular, I would like the authors to distinguish their original contributions and how they differ from the work of Lin et al [36].\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive suggestions and your positive feedback! Please see our detailed responses to your comments and suggestions below.\n\n> W1 Using two markers wrt. confidence interval.\n\nWe revised our table and add shadows to highlight entries with the best performance and a confidence interval excluding the second best results (mean-1*std > mean of the second best result), which considers both the generalizability and stability of the best performed method. \n\nThe revised table is shown in the pdf attached in the \"global\" response, where we can find that the improvements of GALA is non-trivial on many datasets.\n\n> W2 A plausible real word example of Fig. 2.\n\nThe SCM of Fig. 2 is primarily motivated by the real-world drug property prediction task[B], where the property of a drug molecule can usually be characterized by a sub-molecule or functional group (e.g., `-OH` to solubility). We have added the discussion to the introduction of Fig.2.\n\n> W3 Listing contributions and comparing with [36]\n\nWe revised our paper and included a bullet list of our contributions:\n- We identify failure cases of existing invariant graph learning approaches and establish the minimal assumptions for feasible invariant graph learning; \n- We develop a novel framework GALA with provable identifiability of the invariant subgraph for OOD generalization on graphs under the assumptions; \n- We conduct extensive experiments to verify both our theoretical results and the superiority of GALA;\n\nNotably, both our theory and solution differ from Lin et al. [36] fundamentally in terms of **assumptions** and **methodology**. More specifically, we do not rely on the auxiliary information and are compatible with the existing interpretable and generalizable GNN architecture for OOD generalization on graphs.\n\n> Q1 Connection with the explainability literature.\n\nThanks for pointing out the related work [A]. In fact, there is a deep connection between interpretability and generalizability [6,44]. As characterized by the SCMs (Fig. 2), only extracting and using the underlying invariant subgraph for the prediction can generalize under various graph distribution shifts (Proposition E.2 in [6]), meanwhile the extracted subgraph can also serve as the interpretable result for the prediction.\n\nBoth [A] and our work are inspired to incorporate causality in finding faithful interpretable subgraphs, while [A] focuses on *post-hoc explainability* given a trained GNN and ours focuses on *intrinsic interpretable GNN*. We have included the discussion of [A] in our revised draft.\n\nWe’d appreciate it if you could take the above responses into consideration when making the final evaluation of and supporting our work. Please let us know if there are any outstanding questions.\n\n**References**\n\n[6] Learning causally invariant representations for out-of-distribution generalization on graphs, NeurIPS'22.\n\n[36] ZIN: When and how to learn invariance without environment partition? NeurIPS'22.\n\n[A] Cf-gnnexplainer: Counterfactual explanations for graph neural networks, AISTATS’22.\n\n[B] The rise of fragment-based drug discovery, Nature chemistry, 1(3):187, 2009.\n"}, {"review_id": "1dHizkCiG7", "reviewer": "Reviewer_3cig", "summary": "This paper investigates OOD problems on graphs without environmental information. The authors show that it is impossible to learn invariant graph representations without additional assumptions. The authors propose a novel graph OOD method namely GALA based on a new assumption. An auxiliary model is used to make proxy predictions and then the proxy predictions are used to help identify invariant subgraphs. Experiments on 12 datasets show the effectiveness of the proposed method.  ", "questions": "Q1. I am confused about how to get the Oracle performance on these datasets. Can the authors provide more explanations?\n\nQ2. I don't understand why Disc fails on all the datasets except CMNIST-sp. Besides, on CMNIST-sp, only GALA and Disc have some improvement over ERM. \n", "limitations": "The authors didn't discuss the limitations of the proposed method. ", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "S1. The authors clearly demonstrate that the identifiability of causal subgraphs is ill-posed without environment labels. Thereby, the authors give A3.3 and A3.5 and then propose GALA based on these two assumptions.\n\nS2. Both theoretical and experimental analyses are used to support the claim. ", "weaknesses": "W1. This is a math-heavy paper, but the notations are unclear to me, thus making it hard to read. For example, in lines 78-79, should there be something like $e \\in \\mathcal{E}_{all}$? In Equation 1, there is no definition about $\\widehat{G}_c$. There are more examples, I will not list them all.\n\nW2. I think GALA can be used as a general OOD method. Why do the authors focus on graph data? I don't see any specific design only for graphs. Why not test GALA on other data types (ex., images)?\n\nW3. I don't understand why ERM still learns the spurious correlations in invariance dominated cases. Can this be justified theoretically? ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time in reviewing our work. Please see our detailed responses to your comments and suggestions below where we use reference numbers in our manuscript due to the character limit.\n\n> W1 Notations\n\nWe understand it might be difficult for readers who are not familiar with the literature to follow our notations and hence provide a table of key notations in Appendix A for your reference.\n\nWe have revised our draft to include $e\\in\\mathcal{E}\\_{all}$ in line 78. $\\widehat{G}_c$ is the estimated subgraph, and we have added its definition before Eq 1, which is the output of the subgraph extractor. \n\nWe also double checked and corrected the other parts of our draft.\n\n> W2 Why focusing on graph data?\n\nFirst, we need to clarify that, as pointed out in CIGA[6], OOD generalization on graphs is fundamentally more challenging than that in Euclidean data. Given that graph is a generalized form of Euclidean data (e.g., images, languages), **methods developed on graphs usually are not limited to graphs and can be applied directly to Euclidean data, but methods developed in Euclidean data can hardly be applied directly to graphs**.\n\nThe specialization of GALA for graph data mainly lies in three aspects:\n\n- **The study of GALA is more desired by graph data**. Different from Euclidean data, all OOD methods on graph data [6,14,35,38,44,55,56,60,64] do not assume the availability of environment labels nor the auxiliary information [36]. Given the faithfulness of environment augmentation has never been discussed on graphs, the study of faithfulness is more desired by graph data.\n- **Graph data is also a suitable testbed for GALA**. It can be found that, on the two-piece graph datasets that satisfy the theory well, GALA outperforms other methods by a large margin up to $30$% and $13%$% in average, owing to its theoretical guarantee.\n- **GALA inherits the graph specific design from existing works** such as CIGA[6] and GSAT[44] that use information theory to handle the complexity of graph data and an interpretable GNN architecture to address graph distribution shifts modeled by the SCM in Fig. 2. Extracting the invariant subgraph is the central task specific for OOD generalization on graphs. Therefore, all of our theory and solution concentrates on the identifiability of the invariant subgraph.\n\n> W3 Why ERM still learns spurious correlations in invariance-dominated cases?\n\nIt is theoretically proved in [8]. Intuitively, ERM will learn **any correlations presented in the data**, no matter whether the correlation is invariant or spurious. If we look into the first two columns in Table 1, it can be found that **GNNs trained with ERM still have a large gap from the oracle models**, meaning that GNNs trained with ERM still learn the spurious correlations. \n\n> Q1 Oracle performance\n\nThe Oracle performance is obtained via canceling the distribution shifts among training and test data by randomly re-splitting the whole dataset, following [6]. \n\n> Q2 Why DisC fails on all datasets except for CMNIST-sp and why only GALA and DisC succeed in CMNIST-sp?\n\nCMNIST-sp is converted from the ColoredMNIST dataset[3] which contains severe spuriousness dominated shifts in node features. DisC is designed to tackle this kind of severe distribution shifts in CMNIST-sp via an **aggressive strategy** to discard the dominant patterns [14,46], and additionally incorporate environment generation procedure based on the aggressively separated subgraphs. However, the aggressive strategy in DisC may affect learning useful patterns for generalization and fails when the distribution shifts are mild. \n\nOther methods fail in CMNIST-sp because of various drawbacks. For example, GREA, CIGA can only handle invariance dominated cases. In contrast, GALA succeeds because of its theoretical guarantee in handling both invariance and spuriousness dominated cases.\n\n> L1 Limitation discussion.\n\nWe actually provided the discussion of limitations and future works in Appendix B.\n\nPlease let us know if there are any further questions. We’d appreciate it if you could take the above responses into consideration when making the final evaluation of our work. \n"}, {"review_id": "DAwgs2Fr8b", "reviewer": "Reviewer_voKo", "summary": "The paper proposes a new framework called Graph invAriant Learning Assistant (GALA) to address the problem of learning invariance among data from different environments for out-of-distribution generalization on graphs. The authors argue that the performance of graph neural networks (GNNs) can be seriously degraded by graph distribution shifts caused by mismatches between the training and test graph distributions. GALA incorporates an assistant model to differentiate the variations in spurious subgraphs and identify the underlying invariant subgraphs. The authors conduct extensive experiments on 12 datasets with various graph distribution shifts and show that GALA brings improvements in multiple graph datasets. ", "questions": "1. I have difficulty understanding the expression $G_{c}:=f_{gen}^{G_c}(Y^{e}\\cdot Rad(\\alpha_{e}))$ on line 125. It seems to me that $G_{c}$ is determined purely by the label and not influenced by randomness ($Rad(\\alpha_{e})$). Could you please clarify this for me?\n2. In your discussion, you mentioned that when $\\hat{G}_c$ contains parts of $G_s$, assigning the same labels to the generated graph is more likely to reinforce the spurious correlation between $G_s$ and $Y$. I would appreciate it if you could provide a detailed explanation of how your method effectively prevents this pitfall.\n3. On line 179, you stated that \"Any learning algorithm will fail in either $E_{tr}$ or $E_{tr}$\". Could you please explain this sentence using clearer and more precise mathematical language?\n4. In Figure 1, the last sentence mentions \"pulling graphs with the same label but from $\\{G_{p}\\}$ and $\\{G_{n}\\}$ closer in the latent space.\" However, the figure seems to depict \"pulling graphs with different labels but from $\\{G_{p}\\}$ and $\\{G_{n}\\}$ closer in the latent space.\" Is this a typographical error?\n5. I am curious about your final loss function as mentioned in Algorithm 1. Equation 9 does not seem to be optimized directly. Could you please clarify what the final loss function is?\n", "limitations": "The writing in the paper is somewhat ambiguous, leading to occasional confusion. For more specific details, please refer to the sections labeled 'Questions' and 'Weaknesses' mentioned earlier.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. Originality: The paper proposes a novel framework, Graph invAriant Learning Assistant (GALA), that introduces a new assistant model to differentiate the variations in spurious subgraphs and identify the underlying invariant subgraphs. \n2. Quality: The authors do extensive experiments to validate the effectiveness of the proposed approach. They also provide a detailed analysis of the results and compare their approach with several state-of-the-art methods.\n3. Clarity: The paper explanations of the key concepts, the methodology, and detailed experimental results. \n4. Significance: The paper's contributions are significant because they provide a new perspective on the problem of learning invariant graph representations and offer a practical solution that can be applied to real-world problems. \n", "weaknesses": "1.\tIt would be preferable to provide the definition of the problem before presenting your proposed method.\n2.\tThe novel concepts introduced in your proposal would benefit from a more detailed explanation. For instance, the concept of 'two-piece graphs' is somewhat confusing. It would be helpful if you could elaborate on why you have incorporated this concept and provide additional motivation.\n3.\tIt appears that Assumption 3.5 may not appropriately be considered as an assumption. Instead, you could simply explain that $H(C|Y)\\neq H(S|Y)$. Treating \"either $H(C|Y)>H(S|Y)$ holds or $H(C|Y)<H(S|Y)$ holds\" as an assumption may seem unusual in this context.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time in reviewing our work. Please see our detailed responses to your comments and suggestions below where we use reference numbers in our manuscript due to the character limit.\n\n> W1 Problem definition\n\nWe complemented the problem definition in Sec. 2:\n\nThe goal of OOD generalization on graphs is to find a GNN $f$ that minimizes the maximal loss among all environments, i.e., to minimize $\\max_{e\\in \\mathcal{E}_\\text{all}} R^e$, where $R^e$ is the risk of $f$ under environment $e$.\n\n> W2 Motivation of ``two-piece’’ graphs.\n\nWe have revised our paper to include the explanation for why adopting the two-piece graphs. The main reason is that the two-piece graph is the **minimal implementation** of PIIF distribution shifts in Fig. 2, which is more suitable for theoretical discussions. Specifically, PIIF distribution shifts in Fig. 2 model the invariant and spurious correlations via the correlation of subgraphs with the labels. For simplicity, we restrict the types of invariant and spurious subgraphs and use a simple flipping operation to simulate the invariant/spurious correlations, motivated by [25] that is used to discuss the failures of IRM[3].\n\n> W3 Assumption 3.5\n\nWe have revised it following your suggestion.\n\n> Q1 Understanding line 125.\n\nBy definition, the Rad function will flip the labels with a probability $\\alpha$. Considering a $f^{G_c}_{gen}$ that maps +1 to a ``house`` subgraph and -1 to a ``cycle`` subgraph. Let the sampled $Y^e=+1$, and $\\alpha=0.1$, then with a probability of 10%, $G_c$ will be a ``cycle`` and a probability of 90%, $G_c$ will be a ``house``.\n\n> Q2 How GALA addresses the failure of environment generation?\n\nIn fact, the faithfulness of environment generation **relies on the success of partitioning the invariant and spurious subgraphs**, when without additional knowledge. If the model already learns to partition the graph, then it is unnecessary to generate any new environments. Otherwise, environment generation can hardly avoid generating harmful environments for OOD generalization on graphs, as shown in Proposition 3.2.\n\nTherefore, **GALA avoids the failures of environment generation by not performing environment generation**. Instead, GALA aims to identify the maximal invariant subgraph among the partitioned subsets by the assistant model.\n\n> Q3 Explanation of Proposition 3.4.\n\nAny invariant graph learning algorithm accepts the input of the joint distribution $P(Y,G)$ and aims to identify an invariant subgraph $\\widehat{G}_c$. Given the same $P(Y,G)$, the algorithm will stick to the same $\\widehat{G}_c$:\n\nIf $\\mathcal{E}\\_\\text{tr}$ and $\\mathcal{E}\\_\\text{tr}’$ have the same $P(Y,G)$ but different underlying invariant subgraphs ($G_c$, $G_c’$ respectively, with $G_c\\neq G_c’$). If the algorithm succeeds in one environment, for example, $\\mathcal{E}\\_{tr}$, and outputs $\\widehat{G}\\_c=G_c$ given the input $P(Y,G)$. Then, the algorithm will fail in $\\mathcal{E}\\_{tr}’$ as $\\widehat{G}_c=G_c\\neq G_c’$.\n\n> Q4 Caption of Figure 1.\n\nWe need to clarify that **the black colored labels are the predictions yielded by the assistant model instead of the ground truth labels (in orange color)**. According to the correctness of the assistant model, we can partition the whole dataset into positive graphs (i.e., correctly predicted samples in the top left and the bottom right) and negative graphs (i.e., incorrectly predicted samples in the top right and the bottom left). GALA will pull graphs with the same ground truth label but from positive and negative graphs closer. \nWe revised the caption to highlight that the labels are graph labels in the figure.\n\n> Q5 The final loss.\n\nEq. 9 is implemented in a contrastive loss. We provide the specific loss in Eq. 18 of Appendix F. We revised the text around Eq. 9 to clarify this point.\n\nPlease let us know if you have any further questions. We’d be grateful if you could consider our responses above when making final evaluations of our work.\n\n"}]

Author rebuttal:
Dear reviewers,


Thank you for your comments and suggestions. To summarize, all reviewers agree that our theoretical analysis is clear, significant, and interesting. The motivated solution GALA is novel (waKT, m4Xg, voKo, voKo). The empirical study is comprehensive and extensive to demonstrate the superiority of GALA  (W1GV, waKT, m4Xg, voKo). The presentation is clear, well-written and easily understandable (waKT, m4Xg, voKo).

Regarding the reviewers’ concerns, we believe they can be addressed and we have revised our draft according to the reviewers’ valuable suggestions.  In the following, we give our responses to the reviewers’ concerns and suggestions. We will use reference numbers in our manuscript due to the character limit:

Regarding the specialization of GALA for graph data:
- First, we clarify that, as pointed out in CIGA[6], OOD generalization on graphs is fundamentally more challenging than that in Euclidean data. Given that graph is a generalized form of Euclidean data (e.g., images, languages), **methods developed on graphs usually are not limited to graphs and can be applied directly to Euclidean data, but methods developed in Euclidean data can hardly be applied directly to graphs**.
- The necessity and specialization of GALA for graph data originate from the following three aspects:
   - The study of GALA is more desired by graph data. Different from Euclidean data, all OOD methods on graph data [6,14,35,38,44,55,56,60,64] do not assume the availability of environment labels nor the auxiliary information [36]. **Given the faithfulness of environment augmentation has never been discussed on graphs, the study of faithfulness is more desired by graph data**.
   - **Graph data is also a suitable testbed for GALA**. It can be found that, on the two-piece graph datasets that satisfy the theory well, GALA outperforms other methods by a large margin up to $30$% and $13%$% in average, owing to its theoretical guarantee.
  - **GALA inherits the graph specific design from existing works** such as CIGA[6] and GSAT[44] that use information theory to handle the complexity of graph data and an interpretable GNN architecture to address graph distribution shifts modeled by the SCM in Fig. 2. Extracting the invariant subgraph is the central task specific for OOD generalization on graphs. Therefore, all of our theory and solution concentrates on the identifiability of the invariant subgraph.

Regarding the computational overhead, we calculated the average total training time of different methods at various datasets in seconds. It can be found that:
- The training of GALA (no matter with clustering based sampling or prediction based sampling) does not bring much additional overhead than its counterpart CIGA. 
- When considering the additional training time of the assistant model with ERM, GALA cost only a competitive training time as environment generation based methods such as GREA and DisC. Notably, some methods such as DisC and GIL sometimes may be slow to converge even with the same early stop setting, which will cost even more time than the time cost by GALA plus the ERM training.
- Besides, the ERM training time (for a assistant model) is not much long and usually around **5mins** (or 300seconds in the table). 
| Training Time (seconds) | Two-piece graphs | EC50-Assay      | CMNIST-sp         | Graph-SST2       |
|---------------|------------------|-----------------|-------------------|------------------|
| ERM           | 435.85-+-2.14    | 80.45-+-10.27   | 315.84-+-5.55     | 374.31-+-1.28    |
| GREA          | 1128.28-+-34.57  | 210.30-+-21.23  | 902.06-+-8.49     | 979.15-+-18.3114 |
| GSAT          | 1205.67-+-62.54  | 228.88-+-25.04  | 791.55-+-15.67    | 949.57-+-97.68   |
| DisC          | 1244.68-+-4.76   | 207.50-+-17.72  | 932.40-+-76.99    | 1280.77-+-551.97 |
| MoleOOD       | 714.06-+-6.53    | 136.39-+-17.87  | 439.49-+-9.10     | 712.31-+-81.62   |
| GIL           | 533.46-+-11.42   | 279.30-+-25.39  | 919.53-+-14.1487  | 733.36-+-147.08  |
| CIGA          | 873.49-+-16.21   | 167.63-+-1.10   | 650.94-+-5.01     | 792.10-+-59.12   |
| GALA-cluster  | 811.41-+-3.20    | 147.97-+-2.05   | 756.41-+-21.63    | 765.32-+-20.86   |
| GALA-pred     | 793.27-+-8.58    | 149.89-+-2.71   | 644.78-+-53.58    | 764.69-+-30.98   |


We have also revised our manuscript and improved its clarity thanks to the suggestions of all reviewers.

In addition, we have provided a link to our codes for reproducing the results in our paper to the AC.

We hope our previous responses could address the concerns of all reviewers. Otherwise, we are happy to discuss and provide further evidence for any outstanding questions. We would appreciate it if you could take our responses into consideration when making the final evaluation of our work.

Sincerely,
Authors


Meta-review: The paper introduces the Graph invAriant Learning Assistant (GALA), a novel framework for learning invariant graph representations. The paper proves the limitation of existing environment augmentation, and establish minimal assumptions—variation sufficiency and variation consistency—for invariant graph learning. They introduce an assistant model to distinguish variations in spurious subgraphs, thereby enhancing robustness against environmental shifts.  Empirical results, both synthetic and real-world, convincingly validate the efficacy of GALA. The paper is well-structured, clear, and supported by comprehensive experiments and theoretical analyses. Overall, this work advances the understanding of OOD in graph neural networks and should be accepted.
Decision: Accept (poster)