Title: Does Invariant Graph Learning via Environment Augmentation Learn Invariance?
Venue: NeurIPS 2023
Accepted: yes
Keywords: Graph Neural Networks

Out-of-Distribution Generalization

Invariant Learning
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?

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 G raph inv A riant L earning A ssistant ( 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 12 datasets including DrugOOD with various graph distribution shifts confirm the effectiveness of GALA 1 .

1 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 (iid.), 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-Distribution (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 to identify the underlying invariant subgraph of the input graph, which shares an invariant correlation with the target labels across multiple graph

Cycle

House

{

G

n

}

{

G

p

}

'Cycle'

'House'

Training Data

Proxy Predictions

'House'

'Cycle'

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 2 ) of the augmented environment information. For example, if the generated environments or inferred environment labels induce a higher bias or noise, it would make the learning of graph invariance even harder. Although it looks appealing to 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 a 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 a set of minimal assumptions on the underlying environments for feasible invariant graph learning. Based on a 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 the generated environments can even lead to a worse OOD performance. The failure of faithful environment generation implies the necessity of variation sufficiency (Assumption 3.3). Moreover, even 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 at 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 a consistent relationship in the correlation strengths with the labels.

To resolve the OOD generalization challenge under the established assumptions, we propose a new framework G raph inv A riant L earning A ssistant ( 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 not require explicit environment labels but merely proxy predictions to differentiate the variations in the spurious

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 { G p } and a negative { G n } , according to 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 p } and { G n } . Then, only the invariant subgraph holds an invariant correlation with the label among { G p } and { G n } , and hence can be identified by extracting the subgraphs that maximize the intra-class subgraph mutual information among { G p } and { G n } (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 on graphs. This work focuses on graph classification, while the results generalize to 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 e } e ∈E all collected from multiple environments E all . Samples ( G e i , Y e i ) ∈ D e from the environment e are drawn independently from an identical distribution P e . The goal of OOD generalization on graphs is to find a GNN f that

S

C

G

c

G

s

G

(a)

G

E

S

Y

G

C

(b) FIIF SCM

E

S

Y

G

C

(c) PIIF SCM

minimizes the maximal loss among all environments, i.e., to minimize max e ∈E all R e , where R e 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 . C and S control Y and G by controlling the generation of the underlying invariant subgraph G c and spurious subgraph G s , respectively. Since S can be affected by the environment E , the correlation between Y and G s 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 ⊥ ⊥ S | C , and Partially Informative Invariant Features ( PIIF ) when Y ̸⊥ ⊥ S | C .

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 c to predict the label Y [8, 72]. Specifically, the goal of OOD generalization on graphs is to learn an invariant GNN f := f c ◦ g , which is composed of: a) a featurizer g : G → G c that estimates the invariant subgraph ̂ G c ; b) a classifier f c : G c →Y that predicts the label Y based on the extracted ̂ G c , where G c refers to the space of subgraphs of G . The learning objectives of f c and g are formulated as

Since E is not observed, many strategies are proposed to impose the independence of ̂ G c and E . A prevalent approach is to augment the environment information. Based on the estimated invariant

subgraphs ̂ G c and spurious subgraphs ̂ G s , Liu et al. [45], Wu et al. [72, 73] propose to generate new environments, while Li et al. [41], Yang et al. [77] propose to infer the underlying environment labels. However, we show that they all fail to augment faithful environment information in Sec. 3.

Besides, Miao et al. [55, 56], Yu et al. [81, 82, 83] adopt graph information bottleneck to tackle FIIF graph shifts, but they cannot generalize to PIIF shifts, while Our work focuses on PIIF shifts as it is more challenging when without environment labels [43]. Fan et al. [17] generalize [40] to 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 H ( S | Y ) > H ( C | Y ) . In practice, as it 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 are 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 a self-boosting framework based on the estimated invariant and variant features. Deng et al. [14], Liu et al. [44], Pezeshki et al. [61], Zhang et al. [85] propose to 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 a 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 to generate faithful environments or 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 e is defined with two parameters, α e , β e ∈ [0 , 1] , and the dataset ( G e , Y e ) ∈ D e is generated as follows:

(a) Sample Y e ∈ {-1 , 1 } uniformly;

(b) Generate G c and G s via : G c := f G c gen ( Y e · Rad ( α e )) , G s := f G s gen ( Y e · Rad ( β e )) , where f G c gen , f G s gen map the input {-1 , 1 } to a corresponding graph selected from a given set, and Rad ( α ) is a random variable taking value -1 with probability α and +1 with 1 -α ;

(c) Synthesize G e by randomly assembling G c and G s : G e := f G gen ( G c , G s ) .

We denote an environment e with ( α, β e ) for simplicity. Different environments will have a different β e , thus P ( Y | G s ) will change across different environments, while P ( Y | G c ) remains invariant.

3.1 Pitfalls of environment generation

We begin by discussing the cases where there are few environments, and generating new environments is necessary [45, 72, 73]. Environment generation aims to provide some additional 'virtual' environments E v such that the invariant subgraph can be identified via applying an OOD risk to the joint dataset with the augmented data D v tr = {D e | e ∈ E tr ∪ E v } .

The generation of 'virtual' environments is primarily based on the intermediate estimation of the invariant and spurious subgraphs, denoted as ̂ G c and ̂ G s , respectively. Liu et al. [45], Wu et al. [73] propose DIR and GREA to construct new graphs by assembling ̂ G c and ̂ G s from different graphs. Specifically, given n samples { G i , Y i } n i =1 , 3 the new graph samples in E v is generated as follows:

which generates a new environment E v with n 2 samples. Although both DIR and GREA gain some empirical success, the faithfulness of E v remains questionable, as the generation is merely based on

inaccurate estimations of the invariant and spurious subgraphs. Specifically, when ̂ G c contains parts of G s , assigning the same labels to the generated graph is more likely to strengthen the spurious correlation between G s and Y . For example, when the model yields a reversed estimation, i.e., ̂ G c = G s and ̂ G s = G c , the generated environment will destroy the invariant correlations.

Proposition 3.2. Consider the two-piece graph dataset E tr = { ( α, β 1 ) , ( α, β 2 ) } with α ≥ β 1 , β 2 (e.g., E tr = { (0 . 25 , 0 . 1) , (0 . 25 , 0 . 2) } ), and its corresponding mixed environment E mix tr = { ( α, ( β 1 + β 2 ) / 2) } (e.g., E mix tr = { (0 . 25 , 0 . 15) } ). When ̂ G c = G s and ̂ G s = G c , it holds that the augmented environment E v is also a two-piece graph dataset with

The proof is given in Appendix E.1. This also extends to the adversarial augmentation [72, 83], which will destroy the actual ̂ G c . As both DIR and GREA adopt the same environment generation procedure, we verify the failures of environment generation with GREA in Table 2 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, it is fundamentally impossible to identify the underlying invariant graph from the spurious subgraph. More formally, if ∃ G s , such that P e 1 ( Y | G s ) = P e 2 ( Y | G s ) for any e 1 , e 2 ∈ E tr, where P e ( Y | G s ) is the conditional distribution P ( Y | G s ) under environment e ∈ E all, it is impossible for any graph learning algorithm to identify G c . We provide a formal discussion in Appendix E.2. The failure implies a fundamental requirement that E tr should uncover all the potential variations in the spurious subgraph.

̸

Assumption 3.3. (Variation sufficiency) For graphs generated following Fig. 2, for any G s , ∃ e 1 , e 2 ∈ E tr, such that P e 1 ( Y | G s ) = P e 2 ( Y | G s ) , and P e 1 ( Y | G c ) = P e 2 ( Y | G c ) .

Assumption 3.3 aligns with the definition of invariance [8, 30] that the invariant subgraph G c is expected to satisfy P e 1 ( Y | G c ) = P e 2 ( Y | G c ) for e 1 , e 2 ∈ E all . If there exists G s satisfying the invariance condition as well, then it is impossible to tell G c from G s even with environment labels.

3.2 Pitfalls of environment inferring

Although environment sufficiency (Assumption 3.3) relieves the need for generating new environments, is it possible to infer the underlying environment labels via approaches such as MoleOOD [77] and GIL [41], to facilitate invariant graph learning? Unfortunately, we find a negative answer.

Considering the two-piece graph examples E tr = { (0 . 2 , 0 . 1) , (0 . 2 , 0 . 3) } , when given the underlying environment labels, it is easy to identify the invariant subgraphs from spurious subgraphs. However, when the environment labels are not available, we have the mixed data as E tr = { (0 . 2 , 0 . 2) } , where P ( Y | G c ) = P ( Y | G s ) . The identifiability of G s is ill-posed , as it does not affect the E tr even if we swap G c and G s . More formally, considering the environment mixed from two two-piece graph environments { ( α, β 1 ) } and { ( α, β 2 ) } , then we have E tr = { ( α, ( β 1 + β 2 ) / 2 } . For each E tr, we can also find a corresponding E tr ′ = { (( β ′ 1 + β ′ 1 ) / 2 , α ′ ) } with { ( β ′ 1 , α ′ ) } and { ( β ′ 2 , α ′ ) } . Then, let

We now obtain E tr and E tr ′ which share the same joint distribution P ( Y, G ) while the underlying G c is completely different. More generally, we have the following proposition.

Proposition 3.4. There exist 2 two-piece graph training environments E tr and E tr ′ that share the same joint distribution P ( Y, G ) . Any learning algorithm will fail in either E tr or E tr ′ .

The proof is given in Appendix E.3. The experiments in Sec. 5 validate that both MoleOOD 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 α (i.e., H ( Y | G c ) ) and β e (i.e., H ( Y | G s ) ).

̸

Assumption 3.5. (Variation consistency) For all environments in E tr , H ( C | Y ) = 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 a new environment with the same variation strength relationships.

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 a 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 unidentifiable case. However, such auxiliary information is often unavailable and expensive to obtain on 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 one of them while lacking additional inductive biases will result in the 'no free lunch' dilemma [71] and suffer from the unidentifiability issue.

Corollary 3.6. (No Free Graph OOD Lunch) Without Assumption 3.3 or Assumption 3.5, there does not exist a 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 straightforward and given in Appendix E.4. Assumption 3.3 and Assumption 3.5 establish the minimal premises for identifying the underlying invariant subgraphs. However, it also raises new challenges, as shown in Table. 1. Chen et al. [8]

propose CIGA to maximize the intra-class mutual information of the estimated invariant subgraphs to tackle the case when H ( C | Y ) < H ( S | Y ) . While for the case when H ( S | Y ) < H ( C | Y ) , 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 the relations between H ( C | Y ) and H ( S | Y ) . We verify the failures of DisC and CIGA in Table. 2. The failure thus raises a challenging question:

Given the established minimal assumptions, is there a unified framework that tackles both cases when H ( C | Y ) < H ( S | Y ) and H ( C | Y ) > H ( S | Y ) ?

4 Learning Invariant Graph Representations with Environment Assistant

We give an affirmative answer by proposing a new framework, GALA : G raph inv A riant L earning A ssistant, which adopts an assistant model to provide proxy information about the environments.

4.1 Learning with An Environment Assistant

Intuitively, a straightforward approach to 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 to extend (which is empirically verified in Sec. 5), we are motivated to extend the framework of CIGA to resolve the case when H ( S | Y ) < H ( C | Y ) .

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

where ̂ G s c = g ( G s ) and G s ∼ P ( G | Y ) , i.e., ̂ G is 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 c maximizes the mutual information of subgraphs from any two environments, i.e., ∀ e 1 , e 2 ∈ E all ,

where ̂ G e 1 c and ̂ G e 2 c are the estimated invariant subgraphs corresponding to the same latent causal variable C = c under the environments e 1 , e 2 , respectively. Since C is not observable, CIGA adopts Y as a proxy for C , as when H ( S | Y ) > H ( C | Y ) , G c maximizes I ( ̂ G e 1 c ; ̂ G e 2 c | Y ) and thus I ( ̂ G c ; ̂ G s c | Y ) . However, when H ( S | Y ) < H ( C | Y ) , the proxy no longer holds. Given the absence of E , simply maximizing intra-class mutual information favors the spurious subgraph G s instead, i.e.,

Invalidating spuriousness dominance. To mitigate the issue, we are motivated to find a new proxy that samples ̂ G c for Eq. 5, while preserving only the G c as the solution under both cases.

To begin with, we consider the case of H ( S | Y ) < H ( C | Y ) . Although the correlation between G s and Y dominates the intra-class mutual information, Assumption 3.3 implies that there exists a subset of training data where P ( Y | G s ) varies, while P ( Y | G c ) 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 s c invalidates the dominance of G s . Denote the subset as { ̂ G n c } , then

where ̂ G p c ∈ { ̂ G p c } is sampled from the subset { ̂ G p c } dominated by spurious correlations, while ̂ G n c ∈ { ̂ G n c } is sampled from the subset { ̂ G n c } where spurious correlation no long dominates, or is dominated by invariant correlations. We prove the effectiveness of Eq. 6 in Theorem 4.1.

Environment assistant model A . To find the desired subsets { ̂ G p c } and { ̂ G n c } , inspired by the success in tackling spuriousness-dominated OOD generalization via learning from a biased predictors [40, 44, 58, 85], we propose to incorporate an assistant model A that is prone to spurious correlations. Simply training A with ERM using the spuriousness-dominated 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 I ( ̂ A ( G ); Y ) , we have

̸

Reducing to invariance dominance case. After showing that Eq. 6 resolves the spuriousness dominance case, we still need to show that Eq. 6 preserves G c as the only solution when H ( S | Y ) > H ( C | Y ) . Considering training A with ERM using the invariance-dominated data, A will learn both invariant correlations and spurious correlations [10, 17]. Therefore, { ̂ G n c } switches to the subset that is dominated by spurious correlations, while { ̂ G p c } switches to the subset dominated by invariant correlations. Then, Eq. 6 establishes a lower bound for the intra-class mutual information, i.e.,

where ̂ G p c ∈ { ̂ G p c } , ̂ G n c ∈ { ̂ G n c } , and ̂ G c , ̂ G s c are the same as in Eq. 3. The inequality in Eq. 8 holds as any subgraph maximizes the left hand side can also be incorporated in right hand side, while the sampling space of ̂ G c and ̂ G s c in the right hand side (i.e., both ̂ G c and ̂ G s c 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 G c as the solution for the featurizer g . We verify the correctness of Eq. 6 and Eq. 8 in Fig. 3(a).

4.2 Practical implementations.

The detailed algorithm description of

Algorithm 1 GALA : G raph inv A riant L earning A ssistant

1: Input: Training data D tr ; environment assistant A ; featurizer GNN g ; classifier GNN f c ; length of maximum training epochs e ; batch size b ;

2: Initialize environment assistant A ;

3: for p ∈ [1 , . . . , e ] do

4: Sample a batch of data { G i , Y i } b i =1 from D tr ;

6: for each sample G i , y i ∈ { G i , Y i } b i =1 do

5: Obtain Environment Assistant predictions { ˆ y e i } b i =1 ;

7: Find positive graphs with same y i and different ˆ y e i ;

8: Find negative graphs with different y i but same environment assistant prediction ˆ y e i ;

9: Calculate GALA risk via Eq. 9;

10: Update f c , g via gradients from GALA risk;

11: end for

12: end for

13: return final model f c ◦ g ;

GALA is shown as in Algorithm 1. In practice, the environment assistant can have multiple implementation choices so long as it is prone to distribution shifts. As discussed in Sec. 4.1, ERM trained

model can serve as a reliable environment assistant, since ERM tends to learn the dominant features no matter whether the features are invariant or spurious. For example, when H ( S | Y ) < H ( C | Y ) , ERM will first learn to use spurious subgraphs G s to make predictions. Therefore, we can obtain { G p } by finding samples where ERM correctly predicts the labels, and { G n } 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 p } and { G n } [85]. Besides, we can also incorporate models that are easier to overfit to the first dominant features to better differentiate { G p } from { G n } . 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) D tr that satisfies variation sufficiency (Assumption 3.3) and variation consistency (Assumption 3.5); iii) { G p } and { G n } are distinct subsets of D tr such that I ( G p s ; G n s | Y ) = 0 , ∀ G p s = arg max ̂ G p s I ( ̂ G p s ; Y ) under { G p } , and ∀ G n s = arg max ̂ G n s I ( ̂ G n s ; Y ) under { G n } ; suppose | G c | = s c , ∀ G c , resolving the following GALA objective elicits an invariant GNN defined via Eq. 1,

where ̂ G p c ∈ { ̂ G p c = g ( G p ) } and ̂ G n c ∈ { ̂ G n c = g ( G n ) } are the estimated invariant subgraphs via g from { G p } and { G n } , respectively.

The proof is given in Appendix E.5. Essentially, assumption iii) in Theorem 4.1 is an implication of the variation sufficiency (Assumption 3.3). When given the distinct subsets { G p } and { G n } with different relations of H ( C | Y ) and H ( S | Y ) , since H ( C | Y ) 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-art 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.1.

Two-piece graph datasets. We adopt BA-2motifs [50] to implement 4 variants of 3 -class two-piece graph (Def. 3.1) datasets. The datasets contain different relationships of H ( C | Y ) and H ( S | Y ) by controlling the α and β in the mixed environment, respectively. We consider 4 cases of α -β , ranging from { +0 . 2 , +0 . 1 , -0 . 1 , -0 . 2 } , 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 AI-aided drug affinity prediction. The DrugOOD datasets include splits using Assay, Scaffold, and Size from the EC50 category (denoted as EC50-* ) 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

shifts in node attributes (denoted as CMNIST-sp ). 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 Euclidean regime, including IRMv1 [3], VREx [38], EIIL [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 a 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 CIGAv1. 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-concept study. The results in two-piece graph datasets are reported in Table 2. It can be found that the previous environment augmentation approaches fail either in datasets where the invariant correlations dominate or where the spurious correlations dominate, aligned with our discussions in Sec. 3. In particular, GREA , CIGA and GIL achieve high performance when the invariant correlation dominates, but suffer great performance decrease when the spurious correlations are stronger. Although DisC is expected to succeed when spurious correlations dominate, DisC fails to outperform others because of its excessive destruction of the learned information. MoleOOD also yields degraded performance, which

Table 2: OOD generalization performance under various invariant and spurious correlation degrees in the two-piece graph datasets. Each dataset is generated from a variation of two-piece graph model, denoted as { a, b } , where a refers to the invariant correlation strength and b refers to the spurious correlation strength. The blue shadowed entries are the results with the mean-1*std larger than the mean of the corresponding second best results.

could be caused by the failures to infer reliable environment labels. In contrast, GALA achieves 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.

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, MoleOOD in Ki-Assay, GIL in Graph-SST2, or CIGA in EC50-Size, however, inevitably fail to 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 a 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 .

100

75

50

25

0

25

50

75

100

Correlation degree

{0.8,0.6}

{0.8,0.7}

{0.8,0.9}

{0.7,0.9}

Positive graphs    Negative graphs

inv

spu

Ki-Sca

CMNIST-sp

0

20

40

60

80

100

OOD performance

CIGA

CIGAv2

GALA

GALA+CIGAv2

(b) CIGAv2 compatibility

upsampling

0

1

2

3

4

5

penalty weight

0

0.5

1

2

4

8

16

32

performance

0

10

20

30

40

50

60

70

Correlation strengths of { G p } and { G n } . 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 p } and { G n } , no matter whether the dominant correlation is spurious or not. We measure the invariant and spurious correlation strengths in terms of co-occur probability of the invariant and spurious subgraphs with the labels. The results are 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 { G p } and { G n } in GALA succeeds in identifying the underlying invariant subgraph.

CIGAv2 compatibility. Although GALA focuses on the contrastive term in CIGA, both GALA and CIGA are compatible with the additional CIGAv2 term that facilitates constraining the graph sizes. To verify, we 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 can only achieve a similar result.

Hyperparameter sensitivity. We also test the hyperparameter sensitivity of GALA to the contrastive penalty weights as 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 0 , the performance will decrease a lot, which serves as 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 a retrospective study on the faithfulness of the augmented environment information for OOD generalization on graphs. By showing hardness cases and impossibility results of 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 .

References

I. Ahmad and P.-E. 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)

K. Ahuja, E. Caballero, D. Zhang, J.-C. Gagnon-Audet, Y. Bengio, I. Mitliagkas, and I. Rish. Invariance principle meets information bottleneck for out-of-distribution generalization. In Advances in Neural Information Processing Systems , 2021. (Cited on pages 9 and 32)

M. Arjovsky, L. Bottou, I. Gulrajani, and D. Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv:1907.02893 , 2019. (Cited on pages 8, 9, 30 and 32)

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)

B. Bevilacqua, Y. Zhou, and B. Ribeiro. Size-invariant graph representations for graph classification extrapolations. In International Conference on Machine Learning , pages 837-851, 2021. (Cited on page 21)

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)

Y. Chen, H. Yang, Y. Zhang, K. Ma, T. Liu, B. Han, and J. Cheng. Understanding and improving graph injection attack by promoting unnoticeability. In International Conference on Learning Representations , 2022. (Cited on page 21)

Y. Chen, Y. Zhang, Y. Bian, H. Yang, K. Ma, B. Xie, T. Liu, B. Han, and J. Cheng. Learning causally invariant representations for out-of-distribution generalization on graphs. In Advances in Neural Information Processing Systems , 2022. (Cited on pages 1, 2, 3, 4, 5, 6, 9, 19, 20, 21, 22, 25, 27, 31, 32 and 33)

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)

Y. Chen, W. Huang, K. Zhou, Y. Bian, B. Han, and J. Cheng. Towards understanding feature learning in out-of-distribution generalization. arXiv preprint arXiv:2304.11327 , 2023. (Cited on page 7)

S. Chopra, R. Hadsell, and Y. LeCun. Learning a similarity metric discriminatively, with application to face verification. In 2005 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR 2005), 20-26 June 2005, San Diego, CA, USA , pages 539-546, 2005. (Cited on page 27)

E. Creager, J. Jacobsen, and R. S. Zemel. Environment inference for invariant learning. In International Conference on Machine Learning , pages 2189-2200, 2021. (Cited on pages 4, 21 and 32)

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)

Y. Deng, Y. Yang, B. Mirzasoleiman, and Q. Gu. Robust learning with progressive data expansion against spurious correlation. arXiv preprint , arXiv:2306.04949, 2023. (Cited on pages 4 and 21)

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 of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers) , pages 4171-4186, 2019. (Cited on page 30)

M. Ding, K. Kong, J. Chen, J. Kirchenbauer, M. Goldblum, D. Wipf, F. Huang, and T. Goldstein. A closer look at distribution shifts and out-of-distribution generalization on graphs. In NeurIPS 2021 Workshop on Distribution Shifts: Connecting Methods and Applications , 2021. (Cited on page 1)

S. Fan, X. Wang, Y. Mo, C. Shi, and J. Tang. Debiasing graph neural networks via learning disentangled causal substructure. In Advances in Neural Information Processing Systems , 2022. (Cited on pages 1, 4, 6, 7, 9, 21, 22 and 32)

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)

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)

M. Gardner, J. Grus, M. Neumann, O. Tafjord, P. Dasigi, N. F. Liu, M. E. Peters, M. Schmitz, and L. Zettlemoyer. Allennlp: A deep semantic natural language processing platform. arXiv preprint , arXiv:1803.07640, 2018. (Cited on page 30)

S. Gui, X. Li, L. Wang, and S. Ji. GOOD: A graph out-of-distribution benchmark. In Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track , 2022. (Cited on page 1)

S. Gui, M. Liu, X. Li, Y. Luo, and S. Ji. Joint learning of label and environment causal independence for graph out-of-distribution generalization. arXiv preprint , arXiv:2306.01103, 2023. (Cited on page 21)

I. Gulrajani and D. Lopez-Paz. In search of lost domain generalization. In International Conference on Learning Representations , 2021. (Cited on pages 30 and 32)

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)

W. Hu, M. Fey, M. Zitnik, Y. Dong, H. Ren, B. Liu, M. Catasta, and J. Leskovec. Open graph benchmark: Datasets for machine learning on graphs. In Advances in Neural Information Processing Systems , 2020. (Cited on page 1)

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)

S. Ioffe 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)

Y. Ji, L. Zhang, J. Wu, B. Wu, L.-K. Huang, T. Xu, Y. Rong, L. Li, J. Ren, D. Xue, H. Lai, S. Xu, J. Feng, W. Liu, P. Luo, S. Zhou, J. Huang, P. Zhao, and Y . Bian. DrugOOD: Out-of-Distribution (OOD) Dataset Curator and Benchmark for AI-aided Drug Discovery - A Focus on Affinity Prediction Problems with Noise Annotations. arXiv preprint , arXiv:2201.09637, 2022. (Cited on pages 1, 8 and 31)

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)

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)

B. F. Kamhoua, L. Zhang, Y. Chen, H. Yang, M. KAILI, B. Han, B. Li, and J. Cheng. Exact shape correspondence via 2d graph convolution. In Advances in Neural Information Processing Systems , 2022. (Cited on page 21)

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)

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)

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)

T. N. Kipf and M. Welling. Semi-supervised classification with graph convolutional networks. In International Conference on Learning Representations , 2017. (Cited on page 1)

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)

P. W. Koh, S. Sagawa, H. Marklund, S. M. Xie, M. Zhang, A. Balsubramani, W. Hu, M. Yasunaga, R. L. Phillips, I. Gao, T. Lee, E. David, I. Stavness, W. Guo, B. Earnshaw, I. Haque, S. M. Beery, J. Leskovec, A. Kundaje, E. Pierson, S. Levine, C. Finn, and P. Liang. WILDS: A benchmark of in-the-wild distribution shifts. In International Conference on Machine Learning, , pages 5637-5664, 2021. (Cited on page 1)

D. Krueger, E. Caballero, J. Jacobsen, A. Zhang, J. Binas, D. Zhang, R. L. Priol, and A. C. Courville. Out-of-distribution generalization via risk extrapolation (rex). In International Conference on Machine Learning , pages 5815-5826, 2021. (Cited on pages 9 and 32)

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)

J. Lee, E. Kim, J. Lee, J. Lee, and J. Choo. Learning debiased representation via disentangled feature augmentation. In Advances in Neural Information Processing Systems , 2021. (Cited on pages 4, 6, 7 and 21)

H. Li, Z. Zhang, X. Wang, and W. Zhu. Learning invariant graph representations for outof-distribution generalization. In Advances in Neural Information Processing Systems , 2022. (Cited on pages 1, 2, 4, 5, 9, 21, 22, 32 and 33)

X. Li, S. Gui, Y. Luo, and S. Ji. Graph structure and feature extrapolation for out-of-distribution generalization. arXiv preprint , arXiv:2306.08076, 2023. (Cited on page 21)

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)

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, 7 and 21)

G. Liu, T. Zhao, J. Xu, T. Luo, and M. Jiang. Graph rationalization with environment-based augmentations. arXiv preprint arXiv:2206.02886 , 2022. (Cited on pages 1, 2, 4, 9, 21, 22 and 32)

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 4 and 21)

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 21778-21793, 2023. (Cited on page 21)

Y. Liu, X. Ao, F. Feng, Y. Ma, K. Li, T. Chua, and Q. He. FLOOD: A flexible invariant learning framework for out-of-distribution generalization on graphs. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining , pages 1548-1558, 2023. (Cited on page 21)

A. Lucic, 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-4511, 2022. (Cited on page 21)

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)

K. Ma, G. Yang, H. Yang, Y. Chen, and J. Cheng. Calibrating and improving graph contrastive learning. Transactions on Machine Learning Research , 2023. ISSN 2835-8856. (Cited on page 19)

S. Mahdavi, K. Swersky, T. Kipf, M. Hashemi, C. Thrampoulidis, and R. Liao. Towards better out-of-distribution generalization of neural algorithmic reasoning tasks. arXiv preprint arXiv:2211.00692 , 2022. (Cited on page 21)

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)

D. Mendez, A. Gaulton, A. P. Bento, J. Chambers, M. D. Veij, E. Felix, M. P. Magariños, J. F. Mosquera, P. Mutowo-Meullenet, M. Nowotka, M. Gordillo-Marañón, F. M. I. Hunter, L. Junco, G. Mugumbate, M. Rodríguez-López, F. Atkinson, N. Bosc, C. J. Radoux, A. Segura-Cabrera, A. Hersey, and A. R. Leach. Chembl: towards direct deposition of bioassay data. Nucleic Acids Research , 47(Database-Issue):D930-D940, 2019. (Cited on page 31)

S. Miao, M. Liu, and P. Li. Interpretable and generalizable graph learning via stochastic attention mechanism. arXiv preprint arXiv:2201.12987 , 2022. (Cited on pages 1, 4, 9, 21, 31 and 32)

S. Miao, Y. Luo, M. Liu, and P. Li. Interpretable geometric deep learning via learnable randomness injection. In International Conference on Learning Representations , 2023. (Cited on pages 4 and 21)

C. Murray and D. Rees. The rise of fragment-based drug discovery. Nature chemistry , 1:187-92, 06 2009. (Cited on page 3)

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 7)

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. DeVito, 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 8024-8035, 2019. (Cited on page 33)

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 B (Statistical Methodology) , 78(5):947-1012, 2016. (Cited on pages 1 and 2)

M. Pezeshki, D. Bouchacourt, M. Ibrahim, N. Ballas, P. Vincent, and D. Lopez-Paz. Discovering environments with XRM. arXiv preprint , arXiv:2309.16748, 2023. (Cited on pages 4 and 21)

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)

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)

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)

S. 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 on page 21)

D. Ulyanov, A. Vedaldi, and V. S. Lempitsky. Instance normalization: The missing ingredient for fast stylization. arXiv preprint , arXiv:1607.08022, 2016. (Cited on page 31)

A. van den Oord, Y. Li, and O. Vinyals. Representation learning with contrastive predictive coding. arXiv preprint , arXiv:1807.03748, 2018. (Cited on page 27)

P. Veliˇ ckovi´ c, G. Cucurull, A. Casanova, A. Romero, P. Liò, and Y. Bengio. Graph attention networks. In International Conference on Learning Representations , 2018. (Cited on page 1)

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)

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 , arXiv:2310.03152, 2023. (Cited on page 1)

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)

Q. Wu, H. Zhang, J. Yan, and D. Wipf. Handling distribution shifts on graphs: An invariance perspective. In International Conference on Learning Representations , 2022. (Cited on pages 1, 2, 3, 4, 5, 20 and 21)

Y. Wu, X. Wang, A. Zhang, X. He, and T.-S. Chua. Discovering invariant rationales for graph neural networks. In International Conference on Learning Representations , 2022. (Cited on pages 1, 2, 4, 9, 21 and 22)

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-5458, 2018. (Cited on pages 1 and 31)

K. Xu, W. Hu, J. Leskovec, and S. Jegelka. How powerful are graph neural networks? In International Conference on Learning Representations , 2019. (Cited on pages 1, 9, 20, 22 and 31)

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)

N. Yang, K. Zeng, Q. Wu, X. Jia, and J. Yan. Learning substructure invariance for out-ofdistribution molecular representations. In Advances in Neural Information Processing Systems , 2022. (Cited on pages 1, 2, 4, 5, 9, 21, 22, 32 and 33)

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)

R. Yeung. Information Theory and Network Coding . Springer New York, NY, 01 2008. ISBN 978-0-387-79233-0. (Cited on page 25)

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)

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 4, 9 and 21)

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)

J. Yu, J. Liang, and R. He. Mind the label shift of augmentation-based graph OOD generalization. In IEEE/CVF Conference on Computer Vision and Pattern Recognition , 2023. (Cited on pages 1, 4, 5 and 21)

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)

M. Zhang, N. S. Sohoni, H. R. Zhang, C. Finn, and C. Ré. Correct-n-contrast: A contrastive approach for improving robustness to spurious correlations. arXiv preprint , arXiv:2203.01517, 2022. (Cited on pages 4, 7, 8, 21 and 27)

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. Saxton, H. Ling, H. Lawrence, H. Stärk, S. Gui, C. Edwards, N. Gao, A. Ladera, T. Wu, E. F. Hofgard, A. M. Tehrani, 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)

Y. Zhang, M. Gong, T. Liu, G. Niu, X. Tian, B. Han, B. Schölkopf, and K. Zhang. Adversarial robustness through the lens of causality. In International Conference on Learning Representations , 2022. (Cited on page 1)

J. Zhou, B. Bevilacqua, and B. Ribeiro. An ood multi-task perspective for link prediction with new relation types and nodes. arXiv preprint , arXiv:2307.06046, 2023. (Cited on page 21)

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)

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)

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)

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)

Q. Zhu, Y. Jiao, N. Ponomareva, J. Han, and B. Perozzi. Explaining and adapting graph conditional shift. arXiv preprint , arXiv:2306.03256, 2023. (Cited on page 21)

D. Zou, S. Liu, S. Miao, V. Fung, S. Chang, and P. Li. GDL-DS: A benchmark for geometric deep learning under distribution shifts. arXiv preprint , abs/2310.08677, 2023. (Cited on page 1)

Contents

Appendix of GALA

A Notations

Typically, for graphs that appeared in the discussion, we will use the superscript to denote the sampling process (e.g., G p is the positive graph), and the subscript to denote the specific invariant (i.e., G c ) or spurious subgraph (i.e., G s ). Graph symbols with ̂ G are the predicted graphs of a model (i.e., the estimated invariant subgraph ̂ G c . Below, we list some examples of graphs involved in this paper.

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 a suitable framework that one 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.

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

Besides, a 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 G = ( A,X ) denote a graph with n nodes and m edges, where A ∈ { 0 , 1 } n × n is the adjacency matrix, and X ∈ R n × d is the node feature matrix with a node feature dimension of d . In graph classification, we are given a set of N graphs { G i } N i =1 ⊆ G and their labels

{ Y i } N i =1 ⊆ Y = R c from c classes. Then, we train a GNN ρ ◦ h with an encoder h : G → R h that learns a meaningful representation h G for each graph G to help predict their labels y G = ρ ( h G ) with a downstream classifier ρ : R h → Y . The representation h G is typically obtained by performing pooling with a READOUT function on the learned node representations:

where the READOUT is a permutation invariant function (e.g., SUM, MEAN) [75], and h ( K ) u stands for the node representation of u ∈ V at K -th layer that is obtained by neighbor aggregation:

where N ( u ) is the set of neighbors of node u , σ ( · ) is an activation function, e.g., ReLU, and a ( · ) is an aggregation function over neighbors, e.g., MEAN.

S

C

Z

c

X

Z

c

A

Z

s

X

Z

s

A

G

c

G

s

G

(a) Graph generation SCM

E

S

Y

G

C

(b) FIIF SCM

E

S

Y

G

C

(c) PIIF SCM

E

S

1

C

S

2

Y

G

(d) MIIF SCM

Graph generation process. This work focuses on graph classification, while the results generalize to 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 e } e collected from multiple environments E all. Samples ( G e i , Y e i ) ∈ D e from the same environment are considered as drawn independently from an identical distribution P e . 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 et al. [8]. The generation of the observed graph G and labels Y are controlled by a set of latent causal variable C and spurious variable S , i.e.,

C and S control the generation of G by controlling the underlying invariant subgraph G c and spurious subgraph G s , respectively. Since S can be affected by the environment E , the correlation between Y , S and G s can change arbitrarily when the environment changes. C and S control the generation of the underlying invariant subgraph G c and spurious subgraph G s , respectively. Since S can be affected by the environment E , the correlation between Y , S and G s 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 | C and Partially Informative Invariant Features ( PIIF ) when Y ̸⊥ ⊥ S | C . Furthermore, PIIF and FIIF shifts can be mixed together and yield Mixed Informative Invariant Features ( MIIF ), as shown in Fig. 4. We refer interested readers to Chen et al. [8] for a 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 G c to predict the label Y [8, 72]. Specifically, the goal of OOD generalization on graphs is to learn an invariant GNN f := f c ◦ g , which is composed of two modules: a) a featurizer g : G → G c that extracts the invariant subgraph G c ; b) a classifier f c : G c →Y that predicts the label Y based on the extracted G c , where G c refers to the space of subgraphs of G . The learning objectives of f c and g are formulated as

Since E is not observed, many strategies are proposed to impose the independence of ̂ G c and E . A common approach is to augment the environment information. For example, based on the

estimated invariant subgraphs ̂ G c and spurious subgraphs ̂ G s , Liu et al. [45], Wu et al. [72, 73] proposed to 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 al. [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 H ( S | Y ) < H ( C | Y ) . Chen et al. [8] proposed a contrastive framework to tackle both FIIF and PIFF graph shifts, but limited to H ( S | Y ) > H ( C | Y ) . However, in practice it is usually unknown whether H ( S | Y ) < H ( C | Y ) or H ( S | Y ) > H ( C | Y ) without environment information.

More OOD generalization on graphs. In addition to the aforementioned invariant learning approaches, Bevilacqua et al. [5], Li et al. [42], Yehudai et al. [78], Zhou et al. [89] study the OOD generalization as 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 et 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.Gao 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 et al. [46] proposed a self-boosting 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 is 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 augmentation 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 to the two-piece graph generation models. The specific description is given as the following.

Definition D.1 ( 3 -class two-piece graphs) . Each environment is defined with two parameters, α e , β e ∈ [0 , 1] , and the dataset D e is generated as follows:

(a) Sample y e ∈ { 0 , 1 , 2 } uniformly;

(b) Generate G c and G s via :

0.5

0.6

0.7

0.8

0.9

spurious correlation strength

30

40

50

60

70

80

OOD test acc

ERM

DIR

GALA

(a) Failures of env. generation

0.5

0.6

0.7

0.8

0.9

spurious correlation strength

30

40

50

60

70

80

OOD test acc

ERM

MoleOOD

GALA

(b) Failures of env. inferring

0.5

0.6

0.7

0.8

0.9

spurious correlation strength

30

40

50

60

70

80

OOD test acc

ERM

DisC

GALA

(c) Failures of resolving env. consis-

tency

0.5

0.6

0.7

0.8

0.9

spurious correlation strength

40

45

50

55

60

65

70

75

80

OOD test acc

XGNN

GREA

GALA

(d) Failures of env. generation

0.5

0.6

0.7

0.8

0.9

spurious correlation strength

40

45

50

55

60

65

70

75

80

OOD test acc

XGNN

GIL

GALA

(e) Failures of env. inferring

0.5

0.6

0.7

0.8

0.9

spurious correlation strength

40

45

50

55

60

65

70

75

80

OOD test acc

XGNN

CIGA

GALA

(f) Failures of resolving env. consis-

tency

where f G c gen , f G s gen respectively map input { 0 , 1 , 2 } to a specific graph selected from a given set, and Rad ( α ) is a random variable with probability α taking a uniformly random value from { 0 , 1 , 2 } , and a probability of 1 -α taking the value of +1 ;

(c) Sythesize G by randomly concatenating G c and G s :

In experiments, we implement the 3 -class two-piece graphs with the BA-motifs [50] model.

In experiments, we adopt a 3 -layer GIN [75] with a hidden dimension of 32 and a 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 a learning rate of 1 e -3 . All experiments are repeated with 5 different random seeds of { 1 , 2 , 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 1 e -2 for DIR following its original experiment in spurious motif datasets generated similarly using BA-motifs [73]. We use a penalty weight of 1 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 3 . We tune the penalty weights of MoleOOD with values from { 1 e -2 , 1 e -1 , 1 , 10 } but did not observe much performance differences. We tune the penalty weights of GIL with values from { 1 e -5 , 1 e -3 , 1 e -1 } recommended by the authors. For DisC, we tune only the q 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 } as recommended by the authors.

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 E tr = { ( α, β 1 ) , ( α, β 2 ) } with α ≥ β 1 , β 2 (e.g., E tr = { (0 . 25 , 0 . 1) , (0 . 25 , 0 . 2) } ), and its corresponding mixed environment E mix tr = { ( α, ( β 1 + β 2 ) / 2 } (e.g., E mix tr = { (0 . 25 , 0 . 15) } ). When ̂ G c = G s and ̂ G s = G c , it holds that the augmented environment E v is also a two-piece graph dataset with

Proof. From Definition 3.1, we known that for each graph G i ∼ E mix tr = { ( α, ( β 1 + β 2 ) / 2) } , G i is the concatenation of the G i c and G i s defined as

where Rad ( · ) i denotes the i th sample of the random variable Rad ( · ) .

Denote and

Considering applying the augmentation to 2 n samples randomly sampled from E mix tr , since the featurizer g separates each G ∈ E mix tr into ̂ G c = G s and ̂ G s = G c , and the augmented graph G i is obtained by

Then, the new α v , β v in E v can be obtained by summing up the overall numbers of G A , G B , G C , G D concatenated into 2 n 2 samples in E v .

Specifically, we can inspect the changes of the distributions of motifs and labels. Let ¯ β = ( β 1 + β 2 ) / 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 as follows:

Then, new distributions of the motifs and labels are determined by the number of original motifs identified as ̂ G c and ̂ G s , respectively. When ̂ G c = G s and ̂ G s = G c , in the new environment E v , given Y = +1 , G C contributes (1 -¯ β ) n ∗ 2 n samples as the 'invariant' subgraph. More specifically, G C will be concatenated with G A and G B by n times, respectively. Then we have the new distribution tables shown as follows:

Since given the same Y , the spurious subgraph G C and G D will still have the same chance being flipped, we have β v = ¯ β . While as G A and G B appear the same times given the same Y , it suffices to know that α v = 0 . 5 .

E.2 Complementary discussion for Sec. 3.1

Proposition E.2. Given the same graph generation process as in Fig. 2, when there exists spurious subgraph G s such that P e 1 ( Y | G s ) = P e 2 ( Y | G s ) for any two environments e 1 , e 2 ∈ E tr, where P e ( Y | G s ) is the conditional distribution P ( Y | G s ) under environment e ∈ E all, it is impossible for any learning algorithm applied to f c ◦ g to differentiate G c from G s .

Proof. Let G ∗ s be the spurious subgraph such that P e 1 ( Y | G s ) = P e 2 ( Y | G s ) for any two environments e 1 , e 2 ∈ E tr, and G c be the invariant subgraph which P e 1 ( Y | G c ) = P e 2 ( Y | G c ) , ∀ e 1 , e 2 ∈ E tr by definition. Consider a learning algorithm applied to f c ◦ g that accepts the input of E mix tr , and extracts a subgraph ̂ G c = g ( Y ) as an estimation of the invariant subgraph for any G to predict Y via f c ( ̂ G c ) in a deterministic manner. If the algorithm succeed to extract G c from E mix tr , then there always exists a E mix tr ′ with the desired spurious subgraph G ′ s and a underlying invariant subgraph G ′ c , such that G ′ s = G c and G ′ c = G ∗ s . Due to the deterministic nature, the algorithm fails to identify G ′ c in E mix tr ′ .

E.3 Proof of Proposition 3.4

Proposition E.3. (Restatement of Proposition 3.4) There exist 2 two-piece graph training environments E tr and E tr ′ that share the same joint distribution P ( Y, G ) . Any learning algorithm will fail in either E tr or E tr ′ .

Proof. Let the mixed training environment of E tr and E tr ′ be E mix tr = { ( α, β ) } . Based on the definition of two-piece graphs (Definition 3.1), the joint distribution of the mixed training dataset ( G = Concat [ G c , G s ] , Y ) can be computed as

̸

̸

Here we use Bit G c ( G c ) to obtain the input bit of a subgraph G c (or ( f G c gen ) -1 ), and Bit G s ( G s ) for G s , respectively.

Any learning algorithm that tries to identify the invariant subgraph from this training dataset will compute a model that uses subgraph G c , or subgraph G s , or both G c and G s to predict Y deterministically. 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 c as the invariant subgraph for E tr = { ( α, β 1 ) , ( α, β 2 ) } with β = ( β 1 + β 2 ) / 2 .

Now let the other training environment be E tr ′ = { ( α 1 , β ) , ( α 2 , β ) } with α = ( α 1 + α 2 ) / 2 . It is clear that since the mixed training environment of E tr ′ is still { ( α, β ) } , the model keeps regarding G c as the invariant subgraph. However, for E tr ′ , the model fails to identify the invariance since now the invariant subgraph is G s .

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 a 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 f c ◦ g that accepts the input of E mix tr , and extracts a subgraph ̂ G c = g ( Y ) as an estimation of the invariant subgraph for any G to predict Y via f c ( ̂ G c ) in a deterministic manner. Without the holding of Assumption 3.5, due to Proposition 3.4, there exists E mix tr ′ for each E mix tr that have the identical joint distribution but different underlying invariant subgraph. Thus, any learning algorithm that succeeds in either E mix tr or E mix tr ′ 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; ii) D tr that satisfies variation sufficiency (Assumption 3.3) and variation consistency

̸

(Assumption 3.5); iii) { G p } and { G n } are distinct subsets of D tr such that I ( G p s ; G n s | Y ) = 0 , ∀ G p s = arg max ̂ G p s I ( ̂ G p s ; Y ) under { G p } , and ∀ G n s = arg max ̂ G n s I ( ̂ G n s ; Y ) under { G n } ; suppose | G c | = s c , ∀ G c , resolving the following GALA objective elicits an invariant GNN defined via Eq. 12,

where ̂ G p c ∈ { ̂ G p c = g ( G p ) } and ̂ G n c ∈ { ̂ G n c = g ( G n ) } are the estimated invariant subgraphs via g from { G p } and { G n } , respectively.

Proof. Without loss of generality, we assume that { G p } has the same spurious dominance situation as E tr . In other words, when H ( S | Y ) < H ( C | Y ) , the data distribution in { G p } also follows H ( S | Y ) < H ( C | Y ) , while H ( S | Y ) > H ( C | Y ) in { G n } . To proceed, we will use the language of Chen et al. [8].

We begin by discussing the case of H ( S | Y ) < H ( C | Y ) . Given H ( S | Y ) < H ( C | Y ) , we have H ( S | Y ) < H ( C | Y ) in { G p } and H ( S | Y ) > H ( C | Y ) in { G n } . Then, we claim that

Otherwise, consider there exists a subgraph of the spurious subgraph △ ̂ G p s ⊆ G p s in ̂ G p c , which takes up the space of △ ̂ G p c ⊆ G p c from ̂ G p c . Then, let △ ̂ G p c = G p c -△ ̂ G p c we can inspect the changes to I ( ̂ G p c ; ̂ G n c | Y ) led by △ ̂ G p s :

where the last equality is obtained via expanding the conditional entropy. Then, considering the contents in ̂ G n c , without loss of generality, we can divide all of the possible cases into two:

(i) ̂ G n c contains only the corresponding invariant subgraph G n c ;

(ii) ̂ G n c contains subgraph from the corresponding spurious subgraph G n s , denoted as △ ̂ G n s ⊆ G n s ;

For case (i), it is easy to write Eq. 15 as:

since H ( △ ̂ G p s |△ ̂ G p c , Y ) = H ( △ ̂ G p s | ̂ G n c , △ ̂ G p c , Y ) = H ( △ ̂ G p s | Y ) given C ⊥ ⊥ S | Y for PIIF shifts. Then, it suffices to know that △ I ( ̂ G p c ; ̂ G n c | Y ) ≤ 0 as conditioning on new variables will not increase the entropy [79].

For case (ii), we have :

where we claim that H ( △ ̂ G p s |△ ̂ G p c , Y ) -H ( △ ̂ G p s |△ ̂ G p c , ̂ G n c , Y ) = 0 , and similarly conclude that △ I ( ̂ G p c ; ̂ G n c | Y ) ≤ 0 . More specifically, we can rewrite the first term in Eq. 17 as

using the variation condition (i.e., assumption iii)) for △ ̂ G p s under { G p } , and △ ̂ G n s under { G n } . After showing the success of GALA in tackling H ( S | Y ) < H ( C | Y ) , it also suffices to know that the aforementioned discussion also generalizes to the other case, i.e., when H ( S | Y ) > H ( C | Y ) in { G p } and H ( S | Y ) < H ( C | Y ) in { G n } .

F More Discussions on Practical Implementations of GALA

In this section, we provide more implementation discussions about GALA in complementary to Sec. 4.

Objective implementation. As the estimation of mutual information could be highly expensive [4, 67], inspired by Chen et al. [8], we adopt the contrastive learning to approximates the mutual information between subgraphs in Eq. 9 [4, 11, 33, 62, 67]:

where ( ̂ G p c , ̂ G n c ) are subgraphs extracted by g from { G p } , { G n } that share the same label, respectively. { G i c } i M =1 are subgraphs extracted by g from G that has a different label. P g ( G |Y = Y ) is the pushforward distribution of P ( G |Y = Y ) by featurizer g , P ( G |Y = Y ) refers to the distribution of G given the label Y , P ( G |Y ̸ = Y ) refers to the distribution of G given the label that is different from Y , ̂ G c = g ( ̂ G ) , ̂ G c = g ( ̂ G ) , G i c = g ( G i ) are the estimated subgraphs, h ̂ G p c , h ̂ G n c , h G i c are the graph presentations of the extracted subgraphs. ϕ is a similarity measure. As M →∞ , Eq. 18 approximates I ( ̂ G p c ; ̂ G n c | Y ) [1, 32, 69].

Environment assistant implementation. Theorem 4.1 shows the effectiveness of GALA when given proper subsets of { G p } and { G n } . In practice, we can implement the environment assistant into 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 H ( S | Y ) < H ( C | Y ) , ERM will firstly learn to extract spurious subgraphs G s to make predictions. Therefore, we can obtain { G p } by finding samples where ERM correctly predicts the labels, while { G n } 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.

Algorithm 2 GALA : Clustering based G raph inv A riant L earning A ssistant

1: Input: Training data D tr ; environment assistant A ; featurizer g ; classifier f c ; length of maximum training epochs e ; batch size b ;

2: Initialize environment assistant A ;

3: for p ∈ [1 , . . . , e ] do

4: Sample a batch of data { G i , Y i } b i =1 from D tr ;

5: Obtain Environment Assistant predictions { ˆ c e i } b i =1 using k -means clustering on the graph representations yielded by A ;

6: for each sample G i , y i ∈ { G i , Y i } b i =1 do

8: Find negative graphs with different y i but same environment assistant prediction ˆ c e i ;

7: Find postive graphs with same y i and different ˆ c e i ;

9: Calculate GALA risk via Eq. 18;

10: Update f c , g via gradients from GALA risk;

11: end for

12: end for

13: return final model f c ◦ g ;

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 the first dominant features to better differentiate { G p } from { G n } . 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 a CIGAv1 penalty of 4 .

5.0

2.5

0.0

2.5

5.0

7.5

10.0

12.5

2

0

2

4

6

8

10

12

colored by group labels

5.0

2.5

0.0

2.5

5.0

7.5

10.0

12.5

15.0

2

0

2

4

6

8

10

12

colored by prediction

5.0

2.5

0.0

2.5

5.0

7.5

10.0

12.5

2

0

2

4

6

8

10

12

colored by cluster prediction

(a) Colored by environment labels.

(b) Colored by label predictions.

(c) Colored by cluster predictions.

Figure 6: Umap visualizations of learned graph representations in ERM trained environment assistant model based on the 3 -class two-piece graph { 0 . 7 , 0 . 9 } .

5

0

5

10

15

5.0

2.5

0.0

2.5

5.0

7.5

10.0

12.5

15.0

colored by group labels

5

0

5

10

15

5.0

2.5

0.0

2.5

5.0

7.5

10.0

12.5

15.0

colored by prediction

5

0

5

10

15

5.0

2.5

0.0

2.5

5.0

7.5

10.0

12.5

15.0

colored by cluster prediction

(a) Colored by environment labels.

(b) Colored by label predictions.

(c) Colored by cluster predictions.

5

0

5

10

15

20

5

0

5

10

15

colored by group labels

5

0

5

10

15

20

5

0

5

10

15

colored by prediction

5

0

5

10

15

20

5

0

5

10

15

colored by cluster prediction

(a) Colored by environment labels.

(b) Colored by label predictions.

(c) Colored by cluster predictions.

10

0

10

20

30

15

10

5

0

5

10

15

colored by group labels

10

0

10

20

30

15

10

5

0

5

10

15

colored by prediction

10

0

10

20

30

10

5

0

5

10

15

20

25

colored by cluster prediction

(a) Colored by environment labels.

(b) Colored by label predictions.

(c) Colored by cluster predictions.

20

10

0

10

20

30

10

5

0

5

10

15

20

25

colored by group labels

20

10

0

10

20

30

10

5

0

5

10

15

20

25

colored by prediction

20

10

0

10

20

30

10

5

0

5

10

15

20

25

colored by cluster prediction

20

10

0

10

20

10

5

0

5

10

15

20

colored by group labels

20

10

0

10

20

10

5

0

5

10

15

20

colored by prediction

20

10

0

10

20

10

5

0

5

10

15

20

colored by cluster prediction

(a) Colored by environment labels.

(b) Colored by label predictions.

(c) Colored by cluster predictions.

In Fig. 7 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 a interpretable GNN as well as a CIGAv1 penalty, which facilitates the overfitting to the spurious correlations, then the vanilla label predictions can also yield a 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. 9 to Fig. 11, the learned representations have poor quality for approximating the underlying environment labels. Empirically, we 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 as anchor points. However, when it is used to contrast samples from ̂ G p c and ˜ G n c , there could be undesired behaviors. First, it can often happen that there are few to no negative cases when the spurious correlations are too strong. The samples from { G p } in a batch may pull the representations of samples from { G n } to even closer, which makes the model further overfitted to the spurious correlations. Second, the sampling over ̂ G p c and ˜ G n c , can be seen as hard positive and negative samples, that may impose a 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 as anchor points. We empirically observe one-side sampling could yield better performance in two-piece graphs.

Upsampling of 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-dominated 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 to trivial solutions. Therefore, we propose a simple strategy to mitigate the issue. We directly upsample the minority group samples. The minority group of samples will be repeated k 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 as well as the evaluation protocols.

G.1 Datasets

We provide more details about the motivation and construction method of the datasets that are used in our experiments. Statistics of the datasets are presented in Table 5.

Two-piece graph datasets. We construct 3-class synthetic datasets based on BAMotif [50] following Def. D.1, 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.1. While for the generation of validation sets, we use a b v = max(1 / 3 , b -0 . 2) that facilitates the model selection for OOD generalization [9, 23]. While for the generation of test datasets, we merely use a b = 0 . 33 that contains no distribution shifts, to fully examine to what extent the model learns the invariant correlations. During the construction, we 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 ColoredMnist 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 -4 are assigned to y = 0 and those with digits 5 -9 are assigned to y = 1 . Then, y will be flipped with a probability of 0 . 25 . Thirdly, green and red colors will be respectively assigned to images with labels 0 and 1 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 a 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

50 -th percentile while smaller than 80 -th percentile are assigned to the validation set, and the left are assigned to test set.

DrugOOD datasets. To evaluate the OOD performance in realistic scenarios with realistic distribution shifts, we also include three datasets from DrugOOD benchmark [28]. DrugOOD is a systematic OOD benchmark for AI-aided drug discovery, focusing on the task of drug target binding affinity prediction for both macromolecule (protein target) and smallmolecule (drug compound). The molecule data and the notations are curated from realistic ChEMBL database [54]. Complicated distribution shifts can happen on different assays, scaffolds and molecule sizes. In particular, we select DrugOOD-lbap-core-ec50-assay , DrugOOD-lbap-core-ec50-scaffold , DrugOOD-lbap-core-ec50-size , DrugOOD-lbap-core-ki-assay , DrugOOD-lbap-core-ki-scaffold , and DrugOOD-lbap-core-ki-size , 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. 4 For more details, we refer interested readers to Ji et al. [28].

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 use 3 -layer GIN [75] with Batch Normalization [27] between layers and JK residual connections at the last layer [74]. The hidden dimension is set to 32 for Two-piece graphs, CMNIST-sp, and 128 for SST2, and DrugOOD datasets. The pooling is by default a 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 as the basic backbone model for the whole predictor f = f c ◦ g , where g : G → G c is a featurizer GNN and f c : G c → Y is a classifier GNN. g first calculates the sampling weights as in ̂ G c for each edge. More formally, given a graph G containing n nodes, a soft mask is predicted through the following equation:

where a calculates the sampling weights for each edge using a MLP: M ij = MLP ([ Z i , Z j ]) . 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 use the score to reweight the messaging passing process along the edge, as we empirically find it yields more stable performance. While for CMNIST-sp and Graph-SST2, we 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 a layernorm [66] following the practice of [55].

Besides, we also have various implementation options for obtaining the features in ̂ G c , for further obtaining h ̂ G c , as well as for obtaining predictions based on ̂ G s . 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 G based on ̂ G s , we use a separate MLP downstream classifier in the classifier GNN f c .

Optimization and model selection. By default, we use Adam optimizer [34] with a learning rate of 1 e -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].

The final model is selected according to the performance at the validation set. All experiments are repeated with 5 different random seeds of { 1 , 2 , 3 , 4 , 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 is not available, we perform random partitions on the training data to obtain two equally large environments for these objectives following previous works [8, 12]. Moreover, we select the weights for the corresponding regularization from { 0 . 01 , 0 . 1 , 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], we use the author-released 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], MoleOOD [77], GIL [41], DisC [17], and CIGA [8], according to the author provided codes (if available).

GREA [45]: We use a penalty weight of 1 for GREA as 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 implementations by the authors.

-Interpretable ratio: same as others;

-Penalty weight: { 0 . 1 , 0 . 5 , 1 . 0 } ;;

-Number of environments: N/A;

MoleOOD [77]: We tune the penalty weights of MoleOOD with values from { 1 e -2 , 1 e -1 , 1 , 10 } but did not observe much performance differences. Hence we stick the penalty weight as 1 for all datasets.

-Interpretable ratio: N/A;

-Penalty weight: 1 ;

-Number of environments: same as others;

GIL [41]: We follow the recommendations of the paper.

-Interpretable ratio: same as others;

-Penalty weight: { 1 e -5 , 1 e -3 , 1 e -1 } ;

-Number of environments: same as others;

DisC [17]: We tune only the q 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 as others;

-q weight: { 0 . 9 , 0 . 7 , 0 . 5 } ;

-Number of environments: same as others;

CIGA [8]: We follow the recommendations of the released implementations by the authors..

-Interpretable ratio: same as others;

-Penalty weight: { 0 . 5 , 1 , 2 , 4 , 8 , 16 , 32 } ;

-Number of environments: N/A;

• GALA:

-Interpretable ratio: same as others;

-Penalty weight: { 0 . 5 , 1 , 2 , 4 , 8 , 16 , 32 } ;

-Environment assistant: { vanilla GNN , XGNN } ;

-Sampling proxy: { label predictions , cluster predictions } ;

-Number of environments: same as others;

All of the graph learning methods adopt an interpretable GNN as the backbone by default. The only exception is MoleOOD, we follow the original implementation while using a shared GNN encoder for the variational losses to 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 a ratio of 25% for the interpretable GNN backbone.

For environment inferring methods, we search the number of environments

Two-piece graphs: fixed as 3 (since there are 3 spurious graphs);

CMNIST-sp: 2 (since there are 2 environments);

Graph-SST2: { 2 , 3 , 4 } following previous practice [41];

DrugOOD datasets: { 2 , 3 , 5 , 10 , 20 } following previous practice [77].;

Implementations of GALA. For a fair comparison, GALA uses the same GNN architecture for GNN encoders as the baseline methods. By default, we fix the temperature to be 1 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 use 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 is 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 it 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, EC50-Scaffold, EC50-Size, Ki-Assay, Ki-Scaffold, CMNIST-sp and Graph-SST2, we find implementing the environment assistant as a ERM model already yield impressive improvements. While for the other DrugOOD datasets, we implement the environment assistant as an interpretable GNN trained with ERM and cluster the learned graph representations of the model to sample positive and negative pairs.

Since GALA imposes a strong regularization to the data that may hinder the learning of graph representations, we 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 PyTorch [59] and PyTorch Geometric [18]. We ran our experiments on Linux Servers installed with V100 graphics cards and CUDA 10.2.

G.4 Computational analysis

We calculate the average total training time of different methods at various datasets in 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 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).

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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\nWEAKNESSES:\n1.\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\nQUESTIONS:\n1. 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\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

comments: 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.

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.

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

W2. 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)?

W3. I don't understand why ERM still learns the spurious correlations in invariance dominated cases. Can this be justified theoretically?

QUESTIONS:
Q1. I am confused about how to get the Oracle performance on these datasets. Can the authors provide more explanations?

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

LIMITATIONS:
The authors didn't discuss the limitations of the proposed method.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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\nWEAKNESSES:\n- 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\nQUESTIONS:\nI 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\nLIMITATIONS:\n-", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nInvariant 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\nSTRENGTHS:\n1. 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).\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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\nLIMITATIONS:\nThe authors discussed the broader impacts in the appendix.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

comments: 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.

**Post Rebuttal**: I have read the rebuttal of the authors and will keep my rating.

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.

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

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07